diff --git a/src/brad/query_rep.py b/src/brad/query_rep.py index d2d10777..4bdf2916 100644 --- a/src/brad/query_rep.py +++ b/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) diff --git a/src/brad/routing/engine_functionality.yml b/src/brad/routing/engine_functionality.yml index 03a9a979..2a2b2695 100644 --- a/src/brad/routing/engine_functionality.yml +++ b/src/brad/routing/engine_functionality.yml @@ -3,6 +3,7 @@ database_engines: functionalities: - geospatial - transactions + - vector - name: Athena functionalities: diff --git a/src/brad/routing/functionality_catalog.py b/src/brad/routing/functionality_catalog.py index 43abfdc9..9a579cd2 100644 --- a/src/brad/routing/functionality_catalog.py +++ b/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,7 +66,7 @@ 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. """ @@ -73,5 +74,6 @@ def get_engine_functionalities(self) -> List[Tuple[int, int]]: FunctionalityBitmapValues: Dict[str, int] = {} -FunctionalityBitmapValues[Functionality.Geospatial] = 0b01 -FunctionalityBitmapValues[Functionality.Transaction] = 0b10 +FunctionalityBitmapValues[Functionality.Geospatial] = 0b001 +FunctionalityBitmapValues[Functionality.Transaction] = 0b010 +FunctionalityBitmapValues[Functionality.Vector] = 0b100