Functionality catalog

src/brad/front_end/front_end.py

@@ -283,14 +283,8 @@ async def _run_query_impl(
                 return await self._handle_internal_command(session, query, debug_info)
 
             # 2. Select an engine for the query.
-            query_rep = QueryRep(query)
-            transactional_query = (
-                session.in_transaction or query_rep.is_data_modification_query()
-            )
-            if transactional_query:
-                engine_to_use = Engine.Aurora
-            else:
-                engine_to_use = await self._router.engine_for(query_rep)
+            query_rep = QueryRep(query, session)
+            engine_to_use = await self._router.engine_for(query_rep)
 
             log_verbose(
                 logger,
@@ -303,6 +297,9 @@ async def _run_query_impl(
 
             # 3. Actually execute the query.
             try:
+                transactional_query: bool = (
+                    session.in_transaction or query_rep.is_data_modification_query()
+                )
                 if transactional_query:
                     connection = session.engines.get_connection(engine_to_use)
                     cursor = connection.cursor_sync()

src/brad/front_end/session.py

@@ -1,13 +1,16 @@
 import asyncio
 import logging
-from typing import Dict, Tuple, Optional
+from typing import Dict, Tuple, Optional, TYPE_CHECKING
 
 from brad.config.engine import Engine
-from brad.blueprint.manager import BlueprintManager
 from brad.config.file import ConfigFile
 from brad.config.session import SessionId
 from .engine_connections import EngineConnections
 
+if TYPE_CHECKING:
+    from brad.blueprint.manager import BlueprintManager
+
+
 logger = logging.getLogger(__name__)
 
 
@@ -50,7 +53,7 @@ async def close(self):
 
 class SessionManager:
     def __init__(
-        self, config: ConfigFile, blueprint_mgr: BlueprintManager, schema_name: str
+        self, config: ConfigFile, blueprint_mgr: "BlueprintManager", schema_name: str
     ) -> None:
         self._config = config
         self._blueprint_mgr = blueprint_mgr

src/brad/planner/workload/query.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Dict
+from typing import Dict, Optional
 
 from brad.blueprint import Blueprint
 from brad.config.engine import Engine
@@ -9,6 +9,7 @@
     extract_base_cardinalities,
 )
 from brad.front_end.engine_connections import EngineConnections
+from brad.front_end.session import Session
 
 logger = logging.getLogger(__name__)
 
@@ -21,8 +22,10 @@ class Query(QueryRep):
     Objects of this class are logically immutable.
     """
 
-    def __init__(self, sql_query: str, arrival_count: int = 1):
-        super().__init__(sql_query)
+    def __init__(
+        self, sql_query: str, arrival_count: int = 1, session: Optional[Session] = None
+    ):
+        super().__init__(sql_query, session)
         self._arrival_count = arrival_count
 
         # Legacy statistics.

src/brad/query_rep.py

@@ -1,6 +1,10 @@
 import sqlglot
 import sqlglot.expressions as exp
-
+import yaml
+import pathlib
+import os
+from brad.routing.functionality_catalog import Functionality
+from brad.front_end.session import Session
 from typing import List, Optional
 
 _DATA_MODIFICATION_PREFIXES = [
@@ -17,6 +21,14 @@
     "TRUNCATE",
 ]
 
+# Load geospatial keywords used to detect if geospatial query
+_GEOSPATIAL_KEYWORDS_PATH = os.path.join(
+    pathlib.Path(__file__).parent.resolve(), "routing/geospatial_keywords.yml"
+)
+with open(_GEOSPATIAL_KEYWORDS_PATH, "r") as f:
+    _GEOSPATIAL_KEYWORDS = yaml.safe_load(f)
+_GEOSPATIAL_KEYWORDS = [k.upper() for k in _GEOSPATIAL_KEYWORDS]
+
 
 class QueryRep:
     """
@@ -29,13 +41,16 @@ class QueryRep:
     Objects of this class are logically immutable.
     """
 
-    def __init__(self, sql_query: str):
+    def __init__(self, sql_query: str, session: Optional[Session] = None):
         self._raw_sql_query = sql_query
 
         # Lazily computed.
         self._ast: Optional[sqlglot.Expression] = None
         self._is_data_modification: Optional[bool] = None
         self._tables: Optional[List[str]] = None
+        self.in_transaction: bool = False
+        if session is not None:
+            self.in_transaction = session.in_transaction
 
     @property
     def raw_query(self) -> str:
@@ -55,6 +70,25 @@ def is_transaction_end(self) -> bool:
         raw_sql = self._raw_sql_query.upper()
         return raw_sql == "COMMIT" or raw_sql == "ROLLBACK"
 
+    def is_geospatial(self) -> bool:
+        query = self._raw_sql_query.upper()
+        for keyword in _GEOSPATIAL_KEYWORDS:
+            if keyword in query:
+                return True
+        return False
+
+    def is_transaction(self) -> bool:
+        return self.is_data_modification_query() or self.in_transaction
+
+    def get_required_functionality(self) -> int:
+        req_functionality: List[str] = []
+        if self.is_geospatial():
+            req_functionality.append(Functionality.Geospatial)
+        if self.is_transaction():
+            req_functionality.append(Functionality.Transaction)
+
+        return Functionality.to_bitmap(req_functionality)
+
     def tables(self) -> List[str]:
         if self._tables is None:
             if self._ast is None:

src/brad/routing/engine_functionality.yml

@@ -0,0 +1,12 @@
+database_engines:
+  - name: Aurora
+    functionalities:
+      - geospatial
+      - transactions
+
+  - name: Athena
+    functionalities:
+      - geospatial
+
+  - name: Redshift
+    functionalities: []

src/brad/routing/functionality_catalog.py

@@ -0,0 +1,62 @@
+from typing import List
+import operator
+import yaml
+from functools import reduce
+from typing import Dict
+
+
+class Functionality:
+    Geospatial = "geospatial"
+    Transaction = "transactions"
+
+    def __init__(self, functionality_yaml="engine_functionality.yml"):
+        # Read the YAML file
+        with open(functionality_yaml, "r") as yaml_file:
+            data = yaml.load(yaml_file, Loader=yaml.FullLoader)
+
+        # Initialize lists for each database engine's functionalities
+        aurora_functionalities = []
+        athena_functionalities = []
+        redshift_functionalities = []
+
+        # Parse the data into the respective lists
+        for engine in data["database_engines"]:
+            if engine["name"] == "Aurora":
+                aurora_functionalities = engine["functionalities"]
+            elif engine["name"] == "Athena":
+                athena_functionalities = engine["functionalities"]
+            elif engine["name"] == "Redshift":
+                redshift_functionalities = engine["functionalities"]
+
+        # Convert to bitmaps
+        engine_functionality_strings = [
+            athena_functionalities,
+            aurora_functionalities,
+            redshift_functionalities,
+        ]
+        self.engine_functionalities = [
+            Functionality.to_bitmap(f) for f in engine_functionality_strings
+        ]
+
+    @staticmethod
+    def to_bitmap(functionalities: List[str]) -> int:
+        if len(functionalities) == 0:
+            return 0
+        return reduce(
+            # Bitwise OR
+            operator.or_,
+            map(lambda f: FunctionalityBitmapValues[f], functionalities),
+            0,
+        )
+
+    def get_engine_functionalities(self) -> List[int]:
+        """
+        Return a bitmap for each engine that states what functionalities the
+        engine supports
+        """
+        return self.engine_functionalities
+
+
+FunctionalityBitmapValues: Dict[str, int] = {}
+FunctionalityBitmapValues[Functionality.Geospatial] = 0b01
+FunctionalityBitmapValues[Functionality.Transaction] = 0b10