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

datahub_airflow_plugin/_airflow_version_specific.py

@@ -0,0 +1,184 @@
+"""
+Version-specific utilities for Airflow 2 vs 3 compatibility.
+This module provides clean abstractions for version-specific behavior.
+"""
+
+import logging
+from typing import TYPE_CHECKING, Dict
+
+import airflow
+import packaging.version
+
+if TYPE_CHECKING:
+    from airflow.models import TaskInstance
+
+logger = logging.getLogger(__name__)
+
+# Version detection
+AIRFLOW_VERSION = packaging.version.parse(airflow.__version__)
+IS_AIRFLOW_3_OR_HIGHER = AIRFLOW_VERSION >= packaging.version.parse("3.0.0")
+
+
+def _get_duration_attribute(ti: "TaskInstance") -> Dict[str, str]:
+    """
+    Extract duration attribute, calculating it if necessary.
+
+    Airflow 2.x has duration as a direct attribute.
+    Airflow 3.x requires calculation from end_date - start_date.
+    """
+    if hasattr(ti, "duration"):
+        return {"duration": str(ti.duration)}
+
+    if (
+        hasattr(ti, "end_date")
+        and ti.end_date
+        and hasattr(ti, "start_date")
+        and ti.start_date
+    ):
+        try:
+            duration_seconds = (ti.end_date - ti.start_date).total_seconds()
+            return {"duration": str(duration_seconds)}
+        except Exception as e:
+            logger.debug(f"Could not calculate duration: {e}")
+
+    return {}
+
+
+def _get_operator_attribute(ti: "TaskInstance") -> Dict[str, str]:
+    """
+    Extract operator name in a version-compatible way.
+
+    In Airflow 2.x: Available as database column attribute ti.operator
+    In Airflow 3.x (RuntimeTaskInstance): Must extract from ti.task.__class__.__name__
+    """
+    if hasattr(ti, "operator"):
+        operator_from_db = str(ti.operator)
+        logger.debug(
+            f"Operator from ti.operator (DB): {operator_from_db}, "
+            f"hasattr task: {hasattr(ti, 'task')}, "
+            f"task class: {ti.task.__class__.__name__ if hasattr(ti, 'task') and ti.task else 'N/A'}"
+        )
+        return {"operator": operator_from_db}
+
+    if hasattr(ti, "task") and ti.task is not None:
+        try:
+            return {"operator": ti.task.__class__.__name__}
+        except Exception as e:
+            logger.debug(f"Could not get operator name from task: {e}")
+
+    return {}
+
+
+def _get_date_attributes(ti: "TaskInstance") -> Dict[str, str]:
+    """
+    Extract date-related attributes.
+
+    Handles execution_date -> logical_date rename in Airflow 3.0.
+    """
+    attributes = {}
+
+    if hasattr(ti, "end_date"):
+        attributes["end_date"] = str(ti.end_date)
+
+    if hasattr(ti, "execution_date"):
+        attributes["execution_date"] = str(ti.execution_date)
+    elif hasattr(ti, "logical_date"):
+        attributes["logical_date"] = str(ti.logical_date)
+
+    return attributes
+
+
+def get_task_instance_attributes(ti: "TaskInstance") -> Dict[str, str]:
+    """
+    Extract attributes from a TaskInstance in a version-compatible way.
+
+    Airflow 3.0 introduced RuntimeTaskInstance which has different attributes
+    than Airflow 2.x TaskInstance.
+
+    Returns a dict of attribute name -> string value.
+    """
+    attributes = {}
+
+    # Common attributes (both Airflow 2 and 3)
+    if hasattr(ti, "run_id"):
+        attributes["run_id"] = str(ti.run_id)
+    if hasattr(ti, "start_date") and ti.start_date:
+        attributes["start_date"] = str(ti.start_date)
+    if hasattr(ti, "try_number"):
+        attributes["try_number"] = str(ti.try_number - 1)
+    if hasattr(ti, "state"):
+        attributes["state"] = str(ti.state)
+    if hasattr(ti, "task_id"):
+        attributes["task_id"] = str(ti.task_id)
+    if hasattr(ti, "dag_id"):
+        attributes["dag_id"] = str(ti.dag_id)
+
+    # Complex extractions via helper functions
+    attributes.update(_get_duration_attribute(ti))
+    attributes.update(_get_date_attributes(ti))
+    attributes.update(_get_operator_attribute(ti))
+
+    # Optional attributes
+    if hasattr(ti, "max_tries"):
+        attributes["max_tries"] = str(ti.max_tries)
+    if hasattr(ti, "external_executor_id"):
+        attributes["external_executor_id"] = str(ti.external_executor_id)
+    if hasattr(ti, "priority_weight"):
+        attributes["priority_weight"] = str(ti.priority_weight)
+    if hasattr(ti, "log_url"):
+        attributes["log_url"] = ti.log_url
+
+    return attributes
+
+
+def get_airflow_compatible_dag_kwargs(**kwargs):  # type: ignore[no-untyped-def]
+    """
+    Get DAG kwargs that are compatible with current Airflow version.
+
+    Handles differences between Airflow 2.x and 3.x:
+    - schedule_interval -> schedule in Airflow 3.0
+    - default_view removed in Airflow 3.0
+    - start_date handling
+    """
+    compatible_kwargs = kwargs.copy()
+
+    if IS_AIRFLOW_3_OR_HIGHER:
+        # Airflow 3.0 renamed schedule_interval to schedule
+        if "schedule_interval" in compatible_kwargs:
+            compatible_kwargs["schedule"] = compatible_kwargs.pop("schedule_interval")
+
+        # Airflow 3.0 removed default_view
+        if "default_view" in compatible_kwargs:
+            del compatible_kwargs["default_view"]
+
+    return compatible_kwargs  # type: ignore[no-any-return]
+
+
+def days_ago(n: int):  # type: ignore[no-untyped-def]
+    """
+    Compatibility helper for days_ago which was removed in Airflow 3.0.
+
+    In Airflow 2.x, use airflow.utils.dates.days_ago()
+    In Airflow 3.0, use datetime.datetime - datetime.timedelta
+    """
+    from datetime import datetime, timedelta, timezone
+
+    if IS_AIRFLOW_3_OR_HIGHER:
+        # Airflow 3.0: use datetime directly
+        return datetime.now(timezone.utc) - timedelta(days=n)
+    else:
+        # Airflow 2.x: use the official helper
+        from airflow.utils.dates import (  # type: ignore[attr-defined]
+            days_ago as airflow_days_ago,
+        )
+
+        return airflow_days_ago(n)  # type: ignore[no-any-return]
+
+
+__all__ = [
+    "AIRFLOW_VERSION",
+    "IS_AIRFLOW_3_OR_HIGHER",
+    "get_task_instance_attributes",
+    "get_airflow_compatible_dag_kwargs",
+    "days_ago",
+]

