JDSS/mem0
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: prevent SQL injection in Databricks vector store (#4558)

Browse Source 
Co-authored-by: utkarsh240799 <utkarsh240799@users.noreply.github.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Utkarsh
2026-03-26 20:22:34 +05:30
committed by GitHub
co-authored by utkarsh240799 Claude Opus 4.6
parent 3a9fcdbec2
commit a2ffca3266
2 changed files with 286 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mem0/vector_stores/databricks.py
+51 -12
View File
@@ -1,11 +1,13 @@
import json
import logging
import re
import uuid
from typing import Optional, List
from datetime import datetime, date
from databricks.sdk.service.catalog import ColumnInfo, ColumnTypeName, TableType, DataSourceFormat
from databricks.sdk.service.catalog import TableConstraint, PrimaryKeyConstraint
from databricks.sdk import WorkspaceClient
from databricks.sdk.service.sql import StatementParameterListItem
from databricks.sdk.service.vectorsearch import (
VectorIndexType,
DeltaSyncVectorIndexSpecRequest,
@@ -28,6 +30,9 @@ 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
_VALID_SQL_IDENTIFIER = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")
class Databricks(VectorStoreBase):
def __init__(
@@ -388,28 +393,46 @@ class Databricks(VectorStoreBase):
# Determine the number of items to process
num_items = len(payloads) if payloads else len(vectors) if vectors else 0
params = []
value_tuples = []
for i in range(num_items):
values = []
placeholders = []
for col in self.columns:
param_name = f"{col.name}_{i}"
if col.name == "memory_id":
val = ids[i] if ids and i < len(ids) else str(uuid.uuid4())
elif col.name == "embedding":
val = vectors[i] if vectors and i < len(vectors) else []
# Vectors are numeric arrays — ARRAY type not supported by StatementParameterListItem,
# so we inline using _format_sql_value (values are floats from the embedding model).
placeholders.append(self._format_sql_value(val))
continue
elif col.name == "memory":
val = payloads[i].get("data") if payloads and i < len(payloads) else None
else:
val = payloads[i].get(col.name) if payloads and i < len(payloads) else None
values.append(val)
formatted = [self._format_sql_value(v) for v in values]
value_tuples.append(f"({', '.join(formatted)})")
if val is None:
placeholders.append("NULL")
else:
placeholders.append(f":{param_name}")
if isinstance(val, dict):
val = json.dumps(val)
# Use explicit type for TIMESTAMP columns so Databricks doesn't
# rely on implicit STRING→TIMESTAMP casting.
param_type = "TIMESTAMP" if col.type_name == ColumnTypeName.TIMESTAMP else None
params.append(StatementParameterListItem(name=param_name, value=str(val), type=param_type))
value_tuples.append(f"({', '.join(placeholders)})")
insert_sql = f"INSERT INTO {self.fully_qualified_table_name} ({', '.join(self.column_names)}) VALUES {', '.join(value_tuples)}"
# Execute the insert
try:
response = self.client.statement_execution.execute_statement(
statement=insert_sql, warehouse_id=self.warehouse_id, wait_timeout="30s"
statement=insert_sql,
warehouse_id=self.warehouse_id,
wait_timeout="30s",
parameters=params,
)
if response.status.state.value == "SUCCEEDED":
logger.info(
@@ -494,10 +517,13 @@ class Databricks(VectorStoreBase):
try:
logger.info(f"Deleting vector with ID {vector_id} from Delta table {self.fully_qualified_table_name}")
delete_sql = f"DELETE FROM {self.fully_qualified_table_name} WHERE memory_id = '{vector_id}'"
delete_sql = f"DELETE FROM {self.fully_qualified_table_name} WHERE memory_id = :vector_id"
response = self.client.statement_execution.execute_statement(
statement=delete_sql, warehouse_id=self.warehouse_id, wait_timeout="30s"
statement=delete_sql,
warehouse_id=self.warehouse_id,
wait_timeout="30s",
parameters=[StatementParameterListItem(name="vector_id", value=str(vector_id))],
)
if response.status.state.value == "SUCCEEDED":
@@ -519,8 +545,8 @@ class Databricks(VectorStoreBase):
payload (dict, optional): New payload data.
"""
update_sql = f"UPDATE {self.fully_qualified_table_name} SET "
set_clauses = []
params = []
if not vector_id:
logger.error("vector_id is required for update operation")
return
@@ -528,25 +554,38 @@ class Databricks(VectorStoreBase):
if not isinstance(vector, list):
logger.error("vector must be a list of float values")
return
set_clauses.append(f"embedding = {vector}")
# Vectors are numeric arrays — safe to inline since StatementParameterListItem
# doesn't support ARRAY types, and values are validated as list of floats above.
# Use array() SQL syntax, not Python list repr which is invalid Databricks SQL.
set_clauses.append(f"embedding = {self._format_sql_value(vector)}")
if payload:
if not isinstance(payload, dict):
logger.error("payload must be a dictionary")
return
for key, value in payload.items():
if key not in excluded_keys:
set_clauses.append(f"{key} = '{value}'")
if not _VALID_SQL_IDENTIFIER.match(key):
logger.warning(f"Skipping invalid column name in payload: {key!r}")
continue
param_name = f"payload_{key}"
set_clauses.append(f"{key} = :{param_name}")
params.append(StatementParameterListItem(name=param_name, value=str(value)))
if not set_clauses:
logger.error("No fields to update")
return
update_sql = f"UPDATE {self.fully_qualified_table_name} SET "
update_sql += ", ".join(set_clauses)
update_sql += f" WHERE memory_id = '{vector_id}'"
update_sql += " WHERE memory_id = :vector_id"
params.append(StatementParameterListItem(name="vector_id", value=str(vector_id)))
try:
logger.info(f"Updating vector with ID {vector_id} in Delta table {self.fully_qualified_table_name}")
response = self.client.statement_execution.execute_statement(
statement=update_sql, warehouse_id=self.warehouse_id, wait_timeout="30s"
statement=update_sql,
warehouse_id=self.warehouse_id,
wait_timeout="30s",
parameters=params,
)
if response.status.state.value == "SUCCEEDED":