From 559533a838e045e70ac6c61f715d6d81b74339bd Mon Sep 17 00:00:00 2001 From: Yannick Utard Date: Thu, 5 Mar 2026 21:57:39 +0100 Subject: [PATCH 1/3] Fix execute result --- src/altertable_flightsql/client.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/src/altertable_flightsql/client.py b/src/altertable_flightsql/client.py index 9735fe9..f6b7951 100644 --- a/src/altertable_flightsql/client.py +++ b/src/altertable_flightsql/client.py @@ -242,14 +242,18 @@ def execute( # Execute via DoPut writer, reader = self._client.do_put(descriptor, pa.schema([])) - writer.close() + # Signal end of upload while keeping the read side open to receive the + # server's DoPutUpdateResult metadata. writer.close() would close both + # sides prematurely, causing reader.read() to return None. + writer.done_writing() # Read result from metadata result = sql_pb2.DoPutUpdateResult() metadata = reader.read() if metadata: - result.ParseFromString(metadata) + result.ParseFromString(bytes(metadata)) + writer.close() return result.record_count def ingest( From 64ad166e8051756d2d979648abf2278b9d7e07e2 Mon Sep 17 00:00:00 2001 From: Yannick Utard Date: Thu, 5 Mar 2026 22:07:42 +0100 Subject: [PATCH 2/3] Add tests --- tests/test_execute.py | 48 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 48 insertions(+) create mode 100644 tests/test_execute.py diff --git a/tests/test_execute.py b/tests/test_execute.py new file mode 100644 index 0000000..4e0ad2b --- /dev/null +++ b/tests/test_execute.py @@ -0,0 +1,48 @@ +""" +Integration tests for client.execute() DML row count reporting. + +Verifies that execute() correctly returns the number of rows +affected by INSERT, UPDATE, and DELETE statements. +""" + +from altertable_flightsql import Client +from tests.conftest import TableInfo + + +class TestExecuteRowCount: + """Test that execute() returns the correct number of affected rows.""" + + def test_insert_returns_row_count(self, altertable_client: Client, test_table: TableInfo): + """Test that INSERT returns the number of inserted rows.""" + rows = altertable_client.execute( + f"INSERT INTO {test_table.full_name} (id, name, value) VALUES (4, 'Dave', 400), (5, 'Eve', 500)" + ) + assert rows == 2 + + def test_update_returns_row_count(self, altertable_client: Client, test_table: TableInfo): + """Test that UPDATE returns the number of updated rows.""" + rows = altertable_client.execute( + f"UPDATE {test_table.full_name} SET value = 999 WHERE value >= 200" + ) + assert rows == 2 + + def test_delete_returns_row_count(self, altertable_client: Client, test_table: TableInfo): + """Test that DELETE returns the number of deleted rows.""" + rows = altertable_client.execute( + f"DELETE FROM {test_table.full_name} WHERE id IN (1, 2)" + ) + assert rows == 2 + + def test_delete_no_match_returns_zero(self, altertable_client: Client, test_table: TableInfo): + """Test that DELETE with no matching rows returns 0.""" + rows = altertable_client.execute( + f"DELETE FROM {test_table.full_name} WHERE id = 9999" + ) + assert rows == 0 + + def test_update_no_match_returns_zero(self, altertable_client: Client, test_table: TableInfo): + """Test that UPDATE with no matching rows returns 0.""" + rows = altertable_client.execute( + f"UPDATE {test_table.full_name} SET value = 0 WHERE id = 9999" + ) + assert rows == 0 From 6d3a3da4f9568c2c73347c504b263c5e5f160494 Mon Sep 17 00:00:00 2001 From: Yannick Utard Date: Thu, 5 Mar 2026 22:08:59 +0100 Subject: [PATCH 3/3] Format --- tests/test_execute.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/tests/test_execute.py b/tests/test_execute.py index 4e0ad2b..17b37f3 100644 --- a/tests/test_execute.py +++ b/tests/test_execute.py @@ -28,16 +28,12 @@ def test_update_returns_row_count(self, altertable_client: Client, test_table: T def test_delete_returns_row_count(self, altertable_client: Client, test_table: TableInfo): """Test that DELETE returns the number of deleted rows.""" - rows = altertable_client.execute( - f"DELETE FROM {test_table.full_name} WHERE id IN (1, 2)" - ) + rows = altertable_client.execute(f"DELETE FROM {test_table.full_name} WHERE id IN (1, 2)") assert rows == 2 def test_delete_no_match_returns_zero(self, altertable_client: Client, test_table: TableInfo): """Test that DELETE with no matching rows returns 0.""" - rows = altertable_client.execute( - f"DELETE FROM {test_table.full_name} WHERE id = 9999" - ) + rows = altertable_client.execute(f"DELETE FROM {test_table.full_name} WHERE id = 9999") assert rows == 0 def test_update_no_match_returns_zero(self, altertable_client: Client, test_table: TableInfo):