diff --git a/hq_superset/models.py b/hq_superset/models.py index 175a314..b704909 100644 --- a/hq_superset/models.py +++ b/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'