From 2dcb3542f8e9490b5aab9cba55d29e119b2684f0 Mon Sep 17 00:00:00 2001
From: Hrushikesh Yadav <136978914+HrushiYadav@users.noreply.github.com>
Date: Mon, 29 Jun 2026 20:30:27 +0530
Subject: [PATCH] fix(databricks): validate catalog/schema/table identifiers to
 prevent SQL injection (#5988)

---
 mem0/vector_stores/databricks.py       | 16 ++++--
 tests/vector_stores/test_databricks.py | 70 ++++++++++++++++++++++++++
 2 files changed, 81 insertions(+), 5 deletions(-)

diff --git a/mem0/vector_stores/databricks.py b/mem0/vector_stores/databricks.py
index afc13d7bf..6566123b8 100644
--- a/mem0/vector_stores/databricks.py
+++ b/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.schema = schema
-        self.table_name = table_name
+        self.catalog = _validate_identifier(catalog, "catalog")
+        self.schema = _validate_identifier(schema, "schema")
+        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
diff --git a/tests/vector_stores/test_databricks.py b/tests/vector_stores/test_databricks.py
index 1a970d7a3..cd02af60e 100644
--- a/tests/vector_stores/test_databricks.py
+++ b/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"