openlineage, aws: Add OpenLineage support for AthenaOperator.

airflow/providers/amazon/aws/hooks/athena.py

@@ -82,6 +82,7 @@ def __init__(
         else:
             self.sleep_time = 30  # previous default value
         self.log_query = log_query
+        self.__query_results: dict[str, Any] = {}
 
     def run_query(
         self,
@@ -120,7 +121,23 @@ def run_query(
         self.log.info("Query execution id: %s", query_execution_id)
         return query_execution_id
 
-    def check_query_status(self, query_execution_id: str) -> str | None:
+    def get_query_info(self, query_execution_id: str, use_cache: bool = False) -> dict:
+        """Get information about a single execution of a query.
+
+        .. seealso::
+            - :external+boto3:py:meth:`Athena.Client.get_query_execution`
+
+        :param query_execution_id: Id of submitted athena query
+        :param use_cache: If True, use execution information cache
+        """
+        if use_cache and query_execution_id in self.__query_results:
+            return self.__query_results[query_execution_id]
+        response = self.get_conn().get_query_execution(QueryExecutionId=query_execution_id)
+        if use_cache:
+            self.__query_results[query_execution_id] = response
+        return response
+
+    def check_query_status(self, query_execution_id: str, use_cache: bool = False) -> str | None:
         """Fetch the state of a submitted query.
 
         .. seealso::
@@ -130,7 +147,7 @@ def check_query_status(self, query_execution_id: str) -> str | None:
         :return: One of valid query states, or *None* if the response is
             malformed.
         """
-        response = self.get_conn().get_query_execution(QueryExecutionId=query_execution_id)
+        response = self.get_query_info(query_execution_id=query_execution_id, use_cache=use_cache)
         state = None
         try:
             state = response["QueryExecution"]["Status"]["State"]
@@ -143,7 +160,7 @@ def check_query_status(self, query_execution_id: str) -> str | None:
             # The error is being absorbed to implement retries.
             return state
 
-    def get_state_change_reason(self, query_execution_id: str) -> str | None:
+    def get_state_change_reason(self, query_execution_id: str, use_cache: bool = False) -> str | None:
         """
         Fetch the reason for a state change (e.g. error message). Returns None or reason string.
 
@@ -152,7 +169,7 @@ def get_state_change_reason(self, query_execution_id: str) -> str | None:
 
         :param query_execution_id: Id of submitted athena query
         """
-        response = self.get_conn().get_query_execution(QueryExecutionId=query_execution_id)
+        response = self.get_query_info(query_execution_id=query_execution_id, use_cache=use_cache)
         reason = None
         try:
             reason = response["QueryExecution"]["Status"]["StateChangeReason"]
@@ -277,7 +294,7 @@ def get_output_location(self, query_execution_id: str) -> str:
         """
         output_location = None
         if query_execution_id:
-            response = self.get_conn().get_query_execution(QueryExecutionId=query_execution_id)
+            response = self.get_query_info(query_execution_id=query_execution_id, use_cache=True)
 
             if response:
                 try:

airflow/providers/amazon/aws/operators/athena.py

@@ -18,6 +18,7 @@
 from __future__ import annotations
 
 from typing import TYPE_CHECKING, Any, Sequence
+from urllib.parse import urlparse
 
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException
@@ -27,6 +28,10 @@
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
 
 if TYPE_CHECKING:
+    from openlineage.client.facet import BaseFacet
+    from openlineage.client.run import Dataset
+
+    from airflow.providers.openlineage.extractors.base import OperatorLineage
     from airflow.utils.context import Context
 
 
@@ -160,6 +165,9 @@ def execute(self, context: Context) -> str | None:
                 f"query_execution_id is {self.query_execution_id}."
             )
 
+        # Save output location from API response for later use in OpenLineage.
+        self.output_location = self.hook.get_output_location(self.query_execution_id)
+
         return self.query_execution_id
 
     def execute_complete(self, context, event=None):
@@ -187,3 +195,110 @@ def on_kill(self) -> None:
                         "Polling Athena for query with id %s to reach final state", self.query_execution_id
                     )
                     self.hook.poll_query_status(self.query_execution_id, sleep_time=self.sleep_time)
+
+    def get_openlineage_facets_on_start(self) -> OperatorLineage:
+        """Retrieve OpenLineage data by parsing SQL queries and enriching them with Athena API.
+
+        In addition to CTAS query, query and calculation results are stored in S3 location.
+        For that reason additional output is attached with this location.
+        """
+        from openlineage.client.facet import ExtractionError, ExtractionErrorRunFacet, SqlJobFacet
+        from openlineage.client.run import Dataset
+
+        from airflow.providers.openlineage.extractors.base import OperatorLineage
+        from airflow.providers.openlineage.sqlparser import SQLParser
+
+        sql_parser = SQLParser(dialect="generic")
+
+        job_facets: dict[str, BaseFacet] = {"sql": SqlJobFacet(query=sql_parser.normalize_sql(self.query))}
+        parse_result = sql_parser.parse(sql=self.query)
+
+        if not parse_result:
+            return OperatorLineage(job_facets=job_facets)
+
+        run_facets: dict[str, BaseFacet] = {}
+        if parse_result.errors:
+            run_facets["extractionError"] = ExtractionErrorRunFacet(
+                totalTasks=len(self.query) if isinstance(self.query, list) else 1,
+                failedTasks=len(parse_result.errors),
+                errors=[
+                    ExtractionError(
+                        errorMessage=error.message,
+                        stackTrace=None,
+                        task=error.origin_statement,
+                        taskNumber=error.index,
+                    )
+                    for error in parse_result.errors
+                ],
+            )
+
+        inputs: list[Dataset] = list(
+            filter(
+                None,
+                [
+                    self.get_openlineage_dataset(table.schema or self.database, table.name)
+                    for table in parse_result.in_tables
+                ],
+            )
+        )
+
+        outputs: list[Dataset] = list(
+            filter(
+                None,
+                [
+                    self.get_openlineage_dataset(table.schema or self.database, table.name)
+                    for table in parse_result.out_tables
+                ],
+            )
+        )
+
+        if self.output_location:
+            parsed = urlparse(self.output_location)
+            outputs.append(Dataset(namespace=f"{parsed.scheme}://{parsed.netloc}", name=parsed.path))
+
+        return OperatorLineage(job_facets=job_facets, run_facets=run_facets, inputs=inputs, outputs=outputs)
+
+    def get_openlineage_dataset(self, database, table) -> Dataset | None:
+        from openlineage.client.facet import (
+            SchemaDatasetFacet,
+            SchemaField,
+            SymlinksDatasetFacet,
+            SymlinksDatasetFacetIdentifiers,
+        )
+        from openlineage.client.run import Dataset
+
+        client = self.hook.get_conn()
+        try:
+            table_metadata = client.get_table_metadata(
+                CatalogName=self.catalog, DatabaseName=database, TableName=table
+            )
+
+            # Dataset has also its' physical location which we can add in symlink facet.
+            s3_location = table_metadata["TableMetadata"]["Parameters"]["location"]
+            parsed_path = urlparse(s3_location)
+            facets: dict[str, BaseFacet] = {
+                "symlinks": SymlinksDatasetFacet(
+                    identifiers=[
+                        SymlinksDatasetFacetIdentifiers(
+                            namespace=f"{parsed_path.scheme}://{parsed_path.netloc}",
+                            name=str(parsed_path.path),
+                            type="TABLE",
+                        )
+                    ]
+                )
+            }
+            fields = [
+                SchemaField(name=column["Name"], type=column["Type"], description=column["Comment"])
+                for column in table_metadata["TableMetadata"]["Columns"]
+            ]
+            if fields:
+                facets["schema"] = SchemaDatasetFacet(fields=fields)
+            return Dataset(
+                namespace=f"awsathena://athena.{self.hook.region_name}.amazonaws.com",
+                name=".".join(filter(None, (self.catalog, database, table))),
+                facets=facets,
+            )
+
+        except Exception as e:
+            self.log.error("Cannot retrieve table metadata from Athena.Client. %s", e)
+            return None

tests/providers/amazon/aws/hooks/test_athena.py

@@ -43,6 +43,7 @@
     "QueryExecution": {
         "QueryExecutionId": MOCK_DATA["query_execution_id"],
         "ResultConfiguration": {"OutputLocation": "s3://test_bucket/test.csv"},
+        "Status": {"StateChangeReason": "Terminated by user."},
     }
 }
 
@@ -195,3 +196,21 @@ def test_hook_get_output_location(self, mock_conn):
         mock_conn.return_value.get_query_execution.return_value = MOCK_QUERY_EXECUTION_OUTPUT
         result = self.athena.get_output_location(query_execution_id=MOCK_DATA["query_execution_id"])
         assert result == "s3://test_bucket/test.csv"
+
+    @mock.patch.object(AthenaHook, "get_conn")
+    def test_hook_get_query_info_caching(self, mock_conn):
+        mock_conn.return_value.get_query_execution.return_value = MOCK_QUERY_EXECUTION_OUTPUT
+        self.athena.get_state_change_reason(query_execution_id=MOCK_DATA["query_execution_id"])
+        assert not self.athena._AthenaHook__query_results
+        # get_output_location uses cache
+        self.athena.get_output_location(query_execution_id=MOCK_DATA["query_execution_id"])
+        assert MOCK_DATA["query_execution_id"] in self.athena._AthenaHook__query_results
+        mock_conn.return_value.get_query_execution.assert_called_with(
+            QueryExecutionId=MOCK_DATA["query_execution_id"]
+        )
+        self.athena.get_state_change_reason(
+            query_execution_id=MOCK_DATA["query_execution_id"], use_cache=False
+        )
+        mock_conn.return_value.get_query_execution.assert_called_with(
+            QueryExecutionId=MOCK_DATA["query_execution_id"]
+        )

tests/providers/amazon/aws/operators/athena_metadata.json

@@ -0,0 +1,72 @@
+{
+    "DISCOUNTS": {
+      "TableMetadata": {
+        "Name": "DISCOUNTS",
+        "CreateTime": 1593559968.0,
+        "LastAccessTime": 0.0,
+        "TableType": "EXTERNAL_TABLE",
+        "Columns": [
+          {
+            "Name": "ID",
+            "Type": "int",
+            "Comment": "from deserializer"
+          },
+          {
+            "Name": "AMOUNT_OFF",
+            "Type": "int",
+            "Comment": "from deserializer"
+          },
+          {
+            "Name": "CUSTOMER_EMAIL",
+            "Type": "varchar",
+            "Comment": "from deserializer"
+          },
+          {
+            "Name": "STARTS_ON",
+            "Type": "timestamp",
+            "Comment": "from deserializer"
+          },
+          {
+            "Name": "ENDS_ON",
+            "Type": "timestamp",
+            "Comment": "from deserializer"
+          }
+        ],
+        "PartitionKeys": [],
+        "Parameters": {
+          "EXTERNAL": "TRUE",
+          "inputformat": "com.esri.json.hadoop.EnclosedJsonInputFormat",
+          "location": "s3://bucket/discount/data/path/",
+          "outputformat": "org.apache.hadoop.hive.ql.io.HiveIgnoreKeyTextOutputFormat",
+          "serde.param.serialization.format": "1",
+          "serde.serialization.lib": "com.esri.hadoop.hive.serde.JsonSerde",
+          "transient_lastDdlTime": "1593559968"
+        }
+      }
+    },
+    "TEST_TABLE": {
+      "TableMetadata": {
+        "Name": "TEST_TABLE",
+        "CreateTime": 1593559968.0,
+        "LastAccessTime": 0.0,
+        "TableType": "EXTERNAL_TABLE",
+        "Columns": [
+          {
+            "Name": "column",
+            "Type": "string",
+            "Comment": "from deserializer"
+          }
+        ],
+        "PartitionKeys": [],
+        "Parameters": {
+          "EXTERNAL": "TRUE",
+          "inputformat": "com.esri.json.hadoop.EnclosedJsonInputFormat",
+          "location": "s3://bucket/data/test_table/data/path",
+          "outputformat": "org.apache.hadoop.hive.ql.io.HiveIgnoreKeyTextOutputFormat",
+          "serde.param.serialization.format": "1",
+          "serde.serialization.lib": "com.esri.hadoop.hive.serde.JsonSerde",
+          "transient_lastDdlTime": "1593559968"
+        }
+      }
+    }
+  }