dabla commented on code in PR #52330: URL: https://github.com/apache/airflow/pull/52330#discussion_r3836153710
########## providers/apache/arrow/src/airflow/providers/apache/arrow/hooks/adbc.py: ########## @@ -0,0 +1,345 @@ +# +# 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: + table = cursor.fetch_arrow_table() + return list(zip(*(column.to_pylist() for column in table.columns))) + 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: 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** + recognised 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 + """ + + 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 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._conn.set_options(**{"adbc.connection.autocommit": "true" if autocommit else "false"}) Review Comment: You're right, there is a public API available for this. -- 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]
