feat: test rows_affected

Author: lidavidm

adbc_drivers_validation/model.py

@@ -93,6 +93,7 @@ class DriverFeatures:
     statement_execute_schema: bool = False
     statement_get_parameter_schema: bool = False
     statement_prepare: bool = True
+    statement_rows_affected: bool = False
     _current_catalog: str | FromEnv | None = None
     _current_schema: str | FromEnv | None = None
     _secondary_schema: str | FromEnv | None = None

adbc_drivers_validation/tests/ingest.py

@@ -115,7 +115,11 @@ def test_create(
                     raise
 
             with setup_statement(query, cursor):
-                cursor.adbc_ingest(table_name, data, mode="create")
+                modified = cursor.adbc_ingest(table_name, data, mode="create")
+                if driver.features.statement_rows_affected:
+                    assert modified == len(data)
+                else:
+                    assert modified == -1
 
         idx = driver.quote_identifier("idx")
         value = driver.quote_identifier("value")
@@ -158,8 +162,17 @@ def test_append(
                 # Some databases have no way to do DROP IF EXISTS
                 if not driver.is_table_not_found(table_name=table_name, error=e):
                     raise
-            cursor.adbc_ingest(table_name, data, mode="create")
-            cursor.adbc_ingest(table_name, data2, mode="append")
+            modified = cursor.adbc_ingest(table_name, data, mode="create")
+            if driver.features.statement_rows_affected:
+                assert modified == len(data)
+            else:
+                assert modified == -1
+
+            modified = cursor.adbc_ingest(table_name, data2, mode="append")
+            if driver.features.statement_rows_affected:
+                assert modified == len(data2)
+            else:
+                assert modified == -1
 
         idx = driver.quote_identifier("idx")
         value = driver.quote_identifier("value")
@@ -234,8 +247,17 @@ def test_createappend(
                 # Some databases have no way to do DROP IF EXISTS
                 if not driver.is_table_not_found(table_name=table_name, error=e):
                     raise
-            cursor.adbc_ingest(table_name, data, mode="create_append")
-            cursor.adbc_ingest(table_name, data2, mode="create_append")
+            modified = cursor.adbc_ingest(table_name, data, mode="create_append")
+            if driver.features.statement_rows_affected:
+                assert modified == len(data)
+            else:
+                assert modified == -1
+
+            modified = cursor.adbc_ingest(table_name, data2, mode="create_append")
+            if driver.features.statement_rows_affected:
+                assert modified == len(data2)
+            else:
+                assert modified == -1
 
         idx = driver.quote_identifier("idx")
         value = driver.quote_identifier("value")
@@ -280,10 +302,20 @@ def test_replace(
                 if not driver.is_table_not_found(table_name=table_name, error=e):
                     raise
             cursor.adbc_ingest(table_name, data, mode="replace")
+            modified = cursor.adbc_ingest(table_name, data, mode="replace")
+            if driver.features.statement_rows_affected:
+                assert modified == len(data)
+            else:
+                assert modified == -1
+
             if driver.name == "bigquery":
                 # BigQuery rate-limits metadata operations
                 time.sleep(5)
-            cursor.adbc_ingest(table_name, data2, mode="replace")
+            modified = cursor.adbc_ingest(table_name, data2, mode="replace")
+            if driver.features.statement_rows_affected:
+                assert modified == len(data2)
+            else:
+                assert modified == -1
 
         idx = driver.quote_identifier("idx")
         value = driver.quote_identifier("value")
@@ -316,7 +348,11 @@ def test_replace_noop(
                 # Some databases have no way to do DROP IF EXISTS
                 if not driver.is_table_not_found(table_name=table_name, error=e):
                     raise
-            cursor.adbc_ingest(table_name, data, mode="replace")
+            modified = cursor.adbc_ingest(table_name, data, mode="replace")
+            if driver.features.statement_rows_affected:
+                assert modified == len(data)
+            else:
+                assert modified == -1
 
         idx = driver.quote_identifier("idx")
         value = driver.quote_identifier("value")

adbc_drivers_validation/tests/statement.py

@@ -144,3 +144,61 @@ def test_transaction_toggle(
 
         conn.adbc_connection.set_options(**{"adbc.connection.autocommit": True})
         assert conn.adbc_connection.get_option("adbc.connection.autocommit") == "true"
+
+    def test_rows_affected(
+        self,
+        driver: model.DriverQuirks,
+        conn: adbc_driver_manager.dbapi.Connection,
+    ) -> None:
+        table_name = "test_rows_affected"
+        with conn.cursor() as cursor:
+            cursor.adbc_statement.set_sql_query(
+                driver.drop_table(table_name="test_rows_affected")
+            )
+            try:
+                cursor.adbc_statement.execute_update()
+            except adbc_driver_manager.Error as e:
+                # Some databases have no way to do DROP IF EXISTS
+                if not driver.is_table_not_found(table_name=table_name, error=e):
+                    raise
+
+            cursor.adbc_statement.set_sql_query(f"CREATE TABLE {table_name} (id INT)")
+            rows_affected = cursor.adbc_statement.execute_update()
+
+            if driver.features.statement_rows_affected:
+                assert rows_affected == 0
+            else:
+                assert rows_affected == -1
+
+            cursor.adbc_statement.set_sql_query(
+                f"INSERT INTO {table_name} (id) VALUES (1)"
+            )
+            rows_affected = cursor.adbc_statement.execute_update()
+            if driver.features.statement_rows_affected:
+                assert rows_affected == 1
+            else:
+                assert rows_affected == -1
+
+            cursor.adbc_statement.set_sql_query(f"UPDATE {table_name} SET id = id + 1")
+            rows_affected = cursor.adbc_statement.execute_update()
+            if driver.features.statement_rows_affected:
+                assert rows_affected == 1
+            else:
+                assert rows_affected == -1
+
+            cursor.adbc_statement.set_sql_query(f"DELETE FROM {table_name}")
+            rows_affected = cursor.adbc_statement.execute_update()
+            if driver.features.statement_rows_affected:
+                assert rows_affected == 1
+            else:
+                assert rows_affected == -1
+
+            cursor.adbc_statement.set_sql_query(
+                driver.drop_table(table_name="test_rows_affected")
+            )
+            try:
+                cursor.adbc_statement.execute_update()
+            except adbc_driver_manager.Error as e:
+                # Some databases have no way to do DROP IF EXISTS
+                if not driver.is_table_not_found(table_name=table_name, error=e):
+                    raise