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

acryl_datahub_airflow_plugin-1.3.1.5rc1.dist-info/METADATA

@@ -0,0 +1,91 @@
+Metadata-Version: 2.4
+Name: acryl-datahub-airflow-plugin
+Version: 1.3.1.5rc1
+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: openlineage-airflow<=1.30.1,>=1.2.0
+Requires-Dist: acryl-datahub[datahub-rest]==1.3.1.5rc1
+Requires-Dist: pydantic>=2.4.0
+Requires-Dist: apache-airflow<3,>=2.7.0
+Requires-Dist: acryl-datahub[datahub-rest,sql-parser]==1.3.1.5rc1
+Provides-Extra: ignore
+Provides-Extra: datahub-rest
+Requires-Dist: acryl-datahub[datahub-rest]==1.3.1.5rc1; extra == "datahub-rest"
+Provides-Extra: datahub-kafka
+Requires-Dist: acryl-datahub[datahub-kafka]==1.3.1.5rc1; extra == "datahub-kafka"
+Provides-Extra: datahub-file
+Requires-Dist: acryl-datahub[sync-file-emitter]==1.3.1.5rc1; extra == "datahub-file"
+Provides-Extra: dev
+Requires-Dist: openlineage-airflow<=1.30.1,>=1.2.0; extra == "dev"
+Requires-Dist: types-PyYAML; extra == "dev"
+Requires-Dist: pydantic>=2.4.0; extra == "dev"
+Requires-Dist: mypy==1.17.1; extra == "dev"
+Requires-Dist: tox-uv; extra == "dev"
+Requires-Dist: packaging; extra == "dev"
+Requires-Dist: sqlalchemy-stubs; extra == "dev"
+Requires-Dist: acryl-datahub[datahub-rest]==1.3.1.5rc1; extra == "dev"
+Requires-Dist: types-cachetools; extra == "dev"
+Requires-Dist: coverage>=5.1; extra == "dev"
+Requires-Dist: types-six; extra == "dev"
+Requires-Dist: types-requests; extra == "dev"
+Requires-Dist: types-toml; extra == "dev"
+Requires-Dist: pytest-cov>=2.8.1; extra == "dev"
+Requires-Dist: pytest>=6.2.2; extra == "dev"
+Requires-Dist: tenacity; extra == "dev"
+Requires-Dist: types-setuptools; extra == "dev"
+Requires-Dist: build; extra == "dev"
+Requires-Dist: deepdiff!=8.0.0; extra == "dev"
+Requires-Dist: types-click==0.1.12; extra == "dev"
+Requires-Dist: types-dataclasses; extra == "dev"
+Requires-Dist: ruff==0.11.7; extra == "dev"
+Requires-Dist: tox; extra == "dev"
+Requires-Dist: apache-airflow<3,>=2.7.0; extra == "dev"
+Requires-Dist: types-tabulate; extra == "dev"
+Requires-Dist: types-python-dateutil; extra == "dev"
+Requires-Dist: twine; extra == "dev"
+Requires-Dist: acryl-datahub[datahub-rest,sql-parser]==1.3.1.5rc1; extra == "dev"
+Provides-Extra: integration-tests
+Requires-Dist: apache-airflow-providers-teradata; extra == "integration-tests"
+Requires-Dist: snowflake-connector-python>=2.7.10; extra == "integration-tests"
+Requires-Dist: acryl-datahub[testing-utils]==1.3.1.5rc1; extra == "integration-tests"
+Requires-Dist: apache-airflow[amazon,google,snowflake]>=2.0.2; extra == "integration-tests"
+Requires-Dist: acryl-datahub[datahub-kafka]==1.3.1.5rc1; extra == "integration-tests"
+Requires-Dist: acryl-datahub[sync-file-emitter]==1.3.1.5rc1; 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.5rc1.dist-info/RECORD

