fix(databricks): validate catalog/schema/table identifiers to prevent SQL injection (#5988)

mem0/vector_stores/databricks.py

@@ -38,10 +38,16 @@ class MemoryResult(BaseModel):
 
 excluded_keys = {"user_id", "agent_id", "run_id", "hash", "data", "created_at", "updated_at"}
 
-# Pattern for valid SQL identifiers to prevent column name injection
+# Pattern for valid SQL identifiers to prevent column name / table name injection
 _VALID_SQL_IDENTIFIER = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")
 
 
+def _validate_identifier(name: str, label: str = "identifier") -> str:
+    if not isinstance(name, str) or not _VALID_SQL_IDENTIFIER.match(name):
+        raise ValueError(f"Invalid {label}: {name!r}")
+    return name
+
+
 class Databricks(VectorStoreBase):
     def __init__(
         self,
@@ -90,11 +96,11 @@ class Databricks(VectorStoreBase):
         # Basic identifiers
         self.workspace_url = workspace_url
         self.endpoint_name = endpoint_name
-        self.catalog = catalog
+        self.catalog = _validate_identifier(catalog, "catalog")
-        self.schema = schema
+        self.schema = _validate_identifier(schema, "schema")
-        self.table_name = table_name
+        self.table_name = _validate_identifier(table_name, "table_name")
         self.fully_qualified_table_name = f"{self.catalog}.{self.schema}.{self.table_name}"
-        self.index_name = collection_name
+        self.index_name = _validate_identifier(collection_name, "collection_name")
         self.fully_qualified_index_name = f"{self.catalog}.{self.schema}.{self.index_name}"
 
         # Configuration

tests/vector_stores/test_databricks.py

@@ -1065,3 +1065,73 @@ def test_insert_timestamp_params_have_explicit_type(db_instance_delta, mock_work
     non_ts_params = [p for p in params if "created_at" not in p.name and "updated_at" not in p.name]
     for p in non_ts_params:
         assert p.type is None, f"Parameter {p.name} should not have explicit type, got {p.type}"
+
+
+class TestIdentifierValidation:
+    """Validate that catalog/schema/table_name/collection_name are sanitized."""
+
+    def test_rejects_semicolon_in_table_name(self, mock_workspace_client):
+        with pytest.raises(ValueError):
+            Databricks(
+                workspace_url="https://test",
+                access_token="tok",
+                endpoint_name="ep",
+                catalog="catalog",
+                schema="schema",
+                table_name="t; DROP TABLE x; --",
+                collection_name="mem0",
+                warehouse_name="test-warehouse",
+            )
+
+    def test_rejects_dot_in_catalog(self, mock_workspace_client):
+        with pytest.raises(ValueError):
+            Databricks(
+                workspace_url="https://test",
+                access_token="tok",
+                endpoint_name="ep",
+                catalog="cat.evil",
+                schema="schema",
+                table_name="table",
+                collection_name="mem0",
+                warehouse_name="test-warehouse",
+            )
+
+    def test_rejects_space_in_schema(self, mock_workspace_client):
+        with pytest.raises(ValueError):
+            Databricks(
+                workspace_url="https://test",
+                access_token="tok",
+                endpoint_name="ep",
+                catalog="catalog",
+                schema="schema name",
+                table_name="table",
+                collection_name="mem0",
+                warehouse_name="test-warehouse",
+            )
+
+    def test_rejects_dash_in_collection_name(self, mock_workspace_client):
+        with pytest.raises(ValueError):
+            Databricks(
+                workspace_url="https://test",
+                access_token="tok",
+                endpoint_name="ep",
+                catalog="catalog",
+                schema="schema",
+                table_name="table",
+                collection_name="mem-0",
+                warehouse_name="test-warehouse",
+            )
+
+    def test_accepts_valid_identifiers(self, mock_workspace_client):
+        db = Databricks(
+            workspace_url="https://test",
+            access_token="tok",
+            endpoint_name="ep",
+            catalog="my_catalog",
+            schema="my_schema",
+            table_name="my_table_01",
+            collection_name="mem0",
+            warehouse_name="test-warehouse",
+        )
+        assert db.fully_qualified_table_name == "my_catalog.my_schema.my_table_01"
+        assert db.fully_qualified_index_name == "my_catalog.my_schema.mem0"