aoelvp94 commented on code in PR #76363:
URL: https://github.com/apache/airflow/pull/76363#discussion_r4455576170
##########
providers/apache/arrow/tests/unit/apache/arrow/hooks/test_adbc.py:
##########
@@ -97,22 +97,125 @@ def test_get_records_duplicate_column_names(self):
assert result == [(10, 30), (20, 40)]
- def test_insert_rows_uses_executemany_with_record_batch(self):
- """insert_rows must call executemany with a RecordBatch — the
Arrow-native fast path."""
+ def test_insert_rows_uses_adbc_ingest_by_default(self):
+ """By default insert_rows must hand each RecordBatch to the driver's
bulk ingest."""
rows = [("a",), ("b",)]
self.hook.insert_rows("table", rows)
- assert self.cur.executemany.called
- batch = self.cur.executemany.call_args[0][1]
+ self.cur.executemany.assert_not_called()
+ self.cur.adbc_ingest.assert_called_once()
+ args, kwargs = self.cur.adbc_ingest.call_args
+ assert args[0] == "table"
+ assert isinstance(args[1], RecordBatch)
+ assert args[1].num_rows == 2
+ assert kwargs == {"mode": "append", "db_schema_name": None}
+ assert self.conn.commit.call_count == 1
+
+ def test_insert_rows_fast_executemany_false_uses_executemany(self):
+ """fast_executemany=False keeps the previous bind path (one INSERT
bound to the batch)."""
+ rows = [("a",), ("b",)]
+ self.hook.insert_rows("table", rows, fast_executemany=False)
+
+ self.cur.adbc_ingest.assert_not_called()
+ self.cur.executemany.assert_called_once()
+ sql, batch = self.cur.executemany.call_args[0]
+ assert sql.startswith("INSERT INTO table")
assert batch.num_rows == 2
- assert self.conn.commit.call_count >= 1
+
+ def test_insert_rows_replace_binds_even_with_fast_executemany(self):
+ """replace=True cannot go through bulk ingest, so it must bind to the
generated statement."""
+ rows = [("a",), ("b",)]
+ with mock.patch.object(
+ AdbcHook, "_generate_insert_sql", return_value="REPLACE INTO table
(col) VALUES ($1)"
+ ) as generate_insert_sql:
+ self.hook.insert_rows("table", rows, replace=True,
fast_executemany=True)
+
+ assert generate_insert_sql.call_args[0][3] is True
+ self.cur.adbc_ingest.assert_not_called()
+ self.cur.executemany.assert_called_once()
+ assert self.cur.executemany.call_args[0][0] == "REPLACE INTO table
(col) VALUES ($1)"
+
+ def
test_insert_rows_replace_with_fast_executemany_logs_debug_not_info(self):
+ """The caller did not ask for ingest explicitly, so binding for
replace=True is only a debug note."""
+ self.hook._log = mock.MagicMock()
Review Comment:
Done — AGENTS.md nit: MagicMock(spec=logging.Logger) on both _log mocks, and
@mock.patch.object for the two replace=True tests.
https://github.com/aoelvp94/airflow/commit/55f6391cdc41ea2ab0b804cb56ef08cf0b711209
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
