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..592dc78ece185 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 @@ 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 +126,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 @@ -251,11 +256,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 ], ) @@ -265,7 +272,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 4b0b91f0f3f8b..91fcd3ff84899 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") @@ -321,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):