improving query speed (#206)

* improving query speed

* adding session caching to cache

* Update materializationengine/blueprints/client/cache.py

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

---------

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: fcollman

materializationengine/blueprints/client/api2.py

@@ -36,6 +36,7 @@
     sql_query_warning,
     validate_table_args,
 )
+from materializationengine.blueprints.client.cache import get_cached_view_metadata
 from materializationengine.blueprints.client.common import (
     unhandled_exception as common_unhandled_exception,
 )
@@ -3150,8 +3151,8 @@ def get(
             mat_db_name = f"{aligned_volume_name}"
         else:
             mat_db_name = f"{datastack_name}__mat{version}"
-        with request_db_session(mat_db_name) as meta_db:
-            md = meta_db.database.get_view_metadata(datastack_name, view_name)
+        
+        md = get_cached_view_metadata(aligned_volume_name, datastack_name, view_name)
 
         return md
 
@@ -3189,8 +3190,8 @@ def assemble_view_dataframe(datastack_name, version, view_name, data, args):
     get_count = args.get("count", False)
     if get_count:
         limit = None
-    with request_db_session(mat_db_name) as mat_db:
-        md = mat_db.database.get_view_metadata(datastack_name, view_name)
+    
+    md = get_cached_view_metadata(aligned_volume_name, datastack_name, view_name)
 
     if not data.get("desired_resolution", None):
         des_res = [

materializationengine/blueprints/client/cache.py

@@ -0,0 +1,40 @@
+"""
+Cached metadata functions for MaterializationEngine.
+
+This module provides cached versions of frequently accessed metadata functions
+to reduce database load and improve performance.
+"""
+
+from cachetools import TTLCache, cached
+from materializationengine.database import dynamic_annotation_cache
+from materializationengine.request_db import request_db_session
+
+@cached(cache=TTLCache(maxsize=256, ttl=86400))  # 1 day TTL = 86400 seconds
+def get_cached_table_metadata(meta_db_name: str, table_name: str):
+    """Get table metadata with 1-day TTL cache.
+    
+    Args:
+        meta_db_name (str): The name of the metadata database
+        table_name (str): The name of the table
+        
+    Returns:
+        dict: Table metadata dictionary
+    """
+    with request_db_session(meta_db_name) as meta_db:  
+        return meta_db.database.get_table_metadata(table_name)
+
+
+@cached(cache=TTLCache(maxsize=256, ttl=86400))  # 1 day TTL = 86400 seconds
+def get_cached_view_metadata(meta_db_name: str, datastack_name: str, view_name: str):
+    """Get view metadata with 1-day TTL cache.
+    
+    Args:
+        meta_db_name (str): The name of the metadata database
+        datastack_name (str): The name of the datastack
+        view_name (str): The name of the view
+        
+    Returns:
+        dict: View metadata dictionary
+    """
+    with request_db_session(meta_db_name) as meta_db:  
+        return meta_db.database.get_view_metadata(datastack_name, view_name)

materializationengine/blueprints/client/common.py

@@ -2,11 +2,12 @@
 import traceback
 
 import numpy as np
-from cachetools import LRUCache, cached
+from cachetools import LRUCache, TTLCache, cached
 from cachetools.keys import hashkey
 from dynamicannotationdb.models import AnalysisTable, AnalysisVersion
 from flask import abort, current_app, g, request
 
+from materializationengine.blueprints.client.cache import get_cached_table_metadata
 from materializationengine.blueprints.client.query_manager import QueryManager
 from materializationengine.blueprints.client.utils import (
     collect_crud_columns,
@@ -197,8 +198,7 @@ def get_flat_model(datastack_name: str, table_name: str, version: int):
     if not analysis_version_dict.get("valid", False):
         abort(410, "This materialization version is not available")
 
-    db = dynamic_annotation_cache.get_db(aligned_volume_name)
-    metadata = db.database.get_table_metadata(table_name)
+    metadata = get_cached_table_metadata(aligned_volume_name, table_name)
     reference_table = metadata.get("reference_table")
     if reference_table:
         table_metadata = {"reference_table": reference_table}
@@ -209,6 +209,7 @@ def get_flat_model(datastack_name: str, table_name: str, version: int):
     if schema_type is None:
         abort(500, f"Schema not found for table {table_name} in version {version}")
 
+    db = dynamic_annotation_cache.get_db(aligned_volume_name)
     return db.schema.create_flat_model(
         table_name=table_name,
         schema_type=schema_type,
@@ -254,7 +255,7 @@ def generate_simple_query_dataframe(
     db = dynamic_annotation_cache.get_db(aligned_volume_name)
     check_read_permission(db, table_name)
 
-    ann_md = db.database.get_table_metadata(table_name)
+    ann_md = get_cached_table_metadata(aligned_volume_name, table_name)
 
     analysis_version_dict, analysis_table_dict = get_analysis_version_and_table(
         datastack_name, table_name, version, aligned_volume_name

materializationengine/blueprints/client/query.py

@@ -328,9 +328,12 @@ def _execute_query(
         df = pd.DataFrame({"count": [count]})
     else:
         if direct_sql_pandas:
-            with engine.connect() as connection:
-                statement = str(query.statement.compile(engine, compile_kwargs={"literal_binds": True}))
-                df = pd.read_sql(statement, connection.connection, coerce_float=False, index_col=index_col, dtype_backend='pyarrow')
+            statement = str(query.statement.compile(engine, compile_kwargs={"literal_binds": True}))
+            df = pd.read_sql(statement,
+                             session.connection().connection,
+                             coerce_float=False,
+                             index_col=index_col,
+                             dtype_backend='pyarrow')   
         else:
             df = read_sql_tmpfile(
                 query.statement.compile(engine, compile_kwargs={"literal_binds": True}),

materializationengine/blueprints/client/query_manager.py

@@ -17,6 +17,9 @@
 from sqlalchemy.ext.declarative.api import DeclarativeMeta
 import datetime
 
+from materializationengine.blueprints.client.cache import get_cached_table_metadata, get_cached_view_metadata
+from materializationengine.database import dynamic_annotation_cache
+
 DEFAULT_SUFFIX_LIST = ["x", "y", "z", "xx", "yy", "zz", "xxx", "yyy", "zzz"]
 DEFAULT_LIMIT = 250000
 
@@ -39,8 +42,11 @@ def __init__(
         self._db = dynamic_annotation_cache.get_db(db_name)
         if meta_db_name is None:
             self._meta_db = self._db
+            self._meta_db_name = db_name
         else:
             self._meta_db = dynamic_annotation_cache.get_db(meta_db_name)
+            self._meta_db_name = meta_db_name
+
         self._segmentation_source = segmentation_source
         self._split_mode = split_mode
         self._random_sample = random_sample
@@ -77,7 +83,7 @@ def _get_split_model(self, table_name):
         if table_name in self._split_models.keys():
             return self._split_models[table_name]
         else:
-            md = self._meta_db.database.get_table_metadata(table_name)
+            md = get_cached_table_metadata(self._meta_db_name, table_name)
             if md is None:
                 abort(404, f"Table {table_name} not found in metadata database")
             vox_res = np.array(
@@ -93,7 +99,7 @@ def _get_split_model(self, table_name):
             reference_table = md.get("reference_table")
             if reference_table:
                 table_metadata = {"reference_table": reference_table}
-                ref_md = self._meta_db.database.get_table_metadata(reference_table)
+                ref_md = get_cached_table_metadata(self._meta_db_name, reference_table)
                 _ = self._db.schema.get_split_models(
                     reference_table,
                     ref_md["schema_type"],
@@ -119,7 +125,7 @@ def _get_flat_model(self, table_name):
             return self._models[table_name]
         else:
             # schema = self._meta_db.database.get_table_schema(table_name)
-            md = self._meta_db.database.get_table_metadata(table_name)
+            md = get_cached_table_metadata(self._meta_db_name, table_name)
             vox_res = np.array(
                 [
                     md["voxel_resolution_x"],
@@ -148,7 +154,7 @@ def add_view(self, datastack_name, view_name):
         view_table = self._db.database.get_view_table(view_name)
         self._tables.add(view_table)
         self._models[view_name] = view_table
-        md = self._meta_db.database.get_view_metadata(datastack_name, view_name)
+        md = get_cached_view_metadata(self._meta_db_name, datastack_name, view_name)
         vox_res = np.array(
             [
                 md["voxel_resolution_x"],

materializationengine/utils.py

@@ -128,7 +128,10 @@ def check_write_permission(db, table_name):
 
 @cached(cache=TTLCache(maxsize=100, ttl=600))
 def check_read_permission(db, table_name):
-    metadata = db.database.get_table_metadata(table_name)
+    from materializationengine.blueprints.client.cache import get_cached_table_metadata
+    # Get aligned_volume_name from the db object
+    aligned_volume_name = db._aligned_volume
+    metadata = get_cached_table_metadata(aligned_volume_name, table_name)
     if metadata["read_permission"] == "GROUP":
         if not users_share_common_group(metadata["user_id"]):
             abort(
@@ -142,7 +145,10 @@ def check_read_permission(db, table_name):
 
 
 def check_ownership(db, table_name):
-    metadata = db.database.get_table_metadata(table_name)
+    from materializationengine.blueprints.client.cache import get_cached_table_metadata
+    # Get aligned_volume_name from the db object
+    aligned_volume_name = db._aligned_volume
+    metadata = get_cached_table_metadata(aligned_volume_name, table_name)
     if metadata["user_id"] != str(g.auth_user["id"]):
         abort(401, "You cannot do this because you are not the owner of this table")
     return metadata