update model for import

kontor-scripts/import.py

@@ -3,40 +3,45 @@ import data from json file to MariaDB
 """
 from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter
 
-import yaml
+from typing import Generator
+import os
 from sqlalchemy import create_engine
-from sqlalchemy.orm import sessionmaker
-from platformdirs import PlatformDirs
-from pathlib import Path
-
-from schema import Base, KontorDB
+from sqlalchemy.orm import sessionmaker, Session
+from db.models.base import Base
+from db.models import registry
 from config import get_logger
 
 parser = ArgumentParser(formatter_class=ArgumentDefaultsHelpFormatter)
 parser.add_argument('--verbose', '-v', action='count', default=0)
-parser.add_argument('--config', '-c', default='kontor-docker')
 parser.add_argument('--recreate-db', action='store_true')
 parser.add_argument('--file', '-f', default='~/data.json')
 args = parser.parse_args()
 
+DB_USER: str = os.getenv("DB_USER", "kontor")
+DB_PASSWORD: str = os.getenv("DB_PASSWORD", "kontor")
+DB_SERVER: str = os.getenv("DB_SERVER", "127.0.0.1")
+DB_PORT: int = os.getenv("DB_PORT", 5432)
+DB_DBNAME: str = os.getenv("DB_DBNAME", "kontor")
+DATABASE_URL: str = f"postgresql://{DB_USER}:{DB_PASSWORD}@{DB_SERVER}:{DB_PORT}/{DB_DBNAME}"
+
+
+def cleanup_database(db: Session, log):
+    # get tables from registry
+    for table in registry:
+        log.info(f"{table}")
+        model = registry[table]
+        entries = db.query(model).all()
+        for entry in entries:
+            db.delete(entry)
+            db.commit()
 
 if __name__ == '__main__':
-    logger = get_logger(args.verbose, args.config)
+    logger = get_logger(args.verbose, "kontor")
     logger.info('kontor.import started')
-    dirs = PlatformDirs(args.config)
-    database_config = Path(dirs.user_config_dir, 'database-config.yaml')
-    with open(database_config, 'rt') as f:
-        db_config = yaml.safe_load(f.read())
-    connect_string = ('mariadb+mariadbconnector://{}:{}@{}:{}/{}'.format(
-        db_config['mariadb']['user'],
-        db_config['mariadb']['password'],
-        db_config['mariadb']['host'],
-        db_config['mariadb']['port'],
-        db_config['mariadb']['database']
-    ))
-    engine = create_engine(connect_string)
+    engine = create_engine(DATABASE_URL)
     Base.metadata.create_all(bind=engine, checkfirst=True)
-    __session__ = sessionmaker(bind=engine)
-    kontor_db = KontorDB(engine, logger)
-    kontor_db.import_db(args.file)
+    SessionLocal = sessionmaker(bind=engine)
+    with SessionLocal() as db:
+        cleanup_database(db, logger)
     logger.info('kontor.import finished')
+