@@ -0,0 +1,33 @@
+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=o2amnv1ram58zYNiWpleQcpSREthin33LQ4yRLwoGxA,12759
+datahub_airflow_plugin/_version.py,sha256=mFUiAN31A5PUxNLnvOGfHnDXajphGmegPa0MM9whufg,148
+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.5rc1.dist-info/METADATA,sha256=qSybyzWbxjUDvL88MWC5HVumK3DTVtsT4Ezl4fTnyWw,4088
+acryl_datahub_airflow_plugin-1.3.1.5rc1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+acryl_datahub_airflow_plugin-1.3.1.5rc1.dist-info/entry_points.txt,sha256=HqmajDHtrsz0b5Lswe1-eeuObxdtucd9YoxH77jJBA8,179
+acryl_datahub_airflow_plugin-1.3.1.5rc1.dist-info/top_level.txt,sha256=VBzisOQfzqL1WRbNyItaruf3kTigXltjzgqzbheaFp0,23
+acryl_datahub_airflow_plugin-1.3.1.5rc1.dist-info/RECORD,,

datahub_airflow_plugin/_airflow_shims.py

@@ -1,75 +1,42 @@
-"""
-Pure dispatcher for version-specific Airflow shims.
-
-This module automatically imports the correct shims based on the installed
-Airflow version, dispatching to either airflow2 or airflow3 implementations.
-
-No logic lives here - just clean version detection and re-export.
-"""
+from typing import List
 
+import airflow.version
 import packaging.version
+from airflow.models.operator import Operator
+
+try:
+    from airflow.sensors.external_task import ExternalTaskSensor
+except ImportError:
+    from airflow.sensors.external_task_sensor import ExternalTaskSensor  # type: ignore
 
