diff --git a/mindsdb/api/executor/sql_query/steps/fetch_dataframe.py b/mindsdb/api/executor/sql_query/steps/fetch_dataframe.py index d73666e49e..8694fc63fb 100644 --- a/mindsdb/api/executor/sql_query/steps/fetch_dataframe.py +++ b/mindsdb/api/executor/sql_query/steps/fetch_dataframe.py @@ -16,7 +16,7 @@ from mindsdb.api.executor.exceptions import UnknownError from mindsdb.interfaces.query_context.context_controller import query_context_controller from mindsdb.integrations.utilities.query_traversal import query_traversal -from mindsdb.integrations.libs.response import TableResponse +from mindsdb.integrations.libs.response import TableResponse, OkResponse from .base import BaseStepCall @@ -92,7 +92,9 @@ def call(self, step): if query is None: table_alias = (self.context.get("database"), "result", "result") - response: TableResponse = dn.query(step.raw_query, session=self.session) + response: TableResponse | OkResponse = dn.query(step.raw_query, session=self.session) + if isinstance(response, OkResponse): + return ResultSet(affected_rows=response.affected_rows) df = response.data_frame else: if isinstance(step.query, (Union, Intersect)): diff --git a/mindsdb/integrations/handlers/postgres_handler/postgres_handler.py b/mindsdb/integrations/handlers/postgres_handler/postgres_handler.py index a3456a8e95..cbc7034e8d 100644 --- a/mindsdb/integrations/handlers/postgres_handler/postgres_handler.py +++ b/mindsdb/integrations/handlers/postgres_handler/postgres_handler.py @@ -305,8 +305,8 @@ def _cast_dtypes(self, df: DataFrame, description: list) -> DataFrame: def native_query(self, query: str, params=None, stream: bool = True, **kwargs) -> DataHandlerResponse: """Executes a SQL query on the PostgreSQL database and returns the result. - NOTE: 'INSERT' (and may be some else) queries can not be executed on the server side, - but there are fallbackto client side execution. + NOTE: only SELECT and VALUES queries can be executed on the server side, + other queries fall back to client side execution. Args: query (str): The SQL query to be executed. @@ -383,14 +383,10 @@ def _execute_server_side( try: try: cursor.execute(query) - except psycopg.errors.SyntaxError as e: - # NOTE: INSERT queries cannot be executed server-side. When they fail, they produce a syntax error - # that always starts with the text below, regardless of the INSERT query format. - lower_e = str(e).lower() - if not lower_e.startswith('syntax error at or near "insert"') and not lower_e.startswith( - 'syntax error at or near "drop"' - ): - raise + except (psycopg.errors.SyntaxError, psycopg.errors.FeatureNotSupported): + # NOTE: a server-side cursor is created with 'DECLARE ... CURSOR FOR ', which accepts only + # SELECT and VALUES. Any other statement (INSERT, UPDATE, CREATE, SHOW, EXPLAIN, etc) is rejected + # before it runs, so it is safe to execute it again with a client side cursor. connection.rollback() return self._execute_client_side(query=query) diff --git a/tests/unit/executor/test_executor.py b/tests/unit/executor/test_executor.py index 89a4acdfda..4bbb41cad6 100644 --- a/tests/unit/executor/test_executor.py +++ b/tests/unit/executor/test_executor.py @@ -1255,6 +1255,19 @@ def test_integration_native_query(self, mock_handler): assert mock_handler().native_query.call_args[0][0] == "select * from tasks" assert ret.data.to_lists()[0][0] == 3 + @patch("mindsdb.integrations.handlers.postgres_handler.Handler") + def test_integration_native_query_without_result(self, mock_handler): + from mindsdb.integrations.libs.response import OkResponse + + self.set_handler(mock_handler, name="pg", tables={}) + mock_handler().native_query.side_effect = None + mock_handler().native_query.return_value = OkResponse(affected_rows=2) + + ret = self.execute("select * from pg (update tasks set a = 1)") + + assert mock_handler().native_query.call_args[0][0] == "update tasks set a = 1" + assert ret.data.to_lists() == [] + @pytest.mark.slow @patch("mindsdb.integrations.handlers.postgres_handler.Handler") def test_view_native_query(self, mock_handler): diff --git a/tests/unit/handlers/test_postgres.py b/tests/unit/handlers/test_postgres.py index a0e3adc133..9795f51bd2 100644 --- a/tests/unit/handlers/test_postgres.py +++ b/tests/unit/handlers/test_postgres.py @@ -130,6 +130,71 @@ def test_native_query_command_ok_stream(self): assert isinstance(data, OkResponse) self.assertEqual(data.affected_rows, 1) + def test_native_query_stream_falls_back_to_client_side_when_cursor_is_not_allowed(self): + """ + Tests that statements rejected by 'DECLARE ... CURSOR FOR' are executed with a client side cursor + """ + with_update_error = psycopg.errors.FeatureNotSupported( + "DECLARE CURSOR must not contain data-modifying statements in WITH" + ) + cases = [ + ("UPDATE table SET a = 1", ExecStatus.COMMAND_OK, None), + ("DELETE FROM table", ExecStatus.COMMAND_OK, None), + ("CREATE TABLE table (a int)", ExecStatus.COMMAND_OK, None), + ("ALTER TABLE table ADD COLUMN b int", ExecStatus.COMMAND_OK, None), + ("TRUNCATE table", ExecStatus.COMMAND_OK, None), + ("SHOW search_path", ExecStatus.TUPLES_OK, None), + ("EXPLAIN SELECT 1", ExecStatus.TUPLES_OK, None), + ("WITH x AS (UPDATE table SET a = 1 RETURNING a) SELECT * FROM x", ExecStatus.TUPLES_OK, with_update_error), + ] + for query_str, status, server_error in cases: + with self.subTest(query=query_str): + mock_conn = MagicMock() + mock_cursor_server = MockCursorContextManager() + mock_cursor_client = MockCursorContextManager() + self.handler.connect = MagicMock(return_value=mock_conn) + mock_conn.cursor = MagicMock(side_effect=[mock_cursor_server, mock_cursor_client]) + + if server_error is None: + keyword = query_str.split()[0] + server_error = psycopg.errors.SyntaxError(f'syntax error at or near "{keyword}"') + mock_cursor_server.execute.side_effect = server_error + mock_cursor_client.pgresult = MagicMock(status=status) + mock_cursor_client.rowcount = 1 + mock_cursor_client.description = [ColumnDescription(name="a", type_code=regtype_to_oid["text"])] + mock_cursor_client.fetchall = MagicMock(return_value=[["value"]]) + + data = self.handler.native_query(query_str, stream=True) + + self.assertIn("name", mock_conn.cursor.call_args_list[0].kwargs) + mock_cursor_server.execute.assert_called_once_with(query_str) + mock_conn.rollback.assert_called_once() + mock_cursor_client.execute.assert_called_once_with(query_str) + if status == ExecStatus.COMMAND_OK: + self.assertIsInstance(data, OkResponse) + self.assertEqual(data.affected_rows, 1) + else: + self.assertIsInstance(data, TableResponse) + self.assertEqual(data.data_frame["a"].tolist(), ["value"]) + + def test_native_query_stream_does_not_retry_on_other_errors(self): + """ + Tests that errors not caused by the server-side cursor are returned without running the query again + """ + mock_conn = MagicMock() + mock_cursor = MockCursorContextManager() + self.handler.connect = MagicMock(return_value=mock_conn) + mock_conn.cursor = MagicMock(return_value=mock_cursor) + mock_cursor.execute.side_effect = psycopg.errors.UndefinedTable('relation "missing" does not exist') + + query_str = "SELECT * FROM missing" + data = self.handler.native_query(query_str, stream=True) + + mock_conn.cursor.assert_called_once() + mock_cursor.execute.assert_called_once_with(query_str) + self.assertIsInstance(data, ErrorResponse) + self.assertEqual(data.error_message, 'relation "missing" does not exist') + def test_native_query_command_ok_no_stream(self): """ Tests the `native_query` at client side execution @@ -212,6 +277,7 @@ def test_native_query_with_results_stream(self): query_str = "SELECT * FROM table" data = self.handler.native_query(query_str, stream=True) + mock_conn.cursor.assert_called_once_with(name=f"mindsdb_{id(self.handler)}") mock_cursor.execute.assert_called_once_with(query_str) # Verify the response