github-featured/google--adk-python
2
0
Fork 0
Code Issues Pull Requests Actions 8 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

fix: validate SQL identifiers in Spanner search tool

Browse Source 
Merge https://github.com/google/adk-python/pull/5952

Fixes #5913

PiperOrigin-RevId: 967358801
This commit is contained in:
Ashutosh0x
2026-08-19 12:50:52 -07:00
committed by Copybara-Service
parent 8989aeadce
commit 8d2f2779e6
3 changed files with 550 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/google/adk/tools/spanner/search_tool.py
+196 -5
View File
@@ -16,6 +16,7 @@ from __future__ import annotations
import asyncio
import json
import re
from typing import Any
from typing import Dict
from typing import List
@@ -31,6 +32,115 @@ from .settings import APPROXIMATE_NEAREST_NEIGHBORS
from .settings import EXACT_NEAREST_NEIGHBORS
from .settings import SpannerToolSettings
# Pattern for valid SQL identifiers: alphanumeric, underscores,
# dots (for schema-qualified names), and backtick/double-quote quoting.
# Supports per-part quoting for schema-qualified names.
_IDENTIFIER_PART_PATTERN = r'(?:[A-Za-z_][A-Za-z0-9_]*|`[^`\\]+`|"[^"\\]+")'
_SAFE_IDENTIFIER_RE = re.compile(
rf"^{_IDENTIFIER_PART_PATTERN}(?:\.{_IDENTIFIER_PART_PATTERN})*$"
)
# Operator allowlist for additional_filter
_ALLOWED_OPERATORS = r"(?:=|!=|<=|>=|<|>|(?i:\bLIKE\b|\bIS\s+NOT\b|\bIS\b))"
# Value allowlist for additional_filter: numbers, single-quoted strings (no backslashes), booleans, NULL
_ALLOWED_VALUES = (
r"(?:[+-]?\d+(?:\.\d+)?|'[^'\\]*'|(?i:\bTRUE\b|\bFALSE\b|\bNULL\b))"
)
# IN operator support
_IN_OPERATOR = r"(?i:\bNOT\s+IN\b|\bIN\b)"
_IN_VALUES = rf"\(\s*{_ALLOWED_VALUES}(?:\s*,\s*{_ALLOWED_VALUES})*\s*\)"
# BETWEEN operator support
_BETWEEN_OPERATOR = r"(?i:\bBETWEEN\b)"
_BETWEEN_VALUE = rf"{_ALLOWED_VALUES}\s+(?i:\bAND\b)\s+{_ALLOWED_VALUES}"
# A single condition (without paren)
_BASE_COND = (
rf"(?:"
rf"(?:{_IDENTIFIER_PART_PATTERN}(?:\.{_IDENTIFIER_PART_PATTERN})*)\s*{_ALLOWED_OPERATORS}\s*{_ALLOWED_VALUES}"
rf"|(?:{_IDENTIFIER_PART_PATTERN}(?:\.{_IDENTIFIER_PART_PATTERN})*)\s*{_IN_OPERATOR}\s*{_IN_VALUES}"
rf"|(?:{_IDENTIFIER_PART_PATTERN}(?:\.{_IDENTIFIER_PART_PATTERN})*)\s*{_BETWEEN_OPERATOR}\s*{_BETWEEN_VALUE}"
rf"|{_IDENTIFIER_PART_PATTERN}(?:\.{_IDENTIFIER_PART_PATTERN})*" # Just identifier (e.g. boolean col)
rf"|1\s*=\s*1" # dummy filter
rf")"
)
_BLOCK_0 = rf"{_BASE_COND}(?:\s+(?i:\bAND\b|\bOR\b)\s+{_BASE_COND})*"
_COND_1 = rf"(?:{_BASE_COND}|\(\s*{_BLOCK_0}\s*\))"
_BLOCK_1 = rf"{_COND_1}(?:\s+(?i:\bAND\b|\bOR\b)\s+{_COND_1})*"
_COND_2 = rf"(?:{_BASE_COND}|\(\s*{_BLOCK_1}\s*\))"
# Full filter pattern: conditions joined by AND/OR, supporting up to 2 levels of nested parens
_SAFE_FILTER_RE = re.compile(
rf"^\s*{_COND_2}(?:\s+(?i:\bAND\b|\bOR\b)\s+{_COND_2})*\s*$",
re.IGNORECASE,
)
def _validate_identifier(value: str, param_name: str) -> str:
"""Validate that a value is a safe SQL identifier.
Args:
value: The identifier string to validate.
param_name: Name of the parameter (for error messages).
Returns:
The validated identifier string.
Raises:
ValueError: If the identifier contains unsafe characters.
"""
if not value or not _SAFE_IDENTIFIER_RE.match(value.strip()):
raise ValueError(
f"Invalid SQL identifier for {param_name}: {value!r}. "
"Identifiers must contain only alphanumeric characters, underscores, "
"and dots, or be quoted with backticks or double quotes."
)
return value.strip()
def _validate_column_list(columns: List[str], param_name: str) -> List[str]:
"""Validate that each column name in a list is a safe SQL identifier."""
validated = []
for col in columns:
_validate_identifier(col, param_name)
validated.append(col)
return validated
def _validate_additional_filter(
filter_value: Optional[str],
) -> Optional[str]:
"""Validate that an additional_filter does not contain injection patterns.
This is a defense-in-depth measure. The additional_filter field is
documented as a developer-trusted value, but since it can be populated
by the LLM at runtime via tool calls, we restrict it to an allow-listed
grammar.
Args:
filter_value: The filter string to validate.
Returns:
The validated filter string, or None.
Raises:
ValueError: If the filter contains dangerous patterns.
"""
if filter_value is None:
return None
if not _SAFE_FILTER_RE.match(filter_value):
raise ValueError(
"additional_filter contains unsafe or unsupported patterns: "
f"{filter_value!r}. Only simple filters using =, !=, <=, >=, <, >, "
"LIKE, IS, IS NOT, IN, BETWEEN joined by AND or OR (with up to 2 "
"levels of nested parentheses) are allowed."
)
return filter_value
# Embedding model settings.
# Only for Spanner GoogleSQL dialect database, and use Spanner ML.PREDICT
# function.
@@ -75,6 +185,8 @@ def _generate_postgresql_for_embedding_query(
vertex_ai_embedding_model_endpoint: str,
output_dimensionality: Optional[int],
) -> str:
if output_dimensionality is not None:
output_dimensionality = int(output_dimensionality)
instances_json = f"""
'instances',
JSONB_BUILD_ARRAY(
@@ -166,6 +278,7 @@ def _generate_sql_for_knn(
top_k: int,
) -> str:
"""Generates a SQL query for kNN search."""
top_k = int(top_k)
if dialect == DatabaseDialect.POSTGRESQL:
distance_function = _get_postgresql_distance_function(distance_type)
embedding_parameter = f"${_POSTGRESQL_PARAMETER_QUERY_EMBEDDING}"
@@ -174,7 +287,7 @@ def _generate_sql_for_knn(
distance_type, ann=False
)
embedding_parameter = f"@{_GOOGLESQL_PARAMETER_QUERY_EMBEDDING}"
columns = columns + [f"""{distance_function}(
columns = list(columns) + [f"""{distance_function}(
{embedding_column_to_search},
{embedding_parameter}) AS {_DISTANCE_ALIAS}
"""]
@@ -205,13 +318,15 @@ def _generate_sql_for_ann(
num_leaves_to_search: int,
):
"""Generates a SQL query for ANN search."""
top_k = int(top_k)
num_leaves_to_search = int(num_leaves_to_search)
if dialect == DatabaseDialect.POSTGRESQL:
raise NotImplementedError(
f"{APPROXIMATE_NEAREST_NEIGHBORS} is not supported for PostgreSQL"
" dialect."
)
distance_function = _get_googlesql_distance_function(distance_type, ann=True)
columns = columns + [f"""{distance_function}(
columns = list(columns) + [f"""{distance_function}(
{embedding_column_to_search},
@{_GOOGLESQL_PARAMETER_QUERY_EMBEDDING},
options => JSON '{{"num_leaves_to_search": {num_leaves_to_search}}}'
@@ -297,7 +412,14 @@ async def similarity_search(
credentials (Credentials): The credentials to use for the request.
additional_filter (Optional[str]): An optional filter to apply to the
search query. If provided, this will be added to the WHERE clause of the
final query.
final query. Only simple filters are allowed. Supported grammar:
- Columns and values compared with: =, !=, <, >, <=, >=, LIKE, IS, IS NOT
- Set membership: IN, NOT IN (e.g., col IN (val1, val2))
- Range checks: BETWEEN ... AND ... (e.g., col BETWEEN val1 AND val2)
- Boolean columns (e.g., col_name) or dummy filter '1=1'
- Logical operators: AND, OR (case-insensitive)
- Parentheses: up to 2 levels of nesting (e.g., (col1 = val1 OR col2 = val2) AND col3 = val3)
Values must be numbers, single-quoted strings (without backslashes), booleans, or NULL.
search_options (Optional[Dict[str, Any]]): A dictionary of options to use
for the similarity search. The following options are supported:
- top_k: The number of most similar documents to return. The
@@ -328,7 +450,7 @@ async def similarity_search(
... project_id="my-project",
... instance_id="my-instance",
... database_id="my-database",
... table_name="my-product-table",
... table_name="my_product_table",
... query="Tools that can help me clean my house.",
... embedding_column_to_search="product_description_embedding",
... columns=["product_name", "product_description", "price_in_cents"],
@@ -362,6 +484,67 @@ async def similarity_search(
"""
# fmt: on
try:
# Validate input arguments to prevent SQL injection
_validate_identifier(table_name, "table_name")
_validate_identifier(
embedding_column_to_search, "embedding_column_to_search"
)
_validate_column_list(columns, "columns")
if additional_filter:
_validate_additional_filter(additional_filter)
opts = embedding_options or {}
gsql_model = opts.get(_SPANNER_GSQL_EMBEDDING_MODEL_NAME)
if gsql_model:
_validate_identifier(gsql_model, _SPANNER_GSQL_EMBEDDING_MODEL_NAME)
pg_endpoint = opts.get(_SPANNER_PG_VERTEX_AI_EMBEDDING_MODEL_ENDPOINT)
if pg_endpoint:
if not re.match(
r"^projects/[\w-]+/locations/[\w-]+/publishers/[\w-]+/models/[\w.-]+$",
pg_endpoint,
):
raise ValueError(
"Invalid Vertex AI endpoint format: "
f"{pg_endpoint!r}. Expected format: "
"projects/$project/locations/$location/publishers/google/models/$model"
)
return await _similarity_search_internal(
project_id=project_id,
instance_id=instance_id,
database_id=database_id,
table_name=table_name,
query=query,
embedding_column_to_search=embedding_column_to_search,
columns=columns,
embedding_options=embedding_options,
credentials=credentials,
additional_filter=additional_filter,
search_options=search_options,
)
except Exception as ex:
return {
"status": "ERROR",
"error_details": repr(ex),
}
async def _similarity_search_internal(
project_id: str,
instance_id: str,
database_id: str,
table_name: str,
query: str,
embedding_column_to_search: str,
columns: List[str],
embedding_options: Dict[str, str],
credentials: Credentials,
additional_filter: Optional[str] = None,
search_options: Optional[Dict[str, Any]] = None,
) -> Dict[str, Any]:
try:
# Get Spanner client
spanner_client = client.get_spanner_client(
project=project_id, credentials=credentials
@@ -426,6 +609,14 @@ async def similarity_search(
" must be specified for PostgreSQL dialect Spanner database."
)
output_dimensionality = embedding_options.get(_OUTPUT_DIMENSIONALITY)
if output_dimensionality is not None:
try:
output_dimensionality = int(output_dimensionality)
except (ValueError, TypeError):
raise ValueError(
f"Invalid value for {_OUTPUT_DIMENSIONALITY}:"
f" {output_dimensionality!r}. Must be an integer."
)
if (
output_dimensionality is not None
and spanner_gsql_embedding_model_name is not None
@@ -607,7 +798,7 @@ async def vector_store_similarity_search(
settings.vector_store_settings.num_leaves_to_search
)
return await similarity_search(
return await _similarity_search_internal(
project_id=settings.vector_store_settings.project_id,
instance_id=settings.vector_store_settings.instance_id,
database_id=settings.vector_store_settings.database_id,
tests/unittests/tools/spanner/test_search_tool.py
+64 -13
View File
@@ -33,7 +33,7 @@ def mock_spanner_ids():
"project_id": "test-project",
"instance_id": "test-instance",
"database_id": "test-database",
"table_name": "test-table",
"table_name": "test_table",
}
@@ -264,7 +264,7 @@ async def test_similarity_search_postgresql_knn_success(
columns=["col1"],
embedding_options={
"spanner_postgresql_vertex_ai_embedding_model_endpoint": (
"test_endpoint"
"projects/test-project/locations/us-central1/publishers/google/models/text-embedding-005"
)
},
credentials=mock_credentials,
@@ -302,7 +302,7 @@ async def test_similarity_search_postgresql_ann_unsupported(
columns=["col1"],
embedding_options={
"spanner_postgresql_vertex_ai_embedding_model_endpoint": (
"test_endpoint"
"projects/test-project/locations/us-central1/publishers/google/models/text-embedding-005"
)
},
credentials=mock_credentials,
@@ -341,7 +341,7 @@ async def test_similarity_search_gsql_missing_embedding_model_error(
columns=["col1"],
embedding_options={
"spanner_postgresql_vertex_ai_embedding_model_endpoint": (
"test_endpoint"
"projects/p/locations/l/publishers/google/models/m"
)
},
credentials=mock_credentials,
@@ -396,35 +396,35 @@ async def test_similarity_search_pg_missing_embedding_model_error(
[
pytest.param(
{
"vertex_ai_embedding_model_name": "test-model",
"spanner_googlesql_embedding_model_name": "test-model-2",
"vertex_ai_embedding_model_name": "test_model",
"spanner_googlesql_embedding_model_name": "test_model_2",
},
id="vertex_ai_and_googlesql",
),
pytest.param(
{
"vertex_ai_embedding_model_name": "test-model",
"vertex_ai_embedding_model_name": "test_model",
"spanner_postgresql_vertex_ai_embedding_model_endpoint": (
"test-endpoint"
"projects/p/locations/l/publishers/google/models/m"
),
},
id="vertex_ai_and_postgresql",
),
pytest.param(
{
"spanner_googlesql_embedding_model_name": "test-model",
"spanner_googlesql_embedding_model_name": "test_model",
"spanner_postgresql_vertex_ai_embedding_model_endpoint": (
"test-endpoint"
"projects/p/locations/l/publishers/google/models/m"
),
},
id="googlesql_and_postgresql",
),
pytest.param(
{
"vertex_ai_embedding_model_name": "test-model",
"spanner_googlesql_embedding_model_name": "test-model-2",
"vertex_ai_embedding_model_name": "test_model",
"spanner_googlesql_embedding_model_name": "test_model_2",
"spanner_postgresql_vertex_ai_embedding_model_endpoint": (
"test-endpoint"
"projects/p/locations/l/publishers/google/models/m"
),
},
id="all_three_models",
@@ -530,3 +530,54 @@ async def test_similarity_search_unsupported_algorithm_error(
)
assert result["status"] == "ERROR"
assert "Unsupported search_options" in result["error_details"]
@pytest.mark.asyncio
@mock.patch.object(utils, "embed_contents_async", autospec=True)
@mock.patch.object(client, "get_spanner_client")
async def test_vector_store_similarity_search_bypass_validation(
mock_get_spanner_client,
mock_embed_contents_async,
mock_credentials,
):
"""Test that vector_store_similarity_search bypasses validation for complex filter."""
from google.adk.tools.spanner.settings import SpannerToolSettings
from google.adk.tools.spanner.settings import SpannerVectorStoreSettings
mock_spanner_client = MagicMock()
mock_instance = MagicMock()
mock_database = MagicMock()
mock_snapshot = MagicMock()
mock_database.snapshot.return_value.__enter__.return_value = mock_snapshot
mock_database.database_dialect = DatabaseDialect.GOOGLE_STANDARD_SQL
mock_instance.database.return_value = mock_database
mock_spanner_client.instance.return_value = mock_instance
mock_get_spanner_client.return_value = mock_spanner_client
mock_embed_contents_async.return_value = [[0.1, 0.2, 0.3]]
mock_snapshot.execute_sql.return_value = iter([("result1",), ("result2",)])
# col1 = col2 is rejected by _validate_additional_filter, but should pass here
vector_store_settings = SpannerVectorStoreSettings(
project_id="test-project",
instance_id="test-instance",
database_id="test-database",
table_name="test_table",
content_column="content",
embedding_column="embedding",
vector_length=3,
vertex_ai_embedding_model_name="text-embedding-005",
selected_columns=["col1"],
additional_filter="col1 = col2",
)
tool_settings = SpannerToolSettings(
vector_store_settings=vector_store_settings,
)
result = await search_tool.vector_store_similarity_search(
query="test query",
credentials=mock_credentials,
settings=tool_settings,
)
assert result["status"] == "SUCCESS", result
assert result["rows"] == [("result1",), ("result2",)]
tests/unittests/tools/spanner/test_spanner_sql_validation.py
+290
View File
@@ -0,0 +1,290 @@
# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Tests for SQL identifier validation in Spanner search tool.
Verifies that malicious SQL identifiers and filter patterns are rejected
before being interpolated into SQL queries (defense against SQL injection
via LLM-populated tool parameters).
"""
from google.adk.tools.spanner.search_tool import _generate_postgresql_for_embedding_query
from google.adk.tools.spanner.search_tool import _generate_sql_for_ann
from google.adk.tools.spanner.search_tool import _generate_sql_for_knn
from google.adk.tools.spanner.search_tool import _validate_additional_filter
from google.adk.tools.spanner.search_tool import _validate_column_list
from google.adk.tools.spanner.search_tool import _validate_identifier
from google.adk.tools.spanner.search_tool import similarity_search
from google.cloud.spanner_admin_database_v1.types import DatabaseDialect
import pytest
class TestValidateIdentifier:
"""Tests for _validate_identifier."""
def test_simple_identifier(self):
assert _validate_identifier("documents", "test") == "documents"
def test_schema_qualified_identifier(self):
assert (
_validate_identifier("my_schema.my_table", "test")
== "my_schema.my_table"
)
def test_per_part_quoted_schema_qualified_identifier(self):
assert (
_validate_identifier("`my_schema`.`my_table`", "test")
== "`my_schema`.`my_table`"
)
assert (
_validate_identifier('"my_schema"."my_table"', "test")
== '"my_schema"."my_table"'
)
assert (
_validate_identifier("`my_schema`.my_table", "test")
== "`my_schema`.my_table"
)
assert (
_validate_identifier('my_schema."my_table"', "test")
== 'my_schema."my_table"'
)
def test_identifier_with_underscores(self):
assert _validate_identifier("embedding_col_1", "test") == "embedding_col_1"
def test_backtick_quoted_identifier(self):
assert _validate_identifier("`my table`", "test") == "`my table`"
def test_double_quote_quoted_identifier(self):
assert _validate_identifier('"my column"', "test") == '"my column"'
def test_rejects_join_injection(self):
with pytest.raises(ValueError, match="Invalid SQL identifier"):
_validate_identifier(
"documents JOIN admin_credentials ac ON TRUE", "table_name"
)
def test_rejects_subquery_in_column(self):
with pytest.raises(ValueError, match="Invalid SQL identifier"):
_validate_identifier(
"(SELECT STRING_AGG(table_name, ',') FROM INFORMATION_SCHEMA.TABLES)"
" AS schema_dump",
"columns",
)
def test_rejects_semicolon(self):
with pytest.raises(ValueError, match="Invalid SQL identifier"):
_validate_identifier("table; DROP TABLE users", "table_name")
def test_rejects_empty(self):
with pytest.raises(ValueError, match="Invalid SQL identifier"):
_validate_identifier("", "table_name")
def test_rejects_sql_comment(self):
with pytest.raises(ValueError, match="Invalid SQL identifier"):
_validate_identifier("table -- comment", "table_name")
def test_rejects_hyphen(self):
with pytest.raises(ValueError, match="Invalid SQL identifier"):
_validate_identifier("my-table", "table_name")
def test_rejects_backslash_escaped_identifier(self):
with pytest.raises(ValueError, match="Invalid SQL identifier"):
_validate_identifier("`a\\`", "table_name")
with pytest.raises(ValueError, match="Invalid SQL identifier"):
_validate_identifier('"a\\"', "table_name")
class TestValidateColumnList:
"""Tests for _validate_column_list."""
def test_valid_columns(self):
result = _validate_column_list(["col1", "col2", "col3"], "columns")
assert result == ["col1", "col2", "col3"]
def test_rejects_subquery_column(self):
with pytest.raises(ValueError, match="Invalid SQL identifier"):
_validate_column_list(
[
(
"(SELECT STRING_AGG(table_name, ',') FROM"
" INFORMATION_SCHEMA.TABLES) AS dump"
),
"content",
],
"columns",
)
class TestValidateAdditionalFilter:
"""Tests for _validate_additional_filter."""
def test_none_filter(self):
assert _validate_additional_filter(None) is None
def test_simple_filter(self):
assert (
_validate_additional_filter("price_in_cents < 100000")
== "price_in_cents < 100000"
)
def test_multiple_conditions_and(self):
assert (
_validate_additional_filter(
"price_in_cents < 100000 AND category = 'books'"
)
== "price_in_cents < 100000 AND category = 'books'"
)
def test_rejects_union(self):
with pytest.raises(ValueError, match="unsafe or unsupported patterns"):
_validate_additional_filter(
"1=1 UNION ALL SELECT password, 0.0 FROM admin_credentials"
)
def test_rejects_semicolon(self):
with pytest.raises(ValueError, match="unsafe or unsupported patterns"):
_validate_additional_filter("1=1; SELECT * FROM secrets")
def test_rejects_line_comment(self):
with pytest.raises(ValueError, match="unsafe or unsupported patterns"):
_validate_additional_filter("1=1 -- bypass")
def test_rejects_block_comment(self):
with pytest.raises(ValueError, match="unsafe or unsupported patterns"):
_validate_additional_filter("1=1 /* bypass */")
def test_rejects_hash_comment(self):
with pytest.raises(ValueError, match="unsafe or unsupported patterns"):
_validate_additional_filter("1=1 # bypass")
def test_rejects_subquery_exfiltration(self):
with pytest.raises(ValueError, match="unsafe or unsupported patterns"):
_validate_additional_filter(
"1=1 OR (SELECT password FROM admin_credentials) = 'x'"
)
def test_allows_or_condition(self):
assert (
_validate_additional_filter(
"price_in_cents < 100 OR category = 'books'"
)
== "price_in_cents < 100 OR category = 'books'"
)
def test_allows_in_condition(self):
assert (
_validate_additional_filter("category IN ('books', 'movies')")
== "category IN ('books', 'movies')"
)
def test_allows_between_condition(self):
assert (
_validate_additional_filter("price_in_cents BETWEEN 100 AND 500")
== "price_in_cents BETWEEN 100 AND 500"
)
def test_allows_nested_parentheses(self):
assert (
_validate_additional_filter(
"((price_in_cents < 100 OR category = 'books') AND status ="
" 'active') OR price_in_cents > 1000"
)
== "((price_in_cents < 100 OR category = 'books') AND status ="
" 'active') OR price_in_cents > 1000"
)
class TestGenerateSqlForKnn:
"""Tests for _generate_sql_for_knn."""
def test_valid_query_googlesql(self):
sql = _generate_sql_for_knn(
dialect=DatabaseDialect.GOOGLE_STANDARD_SQL,
table_name="documents",
embedding_column_to_search="embedding",
columns=["content"],
additional_filter=None,
distance_type="COSINE",
top_k=10,
)
assert "FROM documents" in sql
assert "COSINE_DISTANCE" in sql
def test_top_k_string_coerced_to_int(self):
sql = _generate_sql_for_knn(
dialect=DatabaseDialect.GOOGLE_STANDARD_SQL,
table_name="documents",
embedding_column_to_search="embedding",
columns=["content"],
additional_filter=None,
distance_type="COSINE",
top_k="10", # String input
)
assert "LIMIT 10" in sql
class TestSimilaritySearchInternalValidation:
"""Tests that similarity_search itself validates arguments."""
@pytest.mark.asyncio
async def test_rejects_invalid_table_name(self):
result = await similarity_search(
project_id="proj",
instance_id="inst",
database_id="db",
table_name="documents JOIN admin_credentials ac ON TRUE",
query="test query",
embedding_column_to_search="embedding_col",
columns=["col1"],
embedding_options={"vertex_ai_embedding_model_name": "model"},
credentials=None,
)
assert result["status"] == "ERROR"
assert "Invalid SQL identifier" in result["error_details"]
@pytest.mark.asyncio
async def test_rejects_invalid_additional_filter(self):
result = await similarity_search(
project_id="proj",
instance_id="inst",
database_id="db",
table_name="documents",
query="test query",
embedding_column_to_search="embedding_col",
columns=["col1"],
embedding_options={"vertex_ai_embedding_model_name": "model"},
credentials=None,
additional_filter=(
"1=1 UNION ALL SELECT password FROM admin_credentials"
),
)
assert result["status"] == "ERROR"
assert "unsafe or unsupported patterns" in result["error_details"]
class TestGeneratePostgresqlForEmbeddingQuery:
"""Tests for _generate_postgresql_for_embedding_query."""
def test_output_dimensionality_coerced(self):
sql = _generate_postgresql_for_embedding_query(
"projects/p/locations/l/publishers/g/models/m", "128"
)
assert "'outputDimensionality',\n 128" in sql
def test_output_dimensionality_invalid_raises(self):
with pytest.raises(ValueError):
_generate_postgresql_for_embedding_query(
"projects/p/locations/l/publishers/g/models/m", "invalid"
)