This is an automated email from the ASF dual-hosted git repository.
o-nikolas pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new cc19675d24f Add DuckDB support to the Amazon provider (#73301)
cc19675d24f is described below
commit cc19675d24f04b72ec2443d385b9ba2c9be18047
Author: Niko Oliveira <[email protected]>
AuthorDate: Tue Sep 29 12:51:30 2026 -0700
Add DuckDB support to the Amazon provider (#73301)
Run DuckDB queries against S3 without wiring up credentials in the Dag.
Point
DuckDBExecuteQueryOperator at a duckdb_aws connection and the Amazon
provider
brokers credentials from your AWS connection into the DuckDB session.
The connection type is what selects the AWS hook, so the same operator
serves local
files and S3. This follows the precedent of redshift and athena, where the
generic SQL operator picks up the right hook based on the redshift and
athena connections.
DuckDB is an optional extra of this provider.
---
dev/breeze/tests/test_selective_checks.py | 10 +-
providers/amazon/docs/index.rst | 2 +
providers/amazon/docs/operators/duckdb.rst | 180 ++++++++++++
providers/amazon/provider.yaml | 31 +++
providers/amazon/pyproject.toml | 9 +-
.../airflow/providers/amazon/aws/hooks/duckdb.py | 277 +++++++++++++++++++
.../airflow/providers/amazon/get_provider_info.py | 23 ++
.../tests/system/amazon/aws/example_duckdb.py | 158 +++++++++++
.../tests/unit/amazon/aws/hooks/test_duckdb.py | 307 +++++++++++++++++++++
providers/duckdb/docs/operators.rst | 6 +
.../src/airflow/providers/duckdb/hooks/duckdb.py | 6 +-
.../airflow/providers/duckdb/operators/duckdb.py | 48 +++-
.../tests/unit/duckdb/operators/test_duckdb.py | 33 ++-
uv.lock | 8 +-
14 files changed, 1078 insertions(+), 20 deletions(-)
diff --git a/dev/breeze/tests/test_selective_checks.py
b/dev/breeze/tests/test_selective_checks.py
index 7f64795d061..f496912f6b3 100644
--- a/dev/breeze/tests/test_selective_checks.py
+++ b/dev/breeze/tests/test_selective_checks.py
@@ -1130,7 +1130,8 @@ def assert_outputs_are_printed(expected_outputs:
dict[str, str], stderr: str):
("providers/amazon/src/airflow/providers/amazon/provider.yaml",),
{
"selected-providers-list-as-string": "amazon apache.hive
cncf.kubernetes "
- "common.compat common.messaging common.sql databricks exasol
ftp google http imap microsoft.azure "
+ "common.compat common.messaging common.sql databricks duckdb
exasol ftp google http imap "
+ "microsoft.azure "
"mongo mysql openlineage postgres salesforce ssh teradata",
"all-python-versions":
f"['{DEFAULT_PYTHON_MAJOR_MINOR_VERSION}']",
"all-python-versions-list-as-string":
DEFAULT_PYTHON_MAJOR_MINOR_VERSION,
@@ -1156,7 +1157,7 @@ def assert_outputs_are_printed(expected_outputs:
dict[str, str], stderr: str):
{
"description": "amazon...google",
"test_types": "Providers[amazon]
Providers[apache.hive,cncf.kubernetes,"
-
"common.compat,common.messaging,common.sql,databricks,exasol,ftp,http,imap,"
+
"common.compat,common.messaging,common.sql,databricks,duckdb,exasol,ftp,http,imap,"
"microsoft.azure,mongo,mysql,openlineage,postgres,salesforce,ssh,teradata] "
"Providers[google]",
}
@@ -1202,7 +1203,8 @@ def assert_outputs_are_printed(expected_outputs:
dict[str, str], stderr: str):
("providers/amazon/src/airflow/providers/amazon/file.py",),
{
"selected-providers-list-as-string": "amazon apache.hive
cncf.kubernetes "
- "common.compat common.messaging common.sql databricks exasol
ftp google http imap microsoft.azure "
+ "common.compat common.messaging common.sql databricks duckdb
exasol ftp google http imap "
+ "microsoft.azure "
"mongo mysql openlineage postgres salesforce ssh teradata",
"all-python-versions":
f"['{DEFAULT_PYTHON_MAJOR_MINOR_VERSION}']",
"all-python-versions-list-as-string":
DEFAULT_PYTHON_MAJOR_MINOR_VERSION,
@@ -1225,7 +1227,7 @@ def assert_outputs_are_printed(expected_outputs:
dict[str, str], stderr: str):
{
"description": "amazon...google",
"test_types": "Providers[amazon]
Providers[apache.hive,cncf.kubernetes,"
-
"common.compat,common.messaging,common.sql,databricks,exasol,ftp,http,imap,"
+
"common.compat,common.messaging,common.sql,databricks,duckdb,exasol,ftp,http,imap,"
"microsoft.azure,mongo,mysql,openlineage,postgres,salesforce,ssh,teradata] "
"Providers[google]",
}
diff --git a/providers/amazon/docs/index.rst b/providers/amazon/docs/index.rst
index 552f1be4a55..970d7f5cd32 100644
--- a/providers/amazon/docs/index.rst
+++ b/providers/amazon/docs/index.rst
@@ -150,6 +150,7 @@ Dependent package
`apache-airflow-providers-apache-hive
<https://airflow.apache.org/docs/apache-airflow-providers-apache-hive>`_
``apache.hive``
`apache-airflow-providers-cncf-kubernetes
<https://airflow.apache.org/docs/apache-airflow-providers-cncf-kubernetes>`_
``cncf.kubernetes``
`apache-airflow-providers-common-messaging
<https://airflow.apache.org/docs/apache-airflow-providers-common-messaging>`_
``common.messaging``
+`apache-airflow-providers-duckdb
<https://airflow.apache.org/docs/apache-airflow-providers-duckdb>`_
``duckdb``
`apache-airflow-providers-exasol
<https://airflow.apache.org/docs/apache-airflow-providers-exasol>`_
``exasol``
`apache-airflow-providers-ftp
<https://airflow.apache.org/docs/apache-airflow-providers-ftp>`_
``ftp``
`apache-airflow-providers-google
<https://airflow.apache.org/docs/apache-airflow-providers-google>`_
``google``
@@ -194,6 +195,7 @@ Extra Dependencies
``ssh`` ``apache-airflow-providers-ssh``
``standard`` ``apache-airflow-providers-standard``
``common.messaging`` ``apache-airflow-providers-common-messaging>=2.0.0``
+``duckdb`` ``apache-airflow-providers-duckdb>=0.1.0``
``sqlalchemy`` ``sqlalchemy>=1.4.54,!=2.1.0``
====================
============================================================================================================================================================
diff --git a/providers/amazon/docs/operators/duckdb.rst
b/providers/amazon/docs/operators/duckdb.rst
new file mode 100644
index 00000000000..07a75cd35f0
--- /dev/null
+++ b/providers/amazon/docs/operators/duckdb.rst
@@ -0,0 +1,180 @@
+ .. 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.
+
+=============
+DuckDB on AWS
+=============
+
+`DuckDB <https://duckdb.org/>`__ is an in-process analytical database. It can
read and write data in
+Amazon S3 directly, which makes it an affordable way to run transforms that do
not justify the scale
+of a distributed cluster.
+
+DuckDB is not an AWS service, so the generic hook, operator and ``duckdb``
connection type live in the
+:doc:`DuckDB provider <apache-airflow-providers-duckdb:index>`. What this
provider adds is the AWS
+auth. Using an Airflow AWS connection to create the DuckDB secret that grants
access to S3.
+This means there is no credential wiring needed in the Dag.
+
+You run
:class:`~airflow.providers.duckdb.operators.duckdb.DuckDBExecuteQueryOperator`
and point it
+at a ``duckdb_aws`` connection when the SQL should reach S3.
+
+Prerequisite Tasks
+------------------
+
+.. include:: ../_partials/prerequisite_tasks.rst
+
+DuckDB support is an optional extra
+
+.. code-block:: bash
+
+ pip install 'apache-airflow-providers-amazon[duckdb]'
+
+.. note::
+
+ S3 access needs DuckDB's ``httpfs`` and ``aws`` extensions. This provider
installs them if they
+ are missing by default. If your workers should not download anything at
runtime, set
+ ``autoinstall_extensions=False`` and make the extensions available
yourself, either
+ in an ``extension_directory`` or baked into your image.
+
+.. _howto/connection:duckdb_aws:
+
+Run a DuckDB query against Amazon S3
+====================================
+
+Two connections are involved:
+
+* A **DuckDB on AWS** (``duckdb_aws``) connection which configures DuckDB
itself, i.e.: *which DuckDB database and with what settings*.
+ Its id is what you pass to
+
:class:`~airflow.providers.duckdb.operators.duckdb.DuckDBExecuteQueryOperator`
as ``conn_id``, and its
+ type is also what selects
:class:`~airflow.providers.amazon.aws.hooks.duckdb.AwsDuckDBHook`.
+* An **AWS** (``aws``) connection says *which AWS identity*, i.e.: keys or
role, region, endpoint, etc. The DuckDB
+ connection points at it through ``aws_conn_id``, which defaults to
``aws_default``.
+
+Credentials stay in the usual AWS connection. The ``duckdb_aws`` connection
can be empty (though it must exist) and often
+will be with DuckDB's ability to use in-memory databases. Create it, name it
whatever you like, and leave every field blank
+to get an in-memory database.
+
+Given that connection, the hook loads the ``httpfs`` and ``aws`` extensions
and creates the S3 secret,
+so the task supplies only SQL:
+
+.. exampleinclude:: /../../amazon/tests/system/amazon/aws/example_duckdb.py
+ :language: python
+ :dedent: 4
+ :start-after: [START howto_operator_aws_duckdb]
+ :end-before: [END howto_operator_aws_duckdb]
+
+Pass the id explicitly. With no ``conn_id`` the operator uses
``duckdb_default`` and the vendor-neutral
+hook, which has no S3 access.
+
+The DuckDB connection carries the database and the engine settings; leave its
``Database path`` empty for
+an in-memory database, which suits a task that reads from S3 and writes back
to it. Everything
+:class:`~airflow.providers.amazon.aws.hooks.duckdb.AwsDuckDBHook` accepts can
be set in the connection
+``extra``, and ``hook_params`` on the operator overrides it per task. Set
``memory_limit`` and
+``threads`` explicitly: DuckDB otherwise sizes itself from the resources it
detects on the host.
+
+.. code-block:: python
+
+ hook_params = {"memory_limit": "2GB", "threads": 4}
+
+.. note::
+
+ ``aws_conn_id`` resolves on whether the key is present, not on its value,
because ``None`` is a
+ valid choice for it rather than the absence of one. Set it to ``null`` in
the ``extra`` (or
+ ``None`` in ``hook_params``) to use no Airflow AWS connection at all and
let the AWS SDK find
+ credentials in the environment, such as environment variables or an
instance role.
+
+Using the hook directly
+=======================
+
+For work that is not a single statement (e.g. chaining several queries against
one database, or returning
+results to Python) you may use the hook directly. Opening one connection for
several statements also
+avoids paying connection setup and extension loading more than once per task:
+
+.. code-block:: python
+
+ @task
+ def compare_regions():
+ hook = AwsDuckDBHook(aws_conn_id="aws_default")
+ with hook.get_conn() as conn:
+ conn.execute("CREATE TABLE sales AS SELECT * FROM
read_parquet('s3://my-bucket/sales.parquet')")
+ return conn.execute("SELECT region, SUM(revenue) FROM sales GROUP
BY region").fetchall()
+
+How credentials reach DuckDB
+============================
+
+DuckDB does not use ``boto3``. The hook therefore issues a ``CREATE SECRET``
statement that tells
+DuckDB how to authenticate, and the ``credential_strategy`` parameter selects
how:
+
+``credential_chain`` (default)
+ Issues ``CREATE SECRET (TYPE s3, PROVIDER credential_chain)``, which
resolves credentials through
+ the AWS SDK inside DuckDB's ``aws`` extension. Nothing secret is written
into SQL text and the SDK
+ refreshes expiring credentials itself, so a long query does not fail
partway through when a set of
+ temporary credentials expires.
+
+ From DuckDB 1.4 onward the chain is resolved when the secret is created
rather than when S3 is
+ first read. If nothing can be resolved, the hook logs a warning and opens
the connection anyway,
+ so a query that does not use S3 still runs; one that does will fail on its
own. A ``config``
+ secret that cannot be created raises instead, since explicit credentials
failing is a
+ misconfiguration.
+
+``config``
+ Resolves credentials through
+ :class:`~airflow.providers.amazon.aws.hooks.base_aws.AwsBaseHook` and
writes them into the secret
+ explicitly. This also works with a container-assigned role, because boto3
resolves that on the
+ Airflow side, but the credentials are frozen when the connection is opened
rather than refreshed.
+ Needed when the Airflow connection carries static keys that the AWS SDK
running inside
+ DuckDB cannot see. The credentials become part of the SQL statement, so
prefer
+ ``credential_chain`` where it works.
+
+``none``
+ Creates no secret. Use for a DuckDB database that never touches S3.
+
+.. code-block:: python
+
+ DuckDBExecuteQueryOperator(
+ task_id="query",
+ conn_id="duckdb_aws_default",
+ sql="SELECT COUNT(*) FROM read_parquet('s3://my-bucket/data.parquet')",
+ hook_params={"aws_conn_id": "aws_static_keys", "credential_strategy":
"config"},
+ )
+
+Non-default S3 endpoints
+========================
+
+Set ``s3_endpoint_url``, or ``endpoint_url`` in the AWS connection extra, to
reach an S3-compatible
+service or an S3 interface VPC endpoint. The hook translates it into the
secret's ``ENDPOINT``,
+``URL_STYLE`` and ``USE_SSL`` clauses, selecting path-style addressing because
virtual-host-style
+bucket names generally do not resolve against such endpoints.
+
+.. note::
+
+ The S3 secret is the only channel through which Airflow configuration
reaches DuckDB. A DuckDB
+ session with no secret ignores the Airflow connection entirely — including
its region and endpoint
+ falling back to whatever DuckDB's own settings and environment provide.
+
+ DuckDB's ``httpfs`` extension reads credentials from the standard
+ ``AWS_ACCESS_KEY_ID`` / ``AWS_SECRET_ACCESS_KEY`` / ``AWS_SESSION_TOKEN``
environment
+ variables, so it can authenticate where those are set. What its HTTP
client does not implement is
+ the ECS container credential provider, so a container-assigned task role
(often how a managed Airflow
+ worker receives credentials) is invisible to it and S3 access fails with
an opaque
+ ``HTTP 403``. Both ``credential_chain`` and ``config`` handle that case;
the difference is whether
+ DuckDB or Airflow does the resolving, and therefore whether the
credentials refresh.
+
+Reference
+=========
+
+* `DuckDB S3 API support
<https://duckdb.org/docs/stable/extensions/httpfs/s3api>`__
+* `DuckDB secrets manager
<https://duckdb.org/docs/stable/configuration/secrets_manager>`__
diff --git a/providers/amazon/provider.yaml b/providers/amazon/provider.yaml
index e9788a4fa81..f6d91b78fcf 100644
--- a/providers/amazon/provider.yaml
+++ b/providers/amazon/provider.yaml
@@ -143,6 +143,14 @@ integrations:
- /docs/apache-airflow-providers-amazon/operators/athena/athena_sql.rst
- /docs/apache-airflow-providers-amazon/operators/athena/athena_spark.rst
tags: [aws]
+ # DuckDB itself is not an AWS service: the generic hook, operator and
``duckdb`` connection type
+ # live in the ``duckdb`` provider. What lives here is only the AWS-specific
part, brokering AWS
+ # credentials and endpoints into a DuckDB session configured by its own
connection type
+ - integration-name: DuckDB on AWS
+ external-doc-url: https://duckdb.org/docs/stable/extensions/httpfs/s3api
+ how-to-guide:
+ - /docs/apache-airflow-providers-amazon/operators/duckdb.rst
+ tags: [aws]
- integration-name: Amazon Bedrock
external-doc-url: https://aws.amazon.com/bedrock/
logo: /docs/integration-logos/[email protected]
@@ -675,6 +683,9 @@ hooks:
python-modules:
- airflow.providers.amazon.aws.hooks.athena
- airflow.providers.amazon.aws.hooks.athena_sql
+ - integration-name: DuckDB on AWS
+ python-modules:
+ - airflow.providers.amazon.aws.hooks.duckdb
- integration-name: Amazon Bedrock
python-modules:
- airflow.providers.amazon.aws.hooks.bedrock
@@ -1069,6 +1080,26 @@ connection-types:
schema: https
host: hooks.chime.aws/incomingwebhook/
password: T00000000?token=XXXXXXXXXXXXXXXXXXXXXXXX
+ - hook-class-name: airflow.providers.amazon.aws.hooks.duckdb.AwsDuckDBHook
+ hook-name: "DuckDB on AWS"
+ connection-type: duckdb_aws
+ ui-field-behaviour:
+ hidden-fields:
+ - login
+ - password
+ - port
+ - schema
+ relabeling:
+ host: Database path
+ placeholders:
+ host: /tmp/analytics.duckdb (leave empty for an in-memory database)
+ extra: |
+ {
+ "credential_strategy": "credential_chain",
+ "region_name": "us-east-1",
+ "memory_limit": "2GB",
+ "threads": 4
+ }
- hook-class-name: airflow.providers.amazon.aws.hooks.emr.EmrHook
hook-name: "Amazon Elastic MapReduce"
connection-type: emr
diff --git a/providers/amazon/pyproject.toml b/providers/amazon/pyproject.toml
index 06a7032587f..7a8b3e9c5d6 100644
--- a/providers/amazon/pyproject.toml
+++ b/providers/amazon/pyproject.toml
@@ -152,6 +152,11 @@ dependencies = [
"common.messaging" = [
"apache-airflow-providers-common-messaging>=2.0.0"
]
+# DuckDB support is opt-in: the duckdb wheel is around 20MB and irrelevant to
most users of this
+# provider, so it must never become a hard dependency here.
+"duckdb" = [
+ "apache-airflow-providers-duckdb>=0.1.0"
+]
"sqlalchemy" = [
"sqlalchemy>=1.4.54,!=2.1.0",
]
@@ -166,6 +171,7 @@ dev = [
"apache-airflow-providers-common-compat",
"apache-airflow-providers-common-messaging",
"apache-airflow-providers-common-sql",
+ "apache-airflow-providers-duckdb",
"apache-airflow-providers-exasol",
"apache-airflow-providers-ftp",
"apache-airflow-providers-google",
@@ -191,7 +197,8 @@ dev = [
"opensearch-py>=2.2.0",
"responses>=0.25.0",
"apache-airflow-providers-common-sql[pandas,polars]",
- "apache-airflow-providers-amazon[sqlalchemy]"
+ "apache-airflow-providers-amazon[sqlalchemy]",
+ "apache-airflow-providers-duckdb"
]
# To build docs:
diff --git a/providers/amazon/src/airflow/providers/amazon/aws/hooks/duckdb.py
b/providers/amazon/src/airflow/providers/amazon/aws/hooks/duckdb.py
new file mode 100644
index 00000000000..3f6e800ae7f
--- /dev/null
+++ b/providers/amazon/src/airflow/providers/amazon/aws/hooks/duckdb.py
@@ -0,0 +1,277 @@
+# 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
+
+from functools import cached_property
+from typing import TYPE_CHECKING, Any, Literal
+from urllib.parse import urlsplit
+
+from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
+from airflow.providers.amazon.version_compat import NOTSET, ArgNotSet
+from airflow.providers.common.compat.sdk import
AirflowOptionalProviderFeatureException
+
+try:
+ import duckdb
+
+ from airflow.providers.duckdb.hooks.duckdb import DuckDBHook
+except ImportError:
+ raise AirflowOptionalProviderFeatureException(
+ "This feature requires the 'duckdb' provider to be installed. Install
it with: "
+ "pip install 'apache-airflow-providers-amazon[duckdb]'"
+ )
+
+if TYPE_CHECKING:
+ from duckdb import DuckDBPyConnection
+
+CredentialStrategy = Literal["credential_chain", "config", "none"]
+
+#: Used when neither an argument nor the connection extra names an AWS
connection.
+DEFAULT_AWS_CONN_ID = "aws_default"
+
+
+class AwsDuckDBHook(DuckDBHook):
+ """
+ Interact with DuckDB using AWS credentials brokered by Airflow.
+
+ Extends :class:`~airflow.providers.duckdb.hooks.duckdb.DuckDBHook` with an
S3 secret built from
+ an Airflow AWS connection, so ``read_parquet('s3://...')`` and ``COPY ...
TO 's3://...'`` work
+ with no credential wiring in the Dag.
+
+ Two credential strategies are supported:
+
+ ``credential_chain`` (the default)
+ Issues ``CREATE SECRET (TYPE s3, PROVIDER credential_chain)``, which
resolves credentials
+ through the AWS SDK inside DuckDB's ``aws`` extension. Nothing secret
is written into SQL
+ text, and because DuckDB holds the resolution itself the SDK refreshes
expiring credentials,
+ so a query that outlives a set of temporary credentials does not fail
partway through.
+
+ ``config``
+ Resolves credentials through
:class:`~airflow.providers.amazon.aws.hooks.base_aws.AwsBaseHook`
+ and writes them into the secret explicitly. Also works with a
container-assigned role, since
+ boto3 resolves that on the Airflow side, but the credentials are
frozen at connect time and
+ become part of the SQL statement. Needed when the connection carries
static keys the AWS SDK
+ running inside DuckDB cannot see, or when talking to an S3-compatible
endpoint with its own
+ credentials. Prefer ``credential_chain`` where it works.
+
+ Either strategy is better than none: DuckDB's ``httpfs`` extension has its
own HTTP client which
+ reads credentials from the standard ``AWS_*`` environment variables but
does **not** implement the
+ ECS container credential provider, so an ambient task role is invisible to
it and S3 access fails
+ with an opaque ``HTTP 403``. A session with no secret at all also ignores
the Airflow connection
+ entirely, including its region and endpoint.
+
+ This hook needs the ``httpfs`` and ``aws`` extensions and installs them if
they are missing, which
+ is where it diverges from
:class:`~airflow.providers.duckdb.hooks.duckdb.DuckDBHook`. Using this
+ hook says the task runs against AWS, so those two extensions should just
work. Set
+ ``autoinstall_extensions=False`` to forbid the download, in which case the
extensions have to be
+ present already, either in an ``extension_directory`` or baked into the
image.
+
+ :param aws_conn_id: the :ref:`AWS connection <howto/connection:aws>` used
to reach S3. Defaults to
+ ``aws_default``. Set to ``None`` to use no Airflow connection and let
the AWS SDK resolve
+ credentials from the environment instead.
+ :param region_name: region for the S3 secret. Defaults to the region from
``aws_conn_id``.
+ :param credential_strategy: ``credential_chain``, ``config``, or ``none``
to create no secret at
+ all (for a DuckDB database that never touches S3).
+ :param credential_chain: providers DuckDB's credential chain consults, for
example
+ ``"env;config;sts;instance"``. Only used with the ``credential_chain``
strategy; DuckDB's own
+ default order applies when it is not set.
+ :param s3_endpoint_url: override the S3 endpoint, for an S3-compatible
service or a VPC endpoint.
+ Defaults to ``endpoint_url`` from the AWS connection extra.
+ :param secret_name: name of the DuckDB secret to create.
+
+ As with :class:`~airflow.providers.duckdb.hooks.duckdb.DuckDBHook`, every
parameter may also be set
+ in the DuckDB connection ``extra``, and an explicit argument wins.
``aws_conn_id`` resolves on
+ whether the key is present, so ``{"aws_conn_id": null}`` selects the
ambient AWS environment.
+ """
+
+ #: ``httpfs`` provides the S3 filesystem; ``aws`` provides the credential
chain secret provider.
+ required_extensions = ("httpfs", "aws")
+
+ # A connection type is what lets DuckDBExecuteQueryOperator reach this
hook: the
+ # operator resolves the hook from the connection, so a Dag pointing at a
``duckdb_aws``
+ # connection.
+ conn_type = "duckdb_aws"
+ hook_name = "DuckDB on AWS"
+ default_conn_name = "duckdb_aws_default"
+
+ @classmethod
+ def get_ui_field_behaviour(cls) -> dict[str, Any]:
+ """Return custom UI field behaviour for the DuckDB on AWS
connection."""
+ return {
+ "hidden_fields": ["login", "password", "port", "schema"],
+ "relabeling": {"host": "Database path"},
+ "placeholders": {
+ "host": "/tmp/analytics.duckdb (leave empty for an in-memory
database)",
+ "extra": '{"credential_strategy": "credential_chain",
"region_name": "us-east-1", '
+ '"memory_limit": "2GB", "threads": 4}',
+ },
+ }
+
+ def __init__(
+ self,
+ *args,
+ aws_conn_id: str | None | ArgNotSet = NOTSET,
+ region_name: str | None = None,
+ credential_strategy: CredentialStrategy | None = None,
+ credential_chain: str | None = None,
+ s3_endpoint_url: str | None = None,
+ secret_name: str | None = None,
+ **kwargs,
+ ) -> None:
+ super().__init__(*args, **kwargs)
+ self._aws_conn_id = aws_conn_id
+ self._region_name = region_name
+ self._credential_strategy = credential_strategy
+ self._credential_chain = credential_chain
+ self._s3_endpoint_url = s3_endpoint_url
+ self._secret_name = secret_name
+
+ @property
+ def aws_conn_id(self) -> str | None:
+ """
+ Return the AWS connection to take credentials from, or ``None`` for
the ambient environment.
+
+ Resolved by key presence rather than through ``resolve_parameter``,
because ``None`` is a
+ valid choice here and not the absence of one: ``{"aws_conn_id":
null}`` in the ``extra``
+ asks for no Airflow connection at all.
+ """
+ if not isinstance(self._aws_conn_id, ArgNotSet):
+ return self._aws_conn_id
+ if "aws_conn_id" in self.connection_extra:
+ return self.connection_extra["aws_conn_id"]
+ return DEFAULT_AWS_CONN_ID
+
+ @cached_property
+ def aws_hook(self) -> AwsBaseHook:
+ """Return the AWS hook credentials and region are resolved through."""
+ return AwsBaseHook(aws_conn_id=self.aws_conn_id,
region_name=self._region_name)
+
+ @property
+ def autoinstall_extensions(self) -> bool:
+ """
+ Default to installing missing extensions, unlike the generic hook.
+
+ Reaching for the AWS hook is a statement that the task runs against
AWS, so ``httpfs`` and
+ ``aws`` should just work. The generic hook cannot assume that and
leaves downloads off.
+ """
+ return bool(self.resolve_parameter("autoinstall_extensions",
self._autoinstall_extensions, True))
+
+ @property
+ def region_name(self) -> str | None:
+ return self.resolve_parameter("region_name", self._region_name)
+
+ @property
+ def credential_strategy(self) -> CredentialStrategy:
+ return self.resolve_parameter("credential_strategy",
self._credential_strategy, "credential_chain")
+
+ @property
+ def credential_chain(self) -> str | None:
+ return self.resolve_parameter("credential_chain",
self._credential_chain)
+
+ @property
+ def secret_name(self) -> str:
+ return self.resolve_parameter("secret_name", self._secret_name,
"airflow_aws")
+
+ def get_region_name(self) -> str | None:
+ """Return the region for the S3 secret."""
+ return self.region_name or self.aws_hook.region_name
+
+ def get_s3_endpoint_url(self) -> str | None:
+ """Return the S3 endpoint override, if any."""
+ explicit = self.resolve_parameter("s3_endpoint_url",
self._s3_endpoint_url)
+ if explicit:
+ return explicit
+ return self.connection_extra.get("endpoint_url") or
self.aws_hook.conn_config.endpoint_url
+
+ def configure_secrets(self, conn: DuckDBPyConnection) -> None:
+ """Create the DuckDB S3 secret for this connection."""
+ strategy = self.credential_strategy
+ if strategy == "none":
+ self.log.debug("credential_strategy is 'none'; not creating a
DuckDB S3 secret.")
+ return
+
+ self._validate_identifier(self.secret_name, "secret")
+ parts = ["TYPE s3", *self._provider_clauses(strategy)]
+
+ region = self.get_region_name()
+ if region:
+ parts.append(f"REGION '{self._quote(region)}'")
+
+ endpoint_url = self.get_s3_endpoint_url()
+ if endpoint_url:
+ parts.extend(self._endpoint_clauses(endpoint_url))
+
+ self.log.info("Creating DuckDB S3 secret %r using the %r strategy.",
self.secret_name, strategy)
+ try:
+ conn.execute(f"CREATE OR REPLACE SECRET {self.secret_name} ({',
'.join(parts)});")
+ except duckdb.Error as error:
+ if strategy != "credential_chain":
+ raise
+ # DuckDB 1.4 onwards resolves the credential chain when the secret
is created rather than
+ # when S3 is first read, so an environment with no credentials
fails here. A chain that
+ # finds nothing is not a misconfiguration, and a query that never
touches S3 is still
+ # valid, so warn rather than refusing to open the connection.
+ self.log.warning(
+ "Could not create DuckDB S3 secret %r from the credential
chain: %s. S3 access will "
+ "fail. Set credential_strategy='config' to supply credentials
explicitly, or 'none' "
+ "if this database does not use S3.",
+ self.secret_name,
+ error,
+ )
+
+ def _provider_clauses(self, strategy: CredentialStrategy) -> list[str]:
+ if strategy == "credential_chain":
+ clauses = ["PROVIDER credential_chain"]
+ if self.credential_chain:
+ clauses.append(f"CHAIN '{self._quote(self.credential_chain)}'")
+ return clauses
+ if strategy == "config":
+ credentials =
self.aws_hook.get_credentials(region_name=self.get_region_name())
+ clauses = [
+ "PROVIDER config",
+ f"KEY_ID '{self._quote(credentials.access_key)}'",
+ f"SECRET '{self._quote(credentials.secret_key)}'",
+ ]
+ if credentials.token:
+ clauses.append(f"SESSION_TOKEN
'{self._quote(credentials.token)}'")
+ return clauses
+ raise ValueError(
+ f"Unknown credential_strategy {strategy!r}. Expected
'credential_chain', 'config' or 'none'."
+ )
+
+ def _endpoint_clauses(self, endpoint_url: str) -> list[str]:
+ """
+ Translate an endpoint URL into DuckDB secret clauses.
+
+ DuckDB wants a bare ``host[:port]`` and a separate TLS flag rather
than a URL, and an
+ S3-compatible endpoint generally needs path-style addressing because
virtual-host-style
+ buckets do not resolve against it.
+ """
+ split = urlsplit(endpoint_url if "//" in endpoint_url else
f"//{endpoint_url}")
+ host = split.netloc or split.path
+ clauses = [f"ENDPOINT '{self._quote(host)}'", "URL_STYLE 'path'"]
+ if split.scheme == "http":
+ clauses.append("USE_SSL false")
+ return clauses
+
+ @staticmethod
+ def _quote(value: str) -> str:
+ """Escape single quotes so a value cannot break out of its SQL
literal."""
+ return value.replace("'", "''")
+
+ def get_openlineage_database_info(self, connection: Any) -> Any:
+ """Return no lineage metadata; a DuckDB database is task-local and has
no stable namespace."""
+ return None
diff --git a/providers/amazon/src/airflow/providers/amazon/get_provider_info.py
b/providers/amazon/src/airflow/providers/amazon/get_provider_info.py
index 388b4395e4d..74a0f97fb6e 100644
--- a/providers/amazon/src/airflow/providers/amazon/get_provider_info.py
+++ b/providers/amazon/src/airflow/providers/amazon/get_provider_info.py
@@ -38,6 +38,12 @@ def get_provider_info():
],
"tags": ["aws"],
},
+ {
+ "integration-name": "DuckDB on AWS",
+ "external-doc-url":
"https://duckdb.org/docs/stable/extensions/httpfs/s3api";,
+ "how-to-guide":
["/docs/apache-airflow-providers-amazon/operators/duckdb.rst"],
+ "tags": ["aws"],
+ },
{
"integration-name": "Amazon Bedrock",
"external-doc-url": "https://aws.amazon.com/bedrock/";,
@@ -726,6 +732,10 @@ def get_provider_info():
"airflow.providers.amazon.aws.hooks.athena_sql",
],
},
+ {
+ "integration-name": "DuckDB on AWS",
+ "python-modules":
["airflow.providers.amazon.aws.hooks.duckdb"],
+ },
{
"integration-name": "Amazon Bedrock",
"python-modules":
["airflow.providers.amazon.aws.hooks.bedrock"],
@@ -1229,6 +1239,19 @@ def get_provider_info():
},
},
},
+ {
+ "hook-class-name":
"airflow.providers.amazon.aws.hooks.duckdb.AwsDuckDBHook",
+ "hook-name": "DuckDB on AWS",
+ "connection-type": "duckdb_aws",
+ "ui-field-behaviour": {
+ "hidden-fields": ["login", "password", "port", "schema"],
+ "relabeling": {"host": "Database path"},
+ "placeholders": {
+ "host": "/tmp/analytics.duckdb (leave empty for an
in-memory database)",
+ "extra": '{\n "credential_strategy":
"credential_chain",\n "region_name": "us-east-1",\n "memory_limit": "2GB",\n
"threads": 4\n}\n',
+ },
+ },
+ },
{
"hook-class-name":
"airflow.providers.amazon.aws.hooks.emr.EmrHook",
"hook-name": "Amazon Elastic MapReduce",
diff --git a/providers/amazon/tests/system/amazon/aws/example_duckdb.py
b/providers/amazon/tests/system/amazon/aws/example_duckdb.py
new file mode 100644
index 00000000000..81197a566c7
--- /dev/null
+++ b/providers/amazon/tests/system/amazon/aws/example_duckdb.py
@@ -0,0 +1,158 @@
+# 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 json
+import os
+from datetime import datetime
+
+import boto3
+
+from airflow.providers.amazon.aws.hooks.duckdb import AwsDuckDBHook
+from airflow.providers.amazon.aws.operators.s3 import (
+ S3CreateBucketOperator,
+ S3CreateObjectOperator,
+ S3DeleteBucketOperator,
+)
+from airflow.providers.duckdb.operators.duckdb import
DuckDBExecuteQueryOperator
+
+from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS
+
+if AIRFLOW_V_3_0_PLUS:
+ from airflow.sdk import DAG, chain, task
+else:
+ # Airflow 2 path
+ from airflow.decorators import task # type: ignore[attr-defined,no-redef]
+ from airflow.models.baseoperator import chain # type:
ignore[attr-defined,no-redef]
+ from airflow.models.dag import DAG # type:
ignore[attr-defined,no-redef,assignment]
+
+try:
+ from airflow.sdk import TriggerRule
+except ImportError:
+ # Compatibility for Airflow < 3.1
+ from airflow.utils.trigger_rule import TriggerRule # type:
ignore[no-redef,attr-defined]
+
+from system.amazon.aws.utils import SystemTestContextBuilder
+
+sys_test_context_task = SystemTestContextBuilder().build()
+
+DAG_ID = "example_duckdb"
+
+SAMPLE_DATA = """category,price,quantity
+widgets,9.99,10
+widgets,19.99,3
+gadgets,4.50,20
+"""
+SAMPLE_FILENAME = "sales.csv"
+
+# The operator picks its hook from the connection type, so pointing it at a
``duckdb_aws`` connection
+# is the whole of the AWS wiring. Defined here as an env var so every task
process sees it, and in
+# JSON rather than URI form because a URI scheme cannot contain '_' (RFC 3986).
+DUCKDB_AWS_CONN_ID = AwsDuckDBHook.default_conn_name
+os.environ.setdefault(
+ f"AIRFLOW_CONN_{DUCKDB_AWS_CONN_ID.upper()}",
+ json.dumps({"conn_type": AwsDuckDBHook.conn_type}),
+)
+
+
+@task
+def await_bucket(bucket_name):
+ # Avoid a race condition after creating the S3 Bucket.
+ waiter = boto3.client("s3").get_waiter("bucket_exists")
+ waiter.wait(Bucket=bucket_name)
+
+
+@task
+def verify_summary(bucket_name):
+ """Read the summary back through the hook to prove the round trip to S3
worked."""
+ hook = AwsDuckDBHook()
+ with hook.get_conn() as conn:
+ rows = conn.execute(
+ f"SELECT category, revenue FROM
read_parquet('s3://{bucket_name}/summary/revenue.parquet')"
+ " ORDER BY revenue DESC"
+ ).fetchall()
+
+ if [row[0] for row in rows] != ["widgets", "gadgets"]:
+ raise ValueError(f"Unexpected summary contents: {rows}")
+
+
+with DAG(
+ dag_id=DAG_ID,
+ schedule="@once",
+ start_date=datetime(2021, 1, 1),
+ catchup=False,
+) as dag:
+ test_context = sys_test_context_task()
+ env_id = test_context["ENV_ID"]
+
+ s3_bucket = f"{env_id}-duckdb-bucket"
+
+ create_s3_bucket = S3CreateBucketOperator(task_id="create_s3_bucket",
bucket_name=s3_bucket)
+
+ upload_sample_data = S3CreateObjectOperator(
+ task_id="upload_sample_data",
+ s3_bucket=s3_bucket,
+ s3_key=f"sales/{SAMPLE_FILENAME}",
+ data=SAMPLE_DATA,
+ replace=True,
+ )
+
+ # [START howto_operator_aws_duckdb]
+ summarize_sales = DuckDBExecuteQueryOperator(
+ task_id="summarize_sales",
+ conn_id=DUCKDB_AWS_CONN_ID,
+ sql=f"""
+ COPY (
+ SELECT category, SUM(price * quantity) AS revenue
+ FROM read_csv('s3://{s3_bucket}/sales/*.csv')
+ GROUP BY category
+ ORDER BY revenue DESC
+ ) TO 's3://{s3_bucket}/summary/revenue.parquet' (FORMAT PARQUET)
+ """,
+ )
+ # [END howto_operator_aws_duckdb]
+
+ delete_s3_bucket = S3DeleteBucketOperator(
+ task_id="delete_s3_bucket",
+ bucket_name=s3_bucket,
+ force_delete=True,
+ trigger_rule=TriggerRule.ALL_DONE,
+ )
+
+ chain(
+ # TEST SETUP
+ test_context,
+ create_s3_bucket,
+ await_bucket(s3_bucket),
+ upload_sample_data,
+ # TEST BODY
+ summarize_sales,
+ verify_summary(s3_bucket),
+ # TEST TEARDOWN
+ delete_s3_bucket,
+ )
+
+ from tests_common.test_utils.watcher import watcher
+
+ # This test needs watcher in order to properly mark success/failure
+ # when "tearDown" task with trigger rule is part of the DAG
+ list(dag.tasks) >> watcher()
+
+from tests_common.test_utils.system_tests import get_test_run # noqa: E402
+
+# Needed to run the example DAG with pytest (see:
contributing-docs/testing/system_tests.rst)
+test_run = get_test_run(dag)
diff --git a/providers/amazon/tests/unit/amazon/aws/hooks/test_duckdb.py
b/providers/amazon/tests/unit/amazon/aws/hooks/test_duckdb.py
new file mode 100644
index 00000000000..ce72e4313fe
--- /dev/null
+++ b/providers/amazon/tests/unit/amazon/aws/hooks/test_duckdb.py
@@ -0,0 +1,307 @@
+# 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 json
+from unittest import mock
+
+import duckdb
+import pytest
+
+from airflow.models.connection import Connection
+from airflow.providers.amazon.aws.hooks.duckdb import AwsDuckDBHook
+from airflow.providers.duckdb.hooks.duckdb import DuckDBHook
+from airflow.providers.duckdb.operators.duckdb import
DuckDBExecuteQueryOperator
+
+GET_CONNECTION =
"airflow.providers.amazon.aws.hooks.duckdb.AwsDuckDBHook.get_connection"
+AWS_BASE_HOOK = "airflow.providers.amazon.aws.hooks.duckdb.AwsBaseHook"
+
+try:
+ from airflow.sdk.exceptions import AirflowNotFoundException
+except ImportError:
+ from airflow.exceptions import AirflowNotFoundException # type:
ignore[no-redef]
+
+
[email protected]
+def aws_base_hook():
+ """Patch AwsBaseHook so no AWS calls or connections are needed."""
+ with mock.patch(AWS_BASE_HOOK) as patched:
+ instance = patched.return_value
+ instance.region_name = "us-west-2"
+ instance.conn_config.endpoint_url = None
+ credentials = instance.get_credentials.return_value
+ credentials.access_key = "AKIAEXAMPLE"
+ credentials.secret_key = "secret"
+ credentials.token = None
+ yield instance
+
+
[email protected]
+def no_duckdb_connection():
+ with mock.patch(GET_CONNECTION,
side_effect=AirflowNotFoundException("nope")) as patched:
+ yield patched
+
+
+def secret_statement(hook) -> str:
+ """Run configure_secrets against a fake connection and return the
statement it issued."""
+ conn = mock.MagicMock()
+ hook.configure_secrets(conn)
+ if not conn.execute.call_args_list:
+ return ""
+ return conn.execute.call_args_list[-1].args[0]
+
+
+class TestAwsDuckDBHookConnectionType:
+ """The connection type is the integration point: it is how the generic
operator reaches this hook."""
+
+ def test_declares_its_own_connection_type(self):
+ assert AwsDuckDBHook.conn_type == "duckdb_aws"
+ assert AwsDuckDBHook.conn_type != DuckDBHook.conn_type
+ assert AwsDuckDBHook.default_conn_name == "duckdb_aws_default"
+
+ def test_reuses_the_duckdb_connection_id_attribute(self):
+ """``Connection.get_hook`` passes the conn id under this name, so it
has to keep matching."""
+ assert AwsDuckDBHook.conn_name_attr == DuckDBHook.conn_name_attr
+
+ def test_the_generic_operator_resolves_this_hook_from_the_connection(self):
+ """
+ The whole point of the connection type: no AWS-specific operator.
+
+ A Dag author swaps from local DuckDB to DuckDB on S3 by pointing the
same operator at a
+ ``duckdb_aws`` connection.
+ """
+ connection = Connection(conn_id="aws_duck",
conn_type=AwsDuckDBHook.conn_type)
+ operator = DuckDBExecuteQueryOperator(task_id="t", sql="SELECT 1",
conn_id="aws_duck")
+ with (
+ mock.patch(GET_CONNECTION, return_value=connection),
+ mock.patch(
+
"airflow.providers.duckdb.hooks.duckdb.DuckDBHook.get_connection",
return_value=connection
+ ),
+ ):
+ hook = operator.get_db_hook()
+
+ assert isinstance(hook, AwsDuckDBHook)
+ assert hook.get_conn_id() == "aws_duck"
+
+ def test_ui_hides_the_fields_it_does_not_use(self):
+ behaviour = AwsDuckDBHook.get_ui_field_behaviour()
+ assert set(behaviour["hidden_fields"]) == {"login", "password",
"port", "schema"}
+ assert behaviour["relabeling"] == {"host": "Database path"}
+
+
[email protected]("no_duckdb_connection")
+class TestAwsDuckDBHookExtensions:
+ def test_httpfs_and_aws_extensions_are_required(self, aws_base_hook):
+ assert AwsDuckDBHook().get_extensions() == ["httpfs", "aws"]
+
+ def test_required_extensions_are_kept_when_more_are_requested(self,
aws_base_hook):
+ assert AwsDuckDBHook(extensions=["iceberg"]).get_extensions() ==
["httpfs", "aws", "iceberg"]
+
+ def test_extension_downloads_are_enabled_for_aws(self, aws_base_hook):
+ """
+ The AWS hook installs its required extensions, diverging from the
generic hook's default.
+
+ Reaching for this hook implies an AWS environment where httpfs and aws
should just work.
+ """
+ config = AwsDuckDBHook().get_connect_config()
+ assert config["autoinstall_known_extensions"] is True
+ assert config["autoload_known_extensions"] is True
+
+ def test_extension_downloads_can_still_be_forbidden(self, aws_base_hook):
+ config =
AwsDuckDBHook(autoinstall_extensions=False).get_connect_config()
+ assert config["autoinstall_known_extensions"] is False
+
+
[email protected]("no_duckdb_connection")
+class TestAwsDuckDBHookCredentials:
+ def test_credential_chain_is_the_default(self, aws_base_hook):
+ statement = secret_statement(AwsDuckDBHook())
+ assert "PROVIDER credential_chain" in statement
+ assert "KEY_ID" not in statement
+
+ def test_credential_chain_does_not_resolve_credentials_itself(self,
aws_base_hook):
+ """The point of credential_chain: nothing secret is materialised into
the SQL statement."""
+ secret_statement(AwsDuckDBHook())
+ aws_base_hook.get_credentials.assert_not_called()
+
+ def test_explicit_chain_is_passed_through(self, aws_base_hook):
+ statement =
secret_statement(AwsDuckDBHook(credential_chain="env;instance"))
+ assert "CHAIN 'env;instance'" in statement
+
+ def test_config_strategy_writes_the_credentials_into_the_secret(self,
aws_base_hook):
+ statement =
secret_statement(AwsDuckDBHook(credential_strategy="config"))
+ assert "PROVIDER config" in statement
+ assert "KEY_ID 'AKIAEXAMPLE'" in statement
+ assert "SECRET 'secret'" in statement
+ assert "SESSION_TOKEN" not in statement
+
+ def test_config_strategy_includes_a_session_token_when_present(self,
aws_base_hook):
+ aws_base_hook.get_credentials.return_value.token = "token123"
+ statement =
secret_statement(AwsDuckDBHook(credential_strategy="config"))
+ assert "SESSION_TOKEN 'token123'" in statement
+
+ def test_quotes_in_credentials_cannot_break_out_of_the_literal(self,
aws_base_hook):
+ aws_base_hook.get_credentials.return_value.secret_key = "abc' OR
'1'='1"
+ statement =
secret_statement(AwsDuckDBHook(credential_strategy="config"))
+ assert "SECRET 'abc'' OR ''1''=''1'" in statement
+
+ def test_strategy_none_creates_no_secret(self, aws_base_hook):
+ assert secret_statement(AwsDuckDBHook(credential_strategy="none")) ==
""
+
+ def test_unknown_strategy_is_rejected(self, aws_base_hook):
+ with pytest.raises(ValueError, match="Unknown credential_strategy"):
+ secret_statement(AwsDuckDBHook(credential_strategy="magic"))
+
+ def test_secret_name_is_validated(self, aws_base_hook):
+ with pytest.raises(ValueError, match="Invalid DuckDB secret name"):
+ secret_statement(AwsDuckDBHook(secret_name="a); DROP TABLE t; --"))
+
+ def test_secret_name_is_used(self, aws_base_hook):
+ statement = secret_statement(AwsDuckDBHook(secret_name="my_secret"))
+ assert statement.startswith("CREATE OR REPLACE SECRET my_secret (")
+
+ @pytest.mark.parametrize(
+ ("strategy", "tolerated"),
+ [
+ pytest.param("credential_chain", True, id="chain-warns"),
+ pytest.param("config", False, id="config-raises"),
+ ],
+ )
+ def test_secret_creation_failure_depends_on_the_strategy(self,
aws_base_hook, strategy, tolerated):
+ """
+ DuckDB 1.4+ resolves the credential chain when the secret is created,
not when S3 is read.
+
+ A chain that finds nothing is not a misconfiguration, so it must not
stop the connection
+ opening. Explicit keys failing is a misconfiguration and must.
+ """
+ conn = mock.MagicMock()
+ conn.execute.side_effect = duckdb.Error("Secret Validation Failure")
+ hook = AwsDuckDBHook(credential_strategy=strategy)
+ with mock.patch.object(AwsDuckDBHook, "log") as mock_log:
+ if tolerated:
+ hook.configure_secrets(conn)
+ mock_log.warning.assert_called_once()
+ else:
+ with pytest.raises(duckdb.Error):
+ hook.configure_secrets(conn)
+ mock_log.warning.assert_not_called()
+
+
[email protected]("no_duckdb_connection")
+class TestAwsDuckDBHookRegion:
+ def test_region_comes_from_the_aws_connection_by_default(self,
aws_base_hook):
+ assert AwsDuckDBHook().get_region_name() == "us-west-2"
+ assert "REGION 'us-west-2'" in secret_statement(AwsDuckDBHook())
+
+ def test_region_argument_wins(self, aws_base_hook):
+ assert AwsDuckDBHook(region_name="eu-central-1").get_region_name() ==
"eu-central-1"
+
+ def test_no_region_clause_when_no_region_is_resolvable(self,
aws_base_hook):
+ aws_base_hook.region_name = None
+ assert "REGION" not in secret_statement(AwsDuckDBHook())
+
+
+class TestAwsDuckDBHookEndpoint:
+ @pytest.mark.parametrize(
+ ("endpoint_url", "expected_endpoint", "expects_plaintext"),
+ [
+ pytest.param("http://localhost:4566";, "localhost:4566", True,
id="http-with-port"),
+ pytest.param(
+ "https://s3.us-west-2.amazonaws.com";,
"s3.us-west-2.amazonaws.com", False, id="https"
+ ),
+ pytest.param("localhost:9000", "localhost:9000", False,
id="bare-host-port"),
+ ],
+ )
+ def test_endpoint_is_translated_into_secret_clauses(
+ self, aws_base_hook, no_duckdb_connection, endpoint_url,
expected_endpoint, expects_plaintext
+ ):
+ statement =
secret_statement(AwsDuckDBHook(s3_endpoint_url=endpoint_url))
+ assert f"ENDPOINT '{expected_endpoint}'" in statement
+ assert "URL_STYLE 'path'" in statement
+ assert ("USE_SSL false" in statement) is expects_plaintext
+
+ def test_endpoint_falls_back_to_the_aws_connection(self, aws_base_hook,
no_duckdb_connection):
+ aws_base_hook.conn_config.endpoint_url = "http://localstack:4566";
+ assert "ENDPOINT 'localstack:4566'" in
secret_statement(AwsDuckDBHook())
+
+ def test_no_endpoint_clauses_by_default(self, aws_base_hook,
no_duckdb_connection):
+ statement = secret_statement(AwsDuckDBHook())
+ assert "ENDPOINT" not in statement
+ assert "URL_STYLE" not in statement
+
+
+class TestAwsDuckDBHookConnectionOverrides:
+ """The DuckDB connection extra may override how AWS access is
configured."""
+
+ @staticmethod
+ def duckdb_connection(**extra) -> Connection:
+ return Connection(conn_id="duckdb_default", conn_type="duckdb",
extra=json.dumps(extra))
+
+ def test_credential_strategy_can_be_set_on_the_connection(self,
aws_base_hook):
+ with mock.patch(GET_CONNECTION,
return_value=self.duckdb_connection(credential_strategy="none")):
+ assert secret_statement(AwsDuckDBHook()) == ""
+
+ def test_region_can_be_set_on_the_connection(self, aws_base_hook):
+ with mock.patch(GET_CONNECTION,
return_value=self.duckdb_connection(region_name="ap-south-1")):
+ assert AwsDuckDBHook().get_region_name() == "ap-south-1"
+
+ def test_endpoint_can_be_set_on_the_connection(self, aws_base_hook):
+ with mock.patch(
+ GET_CONNECTION,
return_value=self.duckdb_connection(endpoint_url="http://minio:9000";)
+ ):
+ assert "ENDPOINT 'minio:9000'" in secret_statement(AwsDuckDBHook())
+
+ def test_secret_name_can_be_set_on_the_connection(self, aws_base_hook):
+ with mock.patch(GET_CONNECTION,
return_value=self.duckdb_connection(secret_name="from_extra")):
+ assert secret_statement(AwsDuckDBHook()).startswith("CREATE OR
REPLACE SECRET from_extra (")
+
+ def test_credential_chain_can_be_set_on_the_connection(self,
aws_base_hook):
+ with mock.patch(GET_CONNECTION,
return_value=self.duckdb_connection(credential_chain="env;instance")):
+ assert "CHAIN 'env;instance'" in secret_statement(AwsDuckDBHook())
+
+ def test_aws_conn_id_defaults_to_aws_default(self, aws_base_hook,
no_duckdb_connection):
+ assert AwsDuckDBHook().aws_conn_id == "aws_default"
+
+ def test_aws_conn_id_can_be_set_on_the_connection(self, aws_base_hook):
+ with mock.patch(GET_CONNECTION,
return_value=self.duckdb_connection(aws_conn_id="from_extra")):
+ assert AwsDuckDBHook().aws_conn_id == "from_extra"
+
+ def
test_aws_conn_id_null_on_the_connection_selects_the_ambient_environment(self,
aws_base_hook):
+ """
+ ``None`` is a valid choice rather than the absence of one, so a null
in the extra has to stick.
+
+ Resolving it like the other parameters would read it as "not set" and
quietly fall back to
+ ``aws_default``, which is the opposite of what was asked for.
+ """
+ with mock.patch(GET_CONNECTION,
return_value=self.duckdb_connection(aws_conn_id=None)):
+ assert AwsDuckDBHook().aws_conn_id is None
+
+ @pytest.mark.parametrize("explicit", ["explicit_conn", None],
ids=["a-name", "none"])
+ def test_explicit_aws_conn_id_wins_over_the_connection(self,
aws_base_hook, explicit):
+ with mock.patch(GET_CONNECTION,
return_value=self.duckdb_connection(aws_conn_id="from_extra")):
+ assert AwsDuckDBHook(aws_conn_id=explicit).aws_conn_id == explicit
+
+ def test_s3_endpoint_url_can_be_set_on_the_connection(self, aws_base_hook):
+ with mock.patch(
+ GET_CONNECTION,
return_value=self.duckdb_connection(s3_endpoint_url="http://minio:9000";)
+ ):
+ assert "ENDPOINT 'minio:9000'" in secret_statement(AwsDuckDBHook())
+
+ def test_explicit_argument_wins_over_the_connection_extra(self,
aws_base_hook):
+ with mock.patch(GET_CONNECTION,
return_value=self.duckdb_connection(secret_name="from_extra")):
+ statement = secret_statement(AwsDuckDBHook(secret_name="explicit"))
+ assert statement.startswith("CREATE OR REPLACE SECRET explicit (")
diff --git a/providers/duckdb/docs/operators.rst
b/providers/duckdb/docs/operators.rst
index fb5802e52bb..bac82a96764 100644
--- a/providers/duckdb/docs/operators.rst
+++ b/providers/duckdb/docs/operators.rst
@@ -21,6 +21,12 @@ DuckDB Operators
`DuckDB <https://duckdb.org/>`__ is an in-process analytical database. Queries
run inside the
Airflow task process, so there is no cluster to provision and no data to move
before querying it.
+.. seealso::
+
+ Other providers can ship a DuckDB connection type that adds their own
behaviour, which this
+ operator resolves from the connection. For Amazon S3, see
+ :doc:`DuckDB on AWS <apache-airflow-providers-amazon:operators/duckdb>`.
+
.. _howto/operator:DuckDBExecuteQueryOperator:
Run a query
diff --git a/providers/duckdb/src/airflow/providers/duckdb/hooks/duckdb.py
b/providers/duckdb/src/airflow/providers/duckdb/hooks/duckdb.py
index 63cc4c16195..8d57e2253f1 100644
--- a/providers/duckdb/src/airflow/providers/duckdb/hooks/duckdb.py
+++ b/providers/duckdb/src/airflow/providers/duckdb/hooks/duckdb.py
@@ -107,7 +107,7 @@ class DuckDBHook(DbApiHook):
def __init__(
self,
*args,
- duckdb_conn_id: str = default_conn_name,
+ duckdb_conn_id: str | None = None,
database: str | None = None,
extensions: Sequence[str] | None = None,
extension_directory: str | None = None,
@@ -121,7 +121,9 @@ class DuckDBHook(DbApiHook):
settings: dict[str, Any] | None = None,
**kwargs,
) -> None:
- kwargs[self.conn_name_attr] = duckdb_conn_id
+ # Resolved here rather than as a parameter default, which would bind
this class's value and
+ # so hand a subclass declaring its own default_conn_name the wrong
connection id.
+ kwargs[self.conn_name_attr] = duckdb_conn_id or self.default_conn_name
super().__init__(*args, **kwargs)
self.database = database
# Stored unresolved: ``None`` means "not set explicitly", so the
connection extra may supply
diff --git a/providers/duckdb/src/airflow/providers/duckdb/operators/duckdb.py
b/providers/duckdb/src/airflow/providers/duckdb/operators/duckdb.py
index 1e29cb114cb..ce3fc9ca695 100644
--- a/providers/duckdb/src/airflow/providers/duckdb/operators/duckdb.py
+++ b/providers/duckdb/src/airflow/providers/duckdb/operators/duckdb.py
@@ -20,6 +20,7 @@ from typing import ClassVar
from airflow.providers.common.sql.operators.sql import SQLExecuteQueryOperator
from airflow.providers.duckdb.hooks.duckdb import DuckDBHook
+from airflow.providers.duckdb.version_compat import AirflowNotFoundException
class DuckDBExecuteQueryOperator(SQLExecuteQueryOperator):
@@ -31,14 +32,18 @@ class DuckDBExecuteQueryOperator(SQLExecuteQueryOperator):
``duckdb_default`` connection configured it runs against an in-memory
database, which is the common
case for a task that reads its input, transforms it and writes its output
back out.
+ When a connection does exist, its type selects the hook. A provider can
therefore ship a DuckDB
+ connection type that adds its own behaviour (the Amazon provider's
``duckdb_aws`` brokers AWS
+ credentials, for example) and a Dag moves between them by changing
connections.
+
.. seealso::
For more information on how to use this operator, take a look at the
guide:
:ref:`howto/operator:DuckDBExecuteQueryOperator`
:param sql: the SQL statement, list of statements, or ``.sql`` template
file to run.
- :param conn_id: DuckDB :ref:`connection <howto/connection:duckdb>` to use.
The default connection
- need not exist, in which case an in-memory database is used. Any other
id must exist, and
- raises if it does not.
+ :param conn_id: DuckDB :ref:`connection <howto/connection:duckdb>` to use,
or a connection of
+ another DuckDB type such as ``duckdb_aws``. The default connection
need not exist, in which
+ case an in-memory database is used. Any other id must exist, and
raises if it does not.
:param database: DuckDB database to open, for example
``/tmp/analytics.duckdb``. Overrides both the
connection and ``hook_params["database"]``.
:param hook_params: extra keyword arguments for
@@ -56,17 +61,40 @@ class DuckDBExecuteQueryOperator(SQLExecuteQueryOperator):
def get_db_hook(self) -> DuckDBHook:
"""
- Build the hook directly instead of resolving it through the connection.
+ Resolve the hook from the connection, and build one directly when
there is no connection.
+
+ ``BaseSQLOperator`` resolves the hook through
``BaseHook.get_connection(conn_id).get_hook()``,
+ which raises when the connection is absent. DuckDB needs no connection
to be useful, so a
+ missing one is handed to ``hook_class`` to interpret rather than being
an error here.
- ``BaseSQLOperator`` resolves the hook class via
``BaseHook.get_connection(conn_id).get_hook()``,
- which raises when the connection is absent. DuckDB needs no connection
to be useful, so the
- hook is constructed here and left to decide what the missing
connection means.
+ When the connection does exist its type chooses the hook. That is what
lets another provider
+ add cloud-specific behaviour, such as brokering credentials, without a
Dag having to swap
+ operator: pointing this operator at a connection of that type is
enough.
"""
hook_params = dict(self.hook_params)
# ``BaseSQLOperator`` applies ``database`` in ``_hook``, which this
override bypasses, and it
# would apply it as ``hook.schema`` — meaningless for DuckDB, where a
database is a file path.
if self.database:
hook_params["database"] = self.database
- return self.hook_class(
- duckdb_conn_id=self.conn_id or self.hook_class.default_conn_name,
**hook_params
- )
+
+ conn_id = self.conn_id or self.hook_class.default_conn_name
+ try:
+ connection = self.hook_class.get_connection(conn_id)
+ except AirflowNotFoundException:
+ # The hook re-raises this exception unless conn_id is its own
default_conn_name, the
+ # one id allowed to be absent so that DuckDB works with no
configuration (default local
+ # DB). Deciding that here as well would put the same policy in two
places, and the hook
+ # needs it anyway because it is can also be used directly.
+ return self.hook_class(duckdb_conn_id=conn_id, **hook_params)
+
+ # Deliberately not BaseSQLOperator.get_hook: it folds every connection
extra into the hook's
+ # constructor arguments, which would turn an unrecognised extra into a
TypeError. DuckDBHook
+ # reads extras itself, so the connection only has to pick the class.
+ hook = connection.get_hook(hook_params=hook_params)
+ if not isinstance(hook, DuckDBHook):
+ raise ValueError(
+ f"Connection {conn_id!r} has type {connection.conn_type!r},
which resolves to "
+ f"{type(hook).__name__}. {type(self).__name__} needs a DuckDB
connection type, so "
+ f"use a connection whose hook derives from DuckDBHook."
+ )
+ return hook
diff --git a/providers/duckdb/tests/unit/duckdb/operators/test_duckdb.py
b/providers/duckdb/tests/unit/duckdb/operators/test_duckdb.py
index 9b8213efee3..3071faeeee9 100644
--- a/providers/duckdb/tests/unit/duckdb/operators/test_duckdb.py
+++ b/providers/duckdb/tests/unit/duckdb/operators/test_duckdb.py
@@ -38,8 +38,8 @@ class TestDuckDBExecuteQueryOperator:
assert isinstance(hook, DuckDBHook)
assert hook.get_conn_id() == DuckDBHook.default_conn_name
- def test_get_db_hook_bypasses_the_connection_lookup(self):
- """The whole point of the override: no Airflow connection row is
required."""
+ def test_a_missing_connection_falls_back_to_a_plain_hook(self):
+ """No Airflow connection row is required: DuckDB is useful with no
configuration at all."""
operator = DuckDBExecuteQueryOperator(
task_id="t", sql="SELECT 1", conn_id=DuckDBHook.default_conn_name
)
@@ -47,6 +47,35 @@ class TestDuckDBExecuteQueryOperator:
hook = operator.get_db_hook()
assert hook.get_database() == ":memory:"
+ def test_connection_type_selects_the_hook(self):
+ """
+ The connection's type chooses the hook class.
+
+ This is what lets another provider ship a DuckDB connection type with
extra behaviour, so a
+ Dag moves between environments by changing connection rather than by
changing operator.
+ """
+
+ class _CloudDuckDBHook(DuckDBHook):
+ pass
+
+ connection = mock.Mock(conn_type="duckdb_cloud")
+ connection.get_hook.return_value = _CloudDuckDBHook()
+ operator = DuckDBExecuteQueryOperator(task_id="t", sql="SELECT 1",
conn_id="cloud")
+ with mock.patch(GET_CONNECTION, return_value=connection):
+ hook = operator.get_db_hook()
+
+ assert isinstance(hook, _CloudDuckDBHook)
+ connection.get_hook.assert_called_once_with(hook_params={})
+
+ def test_a_connection_that_is_not_duckdb_raises(self):
+ """Resolving by connection type means a wrong type would otherwise
return a foreign hook."""
+ connection = mock.Mock(conn_type="postgres")
+ connection.get_hook.return_value = object()
+ operator = DuckDBExecuteQueryOperator(task_id="t", sql="SELECT 1",
conn_id="pg")
+ with mock.patch(GET_CONNECTION, return_value=connection):
+ with pytest.raises(ValueError, match="needs a DuckDB connection
type"):
+ operator.get_db_hook()
+
def test_hook_params_are_forwarded_to_the_hook(self):
operator = DuckDBExecuteQueryOperator(
task_id="t",
diff --git a/uv.lock b/uv.lock
index 38be4767eb9..ecf33af6816 100644
--- a/uv.lock
+++ b/uv.lock
@@ -3162,6 +3162,9 @@ cncf-kubernetes = [
common-messaging = [
{ name = "apache-airflow-providers-common-messaging" },
]
+duckdb = [
+ { name = "apache-airflow-providers-duckdb" },
+]
exasol = [
{ name = "apache-airflow-providers-exasol" },
]
@@ -3224,6 +3227,7 @@ dev = [
{ name = "apache-airflow-providers-common-compat" },
{ name = "apache-airflow-providers-common-messaging" },
{ name = "apache-airflow-providers-common-sql", extra = ["pandas",
"polars"] },
+ { name = "apache-airflow-providers-duckdb" },
{ name = "apache-airflow-providers-exasol" },
{ name = "apache-airflow-providers-ftp" },
{ name = "apache-airflow-providers-google" },
@@ -3261,6 +3265,7 @@ requires-dist = [
{ name = "apache-airflow-providers-common-compat", editable =
"providers/common/compat" },
{ name = "apache-airflow-providers-common-messaging", marker = "extra ==
'common-messaging'", editable = "providers/common/messaging" },
{ name = "apache-airflow-providers-common-sql", editable =
"providers/common/sql" },
+ { name = "apache-airflow-providers-duckdb", marker = "extra == 'duckdb'",
editable = "providers/duckdb" },
{ name = "apache-airflow-providers-exasol", marker = "extra == 'exasol'",
editable = "providers/exasol" },
{ name = "apache-airflow-providers-fab", marker = "extra == 'fab'",
editable = "providers/fab" },
{ name = "apache-airflow-providers-ftp", marker = "extra == 'ftp'",
editable = "providers/ftp" },
@@ -3296,7 +3301,7 @@ requires-dist = [
{ name = "watchtower", specifier = ">=3.3.1,<4" },
{ name = "xmlsec", marker = "python_full_version < '3.13' and extra ==
'python3-saml'", specifier = ">=1.3.14" },
]
-provides-extras = ["aiobotocore", "cncf-kubernetes", "s3fs", "msk",
"python3-saml", "apache-hive", "exasol", "fab", "ftp", "google", "imap",
"microsoft-azure", "mongo", "pandas", "openlineage", "salesforce", "ssh",
"standard", "common-messaging", "sqlalchemy"]
+provides-extras = ["aiobotocore", "cncf-kubernetes", "s3fs", "msk",
"python3-saml", "apache-hive", "exasol", "fab", "ftp", "google", "imap",
"microsoft-azure", "mongo", "pandas", "openlineage", "salesforce", "ssh",
"standard", "common-messaging", "duckdb", "sqlalchemy"]
[package.metadata.requires-dev]
dev = [
@@ -3310,6 +3315,7 @@ dev = [
{ name = "apache-airflow-providers-common-messaging", editable =
"providers/common/messaging" },
{ name = "apache-airflow-providers-common-sql", editable =
"providers/common/sql" },
{ name = "apache-airflow-providers-common-sql", extras = ["pandas",
"polars"], editable = "providers/common/sql" },
+ { name = "apache-airflow-providers-duckdb", editable = "providers/duckdb"
},
{ name = "apache-airflow-providers-exasol", editable = "providers/exasol"
},
{ name = "apache-airflow-providers-ftp", editable = "providers/ftp" },
{ name = "apache-airflow-providers-google", editable = "providers/google"
},
