acryl-datahub-airflow-plugin 1.3.1.4__py3-none-any.whl → 1.3.1.5__py3-none-any.whl

datahub_airflow_plugin/example_dags/airflow3/snowflake_sample_dag.py

@@ -0,0 +1,89 @@
+"""Snowflake DataHub Ingest DAG
+
+This example demonstrates how to ingest metadata from Snowflake into DataHub
+from within an Airflow DAG. In contrast to the MySQL example, this DAG
+pulls the DB connection configuration from Airflow's connection store.
+
+This is the Airflow 3.0+ version.
+"""
+
+from datetime import datetime, timedelta
+
+from airflow import DAG
+from airflow.hooks.base_hook import BaseHook
+from airflow.operators.python import PythonVirtualenvOperator
+
+
+def ingest_from_snowflake(snowflake_credentials, datahub_gms_server):
+    from datahub.ingestion.run.pipeline import Pipeline
+
+    pipeline = Pipeline.create(
+        # This configuration is analogous to a recipe configuration.
+        {
+            "source": {
+                "type": "snowflake",
+                "config": {
+                    **snowflake_credentials,
+                    # Other Snowflake config can be added here.
+                    "profiling": {"enabled": False},
+                },
+            },
+            # Other ingestion features, like transformers, are also supported.
+            # "transformers": [
+            #     {
+            #         "type": "simple_add_dataset_ownership",
+            #         "config": {
+            #             "owner_urns": [
+            #                 "urn:li:corpuser:example",
+            #             ]
+            #         },
+            #     }
+            # ],
+            "sink": {
+                "type": "datahub-rest",
+                "config": {"server": datahub_gms_server},
+            },
+        }
+    )
+    pipeline.run()
+    pipeline.pretty_print_summary()
+    pipeline.raise_from_status()
+
+
+with DAG(
+    "datahub_snowflake_ingest",
+    default_args={
+        "owner": "airflow",
+    },
+    description="An example DAG which ingests metadata from Snowflake to DataHub",
+    start_date=datetime(2022, 1, 1),
+    schedule=timedelta(days=1),
+    catchup=False,
+) as dag:
+    # This example pulls credentials from Airflow's connection store.
+    # For this to work, you must have previously configured these connections in Airflow.
+    # See the Airflow docs for details: https://airflow.apache.org/docs/apache-airflow/stable/howto/connection.html
+    snowflake_conn = BaseHook.get_connection("snowflake_admin_default")
+    datahub_conn = BaseHook.get_connection("datahub_rest_default")
+
+    # While it is also possible to use the PythonOperator, we recommend using
+    # the PythonVirtualenvOperator to ensure that there are no dependency
+    # conflicts between DataHub and the rest of your Airflow environment.
+    ingest_task = PythonVirtualenvOperator(
+        task_id="ingest_from_snowflake",
+        requirements=[
+            "acryl-datahub[snowflake]",
+        ],
+        system_site_packages=False,
+        python_callable=ingest_from_snowflake,
+        op_kwargs={
+            "snowflake_credentials": {
+                "username": snowflake_conn.login,
+                "password": snowflake_conn.password,
+                "account_id": snowflake_conn.extra_dejson["account"],
+                "warehouse": snowflake_conn.extra_dejson.get("warehouse"),
+                "role": snowflake_conn.extra_dejson.get("role"),
+            },
+            "datahub_gms_server": datahub_conn.host,
+        },
+    )

datahub_airflow_plugin/example_dags/graph_usage_sample_dag.py

@@ -3,17 +3,25 @@
 from datetime import timedelta
 
 import pendulum
-from airflow.decorators import dag, task
+from airflow.decorators import (  # type: ignore[attr-defined] # Decorators not available in all Airflow versions
+    dag,
+    task,
+)
 
 from datahub.ingestion.graph.client import DataHubGraph, RemovedStatusFilter
+from datahub_airflow_plugin._airflow_version_specific import (
+    get_airflow_compatible_dag_kwargs,
+)
 from datahub_airflow_plugin.hooks.datahub import DatahubRestHook
 
