This is an automated email from the ASF dual-hosted git repository.
eladkal 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 26c2c123a3c Allow `AthenaOperator` queries without a `database`
argument (#69846)
26c2c123a3c is described below
commit 26c2c123a3cb3ec191a8e1b7267fc56670c1c65c
Author: T. Horie <[email protected]>
AuthorDate: Wed Jul 15 05:33:44 2026 +0900
Allow `AthenaOperator` queries without a `database` argument (#69846)
* Allow `AthenaOperator` queries without a `database` argument
Athena does not require the `Database` field in `QueryExecutionContext`.
Queries with fully qualified table names can therefore run without a
default database.
https://docs.aws.amazon.com/athena/latest/APIReference/API_StartQueryExecution.html#API_StartQueryExecution_RequestSyntax
Requiring the argument forces users to provide a default database even
when the query does not need one.
* Use `Database` in `query_execution_context` for OpenLineage datasets
When the `database` argument of `AthenaOperator` is `None`, the
`Database` field in `QueryExecutionContext` can still provide the
default database for unqualified table names in a query.
* Clarify `AthenaOperator` database fallback behaviour in its docstring
Suggested in
https://github.com/apache/airflow/pull/69846#issuecomment-4966077688
It is better to explain to users who omit `database` or set it to `None`
that `Database` in `query_execution_context` remains effective.
---
.../providers/amazon/aws/operators/athena.py | 17 +++++++----
.../tests/unit/amazon/aws/operators/test_athena.py | 35 ++++++++++++++++++++++
2 files changed, 47 insertions(+), 5 deletions(-)
diff --git
a/providers/amazon/src/airflow/providers/amazon/aws/operators/athena.py
b/providers/amazon/src/airflow/providers/amazon/aws/operators/athena.py
index 3b8e4d4eb4b..592dc78ece1 100644
--- a/providers/amazon/src/airflow/providers/amazon/aws/operators/athena.py
+++ b/providers/amazon/src/airflow/providers/amazon/aws/operators/athena.py
@@ -47,7 +47,11 @@ class AthenaOperator(AwsBaseOperator[AthenaHook]):
:ref:`howto/operator:AthenaOperator`
:param query: Trino/Presto query to be run on Amazon Athena. (templated)
- :param database: Database to select. (templated)
+ :param database: Default database for query execution. (templated)
+ This argument is optional when the query does not require a default
database,
+ such as when all referenced table names are fully qualified.
+ If omitted or set to ``None``, any ``Database`` value set in
+ the ``query_execution_context`` will be used instead.
:param catalog: Catalog to select. (templated)
:param output_location: s3 path to write the query results into.
(templated)
To run the query, you must specify the query results location using
one of the ways:
@@ -88,7 +92,7 @@ class AthenaOperator(AwsBaseOperator[AthenaHook]):
self,
*,
query: str,
- database: str,
+ database: str | None = None,
output_location: str | None = None,
client_request_token: str | None = None,
workgroup: str = "primary",
@@ -122,7 +126,8 @@ class AthenaOperator(AwsBaseOperator[AthenaHook]):
def execute(self, context: Context) -> str | None:
"""Run Trino/Presto Query on Amazon Athena."""
- self.query_execution_context["Database"] = self.database
+ if self.database:
+ self.query_execution_context["Database"] = self.database
self.query_execution_context["Catalog"] = self.catalog
if self.output_location:
self.result_configuration["OutputLocation"] = self.output_location
@@ -251,11 +256,13 @@ class AthenaOperator(AwsBaseOperator[AthenaHook]):
],
)
+ fallback_database = self.database or
self.query_execution_context.get("Database")
+
inputs: list[Dataset] = list(
filter(
None,
[
- self.get_openlineage_dataset(table.schema or
self.database, table.name)
+ self.get_openlineage_dataset(table.schema or
fallback_database, table.name)
for table in parse_result.in_tables
],
)
@@ -265,7 +272,7 @@ class AthenaOperator(AwsBaseOperator[AthenaHook]):
filter(
None,
[
- self.get_openlineage_dataset(table.schema or
self.database, table.name)
+ self.get_openlineage_dataset(table.schema or
fallback_database, table.name)
for table in parse_result.out_tables
],
)
diff --git a/providers/amazon/tests/unit/amazon/aws/operators/test_athena.py
b/providers/amazon/tests/unit/amazon/aws/operators/test_athena.py
index 4b0b91f0f3f..91fcd3ff848 100644
--- a/providers/amazon/tests/unit/amazon/aws/operators/test_athena.py
+++ b/providers/amazon/tests/unit/amazon/aws/operators/test_athena.py
@@ -139,6 +139,26 @@ class TestAthenaOperator:
)
assert mock_check_query_status.call_count == 1
+ @mock.patch.object(AthenaHook, "check_query_status",
side_effect=("SUCCEEDED",))
+ @mock.patch.object(AthenaHook, "run_query", return_value=ATHENA_QUERY_ID)
+ @mock.patch.object(AthenaHook, "get_conn")
+ def test_hook_run_without_database(self, mock_conn, mock_run_query,
mock_check_query_status):
+ op_kwargs = self.default_op_kwargs.copy()
+ op_kwargs["task_id"] = "test_athena_operator_without_database"
+ op_kwargs.pop("database")
+ op = AthenaOperator(
+ **op_kwargs, output_location="s3://test_s3_bucket/",
aws_conn_id=None, dag=self.dag
+ )
+ op.execute({})
+ mock_run_query.assert_called_once_with(
+ MOCK_DATA["query"],
+ {"Catalog": MOCK_DATA["catalog"]},
+ result_configuration,
+ MOCK_DATA["client_request_token"],
+ MOCK_DATA["workgroup"],
+ )
+ assert mock_check_query_status.call_count == 1
+
@mock.patch.object(AthenaHook, "check_query_status",
side_effect=("SUCCEEDED",))
@mock.patch.object(AthenaHook, "run_query", return_value=ATHENA_QUERY_ID)
@mock.patch.object(AthenaHook, "get_conn")
@@ -321,6 +341,21 @@ class TestAthenaOperator:
)
assert operator.query_execution_id == query_execution_id
+ @mock.patch.object(AthenaOperator, "get_openlineage_dataset")
+ def test_openlineage_uses_database_from_query_execution_context(self,
mock_get_dataset):
+ op = AthenaOperator(
+ task_id="test_athena_openlineage",
+ query="INSERT INTO TEST_TABLE SELECT CUSTOMER_EMAIL FROM
DISCOUNTS",
+ database=None,
+ query_execution_context={"Database": "TEST_DATABASE"},
+ dag=self.dag,
+ )
+
+ op.get_openlineage_facets_on_complete(None)
+
+ mock_get_dataset.assert_any_call("TEST_DATABASE", "DISCOUNTS")
+ mock_get_dataset.assert_any_call("TEST_DATABASE", "TEST_TABLE")
+
@mock.patch.object(AthenaHook, "region_name",
new_callable=mock.PropertyMock)
@mock.patch.object(AthenaHook, "get_conn")
def test_operator_openlineage_data(self, mock_conn, mock_region_name):
