#
# 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
[docs]
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:
table = cursor.fetch_arrow_table()
return list(zip(*(column.to_pylist() for column in table.columns)))
return None
[docs]
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/
[docs]
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: Driver-specific database initialization options passed to
``AdbcDatabase``. Keys and value types are defined by the driver
(e.g., ``"username"``, ``"password"`` for the PostgreSQL driver).
Do **not** put ADBC connection options here — they belong in
``conn_kwargs``.
- conn_kwargs: ADBC connection options as string key-value pairs.
Keys must use the canonical dotted ADBC option names, for example:
``"adbc.connection.autocommit": "true"``,
``"adbc.connection.read_only": "true"``,
``"adbc.connection.current_catalog": "my_catalog"``,
``"adbc.connection.current_db_schema": "my_schema"``.
Short names such as ``autocommit`` or ``read_only`` are **not**
recognized by ADBC drivers and will raise ``NotSupportedError``.
- 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
"""
[docs]
conn_name_attr = "adbc_conn_id"
[docs]
default_conn_name = "adbc_default"
[docs]
hook_name = "ADBC Connection"
[docs]
supports_autocommit = True
@classmethod
[docs]
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
[docs]
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
[docs]
def driver(self) -> str:
return self.connection_extra_lower.get("driver") or f"adbc_driver_{self.dialect_name}"
@cached_property
[docs]
def entrypoint(self) -> str | None:
return self.connection_extra_lower.get("entrypoint")
@cached_property
[docs]
def db_kwargs(self) -> dict:
return {**{"uri": self.uri}, **self.connection_extra_lower.get("db_kwargs", {})}
@cached_property
[docs]
def conn_kwargs(self) -> dict:
return self.connection_extra_lower.get("conn_kwargs", {})
@cached_property
[docs]
def dialect_name(self) -> str:
return self.connection_extra_lower.get("dialect", "default")
[docs]
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,
)
[docs]
def set_autocommit(self, conn: Connection, autocommit: bool) -> None:
"""
Set autocommit on the ADBC connection.
The DBAPI attribute ``conn.autocommit`` has no effect on the underlying
ADBC driver; the real lever is the ``adbc.connection.autocommit`` option.
This override applies the option at the driver level, then calls super()
to keep the Python-level bookkeeping attribute in sync so that the base
class ``get_autocommit()`` and ``run()`` commit-guard behave correctly.
"""
conn.adbc_connection.set_autocommit(autocommit)
super().set_autocommit(conn, autocommit)
[docs]
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_executemany(cls, cursor, statement: str, record_batch: RecordBatch) -> None:
"""
Execute a statement using cursor.executemany with a RecordBatch.
ADBC's executemany() accepts Arrow data directly and routes it through
bind_stream internally, making this the Arrow-native fast path.
"""
cursor.executemany(statement, record_batch)
def _resolve_execute_batch(
self, cursor, executemany: bool, fast_executemany: bool
) -> Callable[[Cursor, str, RecordBatch], None]:
"""
Return the batch-execute callable for the given cursor.
ADBC's executemany() natively accepts a RecordBatch, so it is always
the fast path. fast_executemany is passed through to drivers that
expose that knob (silently ignored otherwise).
"""
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(),
)
return self._execute_executemany
[docs]
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:
fields = {field.name: field for field in table_schema}
table_schema = schema([fields[name] for 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 or None):
batch = self._to_record_batch(rows=chunked_rows, schema=table_schema)
execute_batch(cur, sql, batch)
if not autocommit:
conn.commit()
nb_rows += batch.num_rows
self.log.info("Loaded %s rows into %s so far", nb_rows, table)
self.log.info("Done loading. Loaded a total of %s rows into %s", nb_rows, table)