kaxil commented on code in PR #52330: URL: https://github.com/apache/airflow/pull/52330#discussion_r3677858478
########## providers/apache/arrow/src/airflow/providers/apache/arrow/hooks/adbc.py: ########## @@ -0,0 +1,325 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import contextlib +import pathlib +import re +import sys +from collections.abc import Callable, Iterable, Mapping +from contextlib import closing +from functools import cached_property +from importlib import resources as importlib_resources +from typing import TYPE_CHECKING, Any + +from adbc_driver_manager.dbapi import Connection, connect +from more_itertools import chunked +from pyarrow import RecordBatch, Schema, array, schema + +from airflow.providers.common.sql.dialects.dialect import Dialect +from airflow.providers.common.sql.hooks.sql import DbApiHook + +if TYPE_CHECKING: + from adbc_driver_manager.dbapi import Cursor + + +def fetch_all_handler(cursor) -> list[tuple] | None: + """Return results for DbApiHook.run().""" + if not hasattr(cursor, "description"): + raise RuntimeError( + "The database we interact with does not support DBAPI 2.0. Use operator and " + "handlers that are specifically designed for your database." + ) + if cursor.description is not None: + return list(zip(*cursor.fetch_arrow_table().to_pydict().values())) + return None + + +def replace_placeholders(sql: str, placeholder: str) -> str: + # Replace each placeholder with $1, $2, $3 ... in order + counter = [1] + + def replacer(match): + replacement = f"${counter[0]}" + counter[0] += 1 + return replacement + + return re.sub(placeholder, replacer, sql) + + +# https://arrow.apache.org/adbc/current/python/api/adbc_driver_manager.html +# https://arrow.apache.org/docs/python/ +class AdbcHook(DbApiHook): + """ + General-purpose Airflow hook for interacting with databases via the Arrow Database Connectivity (ADBC) standard. + + This hook enables connections to any database supported by an ADBC driver, using the Python ADBC driver manager. + It provides methods for executing SQL queries, inserting rows in bulk, and handling Arrow-native data transfers. + + Key Features: + - Supports chunked and batched inserts using Apache Arrow RecordBatches for efficient data transfer. + - Discovers and loads ADBC drivers dynamically based on connection extras or naming conventions. + - Handles dialect-specific connection URIs and driver entrypoints. + - Integrates with Airflow's connection system (conn_id, extras, etc.). + - Provides custom placeholder replacement for parameterized SQL queries. + - Supports both native Arrow binding and DBAPI ``executemany`` for inserts. + - Exposes configuration via connection extras: driver, entrypoint, db_kwargs, conn_kwargs, dialect. + + Connection Extras: + - driver: Name of the ADBC driver to use (e.g., "adbc_driver_postgresql"). + - entrypoint: Optional Python entrypoint for the driver. + - db_kwargs: Dict of keyword arguments passed to the database connection. + - conn_kwargs: Dict of keyword arguments passed to the driver connect function. + - dialect: SQL dialect name (default: "default"). + + Example usage: + hook = AdbcHook(adbc_conn_id="my_adbc_conn") + records = hook.get_records("SELECT * FROM my_table") + + For more details, see: + - Apache Arrow ADBC Python API: https://arrow.apache.org/adbc/current/python/api/adbc_driver_manager.html + - Airflow SQL hooks: https://airflow.apache.org/docs/apache-airflow/stable/howto/custom-operator.html#hooks + """ + + conn_name_attr = "adbc_conn_id" + default_conn_name = "adbc_default" + conn_type = "adbc" + hook_name = "ADBC Connection" + supports_autocommit = True + + @classmethod + def get_ui_field_behaviour(cls) -> dict[str, Any]: + """Get custom field behaviour.""" + return { + "hidden_fields": ["port", "schema"], + "relabeling": {"host": "Connection URL"}, + } + + @cached_property + def _driver_path(self) -> str: + # Wheels bundle the shared library + root = importlib_resources.files(self.driver) + # The filename is always the same regardless of platform + entrypoint = root.joinpath(f"lib{self.driver}.so") + if entrypoint.is_file(): + return str(entrypoint) + + # Search sys.prefix + '/lib' (Unix, Conda on Unix) + root = pathlib.Path(sys.prefix) + for filename in (f"lib{self.driver}.so", f"lib{self.driver}.dylib"): + entrypoint = root.joinpath("lib", filename) + if entrypoint.is_file(): + return str(entrypoint) + + # Conda on Windows + entrypoint = root.joinpath("bin", f"{self.driver}.dll") + if entrypoint.is_file(): + return str(entrypoint) + + # Let the driver manager fall back to (DY)LD_LIBRARY_PATH/PATH + # (It will insert 'lib', 'so', etc. as needed) + return self.driver + + @cached_property + def uri(self) -> str: + host = self.connection.host + if host and "::" in str(host): + return str(host) + uri = self.get_uri() + return uri.replace( + f"{self.conn_type.lower().replace('_', '-')}://", + f"{self.dialect_name.lower().replace('_', '-')}://", + ) + + @cached_property + def driver(self) -> str: + return self.connection_extra_lower.get("driver") or f"adbc_driver_{self.dialect_name}" + + @cached_property + def entrypoint(self) -> str | None: + return self.connection_extra_lower.get("entrypoint") + + @cached_property + def db_kwargs(self) -> dict: + return {**{"uri": self.uri}, **self.connection_extra_lower.get("db_kwargs", {})} + + @cached_property + def conn_kwargs(self) -> dict: + return self.connection_extra_lower.get("conn_kwargs", {}) + + @cached_property + def dialect_name(self) -> str: + return self.connection_extra_lower.get("dialect", "default") + + def get_conn(self) -> Connection: + return connect( + driver=self._driver_path, + entrypoint=self.entrypoint, + db_kwargs=self.db_kwargs, + conn_kwargs=self.conn_kwargs, + autocommit=False, + ) + + def get_records( + self, + sql: str | list[str], + parameters: Iterable | Mapping[str, Any] | None = None, + ) -> Any: + """ + Execute the sql and return a set of records. + + :param sql: the sql statement to be executed (str) or a list of sql statements to execute + :param parameters: The parameters to render the SQL query with. + """ + return self.run(sql=sql, parameters=parameters, handler=fetch_all_handler) + + def _run_command(self, cur, sql_statement, parameters): + """Run a statement using an already open cursor.""" + if parameters: + sql_statement = replace_placeholders(sql_statement, re.escape(self.dialect.placeholder)) + + super()._run_command(cur, sql_statement, parameters) + + def _generate_insert_sql(self, table, values, target_fields=None, replace: bool = False, **kwargs) -> str: + sql_statement = super()._generate_insert_sql( + table, values, target_fields=target_fields, replace=replace, **kwargs + ) + sql_statement = replace_placeholders(sql_statement, re.escape(self.dialect.placeholder)) + + if self.log_sql: + self.log.info("Running statement: %s", sql_statement) + + return sql_statement + + @classmethod + def _to_record_batch(cls, rows, schema: Schema) -> RecordBatch: + return RecordBatch.from_arrays( + [array([row[index] for row in rows], type=field.type) for index, field in enumerate(schema)], + schema=schema, + ) + + @classmethod + def _execute_native_bind(cls, cursor, statement: str, record_batch: RecordBatch) -> None: + """Execute a statement using native Arrow bind on the cursor.""" + cursor.bind(record_batch) + cursor.execute(statement) + + @classmethod + def _execute_executemany(cls, cursor, statement: str, record_batch: RecordBatch) -> None: + """Execute a statement using cursor.executemany.""" + cursor.executemany(statement, record_batch) + + def _resolve_execute_batch( + self, cursor, executemany: bool, fast_executemany: bool + ) -> Callable[[Cursor, str, RecordBatch], None]: + """ + Return the appropriate batch-execute callable for the given cursor. + + Picks native Arrow bind when the cursor supports it; otherwise falls back + to executemany, optionally enabling fast_executemany on the cursor first. + """ + use_native_bind = hasattr(cursor, "bind") + + if not use_native_bind and (self.supports_executemany or executemany): + if fast_executemany: + with contextlib.suppress(AttributeError): + cursor.fast_executemany = True + self.log.info( + "Fast_executemany is enabled for conn_id '%s'!", + self.get_conn_id(), + ) + + if use_native_bind: + self.log.info("Native Arrow bind supported!") + return self._execute_native_bind + return self._execute_executemany + + def insert_rows( + self, + table, + rows, + target_fields=None, + commit_every=1000, + replace=False, + *, + executemany=False, + fast_executemany=False, + autocommit=False, + **kwargs, + ): + """ + Insert a collection of tuples into a table. + + Rows are inserted in chunks, each chunk (of size ``commit_every``) is + done in a new transaction. + + :param table: Name of the target table + :param rows: The rows to insert into the table + :param target_fields: The names of the columns to fill in the table + :param commit_every: The maximum number of rows to insert in one + transaction. Set to 0 to insert all rows in one transaction. + :param replace: Whether to replace instead of insert + :param executemany: If True, all rows are inserted at once in + chunks defined by the commit_every parameter. This only works if all rows + have same number of column names, but leads to better performance. + :param fast_executemany: If True, the `fast_executemany` parameter will be set on the + cursor used by `executemany` which leads to better performance, if supported by driver. + :param autocommit: What to set the connection's autocommit setting to + before executing the query. + """ + nb_rows = 0 + + with self._create_autocommit_connection(autocommit) as conn: + table_name, schema_name = Dialect.extract_schema_from_table(table) + + table_schema = conn.adbc_get_table_schema( + table_name=table_name, + db_schema_filter=schema_name, + ) + + if not target_fields: + target_fields = table_schema.names + else: + table_schema = schema([field for field in table_schema if field.name in target_fields]) Review Comment: The filtered schema keeps the *table's* column order, but `_generate_insert_sql` below emits the column list in `target_fields` order, and `_to_record_batch` indexes `rows` positionally against this schema. When the two orders differ, values land in the wrong columns. With `target_fields=["last", "first"]` against a table declared `(first, last)`, the filtered schema comes back as `[first, last]`, so the SQL says `INSERT INTO t (last, first)` while the RecordBatch is `{first: 'Smith', last: 'Alice'}`. Same-typed columns silently swap; differently-typed ones blow up with `ArrowInvalid: Could not convert 'Alice' with type str: tried to convert to int64`. Ordering the filtered schema by `target_fields` instead of by the table schema fixes both cases: ```python fields = {field.name: field for field in table_schema} table_schema = schema([fields[name] for name in target_fields]) ``` ########## providers/apache/arrow/src/airflow/providers/apache/arrow/hooks/adbc.py: ########## @@ -0,0 +1,325 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import contextlib +import pathlib +import re +import sys +from collections.abc import Callable, Iterable, Mapping +from contextlib import closing +from functools import cached_property +from importlib import resources as importlib_resources +from typing import TYPE_CHECKING, Any + +from adbc_driver_manager.dbapi import Connection, connect +from more_itertools import chunked +from pyarrow import RecordBatch, Schema, array, schema + +from airflow.providers.common.sql.dialects.dialect import Dialect +from airflow.providers.common.sql.hooks.sql import DbApiHook + +if TYPE_CHECKING: + from adbc_driver_manager.dbapi import Cursor + + +def fetch_all_handler(cursor) -> list[tuple] | None: + """Return results for DbApiHook.run().""" + if not hasattr(cursor, "description"): + raise RuntimeError( + "The database we interact with does not support DBAPI 2.0. Use operator and " + "handlers that are specifically designed for your database." + ) + if cursor.description is not None: + return list(zip(*cursor.fetch_arrow_table().to_pydict().values())) Review Comment: Routing through `to_pydict()` collapses duplicate column names, so `get_records` silently drops columns whenever a query selects the same name twice. `SELECT a.id, b.id FROM a JOIN b ...` produces an Arrow table with two `id` columns, but `to_pydict()` keeps only the last, and the zip then yields 1-tuples where every other `DbApiHook` returns 2-tuples. Iterating the columns positionally avoids it and behaves the same for the normal and empty-result cases: ```python table = cursor.fetch_arrow_table() return list(zip(*(column.to_pylist() for column in table.columns))) ``` ########## providers/apache/arrow/src/airflow/providers/apache/arrow/hooks/adbc.py: ########## @@ -0,0 +1,325 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import contextlib +import pathlib +import re +import sys +from collections.abc import Callable, Iterable, Mapping +from contextlib import closing +from functools import cached_property +from importlib import resources as importlib_resources +from typing import TYPE_CHECKING, Any + +from adbc_driver_manager.dbapi import Connection, connect +from more_itertools import chunked +from pyarrow import RecordBatch, Schema, array, schema + +from airflow.providers.common.sql.dialects.dialect import Dialect +from airflow.providers.common.sql.hooks.sql import DbApiHook + +if TYPE_CHECKING: + from adbc_driver_manager.dbapi import Cursor + + +def fetch_all_handler(cursor) -> list[tuple] | None: + """Return results for DbApiHook.run().""" + if not hasattr(cursor, "description"): + raise RuntimeError( + "The database we interact with does not support DBAPI 2.0. Use operator and " + "handlers that are specifically designed for your database." + ) + if cursor.description is not None: + return list(zip(*cursor.fetch_arrow_table().to_pydict().values())) + return None + + +def replace_placeholders(sql: str, placeholder: str) -> str: + # Replace each placeholder with $1, $2, $3 ... in order + counter = [1] + + def replacer(match): + replacement = f"${counter[0]}" + counter[0] += 1 + return replacement + + return re.sub(placeholder, replacer, sql) + + +# https://arrow.apache.org/adbc/current/python/api/adbc_driver_manager.html +# https://arrow.apache.org/docs/python/ +class AdbcHook(DbApiHook): + """ + General-purpose Airflow hook for interacting with databases via the Arrow Database Connectivity (ADBC) standard. + + This hook enables connections to any database supported by an ADBC driver, using the Python ADBC driver manager. + It provides methods for executing SQL queries, inserting rows in bulk, and handling Arrow-native data transfers. + + Key Features: + - Supports chunked and batched inserts using Apache Arrow RecordBatches for efficient data transfer. + - Discovers and loads ADBC drivers dynamically based on connection extras or naming conventions. + - Handles dialect-specific connection URIs and driver entrypoints. + - Integrates with Airflow's connection system (conn_id, extras, etc.). + - Provides custom placeholder replacement for parameterized SQL queries. + - Supports both native Arrow binding and DBAPI ``executemany`` for inserts. + - Exposes configuration via connection extras: driver, entrypoint, db_kwargs, conn_kwargs, dialect. + + Connection Extras: + - driver: Name of the ADBC driver to use (e.g., "adbc_driver_postgresql"). + - entrypoint: Optional Python entrypoint for the driver. + - db_kwargs: Dict of keyword arguments passed to the database connection. + - conn_kwargs: Dict of keyword arguments passed to the driver connect function. + - dialect: SQL dialect name (default: "default"). + + Example usage: + hook = AdbcHook(adbc_conn_id="my_adbc_conn") + records = hook.get_records("SELECT * FROM my_table") + + For more details, see: + - Apache Arrow ADBC Python API: https://arrow.apache.org/adbc/current/python/api/adbc_driver_manager.html + - Airflow SQL hooks: https://airflow.apache.org/docs/apache-airflow/stable/howto/custom-operator.html#hooks + """ + + conn_name_attr = "adbc_conn_id" + default_conn_name = "adbc_default" + conn_type = "adbc" + hook_name = "ADBC Connection" + supports_autocommit = True + + @classmethod + def get_ui_field_behaviour(cls) -> dict[str, Any]: + """Get custom field behaviour.""" + return { + "hidden_fields": ["port", "schema"], + "relabeling": {"host": "Connection URL"}, + } + + @cached_property + def _driver_path(self) -> str: + # Wheels bundle the shared library + root = importlib_resources.files(self.driver) + # The filename is always the same regardless of platform + entrypoint = root.joinpath(f"lib{self.driver}.so") + if entrypoint.is_file(): + return str(entrypoint) + + # Search sys.prefix + '/lib' (Unix, Conda on Unix) + root = pathlib.Path(sys.prefix) + for filename in (f"lib{self.driver}.so", f"lib{self.driver}.dylib"): + entrypoint = root.joinpath("lib", filename) + if entrypoint.is_file(): + return str(entrypoint) + + # Conda on Windows + entrypoint = root.joinpath("bin", f"{self.driver}.dll") + if entrypoint.is_file(): + return str(entrypoint) + + # Let the driver manager fall back to (DY)LD_LIBRARY_PATH/PATH + # (It will insert 'lib', 'so', etc. as needed) + return self.driver + + @cached_property + def uri(self) -> str: + host = self.connection.host + if host and "::" in str(host): + return str(host) + uri = self.get_uri() + return uri.replace( + f"{self.conn_type.lower().replace('_', '-')}://", + f"{self.dialect_name.lower().replace('_', '-')}://", + ) + + @cached_property + def driver(self) -> str: + return self.connection_extra_lower.get("driver") or f"adbc_driver_{self.dialect_name}" + + @cached_property + def entrypoint(self) -> str | None: + return self.connection_extra_lower.get("entrypoint") + + @cached_property + def db_kwargs(self) -> dict: + return {**{"uri": self.uri}, **self.connection_extra_lower.get("db_kwargs", {})} + + @cached_property + def conn_kwargs(self) -> dict: + return self.connection_extra_lower.get("conn_kwargs", {}) + + @cached_property + def dialect_name(self) -> str: + return self.connection_extra_lower.get("dialect", "default") + + def get_conn(self) -> Connection: + return connect( + driver=self._driver_path, + entrypoint=self.entrypoint, + db_kwargs=self.db_kwargs, + conn_kwargs=self.conn_kwargs, + autocommit=False, + ) + + def get_records( + self, + sql: str | list[str], + parameters: Iterable | Mapping[str, Any] | None = None, + ) -> Any: + """ + Execute the sql and return a set of records. + + :param sql: the sql statement to be executed (str) or a list of sql statements to execute + :param parameters: The parameters to render the SQL query with. + """ + return self.run(sql=sql, parameters=parameters, handler=fetch_all_handler) + + def _run_command(self, cur, sql_statement, parameters): + """Run a statement using an already open cursor.""" + if parameters: + sql_statement = replace_placeholders(sql_statement, re.escape(self.dialect.placeholder)) + + super()._run_command(cur, sql_statement, parameters) + + def _generate_insert_sql(self, table, values, target_fields=None, replace: bool = False, **kwargs) -> str: + sql_statement = super()._generate_insert_sql( + table, values, target_fields=target_fields, replace=replace, **kwargs + ) + sql_statement = replace_placeholders(sql_statement, re.escape(self.dialect.placeholder)) + + if self.log_sql: + self.log.info("Running statement: %s", sql_statement) + + return sql_statement + + @classmethod + def _to_record_batch(cls, rows, schema: Schema) -> RecordBatch: + return RecordBatch.from_arrays( + [array([row[index] for row in rows], type=field.type) for index, field in enumerate(schema)], + schema=schema, + ) + + @classmethod + def _execute_native_bind(cls, cursor, statement: str, record_batch: RecordBatch) -> None: + """Execute a statement using native Arrow bind on the cursor.""" + cursor.bind(record_batch) + cursor.execute(statement) + + @classmethod + def _execute_executemany(cls, cursor, statement: str, record_batch: RecordBatch) -> None: + """Execute a statement using cursor.executemany.""" + cursor.executemany(statement, record_batch) + + def _resolve_execute_batch( + self, cursor, executemany: bool, fast_executemany: bool + ) -> Callable[[Cursor, str, RecordBatch], None]: + """ + Return the appropriate batch-execute callable for the given cursor. + + Picks native Arrow bind when the cursor supports it; otherwise falls back + to executemany, optionally enabling fast_executemany on the cursor first. + """ + use_native_bind = hasattr(cursor, "bind") Review Comment: `adbc_driver_manager.dbapi.Cursor` has no `bind` method, neither on 1.11.0 nor on the `>=1.7.0` floor declared in `pyproject.toml` (its public surface is `adbc_ingest`, `adbc_prepare`, `execute`, `executemany`, ...). So `hasattr(cursor, "bind")` is always False against a real driver, `_resolve_execute_batch` always returns `_execute_executemany`, and both `_execute_native_bind` and the "Native Arrow bind supported!" log are unreachable in production. `test_insert_rows_native_bind` only passes because `setup_method` assigns `self.cur.bind = mock.MagicMock()` onto a plain `MagicMock`, so the test greenlights a path that cannot execute. Since the Arrow-native transfer is the headline of the provider, worth pinning down what the intended fast path is: `cursor.adbc_ingest`, or `executemany` with the RecordBatch (which does work, I checked it against the SQLite driver)? ########## providers/apache/arrow/docs/connections/adbc.rst: ########## @@ -0,0 +1,200 @@ + .. Licensed to the Apache Software Foundation (ASF) under one + or more contributor license agreements. See the NOTICE file + distributed with this work for additional information + regarding copyright ownership. The ASF licenses this file + to you under the Apache License, Version 2.0 (the + "License"); you may not use this file except in compliance + with the License. You may obtain a copy of the License at + + .. http://www.apache.org/licenses/LICENSE-2.0 + + .. Unless required by applicable law or agreed to in writing, + software distributed under the License is distributed on an + "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + KIND, either express or implied. See the License for the + specific language governing permissions and limitations + under the License. + +.. _howto/connection:adbc: + +ADBC connection +=============== + +The ADBC connection type enables connection to any database that has an +`Arrow Database Connectivity (ADBC) <https://arrow.apache.org/adbc/>`__ driver, +such as PostgreSQL, SQLite, DuckDB, Snowflake, BigQuery, and Flight SQL servers. + +Connections of this type are used by :class:`~airflow.providers.apache.arrow.hooks.adbc.AdbcHook`, +which builds on top of :class:`~airflow.providers.common.sql.hooks.sql.DbApiHook` and transfers +data as Apache Arrow :class:`~pyarrow.RecordBatch` objects for efficient, zero-copy bulk loads. + +Default Connection ID +--------------------- + +``adbc_default`` + +Configuring the Connection +-------------------------- + +Connection URL (Host field) + The database URI passed to the driver. For drivers that use a standard + connection string (e.g. ``postgresql://user:pass@host:5432/db`` or + ``file::memory:``), put it here. If the value contains ``::`` the hook + passes it through unchanged; otherwise the ``adbc://`` scheme prefix is + replaced with the dialect name. + +Login / Password + Convenience fields. When set, Airflow builds a URI of the form + ``<dialect>://login:password@host/schema`` that is merged into + ``db_kwargs["uri"]``. Driver-specific URIs in the Host field take + precedence. + +Extra (JSON) + A JSON object with the following recognized keys: + + ``driver`` *(string)* + Python package name of the ADBC driver to load, e.g. + ``"adbc_driver_postgresql"`` or ``"adbc_driver_sqlite"``. + When omitted the hook derives the name from the dialect as + ``adbc_driver_<dialect>``. + + ``entrypoint`` *(string, optional)* + Fully-qualified Python symbol used as the driver entrypoint, e.g. + ``"adbc_driver_sqlite.dbapi.connect"``. Most drivers do not need + this; it is only required when the automatic entrypoint discovery + built into ``adbc_driver_manager`` cannot locate the function. + + ``dialect`` *(string, optional)* + SQL dialect name used to build the URI scheme and derive the + default driver name. Defaults to ``"default"``. + + ``db_kwargs`` *(object, optional)* + Keyword arguments forwarded verbatim to the ADBC + ``AdbcDatabase`` constructor (database-level connection settings). + ``uri`` is always injected automatically from the Host field and + must not be repeated here. Common keys: + + .. list-table:: + :header-rows: 1 + :widths: 30 15 55 + + * - Key + - Type + - Description + * - ``username`` + - string + - Database user name (alternative to encoding it in the URI). + * - ``password`` + - string + - Database password (alternative to encoding it in the URI). + * - ``adbc.connection.autocommit`` + - bool + - Enable autocommit at the database level (driver-dependent). + * - ``adbc.sqlite.query.batch_rows`` + - int + - SQLite: number of rows fetched per Arrow batch. + + All other keys are driver-specific. Consult the documentation for + your ADBC driver for the full list. See also the + `ADBC driver documentation <https://arrow.apache.org/adbc/current/driver/>`__. + + ``conn_kwargs`` *(object, optional)* + Keyword arguments forwarded verbatim to the ADBC + ``AdbcConnection`` constructor (connection-level settings within the + already-open database). Common keys defined by the ADBC + specification: + + .. list-table:: + :header-rows: 1 + :widths: 30 15 55 + + * - Key + - Type + - Description + * - ``autocommit`` Review Comment: These four keys aren't ADBC option names. Against the real SQLite driver each one is rejected: `NotSupportedError: NOT_IMPLEMENTED: [SQLite] Unknown connection option autocommit='true'`, and the same for `read_only`, `current_catalog` and `current_db_schema`. The dotted form is the canonical spelling, and of the four only `adbc.connection.autocommit` is actually accepted. The `db_kwargs` table above has the mirror-image problem: `adbc.connection.autocommit` is listed there, but it's a connection option and belongs under `conn_kwargs`. `timeout` (the key `test_get_conn_forwards_db_kwargs` asserts) is rejected as an unknown database option too. The tests can't catch any of this because they patch `connect`, so the passthrough assertions pass while the documented values would fail against a driver. Distinct from the earlier `db_kwargs` threads, which were about whether to forward options at all rather than which names are valid. ########## providers/apache/arrow/src/airflow/providers/apache/arrow/hooks/adbc.py: ########## @@ -0,0 +1,325 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import contextlib +import pathlib +import re +import sys +from collections.abc import Callable, Iterable, Mapping +from contextlib import closing +from functools import cached_property +from importlib import resources as importlib_resources +from typing import TYPE_CHECKING, Any + +from adbc_driver_manager.dbapi import Connection, connect +from more_itertools import chunked +from pyarrow import RecordBatch, Schema, array, schema + +from airflow.providers.common.sql.dialects.dialect import Dialect +from airflow.providers.common.sql.hooks.sql import DbApiHook + +if TYPE_CHECKING: + from adbc_driver_manager.dbapi import Cursor + + +def fetch_all_handler(cursor) -> list[tuple] | None: + """Return results for DbApiHook.run().""" + if not hasattr(cursor, "description"): + raise RuntimeError( + "The database we interact with does not support DBAPI 2.0. Use operator and " + "handlers that are specifically designed for your database." + ) + if cursor.description is not None: + return list(zip(*cursor.fetch_arrow_table().to_pydict().values())) + return None + + +def replace_placeholders(sql: str, placeholder: str) -> str: + # Replace each placeholder with $1, $2, $3 ... in order + counter = [1] + + def replacer(match): + replacement = f"${counter[0]}" + counter[0] += 1 + return replacement + + return re.sub(placeholder, replacer, sql) + + +# https://arrow.apache.org/adbc/current/python/api/adbc_driver_manager.html +# https://arrow.apache.org/docs/python/ +class AdbcHook(DbApiHook): + """ + General-purpose Airflow hook for interacting with databases via the Arrow Database Connectivity (ADBC) standard. + + This hook enables connections to any database supported by an ADBC driver, using the Python ADBC driver manager. + It provides methods for executing SQL queries, inserting rows in bulk, and handling Arrow-native data transfers. + + Key Features: + - Supports chunked and batched inserts using Apache Arrow RecordBatches for efficient data transfer. + - Discovers and loads ADBC drivers dynamically based on connection extras or naming conventions. + - Handles dialect-specific connection URIs and driver entrypoints. + - Integrates with Airflow's connection system (conn_id, extras, etc.). + - Provides custom placeholder replacement for parameterized SQL queries. + - Supports both native Arrow binding and DBAPI ``executemany`` for inserts. + - Exposes configuration via connection extras: driver, entrypoint, db_kwargs, conn_kwargs, dialect. + + Connection Extras: + - driver: Name of the ADBC driver to use (e.g., "adbc_driver_postgresql"). + - entrypoint: Optional Python entrypoint for the driver. + - db_kwargs: Dict of keyword arguments passed to the database connection. + - conn_kwargs: Dict of keyword arguments passed to the driver connect function. + - dialect: SQL dialect name (default: "default"). + + Example usage: + hook = AdbcHook(adbc_conn_id="my_adbc_conn") + records = hook.get_records("SELECT * FROM my_table") + + For more details, see: + - Apache Arrow ADBC Python API: https://arrow.apache.org/adbc/current/python/api/adbc_driver_manager.html + - Airflow SQL hooks: https://airflow.apache.org/docs/apache-airflow/stable/howto/custom-operator.html#hooks + """ + + conn_name_attr = "adbc_conn_id" + default_conn_name = "adbc_default" + conn_type = "adbc" + hook_name = "ADBC Connection" + supports_autocommit = True + + @classmethod + def get_ui_field_behaviour(cls) -> dict[str, Any]: + """Get custom field behaviour.""" + return { + "hidden_fields": ["port", "schema"], + "relabeling": {"host": "Connection URL"}, + } + + @cached_property + def _driver_path(self) -> str: + # Wheels bundle the shared library + root = importlib_resources.files(self.driver) + # The filename is always the same regardless of platform + entrypoint = root.joinpath(f"lib{self.driver}.so") + if entrypoint.is_file(): + return str(entrypoint) + + # Search sys.prefix + '/lib' (Unix, Conda on Unix) + root = pathlib.Path(sys.prefix) + for filename in (f"lib{self.driver}.so", f"lib{self.driver}.dylib"): + entrypoint = root.joinpath("lib", filename) + if entrypoint.is_file(): + return str(entrypoint) + + # Conda on Windows + entrypoint = root.joinpath("bin", f"{self.driver}.dll") + if entrypoint.is_file(): + return str(entrypoint) + + # Let the driver manager fall back to (DY)LD_LIBRARY_PATH/PATH + # (It will insert 'lib', 'so', etc. as needed) + return self.driver + + @cached_property + def uri(self) -> str: + host = self.connection.host + if host and "::" in str(host): + return str(host) + uri = self.get_uri() + return uri.replace( + f"{self.conn_type.lower().replace('_', '-')}://", + f"{self.dialect_name.lower().replace('_', '-')}://", + ) + + @cached_property + def driver(self) -> str: + return self.connection_extra_lower.get("driver") or f"adbc_driver_{self.dialect_name}" + + @cached_property + def entrypoint(self) -> str | None: + return self.connection_extra_lower.get("entrypoint") + + @cached_property + def db_kwargs(self) -> dict: + return {**{"uri": self.uri}, **self.connection_extra_lower.get("db_kwargs", {})} + + @cached_property + def conn_kwargs(self) -> dict: + return self.connection_extra_lower.get("conn_kwargs", {}) + + @cached_property + def dialect_name(self) -> str: + return self.connection_extra_lower.get("dialect", "default") + + def get_conn(self) -> Connection: + return connect( + driver=self._driver_path, + entrypoint=self.entrypoint, + db_kwargs=self.db_kwargs, + conn_kwargs=self.conn_kwargs, + autocommit=False, Review Comment: `autocommit` is pinned to False here and there's no `set_autocommit` override, so with `supports_autocommit = True` the base `_create_autocommit_connection` ends up running `conn.autocommit = autocommit` against an ADBC `dbapi.Connection`. That only sets an unused Python attribute: after `conn.autocommit = True` the connection is still in manual-commit mode, and `conn.commit()` keeps succeeding instead of erroring. `insert_rows(..., autocommit=True)` and `run(..., autocommit=True)` therefore do nothing at the driver level. ADBC's actual lever is the `autocommit=` kwarg on `connect()` or `adbc_connection.set_options(**{"adbc.connection.autocommit": "true"})`. Worth noting that wiring it up properly also needs the `conn.commit()` inside the `insert_rows` chunk loop made conditional: once autocommit is really on, SQLite raises `ProgrammingError: INVALID_STATE: No active transaction, cannot commit` on that call. ########## providers/apache/arrow/src/airflow/providers/apache/arrow/hooks/adbc.py: ########## @@ -0,0 +1,325 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import contextlib +import pathlib +import re +import sys +from collections.abc import Callable, Iterable, Mapping +from contextlib import closing +from functools import cached_property +from importlib import resources as importlib_resources +from typing import TYPE_CHECKING, Any + +from adbc_driver_manager.dbapi import Connection, connect +from more_itertools import chunked +from pyarrow import RecordBatch, Schema, array, schema + +from airflow.providers.common.sql.dialects.dialect import Dialect +from airflow.providers.common.sql.hooks.sql import DbApiHook + +if TYPE_CHECKING: + from adbc_driver_manager.dbapi import Cursor + + +def fetch_all_handler(cursor) -> list[tuple] | None: + """Return results for DbApiHook.run().""" + if not hasattr(cursor, "description"): + raise RuntimeError( + "The database we interact with does not support DBAPI 2.0. Use operator and " + "handlers that are specifically designed for your database." + ) + if cursor.description is not None: + return list(zip(*cursor.fetch_arrow_table().to_pydict().values())) + return None + + +def replace_placeholders(sql: str, placeholder: str) -> str: + # Replace each placeholder with $1, $2, $3 ... in order + counter = [1] + + def replacer(match): + replacement = f"${counter[0]}" + counter[0] += 1 + return replacement + + return re.sub(placeholder, replacer, sql) + + +# https://arrow.apache.org/adbc/current/python/api/adbc_driver_manager.html +# https://arrow.apache.org/docs/python/ +class AdbcHook(DbApiHook): + """ + General-purpose Airflow hook for interacting with databases via the Arrow Database Connectivity (ADBC) standard. + + This hook enables connections to any database supported by an ADBC driver, using the Python ADBC driver manager. + It provides methods for executing SQL queries, inserting rows in bulk, and handling Arrow-native data transfers. + + Key Features: + - Supports chunked and batched inserts using Apache Arrow RecordBatches for efficient data transfer. + - Discovers and loads ADBC drivers dynamically based on connection extras or naming conventions. + - Handles dialect-specific connection URIs and driver entrypoints. + - Integrates with Airflow's connection system (conn_id, extras, etc.). + - Provides custom placeholder replacement for parameterized SQL queries. + - Supports both native Arrow binding and DBAPI ``executemany`` for inserts. + - Exposes configuration via connection extras: driver, entrypoint, db_kwargs, conn_kwargs, dialect. + + Connection Extras: + - driver: Name of the ADBC driver to use (e.g., "adbc_driver_postgresql"). + - entrypoint: Optional Python entrypoint for the driver. + - db_kwargs: Dict of keyword arguments passed to the database connection. + - conn_kwargs: Dict of keyword arguments passed to the driver connect function. + - dialect: SQL dialect name (default: "default"). + + Example usage: + hook = AdbcHook(adbc_conn_id="my_adbc_conn") + records = hook.get_records("SELECT * FROM my_table") + + For more details, see: + - Apache Arrow ADBC Python API: https://arrow.apache.org/adbc/current/python/api/adbc_driver_manager.html + - Airflow SQL hooks: https://airflow.apache.org/docs/apache-airflow/stable/howto/custom-operator.html#hooks + """ + + conn_name_attr = "adbc_conn_id" + default_conn_name = "adbc_default" + conn_type = "adbc" + hook_name = "ADBC Connection" + supports_autocommit = True + + @classmethod + def get_ui_field_behaviour(cls) -> dict[str, Any]: + """Get custom field behaviour.""" + return { + "hidden_fields": ["port", "schema"], + "relabeling": {"host": "Connection URL"}, + } + + @cached_property + def _driver_path(self) -> str: + # Wheels bundle the shared library + root = importlib_resources.files(self.driver) + # The filename is always the same regardless of platform + entrypoint = root.joinpath(f"lib{self.driver}.so") + if entrypoint.is_file(): + return str(entrypoint) + + # Search sys.prefix + '/lib' (Unix, Conda on Unix) + root = pathlib.Path(sys.prefix) + for filename in (f"lib{self.driver}.so", f"lib{self.driver}.dylib"): + entrypoint = root.joinpath("lib", filename) + if entrypoint.is_file(): + return str(entrypoint) + + # Conda on Windows + entrypoint = root.joinpath("bin", f"{self.driver}.dll") + if entrypoint.is_file(): + return str(entrypoint) + + # Let the driver manager fall back to (DY)LD_LIBRARY_PATH/PATH + # (It will insert 'lib', 'so', etc. as needed) + return self.driver + + @cached_property + def uri(self) -> str: + host = self.connection.host + if host and "::" in str(host): + return str(host) + uri = self.get_uri() + return uri.replace( + f"{self.conn_type.lower().replace('_', '-')}://", + f"{self.dialect_name.lower().replace('_', '-')}://", + ) + + @cached_property + def driver(self) -> str: + return self.connection_extra_lower.get("driver") or f"adbc_driver_{self.dialect_name}" + + @cached_property + def entrypoint(self) -> str | None: + return self.connection_extra_lower.get("entrypoint") + + @cached_property + def db_kwargs(self) -> dict: + return {**{"uri": self.uri}, **self.connection_extra_lower.get("db_kwargs", {})} + + @cached_property + def conn_kwargs(self) -> dict: + return self.connection_extra_lower.get("conn_kwargs", {}) + + @cached_property + def dialect_name(self) -> str: + return self.connection_extra_lower.get("dialect", "default") + + def get_conn(self) -> Connection: + return connect( + driver=self._driver_path, + entrypoint=self.entrypoint, + db_kwargs=self.db_kwargs, + conn_kwargs=self.conn_kwargs, + autocommit=False, + ) + + def get_records( + self, + sql: str | list[str], + parameters: Iterable | Mapping[str, Any] | None = None, + ) -> Any: + """ + Execute the sql and return a set of records. + + :param sql: the sql statement to be executed (str) or a list of sql statements to execute + :param parameters: The parameters to render the SQL query with. + """ + return self.run(sql=sql, parameters=parameters, handler=fetch_all_handler) + + def _run_command(self, cur, sql_statement, parameters): + """Run a statement using an already open cursor.""" + if parameters: + sql_statement = replace_placeholders(sql_statement, re.escape(self.dialect.placeholder)) + + super()._run_command(cur, sql_statement, parameters) + + def _generate_insert_sql(self, table, values, target_fields=None, replace: bool = False, **kwargs) -> str: + sql_statement = super()._generate_insert_sql( + table, values, target_fields=target_fields, replace=replace, **kwargs + ) + sql_statement = replace_placeholders(sql_statement, re.escape(self.dialect.placeholder)) + + if self.log_sql: + self.log.info("Running statement: %s", sql_statement) + + return sql_statement + + @classmethod + def _to_record_batch(cls, rows, schema: Schema) -> RecordBatch: + return RecordBatch.from_arrays( + [array([row[index] for row in rows], type=field.type) for index, field in enumerate(schema)], + schema=schema, + ) + + @classmethod + def _execute_native_bind(cls, cursor, statement: str, record_batch: RecordBatch) -> None: + """Execute a statement using native Arrow bind on the cursor.""" + cursor.bind(record_batch) + cursor.execute(statement) + + @classmethod + def _execute_executemany(cls, cursor, statement: str, record_batch: RecordBatch) -> None: + """Execute a statement using cursor.executemany.""" + cursor.executemany(statement, record_batch) + + def _resolve_execute_batch( + self, cursor, executemany: bool, fast_executemany: bool + ) -> Callable[[Cursor, str, RecordBatch], None]: + """ + Return the appropriate batch-execute callable for the given cursor. + + Picks native Arrow bind when the cursor supports it; otherwise falls back + to executemany, optionally enabling fast_executemany on the cursor first. + """ + use_native_bind = hasattr(cursor, "bind") + + if not use_native_bind and (self.supports_executemany or executemany): + if fast_executemany: + with contextlib.suppress(AttributeError): + cursor.fast_executemany = True + self.log.info( + "Fast_executemany is enabled for conn_id '%s'!", + self.get_conn_id(), + ) + + if use_native_bind: + self.log.info("Native Arrow bind supported!") + return self._execute_native_bind + return self._execute_executemany + + def insert_rows( + self, + table, + rows, + target_fields=None, + commit_every=1000, + replace=False, + *, + executemany=False, + fast_executemany=False, + autocommit=False, + **kwargs, + ): + """ + Insert a collection of tuples into a table. + + Rows are inserted in chunks, each chunk (of size ``commit_every``) is + done in a new transaction. + + :param table: Name of the target table + :param rows: The rows to insert into the table + :param target_fields: The names of the columns to fill in the table + :param commit_every: The maximum number of rows to insert in one + transaction. Set to 0 to insert all rows in one transaction. + :param replace: Whether to replace instead of insert + :param executemany: If True, all rows are inserted at once in + chunks defined by the commit_every parameter. This only works if all rows + have same number of column names, but leads to better performance. + :param fast_executemany: If True, the `fast_executemany` parameter will be set on the + cursor used by `executemany` which leads to better performance, if supported by driver. + :param autocommit: What to set the connection's autocommit setting to + before executing the query. + """ + nb_rows = 0 + + with self._create_autocommit_connection(autocommit) as conn: + table_name, schema_name = Dialect.extract_schema_from_table(table) + + table_schema = conn.adbc_get_table_schema( + table_name=table_name, + db_schema_filter=schema_name, + ) + + if not target_fields: + target_fields = table_schema.names + else: + table_schema = schema([field for field in table_schema if field.name in target_fields]) + + self.log.info("target fields: %s", target_fields) + self.log.info("table_schema: %s", table_schema) + + sql = self._generate_insert_sql( + table, + target_fields, # values not needed — parameters will come from RecordBatch + target_fields, + replace, + **kwargs, + ) + + with closing(conn.cursor()) as cur: + execute_batch = self._resolve_execute_batch( + cur, executemany=executemany, fast_executemany=fast_executemany + ) + + for chunked_rows in chunked(rows, commit_every): Review Comment: The docstring says `commit_every=0` inserts all rows in one transaction, but `more_itertools.chunked(rows, 0)` yields no chunks at all, so nothing is inserted and the method logs "Loaded a total of 0 rows" without raising. `DbApiHook.insert_rows` reads 0 as "don't commit mid-way", so this also diverges from the base class. `chunked(rows, commit_every or None)` restores the documented behaviour, since `chunked(rows, None)` yields a single chunk containing everything. -- 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]
