From 5f176aa82f2c8f61d420138ebe4feabc74da9adf Mon Sep 17 00:00:00 2001 From: Maha Rana Yadavalli <271375718+maharanay22@users.noreply.github.com> Date: Tue, 13 Oct 2026 11:57:03 -0400 Subject: [PATCH] Fix PostgreSQL native queries other than SELECT failing with server-side cursor Since the data streaming change, native queries are run through a server-side cursor (DECLARE ... CURSOR FOR ). PostgreSQL only accepts SELECT and VALUES there, so UPDATE, DELETE, CREATE, ALTER, TRUNCATE, SHOW, EXPLAIN and similar statements failed with 'syntax error at or near "UPDATE"'. Only INSERT and DROP were retried on the client side, based on the error text. Fall back to a client side cursor whenever DECLARE rejects the statement with SyntaxError or FeatureNotSupported. DECLARE rejects the statement before running it, so running it again is safe. Also return an empty result for 'SELECT * FROM db ()' when the handler answers with OkResponse, instead of failing with "'OkResponse' object has no attribute 'data_frame'". --- .../sql_query/steps/fetch_dataframe.py | 6 +- .../postgres_handler/postgres_handler.py | 16 ++--- tests/unit/executor/test_executor.py | 13 ++++ tests/unit/handlers/test_postgres.py | 66 +++++++++++++++++++ 4 files changed, 89 insertions(+), 12 deletions(-) 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