fix: ensure Databricks SQL compatibility for parameterized queries

- Add explicit type='TIMESTAMP' on StatementParameterListItem for
  created_at/updated_at columns instead of relying on implicit
  STRING->TIMESTAMP casting
- Fix pre-existing bug: update() used Python list repr [0.1, 0.2]
  for embedding which is invalid Databricks SQL, now uses
  _format_sql_value() to produce array(0.1, 0.2) syntax

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
utkarsh240799
2026-03-26 19:11:13 +05:30
parent c6cc8f14b8
commit 1eb6bd8303
2 changed files with 33 additions and 4 deletions
+6 -2
View File
@@ -418,7 +418,10 @@ class Databricks(VectorStoreBase):
placeholders.append(f":{param_name}")
if isinstance(val, dict):
val = json.dumps(val)
params.append(StatementParameterListItem(name=param_name, value=str(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)}"
@@ -553,7 +556,8 @@ class Databricks(VectorStoreBase):
return
# Vectors are numeric arrays — safe to inline since StatementParameterListItem
# doesn't support ARRAY types, and values are validated as list of floats above.
set_clauses.append(f"embedding = {vector}")
# 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")
+27 -2
View File
@@ -275,7 +275,7 @@ def test_update_vector(db_instance_direct, mock_workspace_client):
args, kwargs = mock_workspace_client.statement_execution.execute_statement.call_args
sql = kwargs.get("statement") or args[0]
assert "UPDATE" in sql
assert "embedding = [0.4, 0.5, 0.6, 0.7]" in sql
assert "embedding = array(0.4, 0.5, 0.6, 0.7)" in sql
assert "custom = :payload_custom" in sql
assert ":vector_id" in sql
assert "id-upd" not in sql # value should be in params, not in SQL
@@ -1030,8 +1030,33 @@ def test_update_vector_only(db_instance_direct, mock_workspace_client):
kwargs = mock_workspace_client.statement_execution.execute_statement.call_args.kwargs
sql = kwargs["statement"]
assert "UPDATE" in sql
assert "embedding = [0.1, 0.2, 0.3, 0.4]" in sql
assert "embedding = array(0.1, 0.2, 0.3, 0.4)" in sql
assert ":vector_id" in sql
param_map = {p.name: p.value for p in kwargs["parameters"]}
assert param_map["vector_id"] == "id-1"
assert len(kwargs["parameters"]) == 1 # only vector_id param
def test_insert_timestamp_params_have_explicit_type(db_instance_delta, mock_workspace_client):
"""Verify insert() sets type='TIMESTAMP' on created_at/updated_at parameters
so Databricks doesn't rely on implicit STRING->TIMESTAMP casting."""
db_instance_delta.insert(
vectors=[[0.1, 0.2]],
payloads=[{
"data": "test",
"hash": "h1",
"created_at": "2024-01-01T00:00:00+00:00",
"updated_at": "2024-01-02T00:00:00+00:00",
}],
ids=["id-1"],
)
kwargs = mock_workspace_client.statement_execution.execute_statement.call_args.kwargs
params = kwargs["parameters"]
ts_params = [p for p in params if "created_at" in p.name or "updated_at" in p.name]
assert len(ts_params) == 2
for p in ts_params:
assert p.type == "TIMESTAMP", f"Parameter {p.name} should have type=TIMESTAMP, got {p.type}"
# Non-timestamp params should not have a type set (defaults to STRING)
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}"