diff --git a/src/google/adk/tools/spanner/search_tool.py b/src/google/adk/tools/spanner/search_tool.py index 6fb4a93f..6b723d60 100644 --- a/src/google/adk/tools/spanner/search_tool.py +++ b/src/google/adk/tools/spanner/search_tool.py @@ -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, diff --git a/tests/unittests/tools/spanner/test_search_tool.py b/tests/unittests/tools/spanner/test_search_tool.py index c6a6c742..c210d69f 100644 --- a/tests/unittests/tools/spanner/test_search_tool.py +++ b/tests/unittests/tools/spanner/test_search_tool.py @@ -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",)] diff --git a/tests/unittests/tools/spanner/test_spanner_sql_validation.py b/tests/unittests/tools/spanner/test_spanner_sql_validation.py new file mode 100644 index 00000000..0a458013 --- /dev/null +++ b/tests/unittests/tools/spanner/test_spanner_sql_validation.py @@ -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" + )