Allow AthenaOperator queries without a database argument

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
                 ],
             )

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):