eladkal commented on code in PR #72845: URL: https://github.com/apache/airflow/pull/72845#discussion_r4010432893
########## providers/duckdb/src/airflow/providers/duckdb/hooks/duckdb.py: ########## @@ -0,0 +1,356 @@ +# 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 ``duckdb_default`` connection configured the hook opens + an in-memory database, which is the right default for a stateless, task-scoped analytical query. + Supplying a connection id other than the default asserts that it exists, a missing connection raises + rather than quietly falling back to an in-memory database. + + Extensions are not downloaded at task runtime by default. DuckDB fetches them from its extension + repository on first use, which a deployment without outbound internet access cannot do and which + costs every worker the download. Pre-populate an extension directory, point ``extension_directory`` + at it, and set ``autoinstall_extensions=True`` only if downloading on demand is acceptable. + + :param duckdb_conn_id: reference to a :ref:`DuckDB connection <howto/connection:duckdb>`. The + default connection need not exist; when it is missing an in-memory database is used. Any other + id must exist, and raises if it does not. + :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 + from DuckDB's extension repository. Defaults to ``False``: a locked-down environment usually + cannot reach that repository, and downloading native code at task runtime is a decision worth + making explicitly. Set to ``True`` to allow it. + :param autoload_extensions: whether DuckDB may load an already-installed extension implicitly, so + that for example querying an ``s3://`` path pulls in ``httpfs`` without it being listed in + ``extensions``. Defaults to ``True``; this involves no download and no new code beyond what is + already present. + :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, + autoload_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._autoload_extensions = autoload_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, False)) + + @property + def autoload_extensions(self) -> bool: + return bool(self.resolve_parameter("autoload_extensions", self._autoload_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 the default one does not exist. + + A missing connection is only tolerated for the default connection id, which is what makes the + hook usable with no configuration at all. An explicitly supplied id asserts that a particular + connection exists, so its absence is a misconfiguration: falling back there would let a task + that meant to write to a real database silently write to one that is discarded when it ends. + """ + conn_id = self.get_conn_id() + try: + return self.get_connection(conn_id) + except AirflowNotFoundException: + if conn_id != self.default_conn_name: + raise + self.log.info("No Airflow connection %r; using an in-memory DuckDB database.", conn_id) + return None + + @cached_property + def connection_extra(self) -> dict[str, Any]: + """Return the connection's ``extra``, or an empty mapping when there is no connection.""" + connection = self.airflow_connection + return connection.extra_dejson if connection else {} + + def get_database(self) -> str: + """ + Return the DuckDB database to open. + + Precedence: the ``database`` argument, the connection extra ``database``, a MotherDuck + database built from the connection's token, the connection host, then an in-memory database. Review Comment: MotherDuck is the paid offering of DuckDB isn't it? I am not sure fi this should be referenced here? Can you provide some details about the special handling here? Normally how this works is the underlying tech is separated from the managed service. Much like opensearch has no aws related code in it but AWS provider have it's own version based on opensearch -- 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]
