potiuk commented on code in PR #52330: URL: https://github.com/apache/airflow/pull/52330#discussion_r3520094055
########## providers/apache/arrow/src/airflow/providers/apache/arrow/hooks/adbc.py: ########## @@ -0,0 +1,316 @@ +# +# 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 re +from collections.abc import Iterable, Mapping +from contextlib import closing +from functools import cached_property +from typing import 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 + + +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: + import pathlib Review Comment: `import pathlib` / `import sys` / `importlib_resources` are imported inside `_driver_path`. Per the coding standards imports go at the top of the file — none of the valid exceptions (circular import, worker-isolation lazy load, `TYPE_CHECKING`) apply here. Please move them to module scope. ########## providers/apache/arrow/tests/unit/apache/arrow/hooks/test_adbc.py: ########## @@ -0,0 +1,165 @@ +# 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 importlib +import json +import logging +from unittest import mock + +import pytest +from adbc_driver_manager import dbapi +from adbc_driver_manager.dbapi import Cursor +from pyarrow import field, schema, string + +from airflow.models import Connection +from airflow.providers.apache.arrow.hooks.adbc import AdbcHook +from airflow.providers.common.sql.dialects.dialect import Dialect + + +class TestAdbcHook: + def setup_method(self): + # Create a MagicMock cursor similar to DbApiHook tests + self.cur = mock.MagicMock(rowcount=0, fast_executemany=False) + self.conn = mock.MagicMock() + self.conn.cursor.return_value = self.cur + # Schema and extras that the hook might read + # Provide a real pyarrow Schema so _to_record_batch can build RecordBatch + self.conn.adbc_get_table_schema.return_value = schema([field("col", string())]) + self.conn.extra_dejson = {} + conn = self.conn + + logging.root.disabled = True + + # Instantiate the hook under test + self.hook = self.make_hook_for_conn(conn) + + # Ensure the cursor used in unit tests has a native bind method + # to simulate native Arrow bind support. + self.cur.bind = mock.MagicMock() + + # Make fetch_arrow_table().to_pydict() return a simple column mapping + arrow_table = mock.MagicMock() + arrow_table.to_pydict.return_value = {"col": [1, 2]} + self.cur.fetch_arrow_table.return_value = arrow_table + + def make_hook_for_conn(self, conn): + """ + Return an AdbcHook subclass instance bound to the provided conn. + + Tests previously redefined this subclass locally in multiple places. + This helper centralizes that logic so tests can simply call + `self.make_hook_for_conn(conn)`. + """ + + class AdbcHookMock(AdbcHook): + conn_name_attr = "adbc_default" + + @classmethod + def get_connection(cls, conn_id: str) -> Connection: + return conn + + def get_conn(self): + return conn + + @property + def dialect(self): + return Dialect(self) + + def get_db_log_messages(self, _conn) -> None: + return _conn.get_messages() + + return AdbcHookMock() + + def test_get_records_fetch_all_handler(self): + result = self.hook.get_records("SELECT 1") + assert result == [(1,), (2,)] + + def test_insert_rows_native_bind(self): + table = "table" + rows = [("a",), ("b",)] + + # Native bind supported (cursor has bind attribute) + self.hook.insert_rows(table, rows) + + assert self.cur.bind.called + assert self.cur.execute.called + assert self.conn.commit.call_count >= 1 + + def test_insert_rows_fast_executemany_not_supported(self): + # Cursor without native bind that doesn't support setting fast_executemany + class NoFastExecCursor(mock.MagicMock): + def __setattr__(self, name, value): + if name == "fast_executemany": + raise AttributeError("fast_executemany not supported") + super().__setattr__(name, value) + + cur = NoFastExecCursor(spec=Cursor) + delattr(cur, "bind") # Remove bind to simulate no native bind support + conn = mock.MagicMock() + conn.cursor.return_value = cur + conn.adbc_get_table_schema.return_value = schema([field("col", string())]) + conn.extra_dejson = {} + hook = self.make_hook_for_conn(conn) + + table = "table" + rows = [("x",), ("y",)] + + hook.insert_rows(table, rows, executemany=True, fast_executemany=True) + + assert cur.executemany.called + assert conn.commit.call_count >= 1 + + def test_insert_rows_fast_executemany_supported(self): + # Cursor without native bind but supports setting fast_executemany + cur = mock.MagicMock(spec=Cursor) + delattr(cur, "bind") # Remove bind to simulate no native bind support + conn = mock.MagicMock() + conn.cursor.return_value = cur + conn.adbc_get_table_schema.return_value = schema([field("col", string())]) + conn.extra_dejson = {} + hook = self.make_hook_for_conn(conn) + + table = "table" + rows = [("x",), ("y",)] + + hook.insert_rows(table, rows, executemany=True, fast_executemany=True) + + assert cur.fast_executemany + assert cur.executemany.called + assert conn.commit.call_count >= 1 + + @pytest.mark.skipif( + importlib.util.find_spec("adbc_driver_sqlite") is None, + reason="adbc_driver_sqlite not installed", + ) + def test_dbapi_connection(self, create_connection_without_db): Review Comment: The extras-resolution surface this PR adds isn't tested. `test_dbapi_connection` only sets `extra={"driver": ...}` and asserts the connection type — nothing exercises the `db_kwargs`/`conn_kwargs` passthrough (the feature the earlier review asked for), nor `uri`/`entrypoint`/`dialect` resolution. Per the testing rule, every added behaviour needs a test that fails without the change. Please add a case that sets `db_kwargs`/`conn_kwargs` in the connection `extra` and asserts they're actually forwarded into the ADBC `connect()` call (mock `connect` and check the kwargs), plus coverage for `entrypoint`/`dialect`. ########## providers/apache/arrow/docs/index.rst: ########## @@ -0,0 +1,135 @@ + .. 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. + +``apache-airflow-providers-apache-arrow`` +========================================= + + +.. toctree:: + :hidden: + :maxdepth: 1 + :caption: Basics + + Home <self> + Changelog <changelog> + Security <security> + +.. toctree:: + :hidden: + :maxdepth: 1 + :caption: References + + Python API <_api/airflow/providers/apache/arrow/index> + +.. toctree:: + :hidden: + :maxdepth: 1 + :caption: System tests + + System Tests <_api/tests/system/apache/arrow/index> + +.. toctree:: + :hidden: + :maxdepth: 1 + :caption: Resources + + PyPI Repository <https://pypi.org/project/apache-airflow-providers-arrow/> Review Comment: Broken PyPI link — the package is `apache-airflow-providers-apache-arrow`, not `apache-airflow-providers-arrow`. Should be `https://pypi.org/project/apache-airflow-providers-apache-arrow/`. ########## providers/apache/arrow/src/airflow/providers/apache/arrow/hooks/adbc.py: ########## @@ -0,0 +1,316 @@ +# +# 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 re +from collections.abc import Iterable, Mapping +from contextlib import closing +from functools import cached_property +from typing import 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 + + +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: + import pathlib + import sys + + import importlib_resources + + # 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: Review Comment: `db_kwargs`/`conn_kwargs` let a Connection editor pass arbitrary keyword arguments into the ADBC driver `connect()`. This is the acceptable named-key form (not a blind `**extra_dejson` spread), which is why the earlier review asked for it — but the *values* are unfiltered. Please enumerate in `docs/connections/adbc.rst` what `db_kwargs`/`conn_kwargs` accept, so a Connection editor's blast radius is explicit and there are no surprising file-path/RCE-style kwargs on the driver. -- 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]
