#
# 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 pathlib
import re
import sys
from collections.abc import Callable, Iterable, Mapping
from contextlib import closing
from functools import cached_property, partial
from importlib import resources as importlib_resources
from typing import TYPE_CHECKING, Any
from urllib.parse import quote
from adbc_driver_manager.dbapi import Connection, NotSupportedError, 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.
- Loads rows with the driver's bulk ingest (``adbc_ingest``, e.g. ``COPY`` on PostgreSQL)
by default, or binds them with DBAPI ``executemany`` (``fast_executemany=False``).
- 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": [],
"relabeling": {"host": "Host or URI", "schema": "Database"},
"placeholders": {
"host": "db.example.com, or postgresql://user:pass@db.example.com:5432/db",
"port": "5432",
"schema": "db",
"extra": '{"dialect": "postgresql"}',
},
}
@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:
"""
Return the URI handed to the driver.
A host that already is a URI or a path (``postgresql://user:pass@host:5432/db``,
``file::memory:``, ``:memory:``, ``/data/app.db``) is passed through unchanged. Otherwise the
URI is built as ``<dialect>://login:password@host:port/schema`` from the connection's fields.
Airflow's own connection URI is of no use here: its query string carries the extras, which
drivers such as libpq refuse as unknown parameters.
"""
host = self.connection.host or ""
if "://" in host or "::" in host or host.startswith(("/", ":", "file:")):
return host
uri = f"{self.dialect_name.lower().replace('_', '-')}://"
if self.connection.login:
uri += quote(self.connection.login, safe="")
if self.connection.password:
uri += f":{quote(self.connection.password, safe='')}"
uri += "@"
uri += host
if self.connection.port:
uri += f":{self.connection.port}"
if self.connection.schema:
uri += f"/{quote(self.connection.schema, safe='')}"
return uri
@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,
)
def _resolve_execute_batch(
self,
cursor: Cursor,
*,
table_name: str,
schema_name: str | None,
sql: str,
replace: bool,
fast_executemany: bool,
) -> Callable[[RecordBatch], object]:
"""
Return the callable that loads one RecordBatch through ``cursor``.
With ``fast_executemany`` the batch goes through the driver's bulk ingest
(:meth:`adbc_driver_manager.dbapi.Cursor.adbc_ingest`, ``COPY`` on PostgreSQL).
Otherwise it is bound to the generated ``INSERT`` statement with
``cursor.executemany``, which most drivers execute once per row.
``replace=True`` always binds, because bulk ingest cannot carry the
dialect's conflict clause.
``fast_executemany`` is a hint: if the driver does not implement bulk ingest
(the first ``adbc_ingest`` call raises ``NotSupportedError``), that batch and
all following ones are bound with ``cursor.executemany`` instead.
"""
bind = partial(cursor.executemany, sql)
if replace or not fast_executemany:
if replace and fast_executemany:
self.log.debug("Bulk ingest cannot replace rows; binding them to the statement instead")
return bind
self.log.info("Using bulk ingest (adbc_ingest) for conn_id '%s'", self.get_conn_id())
ingest = partial(cursor.adbc_ingest, table_name, mode="append", db_schema_name=schema_name)
loader: Callable[[RecordBatch], object] | None = None
def execute_batch(batch: RecordBatch) -> object:
nonlocal loader
if loader is not None:
return loader(batch)
try:
result = ingest(batch)
except NotSupportedError:
self.log.info(
"Driver for conn_id '%s' does not support bulk ingest; "
"binding rows with executemany instead",
self.get_conn_id(),
)
loader = bind
return bind(batch)
loader = ingest
return result
return execute_batch
[docs]
def insert_rows(
self,
table,
rows,
target_fields=None,
commit_every=1000,
replace=False,
*,
executemany=False,
fast_executemany=True,
autocommit=False,
**kwargs,
):
"""
Insert a collection of tuples into a table.
Rows are converted to Arrow ``RecordBatch`` objects and inserted in chunks;
each chunk (of size ``commit_every``) is done in a new transaction.
:param table: Name of the target table, optionally prefixed with its schema
(``schema.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. Replacing always binds
the rows to the generated statement, because bulk ingest cannot express it.
:param executemany: Kept for signature compatibility with
:meth:`~airflow.providers.common.sql.hooks.sql.DbApiHook.insert_rows`;
rows are always sent as one Arrow batch per chunk.
:param fast_executemany: If True (default), load each batch with the driver's
bulk ingest (``cursor.adbc_ingest``, ``COPY`` on PostgreSQL). If the driver
does not implement bulk ingest, falls back to binding. If False, bind
each batch to the ``INSERT`` statement with ``cursor.executemany``.
: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,
table_name=table_name,
schema_name=schema_name,
sql=sql,
replace=replace,
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(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)