Improve query API design

VectorInstitute · Sep 14, 2023 · 72b2ead · 72b2ead
1 parent bda2bb1
commit 72b2ead
Show file tree

Hide file tree

Showing 47 changed files with 1,869 additions and 1,713 deletions.
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -21,7 +21,7 @@ repos:
     - id: black
 
   - repo: https://github.com/charliermarsh/ruff-pre-commit
-    rev: 'v0.0.286'
+    rev: 'v0.0.288'
     hooks:
     - id: ruff
       args: [--fix, --exit-non-zero-on-fix]
@@ -36,6 +36,21 @@ repos:
       types: [python]
       exclude: "apps|use_cases|tests|cyclops/(process|models|tasks|monitor|report/plot)"
 
+  - repo: local
+    hooks:
+    - id: nbstripout
+      name: nbstripout
+      language: python
+      entry: nbstripout
+      exclude: ^docs/source/tutorials/gemini/.*\.ipynb$
+
+  - repo: https://github.com/nbQA-dev/nbQA
+    rev: 1.7.0
+    hooks:
+    - id: nbqa-black
+    - id: nbqa-ruff
+      args: [--fix, --exit-non-zero-on-fix]
+
   - repo: local
     hooks:
     - id: pytest
@@ -44,11 +59,3 @@ repos:
       language: system
       pass_filenames: false
       always_run: true
-
-  - repo: local
-    hooks:
-    - id: nbstripout
-      name: nbstripout
-      language: python
-      entry: nbstripout
-      exclude: ^docs/source/tutorials/gemini/.*\.ipynb$
diff --git a/cyclops/query/base.py b/cyclops/query/base.py
@@ -3,7 +3,7 @@
 import logging
 import os
 from functools import partial
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable, Dict, List, Optional
 
 import yaml
 from hydra import compose, initialize
@@ -12,14 +12,12 @@
 from sqlalchemy.sql.selectable import Subquery
 
 from cyclops.query import ops as qo
-from cyclops.query.interface import QueryInterface, QueryInterfaceProcessed
+from cyclops.query.interface import QueryInterface
 from cyclops.query.orm import Database
 from cyclops.query.util import (
     DBSchema,
-    TableTypes,
     _to_subquery,
     get_attr_name,
-    table_params_to_type,
 )
 from cyclops.utils.file import join as join_path
 from cyclops.utils.log import setup_logging
@@ -35,9 +33,9 @@ def _create_get_table_lambdafn(schema_name: str, table_name: str) -> Callable[..
 
     Parameters
     ----------
-    schema_name: str
+    schema_name
         The schema name.
-    table_name: str
+    table_name
         The table name.
 
     Returns
@@ -54,7 +52,7 @@ def _cast_timestamp_cols(table: Subquery) -> Subquery:
 
     Parameters
     ----------
-    table: sqlalchemy.sql.selectable.Subquery
+    table
         Table to cast.
 
     Returns
@@ -78,9 +76,15 @@ class DatasetQuerier:
 
     Attributes
     ----------
-    db: cyclops.query.orm.Database
+    db
         ORM Database used to run queries.
 
+    Parameters
+    ----------
+    config_overrides
+        Override configuration parameters, specified as kwargs.
+
+
     Notes
     -----
     This class is intended to be subclassed to provide methods for querying tables in
@@ -96,14 +100,6 @@ def __init__(
         self,
         **config_overrides: Dict[str, Any],
     ) -> None:
-        """Initialize.
-
-        Parameters
-        ----------
-        **config_overrides
-             Override configuration parameters, specified as kwargs.
-
-        """
         overrides = []
         if config_overrides:
             config_file = join_path(os.path.dirname(__file__), "configs", "config.yaml")
@@ -139,25 +135,39 @@ def list_schemas(self) -> List[str]:
         """
         return list(self.db.inspector.get_schema_names())
 
-    def list_tables(self) -> List[str]:
+    def list_tables(self, schema_name: Optional[str] = None) -> List[str]:
         """List table methods that can be queried using the database.
 
+        Parameters
+        ----------
+        schema_name
+            Name of schema in the database.
+
         Returns
         -------
         List[str]
             List of table names.
 
         """
-        return self.db.list_tables()
+        if schema_name:
+            table_names = []
+            for table in self.db.list_tables():
+                schema_name_, _ = table.split(".")
+                if schema_name_ == schema_name:
+                    table_names.append(table)
+        else:
+            table_names = self.db.list_tables()
+
+        return table_names
 
     def list_columns(self, schema_name: str, table_name: str) -> List[str]:
         """List columns in a table.
 
         Parameters
         ----------
-        schema_name: str
+        schema_name
             Name of schema in the database.
-        table_name: str
+        table_name
             Name of GEMINI table.
 
         Returns
@@ -196,36 +206,6 @@ def list_custom_tables(self) -> List[str]:
 
         return custom_tables
 
-    @table_params_to_type(Subquery)
-    def get_interface(
-        self,
-        table: TableTypes,
-        ops: Optional[qo.Sequential] = None,
-        process_fn: Optional[Callable[..., Any]] = None,
-    ) -> Union[QueryInterface, QueryInterfaceProcessed]:
-        """Get a query interface for a GEMINI table.
-
-        Parameters
-        ----------
-        table: cyclops.query.util.TableTypes
-            Table to wrap in the interface.
-        ops: cyclops.query.ops.Sequential
-            Operations to perform on the query.
-        process_fn
-            Process function to apply on the Pandas DataFrame returned from the query.
-
-        Returns
-        -------
-        cyclops.query.interface.QueryInterface or
-        cyclops.query.interface.QueryInterfaceProcessed
-            A query interface using the GEMINI database object.
-
-        """
-        if process_fn is None:
-            return QueryInterface(self.db, table, ops=ops)
-
-        return QueryInterfaceProcessed(self.db, table, ops=ops, process_fn=process_fn)
-
     def get_table(
         self,
         schema_name: str,
@@ -239,11 +219,11 @@ def get_table(
 
         Parameters
         ----------
-        schema_name: str
+        schema_name
             Name of schema in the database.
-        table_name: str
+        table_name
             Name of GEMINI table.
-        cast_timestamp_cols: bool
+        cast_timestamp_cols
             Whether to cast timestamp columns to datetime.
 
         Returns
@@ -263,21 +243,15 @@ def _template_table_method(
         self,
         schema_name: str,
         table_name: str,
-        join: Optional[qo.JoinArgs] = None,
-        ops: Optional[qo.Sequential] = None,
     ) -> QueryInterface:
         """Template method for table methods.
 
         Parameters
         ----------
-        schema_name: str
+        schema_name
             Name of schema in the database.
-        table_name: str
+        table_name
             Name of table in the database.
-        join: cyclops.query.ops.JoinArgs
-            Join arguments.
-        ops: cyclops.query.ops.Sequential or cyclops.query.ops.QueryOp, optional
-            Operations to perform on the query.
 
         Returns
         -------
@@ -288,7 +262,7 @@ def _template_table_method(
         table = getattr(getattr(self.db, schema_name), table_name).data
         table = _to_subquery(table)
 
-        return QueryInterface(self.db, table, join=join, ops=ops)
+        return QueryInterface(self.db, table)
 
     def _setup_table_methods(self) -> None:
         """Add table methods.

diff --git a/cyclops/query/gemini.py b/cyclops/query/gemini.py
@@ -1,7 +1,7 @@
 """GEMINI query API."""
 
 import logging
-from typing import Any, Dict, Optional
+from typing import Any, Dict
 
 from sqlalchemy import select
 from sqlalchemy.sql.expression import union_all
@@ -40,18 +40,9 @@ def __init__(self, **config_overrides: Dict[str, Any]) -> None:
 
     def ip_admin(
         self,
-        join: Optional[qo.JoinArgs] = None,
-        ops: Optional[qo.Sequential] = None,
     ) -> QueryInterface:
         """Query GEMINI patient encounters.
 
-        Parameters
-        ----------
-        join: qo.JoinArgs, optional
-            Join arguments.
-        ops: qo.Sequential, optional
-            Additional operations to perform on the table.
-
         Returns
         -------
         cyclops.query.interface.QueryInterface
@@ -78,22 +69,13 @@ def ip_admin(
         table = qo.Rename({"description": "discharge_description"})(table)
         table = qo.Drop("value")(table)
 
-        return QueryInterface(self.db, table, join=join, ops=ops)
+        return QueryInterface(self.db, table)
 
     def diagnoses(
         self,
-        join: Optional[qo.JoinArgs] = None,
-        ops: Optional[qo.Sequential] = None,
     ) -> QueryInterface:
         """Query diagnosis data.
 
-        Parameters
-        ----------
-        join: qo.JoinArgs, optional
-            Join arguments.
-        ops: qo.Sequential, optional
-            Additional operations to perform on the table.
-
         Returns
         -------
         cyclops.query.interface.QueryInterface
@@ -117,22 +99,13 @@ def diagnoses(
         # Trim whitespace from ICD codes.
         table = qo.Trim("diagnosis_code")(table)
 
-        return QueryInterface(self.db, table, join=join, ops=ops)
+        return QueryInterface(self.db, table)
 
     def room_transfer(
         self,
-        join: Optional[qo.JoinArgs] = None,
-        ops: Optional[qo.Sequential] = None,
     ) -> QueryInterface:
         """Query room transfer data.
 
-        Parameters
-        ----------
-        join: qo.JoinArgs, optional
-            Join arguments.
-        ops: qo.Sequential, optional
-            Additional operations to perform on the table.
-
         Returns
         -------
         cyclops.query.interface.QueryInterface
@@ -153,22 +126,13 @@ def room_transfer(
         )(table)
         table = qo.Rename({"description": "transfer_description"})(table)
 
-        return QueryInterface(self.db, table, join=join, ops=ops)
+        return QueryInterface(self.db, table)
 
     def care_units(
         self,
-        join: Optional[qo.JoinArgs] = None,
-        ops: Optional[qo.Sequential] = None,
     ) -> QueryInterface:
         """Query care unit data, fetches transfer info from multiple tables.
 
-        Parameters
-        ----------
-        join: qo.JoinArgs, optional
-            Join arguments.
-        ops: qo.Sequential, optional
-            Additional operations to perform on the table.
-
         Returns
         -------
         cyclops.query.interface.QueryInterface
@@ -236,22 +200,13 @@ def care_units(
             select(rt_table),
         ).subquery()
 
-        return QueryInterface(self.db, table, join=join, ops=ops)
+        return QueryInterface(self.db, table)
 
     def imaging(
         self,
-        join: Optional[qo.JoinArgs] = None,
-        ops: Optional[qo.Sequential] = None,
     ) -> QueryInterface:
         """Query imaging reports data.
 
-        Parameters
-        ----------
-        join: qo.JoinArgs, optional
-            Join arguments.
-        ops: qo.Sequential, optional
-            Additional operations to perform on the table.
-
         Returns
         -------
         cyclops.query.interface.QueryInterface
@@ -278,4 +233,4 @@ def imaging(
             table,
         )
 
-        return QueryInterface(self.db, table, join=join, ops=ops)
+        return QueryInterface(self.db, table)