archived-teslamate-fix-addrs/teslamate_fix_addrs.py

from sqlalchemy.ext.automap import automap_base
from sqlalchemy.orm import Session
from sqlalchemy import create_engine, or_, func
from sqlalchemy.engine.url import URL
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,
    format='%(asctime)s - %(levelname)s - %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'  # 设置时间格式
)



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_url = URL.create(
    drivername="postgresql",
    username=args.user,
    password=args.password,
    host=args.host,
    port=args.port,
    database=args.dbname
)

engine = create_engine(conn_url, 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 = {
        'accept': 'text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7',
        'accept-language': 'zh-CN,zh;q=0.9,en-US;q=0.8,en;q=0.7',
        'cache-control': 'max-age=0',
        'dnt': '1',
        'priority': 'u=0, i',
        'sec-ch-ua': '"Chromium";v="128", "Not;A=Brand";v="24", "Google Chrome";v="128"',
        'sec-ch-ua-mobile': '?0',
        'sec-ch-ua-platform': '"macOS"',
        'sec-fetch-dest': 'document',
        'sec-fetch-mode': 'navigate',
        'sec-fetch-site': 'none',
        'sec-fetch-user': '?1',
        'upgrade-insecure-requests': '1',
        '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()