update event message ids

.gitignore

*.pyc
/.venv/

main.py

import argparse
import json
import logging

from sqlalchemy import create_engine, engine
from sqlalchemy.orm import DeclarativeBase, Mapped, Session, mapped_column

OWNER_KEYS = ('americas', 'chailease', 'demo', 'plus', 'thailand', 'thingnario', 'tsc')
logger = logging.getLogger(__name__)
logger.addHandler(logging.FileHandler('update_event_message_ids.log'))


class Base(DeclarativeBase):
    pass


class SolarEvent(Base):
    __tablename__ = 'SolarEvent'

    id:             Mapped[int] = mapped_column(primary_key=True)
    eventMessageId: Mapped[int]


def main():
    args = get_args()
    mapping = load_mapping()
    for owner_key in OWNER_KEYS:
        update_event_message_ids(args.url, owner_key, mapping, args.dry_run)


def get_args():
    args = argparse.ArgumentParser()
    args.add_argument(
        '--url',
        required=True,
        type=engine.make_url,
    )
    args.add_argument(
        '--dry-run',
        action='store_true',
    )
    return args.parse_args()


def load_mapping():
    with open('mapping.json') as file:
        return {int(key): value for key, value in json.load(file).items()}


def update_event_message_ids(
        url: engine.URL, owner_key: str, mapping: dict[int, int], dry_run=False):
    engine = create_engine(url.set(database=f'solar_PV_{owner_key}'))
    with Session(engine) as session:
        for old_id, new_id in mapping.items():
            updated = 0
            for row in session.query(SolarEvent).filter_by(eventMessageId=old_id).all():
                row.eventMessageId = new_id
                updated += 1
            if not dry_run:
                session.commit()
            logger.info('[%s] %s: %d updated', owner_key, old_id, updated)


if __name__ == '__main__':
    logging.basicConfig(format='%(asctime)s| %(message)s', level=logging.INFO)
    main()

requirements.txt

PyMySQL==1.1.1
SQLAlchemy==2.0.35

setup.cfg

[isort]
line_length = 100

[pycodestyle]
max_line_length = 100