Dependency inversion: Move util to model method

Models use utils. Utils don't use models.
dimagi · Feb 21, 2024 · c7404aa · c7404aa
1 parent 191721d
commit c7404aa
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 46 deletions.
diff --git a/hq_superset/api.py b/hq_superset/api.py
@@ -106,7 +106,7 @@ def post_dataset_change(self) -> FlaskResponse:
         try:
             request_json = json.loads(request.get_data(as_text=True))
             change = DataSetChange(**request_json)
-            update_dataset(change)
+            change.update_dataset()
             return self.json_response(
                 'Request accepted; updating dataset',
                 status=HTTPStatus.ACCEPTED.value,

diff --git a/hq_superset/models.py b/hq_superset/models.py
@@ -9,9 +9,10 @@
 from cryptography.fernet import MultiFernet
 from flask import current_app
 from superset import db
+from superset.connectors.sqla.models import SqlaTable
 
 from .const import HQ_DATA
-from .utils import encoded
+from .utils import encoded, get_explore_database, get_hq_database
 
 
 @dataclass
@@ -24,6 +25,45 @@ def __post_init__(self):
         if 'doc_id' not in self.data:
             raise TypeError("'data' missing required key: 'doc_id'")
 
+    def update_dataset(self):
+        database = get_hq_database()
+        explore_database = get_explore_database(database)  # TODO: Necessary?
+        sqla_table = (
+            db.session.query(SqlaTable)
+            .filter_by(
+                table_name=self.data_source_id,
+                database_id=explore_database.id,
+            )
+            .one_or_none()
+        )
+        if sqla_table is None:
+            raise ValueError(f'{self.data_source_id} table not found.')
+
+        if self.action == 'delete':
+            stmt = (
+                sqla_table
+                .delete()
+                .where(sqla_table.doc_id == self.data['doc_id'])
+            )
+        elif self.action == 'upsert':
+            stmt = (
+                sqla_table
+                .insert()
+                .values(self.data)  # TODO: Do we need to cast anything?
+                .on_conflict_do_update(
+                    index_elements=['doc_id'],
+                    set_=self.data,
+                )
+            )
+        else:
+            raise ValueError(f'Invalid DataSetChange action {self.action!r}')
+        try:
+            db.session.execute(stmt)
+            db.session.commit()
+        except Exception:  # pylint: disable=broad-except
+            db.session.rollback()
+            raise
+
 
 class HQClient(db.Model, OAuth2ClientMixin):
     __bind_key__ = HQ_DATA

diff --git a/hq_superset/utils.py b/hq_superset/utils.py
@@ -10,7 +10,6 @@
 from superset.utils.database import get_or_create_db
 
 from .const import HQ_DATA
-from .models import DataSetChange
 
 DOMAIN_PREFIX = "hqdomain_"
 SESSION_USER_DOMAINS_KEY = "user_hq_domains"
@@ -220,46 +219,3 @@ def get_explore_database(database):
         )
     else:
         return database
-
-
-def update_dataset(change: DataSetChange):
-    from superset import db
-    from superset.connectors.sqla.models import SqlaTable
-
-    database = get_hq_database()
-    explore_database = get_explore_database(database)  # TODO: Necessary?
-    sqla_table = (
-        db.session.query(SqlaTable)
-        .filter_by(
-            table_name=change.data_source_id,
-            database_id=explore_database.id,
-        )
-        .one_or_none()
-    )
-    if sqla_table is None:
-        raise ValueError(f'{change.data_source_id} table not found.')
-
-    if change.action == 'delete':
-        stmt = (
-            sqla_table
-            .delete()
-            .where(sqla_table.doc_id == change.data['doc_id'])
-        )
-    elif change.action == 'upsert':
-        stmt = (
-            sqla_table
-            .insert()
-            .values(change.data)  # TODO: Do we need to cast anything?
-            .on_conflict_do_update(
-                index_elements=['doc_id'],
-                set_=change.data,
-            )
-        )
-    else:
-        raise ValueError(f'Invalid DataSetChange action {change.action!r}')
-    try:
-        db.session.execute(stmt)
-        db.session.commit()
-    except Exception:  # pylint: disable=broad-except
-        db.session.rollback()
-        raise