datahub_airflow_plugin/_config.py

@@ -1,12 +1,12 @@
 from enum import Enum
-from typing import TYPE_CHECKING, Dict, List, Optional, Union
+from typing import TYPE_CHECKING, List, Optional, Union
 
 from airflow.configuration import conf
-from pydantic import root_validator
-from pydantic.fields import Field
+from pydantic 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,16 +18,15 @@ 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.
+    # DataHub hook connection ID (can be comma-separated for multiple connections).
     datahub_conn_id: str
 
-    _datahub_connection_ids: List[str]
-
     # Cluster to associate with the pipelines and tasks. Defaults to "prod".
     cluster: str
 
@@ -53,6 +52,32 @@ 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
@@ -69,6 +94,17 @@ 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 (
@@ -76,18 +112,11 @@ class DatahubLineageConfig(ConfigModel):
             DatahubGenericHook,
         )
 
-        if len(self._datahub_connection_ids) == 1:
-            return DatahubGenericHook(self._datahub_connection_ids[0])
+        connection_ids = self._datahub_connection_ids
+        if len(connection_ids) == 1:
+            return DatahubGenericHook(connection_ids[0])
         else:
-            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
+            return DatahubCompositeHook(connection_ids)
 
 
 def get_lineage_config() -> DatahubLineageConfig:
@@ -107,14 +136,58 @@ 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=True
+        "datahub", "disable_openlineage_plugin", fallback=default_disable_openlineage
     )
     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=DatajobUrl.TASKINSTANCE.value
+        "datahub", "datajob_url_link", fallback=default_datajob_url
     )
     dag_filter_pattern = AllowDenyPattern.model_validate_json(
         conf.get("datahub", "dag_filter_str", fallback='{"allow": [".*"]}')
@@ -132,6 +205,11 @@ 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/_constants.py

@@ -0,0 +1,16 @@
+"""
+Shared constants for the DataHub Airflow plugin.
+
+This module centralizes constant values used across multiple modules
+to avoid duplication and ensure consistency.
+"""
+
+# SQL parsing result keys for storing SQL lineage in OpenLineage facets
+
+# Key for DataHub's enhanced SQL parsing result (with column-level lineage)
+# Used in Airflow 3.x to pass results from SQLParser patch to DataHub listener
+DATAHUB_SQL_PARSING_RESULT_KEY = "datahub_sql_parsing_result"
+
+# Key for DataHub's SQL parsing result in Airflow 2.x extractors
+# Used to pass results from extractors to DataHub listener
+SQL_PARSING_RESULT_KEY = "datahub_sql"

datahub_airflow_plugin/_datahub_ol_adapter.py

@@ -1,6 +1,18 @@
 import logging
+from typing import TYPE_CHECKING
 
-from openlineage.client.run import Dataset as OpenLineageDataset
+# 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
 
 import datahub.emitter.mce_builder as builder
 
@@ -13,7 +25,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
 

datahub_airflow_plugin/_version.py

@@ -1,3 +1,3 @@
 # Published at https://pypi.org/project/acryl-datahub-airflow-plugin/.
 __package_name__ = "acryl-datahub-airflow-plugin"
-__version__ = "1.3.1.3rc2"
+__version__ = "1.3.1.5"

datahub_airflow_plugin/airflow2/__init__.py

@@ -0,0 +1,6 @@
+"""
+DataHub Airflow Plugin v2 for Airflow 2.x.
+
+This module provides the DataHub listener implementation for Airflow 2.x,
+using the legacy openlineage-airflow package and extractor-based lineage.
+"""