feat: support dml returning

test/test_suite_13.py

@@ -2065,6 +2065,49 @@ def test_create_engine_wo_database(self):
             assert connection.connection.database is None
 
 
+class ReturningTest(fixtures.TestBase):
+    def setUp(self):
+        self._engine = create_engine(get_db_url())
+        metadata = MetaData()
+
+        self._table = Table(
+            "returning_test",
+            metadata,
+            Column("id", Integer, primary_key=True),
+            Column("data", String(16), nullable=False),
+        )
+
+        metadata.create_all(self._engine)
+
+    def test_returning_for_insert_and_update(self):
+        random_id = random.randint(1, 1000)
+        with self._engine.begin() as connection:
+            stmt = (
+                self._table.insert()
+                .values(id=random_id, data="some % value")
+                .returning(self._table.c.id)
+            )
+            row = connection.execute(stmt).fetchall()
+            eq_(
+                row,
+                [(random_id,)],
+            )
+
+        with self._engine.begin() as connection:
+            update_text = "some + value"
+            stmt = (
+                self._table.update()
+                .values(data=update_text)
+                .where(self._table.c.id == random_id)
+                .returning(self._table.c.data)
+            )
+            row = connection.execute(stmt).fetchall()
+            eq_(
+                row,
+                [(update_text,)],
+            )
+
+
 @pytest.mark.skipif(
     bool(os.environ.get("SPANNER_EMULATOR_HOST")), reason="Skipped on emulator"
 )

test/test_suite_14.py

@@ -2399,6 +2399,49 @@ def test_create_engine_wo_database(self):
             assert connection.connection.database is None
 
 
+class ReturningTest(fixtures.TestBase):
+    def setUp(self):
+        self._engine = create_engine(get_db_url())
+        metadata = MetaData()
+
+        self._table = Table(
+            "returning_test",
+            metadata,
+            Column("id", Integer, primary_key=True),
+            Column("data", String(16), nullable=False),
+        )
+
+        metadata.create_all(self._engine)
+
+    def test_returning_for_insert_and_update(self):
+        random_id = random.randint(1, 1000)
+        with self._engine.begin() as connection:
+            stmt = (
+                self._table.insert()
+                .values(id=random_id, data="some % value")
+                .returning(self._table.c.id)
+            )
+            row = connection.execute(stmt).fetchall()
+            eq_(
+                row,
+                [(random_id,)],
+            )
+
+        with self._engine.begin() as connection:
+            update_text = "some + value"
+            stmt = (
+                self._table.update()
+                .values(data=update_text)
+                .where(self._table.c.id == random_id)
+                .returning(self._table.c.data)
+            )
+            row = connection.execute(stmt).fetchall()
+            eq_(
+                row,
+                [(update_text,)],
+            )
+
+
 @pytest.mark.skipif(
     bool(os.environ.get("SPANNER_EMULATOR_HOST")), reason="Skipped on emulator"
 )

test/test_suite_20.py

@@ -3113,6 +3113,49 @@ def test_create_engine_wo_database(self):
             assert connection.connection.database is None
 
 
+class ReturningTest(fixtures.TestBase):
+    def setUp(self):
+        self._engine = create_engine(get_db_url())
+        metadata = MetaData()
+
+        self._table = Table(
+            "returning_test",
+            metadata,
+            Column("id", Integer, primary_key=True),
+            Column("data", String(16), nullable=False),
+        )
+
+        metadata.create_all(self._engine)
+
+    def test_returning_for_insert_and_update(self):
+        random_id = random.randint(1, 1000)
+        with self._engine.begin() as connection:
+            stmt = (
+                self._table.insert()
+                .values(id=random_id, data="some % value")
+                .returning(self._table.c.id)
+            )
+            row = connection.execute(stmt).fetchall()
+            eq_(
+                row,
+                [(random_id,)],
+            )
+
+        with self._engine.begin() as connection:
+            update_text = "some + value"
+            stmt = (
+                self._table.update()
+                .values(data=update_text)
+                .where(self._table.c.id == random_id)
+                .returning(self._table.c.data)
+            )
+            row = connection.execute(stmt).fetchall()
+            eq_(
+                row,
+                [(update_text,)],
+            )
+
+
 @pytest.mark.skipif(
     bool(os.environ.get("SPANNER_EMULATOR_HOST")), reason="Skipped on emulator"
 )