-
-@dag(
-    schedule_interval=timedelta(days=1),
+dag_decorator_kwargs = get_airflow_compatible_dag_kwargs(
     start_date=pendulum.datetime(2021, 1, 1, tz="UTC"),
+    schedule_interval=timedelta(days=1),
     catchup=False,
 )
+
+
+@dag(**dag_decorator_kwargs)
 def datahub_graph_usage_sample_dag():
     @task()
     def use_the_graph():

datahub_airflow_plugin/hooks/datahub.py

@@ -1,7 +1,12 @@
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence, Tuple, Union
 
 from airflow.exceptions import AirflowException
-from airflow.hooks.base import BaseHook
+
+# BaseHook import - prefer new location in Airflow 3.x
+try:
+    from airflow.sdk.bases.hook import BaseHook
+except (ModuleNotFoundError, ImportError):
+    from airflow.hooks.base import BaseHook  # type: ignore
 
 from datahub.emitter.composite_emitter import CompositeEmitter
 from datahub.emitter.generic_emitter import Emitter
@@ -10,9 +15,13 @@ from datahub.metadata.com.linkedin.pegasus2avro.mxe import (
     MetadataChangeEvent,
     MetadataChangeProposal,
 )
+from datahub_airflow_plugin._airflow_version_specific import IS_AIRFLOW_3_OR_HIGHER
 
 if TYPE_CHECKING:
-    from airflow.models.connection import Connection
+    if IS_AIRFLOW_3_OR_HIGHER:
+        from airflow.sdk.definitions.connection import Connection
+    else:
+        from airflow.models.connection import Connection  # type: ignore[assignment]
 
     from datahub.emitter.kafka_emitter import DatahubKafkaEmitter
     from datahub.emitter.rest_emitter import DataHubRestEmitter

datahub_airflow_plugin/operators/datahub.py

@@ -1,4 +1,4 @@
-from typing import Any, List, Union
+from typing import TYPE_CHECKING, Any, List, Union
 
 from airflow.models import BaseOperator
 from avrogen.dict_wrapper import DictWrapper
@@ -11,6 +11,21 @@ from datahub_airflow_plugin.hooks.datahub import (
     DatahubRestHook,
 )
 
+if TYPE_CHECKING:
+    from jinja2 import Environment
+
+    # Import Context with version compatibility for type checking
+    # Import to different names to avoid redefinition errors, then assign to Context
+    Context: Any
+    try:
+        from airflow.utils.context import Context as _AirflowContext
+
+        Context = _AirflowContext
+    except ImportError:
+        from airflow.sdk.definitions.context import Context as _Airflow3Context
+
+        Context = _Airflow3Context  # type: ignore[no-redef]
+
 
 class DatahubBaseOperator(BaseOperator):
     """
@@ -56,7 +71,9 @@ class DatahubEmitterOperator(DatahubBaseOperator):
         )
         self.metadata = mces
 
-    def _render_template_fields(self, field_value, context, jinja_env):
+    def _render_template_fields(
+        self, field_value: Any, context: "Context", jinja_env: "Environment"
+    ) -> Any:
         if isinstance(field_value, DictWrapper):
             for key, value in field_value.items():
                 setattr(
@@ -73,7 +90,7 @@ class DatahubEmitterOperator(DatahubBaseOperator):
             return super().render_template(field_value, context, jinja_env)
         return field_value
 
-    def execute(self, context):
+    def execute(self, context: "Context") -> None:
         if context:
             jinja_env = self.get_template_env()
 

acryl_datahub_airflow_plugin-1.3.1.4.dist-info/METADATA

@@ -1,90 +0,0 @@
-Metadata-Version: 2.4
-Name: acryl-datahub-airflow-plugin
-Version: 1.3.1.4
-Summary: Datahub Airflow plugin to capture executions and send to Datahub
-Home-page: https://docs.datahub.com/
-License: Apache-2.0
-Project-URL: Documentation, https://docs.datahub.com/docs/
-Project-URL: Source, https://github.com/datahub-project/datahub
-Project-URL: Changelog, https://github.com/datahub-project/datahub/releases
-Classifier: Development Status :: 5 - Production/Stable
-Classifier: Programming Language :: Python
-Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3 :: Only
-Classifier: Intended Audience :: Developers
-Classifier: Intended Audience :: Information Technology
-Classifier: Intended Audience :: System Administrators
-Classifier: Operating System :: Unix
-Classifier: Operating System :: POSIX :: Linux
-Classifier: Environment :: Console
-Classifier: Environment :: MacOS X
-Classifier: Topic :: Software Development
-Requires-Python: >=3.9
-Description-Content-Type: text/markdown
-Requires-Dist: apache-airflow<3,>=2.7.0
-Requires-Dist: openlineage-airflow<=1.30.1,>=1.2.0
-Requires-Dist: acryl-datahub[datahub-rest]==1.3.1.4
-Requires-Dist: acryl-datahub[datahub-rest,sql-parser]==1.3.1.4
-Requires-Dist: pydantic>=2.4.0
-Provides-Extra: ignore
-Provides-Extra: datahub-rest
-Requires-Dist: acryl-datahub[datahub-rest]==1.3.1.4; extra == "datahub-rest"
-Provides-Extra: datahub-kafka
-Requires-Dist: acryl-datahub[datahub-kafka]==1.3.1.4; extra == "datahub-kafka"
-Provides-Extra: datahub-file
-Requires-Dist: acryl-datahub[sync-file-emitter]==1.3.1.4; extra == "datahub-file"
-Provides-Extra: dev
-Requires-Dist: ruff==0.11.7; extra == "dev"
-Requires-Dist: apache-airflow<3,>=2.7.0; extra == "dev"
-Requires-Dist: packaging; extra == "dev"
-Requires-Dist: openlineage-airflow<=1.30.1,>=1.2.0; extra == "dev"
-Requires-Dist: tenacity; extra == "dev"
-Requires-Dist: tox; extra == "dev"
-Requires-Dist: types-PyYAML; extra == "dev"
-Requires-Dist: mypy==1.17.1; extra == "dev"
-Requires-Dist: tox-uv; extra == "dev"
-Requires-Dist: twine; extra == "dev"
-Requires-Dist: types-click==0.1.12; extra == "dev"
-Requires-Dist: deepdiff!=8.0.0; extra == "dev"
-Requires-Dist: coverage>=5.1; extra == "dev"
-Requires-Dist: pytest-cov>=2.8.1; extra == "dev"
-Requires-Dist: types-setuptools; extra == "dev"
-Requires-Dist: pydantic>=2.4.0; extra == "dev"
-Requires-Dist: types-dataclasses; extra == "dev"
-Requires-Dist: build; extra == "dev"
-Requires-Dist: acryl-datahub[datahub-rest]==1.3.1.4; extra == "dev"
-Requires-Dist: types-tabulate; extra == "dev"
-Requires-Dist: types-six; extra == "dev"
-Requires-Dist: types-requests; extra == "dev"
-Requires-Dist: types-cachetools; extra == "dev"
-Requires-Dist: acryl-datahub[datahub-rest,sql-parser]==1.3.1.4; extra == "dev"
-Requires-Dist: sqlalchemy-stubs; extra == "dev"
-Requires-Dist: pytest>=6.2.2; extra == "dev"
-Requires-Dist: types-toml; extra == "dev"
-Requires-Dist: types-python-dateutil; extra == "dev"
-Provides-Extra: integration-tests
-Requires-Dist: acryl-datahub[sync-file-emitter]==1.3.1.4; extra == "integration-tests"
-Requires-Dist: apache-airflow[amazon,google,snowflake]>=2.0.2; extra == "integration-tests"
-Requires-Dist: acryl-datahub[testing-utils]==1.3.1.4; extra == "integration-tests"
-Requires-Dist: acryl-datahub[datahub-kafka]==1.3.1.4; extra == "integration-tests"
-Requires-Dist: snowflake-connector-python>=2.7.10; extra == "integration-tests"
-Requires-Dist: apache-airflow-providers-sqlite; extra == "integration-tests"
-Requires-Dist: virtualenv; extra == "integration-tests"
-Dynamic: classifier
-Dynamic: description
-Dynamic: description-content-type
-Dynamic: home-page
-Dynamic: license
-Dynamic: project-url
-Dynamic: provides-extra
-Dynamic: requires-dist
-Dynamic: requires-python
-Dynamic: summary
-
-# Datahub Airflow Plugin
-
-See [the DataHub Airflow docs](https://docs.datahub.com/docs/lineage/airflow) for details.
-
-## Developing
-
-See the [developing docs](../../metadata-ingestion/developing.md).

acryl_datahub_airflow_plugin-1.3.1.4.dist-info/RECORD

@@ -1,33 +0,0 @@
-datahub_airflow_plugin/__init__.py,sha256=NScUtA8N-m66Pyg0DO--YbPkrl48PK3UevpdQVW_y6E,1009
-datahub_airflow_plugin/_airflow_shims.py,sha256=hLMTkANJzmH9sEcaiwNb0EZgD11vh-XnDBqFQ9yqjr4,1613
-datahub_airflow_plugin/_config.py,sha256=qNbNC6YUGHf06RfOMsU7jpeNeE2ttxjUQIMEpxIhvyM,5221
-datahub_airflow_plugin/_datahub_ol_adapter.py,sha256=RuzMyWZo7MeJzAFoBfkT4cdDw5g1iWshB_nXG7jLnR0,545
-datahub_airflow_plugin/_extractors.py,sha256=HzYf6Nm8HhRfXGVTQvVgFVRUe54RSOq0CDv8oxtor1A,11875
-datahub_airflow_plugin/_version.py,sha256=RJL3ep1RXk_49LNT5HC_UqXUDu65_KSs6TJ87R1_CA0,145
-datahub_airflow_plugin/datahub_listener.py,sha256=plAU-DZ7EQmC1Ec4gzOaksqHnJLxaCYHkm2jEPTkgvc,31100
-datahub_airflow_plugin/datahub_plugin.py,sha256=rbZhs7s5O3_MlkQw5aZToC2W5mMic_EpI3oybHB0ofw,1224
-datahub_airflow_plugin/entities.py,sha256=xDZ-mZH7hjUkZbatWYUwI43_9B40wGiotlyQhiO8rEM,1987
-datahub_airflow_plugin/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-datahub_airflow_plugin/client/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-datahub_airflow_plugin/client/airflow_generator.py,sha256=zWGX6M7hqbNvOPu11VlutkJ-g149Xv2m5_IC3GqfRJk,22120
-datahub_airflow_plugin/example_dags/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-datahub_airflow_plugin/example_dags/generic_recipe_sample_dag.py,sha256=BbrOErFboKMDFn712RHEKI9T4Vh0q6kYSVet56gPqVk,1319
-datahub_airflow_plugin/example_dags/graph_usage_sample_dag.py,sha256=xYnuXhWXL5b9Tij0BFvFLckjKCobjzPU3xlxLg2_NXc,1015
-datahub_airflow_plugin/example_dags/lineage_backend_demo.py,sha256=Dy6MxwtX7G0mQeALqpLRu4F03IyU9fqIkr-CcKpo2JE,1625
-datahub_airflow_plugin/example_dags/lineage_backend_taskflow_demo.py,sha256=kW2rLFtOnoiMxBJ315GzlmR0Sz1cqQ_wwLbG9UC-u7Y,1499
-datahub_airflow_plugin/example_dags/lineage_emission_dag.py,sha256=LE29DzW51a4ZAl_zrcLrqSyzmy8qElcZagXsIMjaZLU,1946
-datahub_airflow_plugin/example_dags/mysql_sample_dag.py,sha256=Unx9Ger3R9ptEutfV-4NjjEaTIEYJ-tLrZr7OsK608k,1922
-datahub_airflow_plugin/example_dags/snowflake_sample_dag.py,sha256=b9iaE7zChQha9u57F84U6uqavGl7WrUnMNOzXEiZxjE,3234
-datahub_airflow_plugin/hooks/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-datahub_airflow_plugin/hooks/datahub.py,sha256=xa-gsJUza3jSAUP1QLJSNBn4bUHxjXo_FdCAf08IWFo,11155
-datahub_airflow_plugin/operators/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-datahub_airflow_plugin/operators/datahub.py,sha256=3fT_qg1tUpBEne1XBk8zDIUNBMcox7mxoEoN9O4XIPA,3814
-datahub_airflow_plugin/operators/datahub_assertion_operator.py,sha256=j_P9M1a5qME55pKHAfTqZsVVtIslFBO59r8UQOOBvsk,2914
-datahub_airflow_plugin/operators/datahub_assertion_sensor.py,sha256=QJIZZYQhqscj3bhBN5Sei-ABMRRAl2KiQxXTXcZQ51Q,2917
-datahub_airflow_plugin/operators/datahub_operation_operator.py,sha256=KJ8M8jJ7UWW6kNbiS-rELc-kqCPkZ3ck7z51oAXGPSI,3351
-datahub_airflow_plugin/operators/datahub_operation_sensor.py,sha256=U19fi5DpjBRWm_1ljXcjnspUzfa3mqYfOQZHjLk-ufI,3618
-acryl_datahub_airflow_plugin-1.3.1.4.dist-info/METADATA,sha256=USLZazGCYGGkU-KVH1etSzsIvnLyhMuQANkHGhKdzf0,3976
-acryl_datahub_airflow_plugin-1.3.1.4.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-acryl_datahub_airflow_plugin-1.3.1.4.dist-info/entry_points.txt,sha256=HqmajDHtrsz0b5Lswe1-eeuObxdtucd9YoxH77jJBA8,179
-acryl_datahub_airflow_plugin-1.3.1.4.dist-info/top_level.txt,sha256=VBzisOQfzqL1WRbNyItaruf3kTigXltjzgqzbheaFp0,23
-acryl_datahub_airflow_plugin-1.3.1.4.dist-info/RECORD,,

datahub_airflow_plugin/_extractors.py

@@ -1,336 +0,0 @@
-import contextlib
-import logging
-import unittest.mock
-from typing import TYPE_CHECKING, Optional
-
-from airflow.models.operator import Operator
-from openlineage.airflow.extractors import (
-    BaseExtractor,
-    ExtractorManager as OLExtractorManager,
-    TaskMetadata,
-)
-from openlineage.airflow.extractors.snowflake_extractor import SnowflakeExtractor
-from openlineage.airflow.extractors.sql_extractor import SqlExtractor
-from openlineage.airflow.utils import get_operator_class, try_import_from_string
-from openlineage.client.facet import (
-    ExtractionError,
-    ExtractionErrorRunFacet,
-    SqlJobFacet,
-)
-
-import datahub.emitter.mce_builder as builder
-from datahub.ingestion.source.sql.sqlalchemy_uri_mapper import (
-    get_platform_from_sqlalchemy_uri,
-)
-from datahub.sql_parsing.sqlglot_lineage import (
-    SqlParsingResult,
-    create_lineage_sql_parsed_result,
-)
-from datahub_airflow_plugin._datahub_ol_adapter import OL_SCHEME_TWEAKS
-
-if TYPE_CHECKING:
-    from airflow.models import DagRun, TaskInstance
-
-    from datahub.ingestion.graph.client import DataHubGraph
-
-logger = logging.getLogger(__name__)
-_DATAHUB_GRAPH_CONTEXT_KEY = "datahub_graph"
-SQL_PARSING_RESULT_KEY = "datahub_sql"
-
-
-class ExtractorManager(OLExtractorManager):
-    # TODO: On Airflow 2.7, the OLExtractorManager is part of the built-in Airflow API.
-    # When available, we should use that instead. The same goe for most of the OL
-    # extractors.
-
-    def __init__(self):
-        super().__init__()
-
-        _sql_operator_overrides = [
-            # The OL BigQuery extractor has some complex logic to fetch detect
-            # the BigQuery job_id and fetch lineage from there. However, it can't
-            # generate CLL, so we disable it and use our own extractor instead.
-            "BigQueryOperator",
-            "BigQueryExecuteQueryOperator",
-            # Athena also does something similar.
-            "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.
-            "SqliteOperator",
-        ]
-        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
-        )
-
-        self._graph: Optional["DataHubGraph"] = None
-
-    @contextlib.contextmanager
-    def _patch_extractors(self):
-        with contextlib.ExitStack() as stack:
-            # Patch the SqlExtractor.extract() method.
-            stack.enter_context(
-                unittest.mock.patch.object(
-                    SqlExtractor,
-                    "extract",
-                    _sql_extractor_extract,
-                )
-            )
-
-            # Patch the SnowflakeExtractor.default_schema property.
-            stack.enter_context(
-                unittest.mock.patch.object(
-                    SnowflakeExtractor,
-                    "default_schema",
-                    property(_snowflake_default_schema),
-                )
-            )
-
-            # TODO: Override the BigQuery extractor to use the DataHub SQL parser.
-            # self.extractor_manager.add_extractor()
-
-            # TODO: Override the Athena extractor to use the DataHub SQL parser.
-
-            yield
-
-    def extract_metadata(
-        self,
-        dagrun: "DagRun",
-        task: "Operator",
-        complete: bool = False,
-        task_instance: Optional["TaskInstance"] = None,
-        task_uuid: Optional[str] = None,
-        graph: Optional["DataHubGraph"] = None,
-    ) -> TaskMetadata:
-        self._graph = graph
-        with self._patch_extractors():
-            return super().extract_metadata(
-                dagrun, task, complete, task_instance, task_uuid
-            )
-
-    def _get_extractor(self, task: "Operator") -> Optional[BaseExtractor]:
-        # By adding this, we can use the generic extractor as a fallback for
-        # any operator that inherits from SQLExecuteQueryOperator.
-        clazz = get_operator_class(task)
-        SQLExecuteQueryOperator = try_import_from_string(
-            "airflow.providers.common.sql.operators.sql.SQLExecuteQueryOperator"
-        )
-        if SQLExecuteQueryOperator and issubclass(clazz, SQLExecuteQueryOperator):
-            self.task_to_extractor.extractors.setdefault(
-                clazz.__name__, GenericSqlExtractor
-            )
-
-        extractor = super()._get_extractor(task)
-        if extractor:
-            extractor.set_context(_DATAHUB_GRAPH_CONTEXT_KEY, self._graph)
-        return extractor
-
-
-class GenericSqlExtractor(SqlExtractor):
-    # Note that the extract() method is patched elsewhere.
-
-    @property
-    def default_schema(self):
-        return super().default_schema
-
-    def _get_scheme(self) -> Optional[str]:
-        # Best effort conversion to DataHub platform names.
-
-        with contextlib.suppress(Exception):
-            if self.hook:
-                if hasattr(self.hook, "get_uri"):
-                    uri = self.hook.get_uri()
-                    return get_platform_from_sqlalchemy_uri(uri)
-
-        return self.conn.conn_type or super().dialect
-
-    def _get_database(self) -> Optional[str]:
-        if self.conn:
-            # For BigQuery, the "database" is the project name.
-            if hasattr(self.conn, "project_id"):
-                return self.conn.project_id
-
-            return self.conn.schema
-        return None
-
-
-def _sql_extractor_extract(self: "SqlExtractor") -> TaskMetadata:
-    # Why not override the OL sql_parse method directly, instead of overriding
-    # extract()? A few reasons:
-    #
-    # 1. We would want to pass the default_db and graph instance into our sql parser
-    #    method. The OL code doesn't pass the default_db (despite having it available),
-    #    and it's not clear how to get the graph instance into that method.
-    # 2. OL has some janky logic to fetch table schemas as part of the sql extractor.
-    #    We don't want that behavior and this lets us disable it.
-    # 3. Our SqlParsingResult already has DataHub urns, whereas using SqlMeta would
-    #    require us to convert those urns to OL uris, just for them to get converted
-    #    back to urns later on in our processing.
-
-    task_name = f"{self.operator.dag_id}.{self.operator.task_id}"
-    sql = self.operator.sql
-
-    default_database = getattr(self.operator, "database", None)
-    if not default_database:
-        default_database = self.database
-    default_schema = self.default_schema
-
-    # TODO: Add better handling for sql being a list of statements.
-    if isinstance(sql, list):
-        logger.info(f"Got list of SQL statements for {task_name}. Using first one.")
-        sql = sql[0]
-
-    # Run the SQL parser.
-    scheme = self.scheme
-    platform = OL_SCHEME_TWEAKS.get(scheme, scheme)
-
-    return _parse_sql_into_task_metadata(
-        self,
-        sql,
-        platform=platform,
-        default_database=default_database,
-        default_schema=default_schema,
-    )
-
-
-def _parse_sql_into_task_metadata(
-    self: "BaseExtractor",
-    sql: str,
-    platform: str,
-    default_database: Optional[str],
-    default_schema: Optional[str],
-) -> TaskMetadata:
-    task_name = f"{self.operator.dag_id}.{self.operator.task_id}"
-
-    run_facets = {}
-    job_facets = {"sql": SqlJobFacet(query=SqlExtractor._normalize_sql(sql))}
-
-    # Prepare to run the SQL parser.
-    graph = self.context.get(_DATAHUB_GRAPH_CONTEXT_KEY, None)
-
-    self.log.debug(
-        "Running the SQL parser %s (platform=%s, default db=%s, schema=%s): %s",
-        "with graph client" if graph else "in offline mode",
-        platform,
-        default_database,
-        default_schema,
-        sql,
-    )
-    sql_parsing_result: SqlParsingResult = create_lineage_sql_parsed_result(
-        query=sql,
-        graph=graph,
-        platform=platform,
-        platform_instance=None,
-        env=builder.DEFAULT_ENV,
-        default_db=default_database,
-        default_schema=default_schema,
-    )
-    self.log.debug(f"Got sql lineage {sql_parsing_result}")
-
-    if sql_parsing_result.debug_info.error:
-        error = sql_parsing_result.debug_info.error
-        run_facets["extractionError"] = ExtractionErrorRunFacet(
-            totalTasks=1,
-            failedTasks=1,
-            errors=[
-                ExtractionError(
-                    errorMessage=str(error),
-                    stackTrace=None,
-                    task="datahub_sql_parser",
-                    taskNumber=None,
-                )
-            ],
-        )
-
-    # Save sql_parsing_result to the facets dict. It is removed from the
-    # facet dict in the extractor's processing logic.
-    run_facets[SQL_PARSING_RESULT_KEY] = sql_parsing_result  # type: ignore
-
-    return TaskMetadata(
-        name=task_name,
-        inputs=[],
-        outputs=[],
-        run_facets=run_facets,
-        job_facets=job_facets,
-    )
-
-
-class BigQueryInsertJobOperatorExtractor(BaseExtractor):
-    def extract(self) -> Optional[TaskMetadata]:
-        from airflow.providers.google.cloud.operators.bigquery import (
-            BigQueryInsertJobOperator,  # type: ignore
-        )
-
-        operator: "BigQueryInsertJobOperator" = self.operator
-        sql = operator.configuration.get("query", {}).get("query")
-        if not sql:
-            self.log.warning("No query found in BigQueryInsertJobOperator")
-            return None
-
-        destination_table = operator.configuration.get("query", {}).get(
-            "destinationTable"
-        )
-        destination_table_urn = None
-        if destination_table:
-            project_id = destination_table.get("projectId")
-            dataset_id = destination_table.get("datasetId")
-            table_id = destination_table.get("tableId")
-
-            if project_id and dataset_id and table_id:
-                destination_table_urn = builder.make_dataset_urn(
-                    platform="bigquery",
-                    name=f"{project_id}.{dataset_id}.{table_id}",
-                    env=builder.DEFAULT_ENV,
-                )
-
-        task_metadata = _parse_sql_into_task_metadata(
-            self,
-            sql,
-            platform="bigquery",
-            default_database=operator.project_id,
-            default_schema=None,
-        )
-
-        if destination_table_urn and task_metadata:
-            sql_parsing_result = task_metadata.run_facets.get(SQL_PARSING_RESULT_KEY)
-            if sql_parsing_result and isinstance(sql_parsing_result, SqlParsingResult):
-                sql_parsing_result.out_tables.append(destination_table_urn)
-
-        return task_metadata
-
-
-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=self.operator.database,
-        )
-
-
-def _snowflake_default_schema(self: "SnowflakeExtractor") -> Optional[str]:
-    if hasattr(self.operator, "schema") and self.operator.schema is not None:
-        return self.operator.schema
-    return (
-        self.conn.extra_dejson.get("extra__snowflake__schema", "")
-        or self.conn.extra_dejson.get("schema", "")
-        or self.conn.schema
-    )
-    # TODO: Should we try a fallback of:
-    # execute_query_on_hook(self.hook, "SELECT current_schema();")[0][0]