diff --git a/doc/changes/unreleased.md b/doc/changes/unreleased.md index c6bd58b..2a58c6f 100644 --- a/doc/changes/unreleased.md +++ b/doc/changes/unreleased.md @@ -13,4 +13,5 @@ * #164: Changed the return type schema for built-in functions. * #168: Updated tool names and descriptions. * #170: Added tool list to the User Guide. -* #172: Made the slow tests running on only one version of Python - 3.12. +* #172: Made the slow SaaS tests running on only one version of Python - 3.12. +* #174: Added query validation tool. diff --git a/doc/user_guide/tool_list.rst b/doc/user_guide/tool_list.rst index 4c0a0d3..fa73721 100644 --- a/doc/user_guide/tool_list.rst +++ b/doc/user_guide/tool_list.rst @@ -177,8 +177,19 @@ describe_exasol_user_defined_function - ``type``: SQL type - ``dynamic_output``: for emit type UDF, indication that the UDF emits dynamic output -Tools Executing a Query ------------------------ +Tools Validating and Executing a Query +-------------------------------------- + +validate_exasol_query +~~~~~~~~~~~~~~~~~~~~~ + +:Description: + Executes the query, which must be a SELECT statement. In case of + successful execution returns nothing. Otherwise conveys an exception. + Validation doesn't work with DML or DDL queries. + +:Returns: + - **Type**: ``None or Exception`` execute_exasol_query ~~~~~~~~~~~~~~~~~~~~ diff --git a/exasol/ai/mcp/server/main.py b/exasol/ai/mcp/server/main.py index 2a99cae..24197e5 100644 --- a/exasol/ai/mcp/server/main.py +++ b/exasol/ai/mcp/server/main.py @@ -129,6 +129,19 @@ def _register_describe_script(mcp_server: ExasolMCPServer) -> None: ) +def _register_validate_query(mcp_server: ExasolMCPServer) -> None: + mcp_server.tool( + mcp_server.validate_query, + name="validate_exasol_query", + description=( + "Executes the query, which must be a SELECT statement. In case of " + "successful execution returns nothing. Otherwise conveys an exception. " + "Validation doesn't work with DML or DDL queries." + ), + annotations=ToolAnnotations(readOnlyHint=True), + ) + + def _register_execute_query(mcp_server: ExasolMCPServer) -> None: mcp_server.tool( mcp_server.execute_query, @@ -324,6 +337,7 @@ def register_tools(mcp_server: ExasolMCPServer, config: McpServerSettings) -> No if config.parameters.enable: _register_describe_function(mcp_server) _register_describe_script(mcp_server) + _register_validate_query(mcp_server) if config.enable_read_query: _register_execute_query(mcp_server) if config.enable_write_query: diff --git a/exasol/ai/mcp/server/tools/mcp_server.py b/exasol/ai/mcp/server/tools/mcp_server.py index 4c4d0f9..f97b282 100644 --- a/exasol/ai/mcp/server/tools/mcp_server.py +++ b/exasol/ai/mcp/server/tools/mcp_server.py @@ -322,6 +322,11 @@ def describe_script( DBReturnFunction | DBEmitFunction, parser.describe(schema_name, script_name) ) + def validate_query(self, query: QueryArg) -> None: + if not verify_query(query): + raise ValueError("Unable to verify that the query is a SELECT statement.") + self.connection.execute_query(query, snapshot=False) + def execute_query(self, query: QueryArg) -> list[dict[str, Any]]: if not self.config.enable_read_query: raise RuntimeError("Query execution is disabled.") diff --git a/test/integration/tools/test_mcp_server.py b/test/integration/tools/test_mcp_server.py index f668808..572ca7a 100644 --- a/test/integration/tools/test_mcp_server.py +++ b/test/integration/tools/test_mcp_server.py @@ -805,6 +805,51 @@ def test_describe_script( assert result_json == expected_json +def test_validate_query(pyexasol_connection, setup_database, db_schemas, db_tables): + """ + Test the `validate_query` tool. Validates a correct SELECT query. + """ + config = McpServerSettings() + for schema in db_schemas: + for table in db_tables: + query = f'SELECT * FROM "{schema.name}"."{table.name}"' + result = run_tool( + pyexasol_connection, + config, + tool_name="validate_exasol_query", + query=query, + ) + assert not result.content + + +@pytest.mark.parametrize( + "query", + [ + 'SELECT * INTO TABLE "{0}"."ANOTHER_TABLE" FROM "{0}"."{1}"', # DDL + 'SELECT ** FROM "{0}"."{1}"', # Invalid SQL + 'SELECT "NON_EXISTENT_COLUMN" FROM "{0}"."{1}"', # Wrong column + ], + ids=["DDL", "Invalid SQL", "Wrong column"], +) +def test_validate_query_error( + pyexasol_connection, setup_database, db_schemas, db_tables, query +): + """ + The test validates that the `validate_query` tool fails when offered either + disallowed or invalid query. + """ + config = McpServerSettings() + for schema in db_schemas: + for table in db_tables: + with pytest.raises(ToolError): + run_tool( + pyexasol_connection, + config, + tool_name="validate_exasol_query", + query=query.format(schema.name, table.name), + ) + + def test_execute_query(pyexasol_connection, setup_database, db_schemas, db_tables): """ Test the `execute_query` tool. Runs the simplest SELECT query that grabs the entire diff --git a/test/integration/tools/test_tool_hints.py b/test/integration/tools/test_tool_hints.py index f19f342..ee9a0d3 100644 --- a/test/integration/tools/test_tool_hints.py +++ b/test/integration/tools/test_tool_hints.py @@ -43,6 +43,7 @@ def test_tool_hints(pyexasol_connection) -> None: ToolHints(tool_name="describe_exasol_table_or_view", read_only=True), ToolHints(tool_name="describe_exasol_custom_function", read_only=True), ToolHints(tool_name="describe_exasol_user_defined_function", read_only=True), + ToolHints(tool_name="validate_exasol_query", read_only=True), ToolHints(tool_name="execute_exasol_query", read_only=True), ToolHints(tool_name="execute_exasol_write_query", destructive=True), ToolHints(tool_name="list_exasol_sql_types", read_only=True, idempotent=True),