Include vector in functionality catalog (#382)

src/brad/query_rep.py

@@ -27,6 +27,9 @@
         _GEOSPATIAL_KEYWORDS = yaml.safe_load(f)
 _GEOSPATIAL_KEYWORDS = [k.upper() for k in _GEOSPATIAL_KEYWORDS]
 
+# Define vector keywords.
+_VECTOR_KEYWORDS = ["<=>"]
+
 
 class QueryRep:
     """
@@ -80,10 +83,19 @@ def is_geospatial(self) -> bool:
                 return True
         return False
 
+    def is_vector(self) -> bool:
+        query = self._raw_sql_query.upper()
+        for keyword in _VECTOR_KEYWORDS:
+            if keyword in query:
+                return True
+        return False
+
     def get_required_functionality(self) -> int:
         req_functionality: List[str] = []
         if self.is_geospatial():
             req_functionality.append(Functionality.Geospatial)
+        if self.is_vector():
+            req_functionality.append(Functionality.Vector)
 
         return Functionality.to_bitmap(req_functionality)
 

src/brad/routing/engine_functionality.yml

@@ -3,6 +3,7 @@ database_engines:
     functionalities:
       - geospatial
       - transactions
+      - vector
 
   - name: Athena
     functionalities:

src/brad/routing/functionality_catalog.py

@@ -11,6 +11,7 @@
 class Functionality:
     Geospatial = "geospatial"
     Transaction = "transactions"
+    Vector = "vector"
 
     def __init__(self):
         # Read the YAML file
@@ -65,13 +66,14 @@ def get_engine_functionalities(self) -> List[Tuple[int, int]]:
         Return a bitmap for each engine that states what functionalities the
         engine supports.
 
-        The first value in the tuple is the bitmask representing the the engine.
+        The first value in the tuple is the bitmask representing the engine.
         The second value in the tuple is the bitmap representing its supported
         functionalities.
         """
         return self.engine_functionalities
 
 
 FunctionalityBitmapValues: Dict[str, int] = {}
-FunctionalityBitmapValues[Functionality.Geospatial] = 0b01
-FunctionalityBitmapValues[Functionality.Transaction] = 0b10
+FunctionalityBitmapValues[Functionality.Geospatial] = 0b001
+FunctionalityBitmapValues[Functionality.Transaction] = 0b010
+FunctionalityBitmapValues[Functionality.Vector] = 0b100