ColtenOuO commented on code in PR #72845: URL: https://github.com/apache/airflow/pull/72845#discussion_r3986256652
########## providers/duckdb/src/airflow/providers/duckdb/hooks/duckdb.py: ########## @@ -0,0 +1,335 @@ +# 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 re +from functools import cached_property +from typing import TYPE_CHECKING, Any + +import duckdb + +from airflow.providers.common.sql.hooks.sql import DbApiHook +from airflow.providers.duckdb.version_compat import AirflowNotFoundException + +if TYPE_CHECKING: + from collections.abc import Sequence + + from duckdb import DuckDBPyConnection + + try: + from airflow.sdk import Connection + except ImportError: + from airflow.models.connection import Connection # type: ignore[assignment] + +IN_MEMORY_DATABASE = ":memory:" + +# DuckDB has no bind-parameter form for extension names or PRAGMA-style identifiers, so anything +# interpolated into those statements is validated against this instead. +_IDENTIFIER = re.compile(r"\A[A-Za-z_][A-Za-z0-9_]*\Z") + + +class DuckDBHook(DbApiHook): + """ + Interact with an in-process `DuckDB <https://duckdb.org/>`__ database. + + The hook opens a DuckDB database — in memory, backed by a local file, or hosted by MotherDuck — + applies resource limits, and loads the requested extensions, so a Dag author only supplies SQL. + + The Airflow connection is optional. With no connection configured the hook opens an in-memory + database, which is the right default for a stateless, task-scoped analytical query. + + Extensions are loaded with ``LOAD`` first and only installed when that fails. DuckDB downloads + extensions from its extension repository on first use, so deployments without outbound internet + access should pre-populate an extension directory, point ``extension_directory`` at it and set + ``autoinstall_extensions=False`` to turn the network access off entirely. + + :param duckdb_conn_id: reference to a :ref:`DuckDB connection <howto/connection:duckdb>`. The + connection need not exist; when it is missing an in-memory database is used. + :param database: database to open. Overrides the connection. ``:memory:`` (the default) opens a + transient database that is discarded when the task finishes. + :param extensions: extensions to load on connect, for example ``["httpfs", "iceberg"]``. + :param extension_directory: directory DuckDB loads extensions from and installs them into. + Point this at a pre-populated directory to avoid downloading extensions at task runtime. + :param autoinstall_extensions: whether an extension that is not installed locally may be + downloaded and installed. Set to ``False`` in environments without outbound internet access + so a missing extension fails loudly instead of hanging on a network call. + :param allow_community_extensions: whether DuckDB may load community (third-party) extensions. + Community extensions are native code from outside the DuckDB project, so this defaults to + ``False``. + :param memory_limit: memory DuckDB may use, for example ``"2GB"``. DuckDB otherwise sizes itself + from the memory it detects on the host, which over-commits inside a container that has a + smaller limit than the host it runs on. + :param threads: number of threads DuckDB may use. Defaults to the cores DuckDB detects, which is + subject to the same container caveat as ``memory_limit``. + :param temp_directory: directory DuckDB spills to when a query exceeds ``memory_limit``. + :param read_only: open the database read-only. Not valid for an in-memory database. + :param settings: additional DuckDB configuration options, passed through verbatim. + + Every parameter above except ``duckdb_conn_id`` may also be set in the connection ``extra``, in + which case an explicit argument wins. This mirrors what + :class:`~airflow.providers.common.sql.operators.sql.BaseSQLOperator` does when it merges connection + extras into hook keyword arguments — a path the connection-optional operator has to bypass. + """ + + conn_name_attr = "duckdb_conn_id" + default_conn_name = "duckdb_default" + conn_type = "duckdb" + hook_name = "DuckDB" + placeholder = "?" + supports_autocommit = False + + #: Extensions every connection opened by this hook class loads, regardless of configuration. + #: Subclasses that integrate a specific backend declare their requirements here. + required_extensions: tuple[str, ...] = () + + def __init__( + self, + *args, + duckdb_conn_id: str = default_conn_name, + database: str | None = None, + extensions: Sequence[str] | None = None, + extension_directory: str | None = None, + autoinstall_extensions: bool | None = None, + allow_community_extensions: bool | None = None, + memory_limit: str | None = None, + threads: int | None = None, + temp_directory: str | None = None, + read_only: bool | None = None, + settings: dict[str, Any] | None = None, + **kwargs, + ) -> None: + kwargs[self.conn_name_attr] = duckdb_conn_id + super().__init__(*args, **kwargs) + self.database = database + # Stored unresolved: ``None`` means "not set explicitly", so the connection extra may supply + # it. Reading these through the properties below is what makes the two sources consistent. + self._extensions = list(extensions) if extensions is not None else None + self._extension_directory = extension_directory + self._autoinstall_extensions = autoinstall_extensions + self._allow_community_extensions = allow_community_extensions + self._memory_limit = memory_limit + self._threads = threads + self._temp_directory = temp_directory + self._read_only = read_only + self._settings = settings + + def resolve_parameter(self, name: str, explicit: Any, default: Any = None) -> Any: + """Return the explicit argument if given, else the connection extra, else the default.""" + if explicit is not None: + return explicit + from_extra = self.connection_extra.get(name) + return default if from_extra is None else from_extra + + @property + def extension_directory(self) -> str | None: + return self.resolve_parameter("extension_directory", self._extension_directory) + + @property + def autoinstall_extensions(self) -> bool: + return bool(self.resolve_parameter("autoinstall_extensions", self._autoinstall_extensions, True)) + + @property + def allow_community_extensions(self) -> bool: + return bool( + self.resolve_parameter("allow_community_extensions", self._allow_community_extensions, False) + ) + + @property + def memory_limit(self) -> str | None: + return self.resolve_parameter("memory_limit", self._memory_limit) + + @property + def threads(self) -> int | None: + return self.resolve_parameter("threads", self._threads) + + @property + def temp_directory(self) -> str | None: + return self.resolve_parameter("temp_directory", self._temp_directory) + + @property + def read_only(self) -> bool: + return bool(self.resolve_parameter("read_only", self._read_only, False)) + + @property + def settings(self) -> dict[str, Any]: + merged = dict(self.connection_extra.get("settings") or {}) + merged.update(self._settings or {}) + return merged + + @classmethod + def get_ui_field_behaviour(cls) -> dict[str, Any]: + """Return custom UI field behaviour for the DuckDB connection.""" + return { + "hidden_fields": ["login", "port"], + "relabeling": { + "host": "Database path", + "schema": "MotherDuck database", + "password": "MotherDuck token", + }, + "placeholders": { + "host": "/tmp/analytics.duckdb (leave empty for an in-memory database)", + "extra": '{"extensions": ["httpfs"], "memory_limit": "2GB", "threads": 4}', + }, + } + + @cached_property + def airflow_connection(self) -> Connection | None: + """Return the configured Airflow connection, or ``None`` when it does not exist.""" + try: + return self.get_connection(self.get_conn_id()) + except AirflowNotFoundException: + self.log.debug( + "No Airflow connection %r; falling back to an in-memory DuckDB database.", + self.get_conn_id(), + ) Review Comment: Thanks! Looks great -- 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]
