from sqlalchemy.ext.automap import automap_base from sqlalchemy.orm import Session from sqlalchemy import create_engine, or_, func import requests from requests.adapters import HTTPAdapter import json from datetime import datetime import logging import argparse import os import signal from threading import Timer logging.basicConfig(level=logging.INFO) def handler(signum, frame): '''Contrl-C handler.''' logging.info("Ctrl-C pressed, exit.") os._exit(0) signal.signal(signal.SIGINT, handler) class EnvDefault(argparse.Action): '''args priority: cli args -> ENV -> default.''' def __init__(self, envvar, required=True, default=None, **kwargs): if envvar in os.environ: default = os.environ[envvar] if required and default: required = False super(EnvDefault, self).__init__(default=default, required=required, **kwargs) def __call__(self, parser, namespace, values, option_string=None): setattr(namespace, self.dest, values) parser = argparse.ArgumentParser(description='Usage of address fixer.') parser.add_argument("-u", "--user", required=True, type=str, action=EnvDefault, envvar="DB_USER", help="db user name(DB_USER).") parser.add_argument("-p", "--password", required=True, type=str, action=EnvDefault, envvar="DB_PASSWD", help="db password(DB_PASSWD).") parser.add_argument("-H", "--host", required=True, type=str, action=EnvDefault, envvar="DB_HOST", help="db host name or ip address(DB_HOST).") parser.add_argument("-P", "--port", required=True, type=str, action=EnvDefault, envvar="DB_PORT", help="db port(DB_PORT).") parser.add_argument("-d", "--dbname", required=True, type=str, action=EnvDefault, envvar="DB_NAME", help="db name(DB_NAME).") parser.add_argument("-b", "--batch", required=False, type=int, default=10, action=EnvDefault, envvar="BATCH", help="batch size for one loop(BATCH).") parser.add_argument("-t", "--timeout", required=False, type=int, default=5, action=EnvDefault, envvar="HTTP_TIMEOUT", help="http request timeout(s)(HTTP_TIMEOUT).") parser.add_argument("-r", "--retry", required=False, type=int, default=5, action=EnvDefault, envvar="HTTP_RETRY", help="http request max retries(HTTP_RETRY).") parser.add_argument( "-i", "--interval", required=False, type=int, default=0, action=EnvDefault, envvar="INTERVAL", help= "if value not 0, run in infinity mode, fix record in every interval seconds(INTERVAL)." ) parser.add_argument( "-m", "--mode", required=False, type=int, default=0, action=EnvDefault, envvar="MODE", help= "run mode: 0 -> fix empty record; 1 -> update address by amap; 2 -> do both(MODE)." ) parser.add_argument("-k", "--key", required=False, type=str, default='', action=EnvDefault, envvar="KEY", help="API key for calling amap(KEY).") parser.add_argument("-s", "--since", required=False, type=lambda d: datetime.strptime(d, '%Y-%m-%d'), default=datetime.min, action=EnvDefault, envvar="SINCE", help="Update from specified date(YYYY-mm-dd).") parser.add_argument( "-ua", "--user_agent", required=False, type=str, default='teslamate/#v1.29.2', action=EnvDefault, envvar="USER_AGENT", help="Custom User-Agent for HTTP requests(USER_AGENT)." ) args = parser.parse_args() def custom_json_dumps(d): '''do not add backslash in json.''' return d conn_str = "postgresql://%s:%s@%s:%s/%s" % (args.user, args.password, args.host, args.port, args.dbname) engine = create_engine(conn_str, json_serializer=custom_json_dumps, echo=False) # open street map api. osm_resolve_url = "https://nominatim.openstreetmap.org/reverse?lat=%.6f&lon=%.6f&format=jsonv2&addressdetails=1&extratags=1&namedetails=1&zoom=18" # amap api. amap_coordinate_transformation_url = "https://restapi.amap.com/v3/assistant/coordinate/convert?key=%s&coordsys=gps&output=json&locations=%s,%s" amap_resolve_url = "https://restapi.amap.com/v3/geocode/regeo?key=%s&output=json&location=%s,%s&poitype=all&extensions=all" # last updated record id. last_update_id = 0 # reflact Objects from db tables. Base = automap_base() Base.prepare(autoload_with=engine) Drives = Base.classes.drives ChargingProcesses = Base.classes.charging_processes Positions = Base.classes.positions Addresses = Base.classes.addresses # reference to teslamate's source code, get address value from multiple keys. house_number_aliases = ['house_number', 'street_number'] road_aliases = [ "road", "footway", "street", "street_name", "residential", "path", "pedestrian", "road_reference", "road_reference_intl", "square", "place" ] neighborhood_aliases = [ "neighbourhood", "suburb", "city_district", "district", "quarter", "borough", "city_block", "residential", "commercial", "houses", "subdistrict", "subdivision", "ward" ] municipality_aliases = [ "municipality", "local_administrative_area", "subcounty" ] village_aliases = ["village", "municipality", "hamlet", "locality", "croft"] city_aliases = ["city", "town", "township"] city_aliases.extend(village_aliases) city_aliases.extend(municipality_aliases) county_aliases = ["county", "county_code", "department"] state_aliases = ['state', 'province', 'state_code'] country_aliases = ['country', 'country_name'] def get_address_str(address, addr_keys): '''get address value from multiple keys.''' for addr_key in addr_keys: if addr_key in address: return address[addr_key] return None def get_address_name(address): ''' address names comes from multiple places. 1. address.name. 2. address.namedetails.name. 3. address.namedetails.alt_name. 4. first element in address.display_name. ''' name = '' if 'name' in address.keys() and len(address['name']): name = address['name'] if 'namedetails' in address.keys() and address['namedetails'] is not None: if 'name' in address['namedetails'].keys(): name = address['namedetails']['name'] if 'alt_name' in address['namedetails'].keys(): name = address['namedetails']['alt_name'] if len(name) == 0: name = address['display_name'].split(',')[0] return name def get_position(session, position_id): '''get position id from table positions by position_ids.''' position = session.query(Positions).filter( Positions.id == position_id).first() # position_id is foreign key to table positions. position will never be None. if position == None: # fatal error, exit now. logging.fatal("Position with ID %s is not found." % position_id) assert (False) return position def http_request(url): '''get response by calling map api.''' http_session = requests.Session() http_session.mount('http://', HTTPAdapter(max_retries=args.retry)) http_session.mount('https://', HTTPAdapter(max_retries=args.retry)) headers = { 'User-Agent': args.user_agent } try: response = http_session.get(url=url, timeout=args.timeout, headers=headers) if response.status_code != requests.codes.ok: logging.error( "Http request failed by url: %s, code: %d, body: %s" % (url, response.status_code, response.text)) return None raw = response.text return raw except: logging.error("Http request exception by url: %s" % (url)) return None def get_address_in_db(session, osm_id): '''select address from db, get address id which just added.''' return session.query(Addresses).filter(Addresses.osm_id == osm_id).first() def add_osm_address(session, osm_address, raw): '''add osm address to db.''' exist_address = get_address_in_db(session, osm_address['osm_id']) if exist_address is None: logging.info("oms id = %d is not exist!" % osm_address['osm_id']) if exist_address is None: address = Addresses( display_name=osm_address['display_name'], latitude=osm_address['lat'], longitude=osm_address['lon'], name=get_address_name(osm_address), house_number=get_address_str(osm_address['address'], house_number_aliases), road=get_address_str(osm_address['address'], road_aliases), neighbourhood=get_address_str(osm_address['address'], neighborhood_aliases), city=get_address_str(osm_address['address'], city_aliases), county=get_address_str(osm_address['address'], county_aliases), postcode=get_address_str(osm_address['address'], ['postcode']), state=get_address_str(osm_address['address'], state_aliases), state_district=get_address_str(osm_address['address'], ['state_district']), country=get_address_str(osm_address['address'], country_aliases), raw=raw, inserted_at=datetime.now().replace(microsecond=0), updated_at=datetime.now().replace(microsecond=0), osm_id=osm_address['osm_id'], osm_type=osm_address['osm_type']) session.add(address) logging.info("address added: %s." % osm_address['display_name']) else: logging.info("address is already exist: %d, %s." % (osm_address['osm_id'], osm_address['display_name'])) def get_address(session, position): ''' return address id and display_name by position id. Address will add into db if not exists. ''' url = osm_resolve_url % (position.latitude, position.longitude) raw = http_request(url) if raw is None: return None, None osm_address = json.loads(raw) if osm_address == None: return None, None add_osm_address(session, osm_address, raw) added_address = get_address_in_db(session, osm_address['osm_id']) return added_address.id, added_address.display_name def fix_address(session, batch_size, empty_count): processed_count = 0 # get empty records in drives. empty_drive_addresses = session\ .query(Drives)\ .filter(or_(Drives.start_address_id.is_(None), Drives.end_address_id.is_(None)))\ .filter(Drives.start_position_id.is_not(None))\ .filter(Drives.end_position_id.is_not(None))\ .limit(batch_size)\ .all() # get empty records in charging_processes, all records are LE batch_size. empty_charging_addresses = [] if len(empty_drive_addresses) < batch_size: empty_charging_addresses = session\ .query(ChargingProcesses)\ .filter(ChargingProcesses.address_id.is_(None))\ .filter(ChargingProcesses.position_id.is_not(None))\ .limit(batch_size - len(empty_drive_addresses))\ .all() # processing drives. for empty_drive_address in empty_drive_addresses: logging.info("processing drive address (%d left)" % (empty_count - processed_count)) # get positions. start_position_id = empty_drive_address.start_position_id end_position_id = empty_drive_address.end_position_id start_position = get_position(session, start_position_id) end_position = get_position(session, end_position_id) # get addresses. start_address_id, start_address = get_address(session, start_position) end_address_id, end_address = get_address(session, end_position) if start_address_id is None or end_address_id is None: continue # update address ids. empty_drive_address.start_address_id = start_address_id empty_drive_address.end_address_id = end_address_id logging.info("Changing drives(id = %d) start address to %s" % (empty_drive_address.id, start_address)) logging.info("Changing drives(id = %d) end address to %s" % (empty_drive_address.id, end_address)) processed_count += 1 # processing charging. for empty_charging_address in empty_charging_addresses: logging.info("processing charging address (%d left)" % (empty_count - processed_count)) # get position. position_id = empty_charging_address.position_id position = get_position(session, position_id) # get address. address_id, address = get_address(session, position) if address_id is None: continue # update address id. empty_charging_address.address_id = address_id logging.info("Changing charging(id = %d) to %s" % (empty_charging_address.id, address)) processed_count += 1 # records processed. return processed_count def get_empty_record_count(session): '''get all empty records count.''' empty_count = session\ .query(Drives.id)\ .filter(or_(Drives.start_address_id.is_(None), Drives.end_address_id.is_(None)))\ .filter(Drives.start_position_id.is_not(None))\ .filter(Drives.end_position_id.is_not(None))\ .count() empty_count += session\ .query(ChargingProcesses.id)\ .filter(ChargingProcesses.address_id.is_(None))\ .filter(ChargingProcesses.position_id.is_not(None))\ .count() return empty_count def fix_empty_records(): # for low memory devices. while True: with Session(engine) as session: logging.info("checking empty records...") empty_count = get_empty_record_count(session) if fix_address(session, args.batch, empty_count) == 0: # all recoreds are fixed. break else: # commit at end of each batch. logging.info("saving...") session.commit() def get_field(find, keys): '''get field from a dict object''' item = find for key in keys: if isinstance(key, str): if key in item: item = item[key] else: return '' elif isinstance(key, int): if len(item) > 0: item = item[key] else: return '' # we should have find address str here. if not isinstance(item, str): # some address will be empty list. if len(item) == 0: return '' logging.fatal("key error when parse amap response.") assert (False) return item def update_address_in_db(need_update_address, address_details): # parse response. country = get_field(address_details, ['regeocode', 'addressComponent', 'country']) province = get_field(address_details, ['regeocode', 'addressComponent', 'province']) municipality = province in ['北京市', '天津市', '上海市', '重庆市'] if municipality: city = province + get_field( address_details, ['regeocode', 'addressComponent', 'district']) else: city = get_field(address_details, ['regeocode', 'addressComponent', 'city']) township = get_field(address_details, ['regeocode', 'addressComponent', 'township']) display_name = get_field(address_details, ['regeocode', 'formatted_address']) neighborhood = get_field( address_details, ['regeocode', 'addressComponent', 'neighborhood', 'name']) street_number = get_field( address_details, ['regeocode', 'addressComponent', 'streetNumber', 'number']) road = get_field(address_details, ['regeocode', 'roads', 0, 'name']) name = get_field(address_details, ['regeocode', 'aois', 0, 'name']) if len(name) == 0: name = get_field(address_details, ['regeocode', 'pois', 0, 'name']) if len(name) == 0: name = get_field(address_details, ['regeocode', 'roads', 0, 'name']) # update db record. logging.info("update address from %s to %s" % (need_update_address.display_name, display_name)) need_update_address.state = province need_update_address.county = township need_update_address.city = city need_update_address.house_number = street_number need_update_address.display_name = display_name need_update_address.country = country need_update_address.updated_at = datetime.now().replace(microsecond=0) # record last processed record id, skip records which id less than this. # assume that address record will not updated. # if language changed, teslamate will update all addressed, remember to # restart me to re-process all records. global last_update_id last_update_id = need_update_address.id # if some address is empty, do not update them. if len(road) > 0: need_update_address.road = road if len(name) > 0: need_update_address.name = name if len(neighborhood) > 0: need_update_address.neighbourhood = neighborhood def request_amap_api(url): '''request from amap api and loads as dict''' response = http_request(url) if response is None: return None response_dict = json.loads(response) if response_dict is None or response_dict['status'] != '1': logging.error("request amap api error: %s" % response) return None return response_dict def get_update_record_count(session): # record last updated id to save cpu time. return session\ .query(Addresses)\ .filter(Addresses.updated_at >= args.since)\ .filter(Addresses.id > last_update_id)\ .count() def get_need_update_addresses(session, batch_size): # record last update id to save cpu time. return session\ .query(Addresses)\ .filter(Addresses.updated_at >= args.since)\ .filter(Addresses.id > last_update_id)\ .order_by(Addresses.id)\ .limit(batch_size)\ .all() def update_address(session, batch_size, need_update_count): '''update address str by amap api.''' processed_count = 0 if len(args.key) == 0: logging.error("Amap key is not set.") return 0 need_update_addresses = get_need_update_addresses(session, batch_size) for need_update_address in need_update_addresses: logging.info("processing update address (%d left)" % (need_update_count - processed_count)) gps_lat = need_update_address.latitude gps_lon = need_update_address.longitude # transform coordinate url = amap_coordinate_transformation_url % (args.key, gps_lon, gps_lat) transformed_coordinate = request_amap_api(url) if transformed_coordinate is None: continue locations = transformed_coordinate['locations'] amap_lon = round(float(locations.split(',')[0]), 6) amap_lat = round(float(locations.split(',')[1]), 6) # get address details url = amap_resolve_url % (args.key, amap_lon, amap_lat) address_details = request_amap_api(url) if address_details is None: continue # update db update_address_in_db(need_update_address, address_details) processed_count += 1 return processed_count def update_address_by_amap(): while True: with Session(engine) as session: logging.info("updating address by amap...") need_update_count = get_update_record_count(session) if update_address(session, args.batch, need_update_count) == 0: # all recoreds are updated. break else: # commit at end of each batch. logging.info("saving...") session.commit() def main(): if args.mode == 0 or args.mode == 2: fix_empty_records() if args.mode == 1 or args.mode == 2: update_address_by_amap() # wrong mode, do nothing and exit. if args.mode < 0 or args.mode > 2: logging.info("nothing to do, bye.") # if interval is set, run in infinity mode. elif args.interval != 0: loop_timer = Timer(args.interval, main) loop_timer.start() if __name__ == '__main__': main()