From 3b8d87039338276f0cbc930cfd5e4293273c056a Mon Sep 17 00:00:00 2001 From: te-horie <16241822+te-horie@users.noreply.github.com> Date: Sun, 5 Jul 2026 21:34:37 +0900 Subject: [PATCH 1/3] 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. --- .../providers/amazon/aws/operators/athena.py | 9 ++++++--- .../unit/amazon/aws/operators/test_athena.py | 20 +++++++++++++++++++ 2 files changed, 26 insertions(+), 3 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 3b8e4d4eb4b82..edc892108ec2f 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,9 @@ 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. :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 +90,7 @@ def __init__( 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 +124,8 @@ def _hook_parameters(self) -> dict[str, Any]: 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 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 4b0b91f0f3f8b..bc4809cb9c7ca 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 @@ def test_hook_run_override_catalog(self, mock_conn, mock_run_query, mock_check_q ) 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") From d7d2c25b8fc817778e6fff56ad7f5e8021f50724 Mon Sep 17 00:00:00 2001 From: te-horie <16241822+te-horie@users.noreply.github.com> Date: Sun, 5 Jul 2026 21:35:06 +0900 Subject: [PATCH 2/3] 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. --- .../providers/amazon/aws/operators/athena.py | 6 ++++-- .../unit/amazon/aws/operators/test_athena.py | 15 +++++++++++++++ 2 files changed, 19 insertions(+), 2 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 edc892108ec2f..9bdca49562406 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/operators/athena.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/operators/athena.py @@ -254,11 +254,13 @@ def get_openlineage_facets_on_complete(self, _) -> OperatorLineage: ], ) + 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 ], ) @@ -268,7 +270,7 @@ def get_openlineage_facets_on_complete(self, _) -> OperatorLineage: 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 bc4809cb9c7ca..91fcd3ff84899 100644 --- a/providers/amazon/tests/unit/amazon/aws/operators/test_athena.py +++ b/providers/amazon/tests/unit/amazon/aws/operators/test_athena.py @@ -341,6 +341,21 @@ def test_execute_complete_reassigns_query_execution_id_after_deferring(self): ) 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): From f72d787a25417f88c251a3144a7a5f98755b0fd6 Mon Sep 17 00:00:00 2001 From: te-horie <16241822+te-horie@users.noreply.github.com> Date: Tue, 14 Jul 2026 20:52:20 +0900 Subject: [PATCH 3/3] 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. --- .../amazon/src/airflow/providers/amazon/aws/operators/athena.py | 2 ++ 1 file changed, 2 insertions(+) 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 9bdca49562406..592dc78ece185 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/operators/athena.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/operators/athena.py @@ -50,6 +50,8 @@ class AthenaOperator(AwsBaseOperator[AthenaHook]): :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: