From d28bf18ce8e136e79552471453bd33bef863ce67 Mon Sep 17 00:00:00 2001
From: Owen <owen@sr2.uk>
Date: Mon, 14 Nov 2022 12:21:41 +0000
Subject: [PATCH] Removed duplicate list of database models

---
 app/cli/db.py | 68 ++++++++++++++++++++-------------------------------
 1 file changed, 26 insertions(+), 42 deletions(-)

diff --git a/app/cli/db.py b/app/cli/db.py
index fbe9787..32d7e17 100644
--- a/app/cli/db.py
+++ b/app/cli/db.py
@@ -20,25 +20,26 @@ from app.models.alarms import Alarm, AlarmState
 from app.models.onions import Onion, Eotk
 from app.models.tfstate import TerraformState
 
-models = {
-    "activity": Activity,
-    "alarm": Alarm,
-    "automation": Automation,
-    "automation_logs": AutomationLogs,
-    "bridge": Bridge,
-    "bridgeconf": BridgeConf,
-    "eotk": Eotk,
-    "group": Group,
-    "list": MirrorList,
-    "onion": Onion,
-    "origin": Origin,
-    "pool": Pool,
-    "pool_group": PoolGroup,
-    "proxy": Proxy,
-    "smart_proxy": SmartProxy,
-    "terraform_state": TerraformState,
-    "webhook": Webhook
-}
+# order matters due to foreign key constraints
+models = [
+    Group,
+    Activity,
+    Pool,
+    PoolGroup,
+    SmartProxy,
+    Origin,
+    Proxy,
+    Onion,
+    Alarm,
+    Automation,
+    AutomationLogs,
+    BridgeConf,
+    Bridge,
+    Eotk,
+    MirrorList,
+    TerraformState,
+    Webhook
+]
 
 
 class ExportEncoder(json.JSONEncoder):
@@ -70,8 +71,8 @@ def db_export() -> None:
     encoder = ExportEncoder()
     output = defaultdict(list)
     for model in models:
-        for row in models[model].query.all():  # type: ignore[attr-defined]
-            output[model].append(model_to_dict(row))
+        for row in model.query.all():  # type: ignore[attr-defined]
+            output[model.__name__].append(model_to_dict(row))
     print(encoder.encode(output))
 
 
@@ -86,9 +87,9 @@ decoder = {
 }
 
 
-def db_import_model(model: str, data: List[Dict[str, Any]]) -> None:
+def db_import_model(model: db.Model, data: List[Dict[str, Any]]) -> None:  # type: ignore[name-defined]
     for row in data:
-        new = models[model]()
+        new = model()
         for col in row:
             type_name, col_name = col.split("_", 1)
             new.__setattr__(col_name, decoder.get(type_name, lambda x: x)(row[col]))  # type: ignore[no-untyped-call]
@@ -98,25 +99,8 @@ def db_import_model(model: str, data: List[Dict[str, Any]]) -> None:
 def db_import() -> None:
     data = json.load(sys.stdin)
     # import order matters due to foreign key constraints
-    for model in [
-        "group",
-        "pool",
-        "pool_group",
-        "smart_proxy",
-        "origin",
-        "proxy",
-        "onion",
-        "alarm",
-        "automation",
-        "automation_logs",
-        "bridgeconf",
-        "bridge",
-        "eotk",
-        "list",
-        "terraform_state",
-        "webhook"
-    ]:
-        db_import_model(model, data[model])
+    for model in models:
+        db_import_model(model, data[model.__name__])
 
     db.session.commit()