Fix PostgreSQL native queries other than SELECT failing with server-side cursor

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)):

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 <query>', 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)
 

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):

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