fix(ingestion/airflow-plugin): fix AthenaOperator extraction

The GenericSqlExtractor which is currently by the DataHub Airflow plugin
to extract lineage information does not properly support the
AthenaOperator and crashes with "AttributeError: 'AthenaOperator' object
has no attribute 'sql'". This patch introduces a AthenaOperatorExtractor
following the BigQueryInsertJobOperatorExtractor example to fix support
for the AthenaOperator.

metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/_extractors.py

@@ -50,7 +50,6 @@ def __init__(self):
             "BigQueryOperator",
             "BigQueryExecuteQueryOperator",
             # Athena also does something similar.
-            "AthenaOperator",
             "AWSAthenaOperator",
             # Additional types that OL doesn't support. This is only necessary because
             # on older versions of Airflow, these operators don't inherit from SQLExecuteQueryOperator.
@@ -59,6 +58,10 @@ def __init__(self):
         for operator in _sql_operator_overrides:
             self.task_to_extractor.extractors[operator] = GenericSqlExtractor
 
+        self.task_to_extractor.extractors[
+            "AthenaOperator"
+        ] = AthenaOperatorExtractor
+
         self.task_to_extractor.extractors[
             "BigQueryInsertJobOperator"
         ] = BigQueryInsertJobOperatorExtractor
@@ -275,6 +278,26 @@ def extract(self) -> Optional[TaskMetadata]:
             default_schema=None,
         )
 
+class AthenaOperatorExtractor(BaseExtractor):
+    def extract(self) -> Optional[TaskMetadata]:
+        from airflow.providers.amazon.aws.operators.athena import (
+            AthenaOperator,  # type: ignore
+        )
+
+        operator: "AthenaOperator" = self.operator
+        sql = operator.query
+        if not sql:
+            self.log.warning("No query found in AthenaOperator")
+            return None
+
+        return _parse_sql_into_task_metadata(
+            self,
+            sql,
+            platform="athena",
+            default_database=None,
+            default_schema=None,
+        )
+
 
 def _snowflake_default_schema(self: "SnowflakeExtractor") -> Optional[str]:
     if hasattr(self.operator, "schema") and self.operator.schema is not None: