From e1c5092130bdc359ded5f0c060a3e4a680d019bf Mon Sep 17 00:00:00 2001
From: OndrejSedlacek <xsedla1o@stud.fit.vutbr.cz>
Date: Thu, 23 Nov 2023 12:57:49 +0100
Subject: [PATCH] CLI: Added `schema-update` command to `dp3`.

---
 dp3/bin/cli.py           | 11 ++++++
 dp3/bin/schema_update.py | 83 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 94 insertions(+)
 create mode 100644 dp3/bin/schema_update.py

diff --git a/dp3/bin/cli.py b/dp3/bin/cli.py
index db07e172..db592e95 100755
--- a/dp3/bin/cli.py
+++ b/dp3/bin/cli.py
@@ -10,6 +10,8 @@
 from dp3.bin.check import main as check_main
 from dp3.bin.config import init_parser as init_config_parser
 from dp3.bin.config import main as config_main
+from dp3.bin.schema_update import init_parser as init_schema_update_parser
+from dp3.bin.schema_update import main as schema_update_main
 from dp3.bin.setup import init_parser as init_setup_parser
 from dp3.bin.setup import main as setup_main
 from dp3.bin.worker import init_parser as init_worker_parser
@@ -49,6 +51,13 @@ def init_parser():
         description="Setup configuration files for a container-less DP3 application deployment. ",
     )
     init_config_parser(config_parser)
+
+    schema_update_parser = commands.add_parser(
+        "schema-update",
+        help="Update the database schema after making conflicting changes to the model.",
+        description="Update the database schema after making conflicting changes to the model. ",
+    )
+    init_schema_update_parser(schema_update_parser)
     return parser
 
 
@@ -66,6 +75,8 @@ def run():
         check_main(args)
     elif args.command == "config":
         config_main(args)
+    elif args.command == "schema-update":
+        schema_update_main(args)
     else:
         parser.print_help()
         sys.exit(1)
diff --git a/dp3/bin/schema_update.py b/dp3/bin/schema_update.py
new file mode 100644
index 00000000..e522ac23
--- /dev/null
+++ b/dp3/bin/schema_update.py
@@ -0,0 +1,83 @@
+"""
+Update the database schema after making conflicting changes to the model.
+
+Authors: Ondřej Sedláček <xsedla1o@stud.fit.vutbr.cz>
+"""
+
+import logging
+
+from dp3.common.config import ModelSpec, read_config_dir
+from dp3.database.database import EntityDatabase
+
+
+def init_parser(parser):
+    parser.add_argument(
+        "config",
+        metavar="CONFIG_DIR",
+        help="Path to a directory containing configuration files (e.g. /etc/my_app/config)",
+    )
+    parser.add_argument(
+        "--bypass",
+        action="store_true",
+        default=False,
+        help="Bypass the suggested database changes and update the schema regardless.",
+    )
+
+
+def confirm_changes(prompt: str):
+    while True:
+        answer = input(prompt).lower()
+        if answer == "" or answer[0] == "n":
+            return False
+        if answer[0] == "y":
+            return True
+
+
+def main(args):
+    # Load DP3 configuration
+    config = read_config_dir(args.config, recursive=True)
+
+    # Setup logging
+    LOGFORMAT = "%(asctime)-15s,%(name)s,[%(levelname)s] %(message)s"
+    LOGDATEFORMAT = "%Y-%m-%dT%H:%M:%S"
+    logging.basicConfig(level=logging.DEBUG, format=LOGFORMAT, datefmt=LOGDATEFORMAT)
+    log = logging.getLogger("SchemaUpdate")
+
+    # Connect to database
+    connection_conf = config.get("database", {})
+    db = EntityDatabase(
+        connection_conf,
+        ModelSpec(config.get("db_entities")),
+        config.get("processing_core.worker_processes"),
+    )
+
+    prev_schema, config_schema, updates = db.schema_cleaner.get_schema_status()
+    if prev_schema["schema"] == config_schema["schema"]:
+        log.info("Schema is OK!")
+        return
+
+    if not updates:
+        db.schema_cleaner.schemas.insert_one(config_schema)
+        log.info("Updated schema without any changes to master records, OK now!")
+        return
+
+    log.info("Suggested changes to master records:")
+    for entity, entity_updates in updates.items():
+        log.info(f"{entity}: {dict(entity_updates)}")
+
+    if args.bypass:
+        if not confirm_changes(
+            "Are you sure you want update the schema without applying changes? (y/[n]): "
+        ):
+            log.info("Aborted schema update.")
+            return
+        db.schema_cleaner.schemas.insert_one(config_schema)
+        log.info("Bypassing suggested changes, updated schema, OK now!")
+        return
+
+    if not confirm_changes("Are you sure you want to apply these changes? (y/[n]): "):
+        log.info("Aborted schema update.")
+        return
+    db.schema_cleaner.execute_updates(updates)
+    db.schema_cleaner.schemas.insert_one(config_schema)
+    log.info("Applied suggested changes, updated schema, OK now!")