cache table & database

hq_superset/models.py

@@ -7,6 +7,7 @@
 )
 from cryptography.fernet import MultiFernet
 from superset import db
+from superset.extensions import cache_manager
 
 from hq_superset.const import OAUTH2_DATABASE_NAME
 from hq_superset.exceptions import TableMissing
@@ -16,6 +17,7 @@
     get_hq_database,
 )
 
+cache = cache_manager.cache
 
 @dataclass
 class DataSetChange:
@@ -31,16 +33,12 @@ def update_dataset(self):
         for a form or a case, which is identified by ``self.doc_id``. If
         the form or case has been deleted, then the list will be empty.
         """
-        database = get_hq_database()
-        try:
-            sqla_table = next((
-                table for table in database.tables
-                if table.table_name == self.data_source_id
-            ))
-        except StopIteration:
+        sqla_table = _get_data_source_table(self.data_source_id)
+        if not sqla_table:
             raise TableMissing(f'{self.data_source_id} table not found.')
         table = sqla_table.get_sqla_table_object()
 
+        database = _get_cached_hq_database()
         with (
             database.get_sqla_engine_with_context() as engine,
             engine.connect() as connection,
@@ -54,6 +52,35 @@ def update_dataset(self):
                 connection.execute(insert_stmt)
 
 
+@cache.memoize(timeout=24*3600)  # 1 day
+def _get_data_source_table(data_source_id):
+    """
+    Fetch table for datasource.
+    Try again after expiring database cache if not found first
+    """
+    sqla_table = _get_sqla_table(data_source_id)
+    if not sqla_table:
+        _get_cached_hq_database.delete_memoized()
+        sqla_table = _get_sqla_table(data_source_id)
+    return sqla_table
+
+
+def _get_sqla_table(data_source_id):
+    database = _get_cached_hq_database()
+    try:
+        return next((
+            table for table in database.tables
+            if table.table_name == data_source_id
+        ))
+    except StopIteration:
+        return None
+
+
+@cache.memoize(timeout=24*3600)  # 1 day
+def _get_cached_hq_database():
+    return get_hq_database()
+
+
 class OAuth2Client(db.Model, OAuth2ClientMixin):
     __bind_key__ = OAUTH2_DATABASE_NAME
     __tablename__ = 'hq_oauth_client'