-from datahub_airflow_plugin._airflow_version_specific import (
-    AIRFLOW_VERSION,
-    IS_AIRFLOW_3_OR_HIGHER,
+# Approach suggested by https://stackoverflow.com/a/11887885/5004662.
+AIRFLOW_VERSION = packaging.version.parse(airflow.version.version)
+HAS_AIRFLOW_DAG_LISTENER_API = True  # this is in Airflow 2.5+
+HAS_AIRFLOW_DATASET_LISTENER_API = AIRFLOW_VERSION >= packaging.version.parse(
+    "2.8.0.dev0"
 )
 
-# Version feature flags - hardcode based on Airflow version
-# These were previously in the old _airflow_shims but are better kept simple
-HAS_AIRFLOW_STANDALONE_CMD = AIRFLOW_VERSION >= packaging.version.parse("2.2")
-HAS_AIRFLOW_LISTENER_API = AIRFLOW_VERSION >= packaging.version.parse("2.3")
-HAS_AIRFLOW_DAG_LISTENER_API = AIRFLOW_VERSION >= packaging.version.parse("2.5")
-HAS_AIRFLOW_DATASET_LISTENER_API = AIRFLOW_VERSION >= packaging.version.parse("2.5")
 
-if IS_AIRFLOW_3_OR_HIGHER:
-    # Airflow 3.x - use airflow3 shims
-    from datahub_airflow_plugin.airflow3._shims import (
-        BaseOperator,
-        ExternalTaskSensor,
-        MappedOperator,
-        OpenLineagePlugin,
-        Operator,
-        TaskHolder,
-        get_operator_class,
-        get_task_inlets,
-        get_task_outlets,
-        redact_with_exclusions,
-        try_import_from_string,
-    )
-else:
-    # Airflow 2.x - use airflow2 shims
-    from datahub_airflow_plugin.airflow2._shims import (  # type: ignore[assignment]
-        BaseOperator,
-        ExternalTaskSensor,
-        MappedOperator,
-        OpenLineagePlugin,
-        Operator,
-        TaskHolder,
-        get_operator_class,
-        get_task_inlets,
-        get_task_outlets,
-        redact_with_exclusions,
-        try_import_from_string,
-    )
+def get_task_inlets(operator: "Operator") -> List:
+    # From Airflow 2.4 _inlets is dropped and inlets used consistently. Earlier it was not the case, so we have to stick there to _inlets
+    if hasattr(operator, "_inlets"):
+        return operator._inlets  # type: ignore[attr-defined, union-attr]
+    if hasattr(operator, "get_inlet_defs"):
+        return operator.get_inlet_defs()  # type: ignore[attr-defined]
+    return operator.inlets or []
+
+
+def get_task_outlets(operator: "Operator") -> List:
+    # From Airflow 2.4 _outlets is dropped and inlets used consistently. Earlier it was not the case, so we have to stick there to _outlets
+    # We have to use _outlets because outlets is empty in Airflow < 2.4.0
+    if hasattr(operator, "_outlets"):
+        return operator._outlets  # type: ignore[attr-defined, union-attr]
+    if hasattr(operator, "get_outlet_defs"):
+        return operator.get_outlet_defs()
+    return operator.outlets or []
+
 
 __all__ = [
-    # Airflow version and feature flags
     "AIRFLOW_VERSION",
-    "IS_AIRFLOW_3_OR_HIGHER",
-    "HAS_AIRFLOW_STANDALONE_CMD",
-    "HAS_AIRFLOW_LISTENER_API",
-    "HAS_AIRFLOW_DAG_LISTENER_API",
-    "HAS_AIRFLOW_DATASET_LISTENER_API",
-    # Airflow objects
-    "BaseOperator",
-    "Operator",
-    "MappedOperator",
     "ExternalTaskSensor",
-    "TaskHolder",
-    "OpenLineagePlugin",
-    "get_operator_class",
-    "try_import_from_string",
-    "redact_with_exclusions",
-    "get_task_inlets",
-    "get_task_outlets",
 ]

datahub_airflow_plugin/_config.py

@@ -1,12 +1,12 @@
 from enum import Enum
-from typing import TYPE_CHECKING, List, Optional, Union
+from typing import TYPE_CHECKING, Dict, List, Optional, Union
 
 from airflow.configuration import conf
-from pydantic import Field
+from pydantic import root_validator
+from pydantic.fields import Field
 
 import datahub.emitter.mce_builder as builder
 from datahub.configuration.common import AllowDenyPattern, ConfigModel
-from datahub_airflow_plugin._airflow_version_specific import IS_AIRFLOW_3_OR_HIGHER
 
 if TYPE_CHECKING:
     from datahub_airflow_plugin.hooks.datahub import (
@@ -18,15 +18,16 @@ if TYPE_CHECKING:
 class DatajobUrl(Enum):
     GRID = "grid"
     TASKINSTANCE = "taskinstance"
-    TASKS = "tasks"  # Airflow 3.x task URL format: /dags/{dag_id}/tasks/{task_id}
 
 
 class DatahubLineageConfig(ConfigModel):
     enabled: bool
 
-    # DataHub hook connection ID (can be comma-separated for multiple connections).
+    # DataHub hook connection ID.
     datahub_conn_id: str
 
+    _datahub_connection_ids: List[str]
+
     # Cluster to associate with the pipelines and tasks. Defaults to "prod".
     cluster: str
 
@@ -52,32 +53,6 @@ class DatahubLineageConfig(ConfigModel):
 
     enable_extractors: bool
 
-    # OpenLineage extractor patching/override controls (only apply when enable_extractors=True)
-    # These allow fine-grained control over DataHub's enhancements to OpenLineage extractors
-
-    # If true (default), patch SqlExtractor to use DataHub's SQL parser
-    # This enables column-level lineage extraction from SQL queries
-    # Works with both Legacy OpenLineage and OpenLineage Provider
-    patch_sql_parser: bool
-
-    # If true (default), patch SnowflakeExtractor's default_schema property
-    # Fixes schema detection issues in Snowflake operators
-    # Works with both Legacy OpenLineage and OpenLineage Provider
-    patch_snowflake_schema: bool
-
-    # If true (default), use DataHub's custom AthenaOperatorExtractor
-    # Provides better Athena lineage with DataHub's SQL parser
-    # Only applies to Legacy OpenLineage (OpenLineage Provider has its own)
-    extract_athena_operator: bool
-
-    # If true (default), use DataHub's custom BigQueryInsertJobOperatorExtractor
-    # Handles BigQuery job configuration and destination tables
-    # Only applies to Legacy OpenLineage (OpenLineage Provider has its own)
-    extract_bigquery_insert_job_operator: bool
-
-    # If true (default) use DataHub's custom TeradataOperator
-    extract_teradata_operator: bool
-
     # If true, ti.render_templates() will be called in the listener.
     # Makes extraction of jinja-templated fields more accurate.
     render_templates: bool
@@ -94,17 +69,6 @@ class DatahubLineageConfig(ConfigModel):
 
     disable_openlineage_plugin: bool
 
-    @property
-    def _datahub_connection_ids(self) -> List[str]:
-        """
-        Parse comma-separated connection IDs into a list.
-
-        This is implemented as a property to avoid the class variable pollution
-        bug that would occur with validators. Each instance computes its own
-        connection ID list from its datahub_conn_id field.
-        """
-        return [conn_id.strip() for conn_id in self.datahub_conn_id.split(",")]
-
     def make_emitter_hook(self) -> Union["DatahubGenericHook", "DatahubCompositeHook"]:
         # This is necessary to avoid issues with circular imports.
         from datahub_airflow_plugin.hooks.datahub import (
@@ -112,11 +76,18 @@ class DatahubLineageConfig(ConfigModel):
             DatahubGenericHook,
         )
 
-        connection_ids = self._datahub_connection_ids
-        if len(connection_ids) == 1:
-            return DatahubGenericHook(connection_ids[0])
+        if len(self._datahub_connection_ids) == 1:
+            return DatahubGenericHook(self._datahub_connection_ids[0])
         else:
-            return DatahubCompositeHook(connection_ids)
+            return DatahubCompositeHook(self._datahub_connection_ids)
+
+    @root_validator(skip_on_failure=True)
+    def split_conn_ids(cls, values: Dict) -> Dict:
+        if not values.get("datahub_conn_id"):
+            raise ValueError("datahub_conn_id is required")
+        conn_ids = values.get("datahub_conn_id", "").split(",")
+        cls._datahub_connection_ids = [conn_id.strip() for conn_id in conn_ids]
+        return values
 
 
 def get_lineage_config() -> DatahubLineageConfig:
@@ -136,58 +107,14 @@ def get_lineage_config() -> DatahubLineageConfig:
     capture_executions = conf.get("datahub", "capture_executions", fallback=True)
     materialize_iolets = conf.get("datahub", "materialize_iolets", fallback=True)
     enable_extractors = conf.get("datahub", "enable_extractors", fallback=True)
-
-    # OpenLineage extractor patching/override configuration
-    # These only apply when enable_extractors=True
-    patch_sql_parser = conf.get("datahub", "patch_sql_parser", fallback=True)
-    patch_snowflake_schema = conf.get(
-        "datahub", "patch_snowflake_schema", fallback=True
-    )
-    extract_athena_operator = conf.get(
-        "datahub", "extract_athena_operator", fallback=True
-    )
-    extract_bigquery_insert_job_operator = conf.get(
-        "datahub", "extract_bigquery_insert_job_operator", fallback=True
-    )
-    extract_teradata_operator = conf.get(
-        "datahub", "extract_teradata_operator", fallback=True
-    )
-
     log_level = conf.get("datahub", "log_level", fallback=None)
     debug_emitter = conf.get("datahub", "debug_emitter", fallback=False)
-
-    # Disable OpenLineage plugin by default (disable_openlineage_plugin=True) for all versions.
-    # This is the safest default since most DataHub users only want DataHub's lineage.
-    #
-    # When disable_openlineage_plugin=True (default):
-    # - Only DataHub plugin runs (OpenLineagePlugin.listeners are cleared if present)
-    # - In Airflow 3: SQLParser calls only DataHub's enhanced parser
-    # - In Airflow 2: DataHub uses its own extractors
-    # - DataHub gets enhanced parsing with column-level lineage
-    #
-    # When disable_openlineage_plugin=False (opt-in for dual plugin mode):
-    # - Both DataHub and OpenLineage plugins run side-by-side
-    # - In Airflow 3: SQLParser calls BOTH parsers
-    #   - OpenLineage plugin uses its own parsing results (inputs/outputs)
-    #   - DataHub extracts its enhanced parsing (with column-level lineage) from run_facets
-    #   - Both plugins get their expected parsing without interference
-    # - In Airflow 2: Not recommended - may cause conflicts
-    default_disable_openlineage = True
-
     disable_openlineage_plugin = conf.get(
-        "datahub", "disable_openlineage_plugin", fallback=default_disable_openlineage
+        "datahub", "disable_openlineage_plugin", fallback=True
     )
     render_templates = conf.get("datahub", "render_templates", fallback=True)
-
-    # Use new task URL format for Airflow 3.x, old taskinstance format for Airflow 2.x
-    # Airflow 3 changed URL structure: /dags/{dag_id}/tasks/{task_id} instead of /taskinstance/list/...
-    default_datajob_url = (
-        DatajobUrl.TASKS.value
-        if IS_AIRFLOW_3_OR_HIGHER
-        else DatajobUrl.TASKINSTANCE.value
-    )
     datajob_url_link = conf.get(
-        "datahub", "datajob_url_link", fallback=default_datajob_url
+        "datahub", "datajob_url_link", fallback=DatajobUrl.TASKINSTANCE.value
     )
     dag_filter_pattern = AllowDenyPattern.model_validate_json(
         conf.get("datahub", "dag_filter_str", fallback='{"allow": [".*"]}')
@@ -205,11 +132,6 @@ def get_lineage_config() -> DatahubLineageConfig:
         capture_executions=capture_executions,
         materialize_iolets=materialize_iolets,
         enable_extractors=enable_extractors,
-        patch_sql_parser=patch_sql_parser,
-        patch_snowflake_schema=patch_snowflake_schema,
-        extract_athena_operator=extract_athena_operator,
-        extract_bigquery_insert_job_operator=extract_bigquery_insert_job_operator,
-        extract_teradata_operator=extract_teradata_operator,
         log_level=log_level,
         debug_emitter=debug_emitter,
         disable_openlineage_plugin=disable_openlineage_plugin,

datahub_airflow_plugin/_datahub_ol_adapter.py

@@ -1,18 +1,6 @@
 import logging
-from typing import TYPE_CHECKING
 
-# Conditional import for OpenLineage (may not be installed)
-try:
-    from openlineage.client.run import Dataset as OpenLineageDataset
-
-    OPENLINEAGE_AVAILABLE = True
-except ImportError:
-    # Not available when openlineage packages aren't installed
-    OpenLineageDataset = None  # type: ignore[assignment,misc]
-    OPENLINEAGE_AVAILABLE = False
-
-if TYPE_CHECKING:
-    from openlineage.client.run import Dataset as OpenLineageDataset
+from openlineage.client.run import Dataset as OpenLineageDataset
 
 import datahub.emitter.mce_builder as builder
 
@@ -25,7 +13,7 @@ OL_SCHEME_TWEAKS = {
 }
 
 
-def translate_ol_to_datahub_urn(ol_uri: "OpenLineageDataset") -> str:
+def translate_ol_to_datahub_urn(ol_uri: OpenLineageDataset) -> str:
     namespace = ol_uri.namespace
     name = ol_uri.name