Fix RETURNING

Author: antoineleclair

src/dqlitedbapi/__pycache__/cursor.cpython-313.pyc

@@ -64,8 +64,11 @@ async def execute(
 
         params = list(parameters) if parameters else None
 
-        # Determine if this is a SELECT query
-        is_query = operation.strip().upper().startswith(("SELECT", "PRAGMA", "EXPLAIN"))
+        # Determine if this is a query that returns rows
+        normalized = operation.strip().upper()
+        is_query = normalized.startswith(("SELECT", "PRAGMA", "EXPLAIN")) or (
+            " RETURNING " in normalized or normalized.endswith(" RETURNING")
+        )
 
         if is_query:
             assert conn._protocol is not None and conn._db_id is not None

src/dqlitedbapi/aio/cursor.py

@@ -79,8 +79,11 @@ async def _execute_async(self, operation: str, parameters: Sequence[Any] | None
         conn = await self._connection._get_async_connection()
         params = list(parameters) if parameters else None
 
-        # Determine if this is a SELECT query
-        is_query = operation.strip().upper().startswith(("SELECT", "PRAGMA", "EXPLAIN"))
+        # Determine if this is a query that returns rows
+        normalized = operation.strip().upper()
+        is_query = normalized.startswith(("SELECT", "PRAGMA", "EXPLAIN")) or (
+            " RETURNING " in normalized or normalized.endswith(" RETURNING")
+        )
 
         if is_query:
             assert conn._protocol is not None and conn._db_id is not None

src/dqlitedbapi/cursor.py

@@ -0,0 +1,64 @@
+"""Integration tests for INSERT/UPDATE/DELETE ... RETURNING support."""
+
+import pytest
+
+from dqlitedbapi import connect
+
+
+@pytest.mark.integration
+class TestReturning:
+    def test_insert_returning(self, cluster_address: str) -> None:
+        """INSERT ... RETURNING should return rows via query path."""
+        conn = connect(cluster_address, database="test_returning")
+        cursor = conn.cursor()
+
+        cursor.execute("CREATE TABLE ret_test (id INTEGER PRIMARY KEY, name TEXT)")
+        cursor.execute("INSERT INTO ret_test (id, name) VALUES (1, 'alice') RETURNING id, name")
+        row = cursor.fetchone()
+        assert row == (1, "alice")
+
+        cursor.execute("DROP TABLE ret_test")
+        conn.close()
+
+    def test_insert_returning_multiple(self, cluster_address: str) -> None:
+        """INSERT ... RETURNING should support fetching all returned rows."""
+        conn = connect(cluster_address, database="test_returning_multi")
+        cursor = conn.cursor()
+
+        cursor.execute("CREATE TABLE ret_multi (id INTEGER PRIMARY KEY, val TEXT)")
+        cursor.execute("INSERT INTO ret_multi (id, val) VALUES (1, 'a') RETURNING id, val")
+        assert cursor.fetchone() == (1, "a")
+
+        cursor.execute("INSERT INTO ret_multi (id, val) VALUES (2, 'b') RETURNING val")
+        assert cursor.fetchone() == ("b",)
+
+        cursor.execute("DROP TABLE ret_multi")
+        conn.close()
+
+    def test_delete_returning(self, cluster_address: str) -> None:
+        """DELETE ... RETURNING should return deleted rows."""
+        conn = connect(cluster_address, database="test_del_returning")
+        cursor = conn.cursor()
+
+        cursor.execute("CREATE TABLE del_ret (id INTEGER PRIMARY KEY, name TEXT)")
+        cursor.execute("INSERT INTO del_ret (id, name) VALUES (1, 'alice')")
+        cursor.execute("DELETE FROM del_ret WHERE id = 1 RETURNING id, name")
+        row = cursor.fetchone()
+        assert row == (1, "alice")
+
+        cursor.execute("DROP TABLE del_ret")
+        conn.close()
+
+    def test_update_returning(self, cluster_address: str) -> None:
+        """UPDATE ... RETURNING should return updated rows."""
+        conn = connect(cluster_address, database="test_upd_returning")
+        cursor = conn.cursor()
+
+        cursor.execute("CREATE TABLE upd_ret (id INTEGER PRIMARY KEY, name TEXT)")
+        cursor.execute("INSERT INTO upd_ret (id, name) VALUES (1, 'alice')")
+        cursor.execute("UPDATE upd_ret SET name = 'bob' WHERE id = 1 RETURNING id, name")
+        row = cursor.fetchone()
+        assert row == (1, "bob")
+
+        cursor.execute("DROP TABLE upd_ret")
+        conn.close()