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

datahub_airflow_plugin/airflow2/_airflow2_sql_parser_patch.py

@@ -0,0 +1,402 @@
+"""
+Patch for Airflow 2.10+ with apache-airflow-providers-openlineage SQLParser.
+
+When using apache-airflow-providers-openlineage with Airflow 2.10+, SQL operators call
+SQLParser.generate_openlineage_metadata_from_sql() directly rather than using extractors.
+This module patches that method to use DataHub's SQL parser, which provides better
+column-level lineage support.
+
+This is analogous to the Airflow 3 SQL parser patch, but adapted for Airflow 2.10+
+when the provider package is installed.
+"""
+
+import logging
+from types import TracebackType
+from typing import TYPE_CHECKING, Any, Callable, Optional
+
+# Try importing from provider package (Airflow 2.10+ with apache-airflow-providers-openlineage)
+try:
+    from airflow.providers.openlineage.extractors import OperatorLineage
+    from airflow.providers.openlineage.sqlparser import DatabaseInfo
+    from openlineage.client.event_v2 import Dataset as OpenLineageDataset
+    from openlineage.client.facet import SqlJobFacet
+
+    PROVIDER_IMPORTS_AVAILABLE = True
+except ImportError:
+    # Provider package not available
+    OperatorLineage = None  # type: ignore[assignment,misc]
+    DatabaseInfo = None  # type: ignore[assignment,misc]
+    OpenLineageDataset = None  # type: ignore[assignment,misc]
+    SqlJobFacet = None  # type: ignore[assignment,misc]
+    PROVIDER_IMPORTS_AVAILABLE = False
+
+# DataHub imports (always available)
+import datahub.emitter.mce_builder as builder
+from datahub.sql_parsing.sqlglot_lineage import create_lineage_sql_parsed_result
+from datahub_airflow_plugin._constants import DATAHUB_SQL_PARSING_RESULT_KEY
+from datahub_airflow_plugin._datahub_ol_adapter import OL_SCHEME_TWEAKS
+
+if TYPE_CHECKING:
+    from airflow.providers.openlineage.extractors import OperatorLineage
+    from airflow.providers.openlineage.sqlparser import DatabaseInfo
+    from openlineage.client.event_v2 import Dataset as OpenLineageDataset
+    from openlineage.client.facet import SqlJobFacet
+
+logger = logging.getLogger(__name__)
+
+# Store the original SQLParser method for fallback
+_original_sql_parser_method: Optional[Callable[..., Any]] = None
+
+
+def _datahub_generate_openlineage_metadata_from_sql(
+    self: Any,
+    sql: Any,
+    hook: Any,
+    database_info: dict,
+    database: Optional[str] = None,
+    sqlalchemy_engine: Optional[Any] = None,
+    use_connection: bool = True,
+) -> Optional["OperatorLineage"]:
+    """
+    Override SQLParser.generate_openlineage_metadata_from_sql to use DataHub's SQL parser.
+
+    This is necessary because in Airflow 2.10+ with provider package, SQL operators call
+    SQLParser directly rather than using extractors. We intercept this call and use
+    DataHub's SQL parser to generate lineage with column-level lineage support.
+
+    When OpenLineage plugin is enabled (disable_openlineage_plugin=False), we call both
+    parsers: OpenLineage gets its own parsing results, while DataHub's enhanced parsing
+    is stored in a custom facet for the DataHub listener to extract.
+    """
+    try:
+        # Import here to avoid circular dependency
+        from datahub_airflow_plugin._config import get_lineage_config
+        from datahub_airflow_plugin.datahub_listener import get_airflow_plugin_listener
+
+        # Check if OpenLineage plugin is enabled
+        try:
+            config = get_lineage_config()
+            openlineage_enabled = not config.disable_openlineage_plugin
+        except Exception as e:
+            logger.warning(
+                f"Could not load config to check disable_openlineage_plugin: {e}"
+            )
+            openlineage_enabled = False
+
+        # If OpenLineage is enabled, call the original parser first to get its results
+        ol_result = None
+        if openlineage_enabled and _original_sql_parser_method is not None:
+            try:
+                logger.debug(
+                    "OpenLineage plugin enabled - calling original parser for OpenLineage"
+                )
+                ol_result = _original_sql_parser_method(
+                    self,
+                    sql,
+                    hook,
+                    database_info,
+                    database,
+                    sqlalchemy_engine,
+                    use_connection,
+                )
+                logger.debug(f"OpenLineage parser result: {ol_result}")
+            except Exception as e:
+                logger.warning(
+                    f"Error calling original OpenLineage parser, will use only DataHub parser: {e}",
+                    exc_info=True,
+                )
+
+        # Handle missing database_info by creating a minimal one from connection
+        if database_info is None:
+            # Get basic properties from hook's connection
+            conn = getattr(hook, "get_connection", lambda: None)()
+            scheme = getattr(conn, "conn_type", None) if conn else None
+            db_name = getattr(conn, "schema", None) if conn else None
+
+            database_info = DatabaseInfo(
+                scheme=scheme,
+                authority=None,
+                database=db_name,
+                information_schema_columns=[],
+                information_schema_table_name="",
+                use_flat_cross_db_query=False,
+                is_information_schema_cross_db=False,
+                is_uppercase_names=False,
+                normalize_name_method=lambda x: x.lower(),
+            )
+            logger.debug(
+                f"Created minimal DatabaseInfo from connection: scheme={scheme}, database={db_name}"
+            )
+
+        # Get platform from dialect or from database_info scheme
+        # If dialect is "generic", prefer database_info.scheme (connection type)
+        platform = self.dialect or "sql"
+        if platform == "generic" and database_info:
+            # Use the actual connection type instead of "generic"
+            platform = getattr(database_info, "scheme", platform) or platform
+            if platform == "generic":
+                raise ValueError(
+                    "Could not determine platform from generic dialect or database_info"
+                )
+
+        platform = OL_SCHEME_TWEAKS.get(platform, platform)
+
+        # Get default database and schema
+        default_database = database or getattr(database_info, "database", None)
+        default_schema = self.default_schema
+
+        # Handle list of SQL statements
+        if isinstance(sql, list):
+            logger.debug("Got list of SQL statements. Using first one for parsing.")
+            sql = sql[0] if sql else ""
+
+        # Run DataHub's SQL parser
+        listener = get_airflow_plugin_listener()
+        graph = listener.graph if listener else None
+
+        logger.debug(
+            "Running DataHub 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 = 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,
+        )
+
+        logger.debug(f"DataHub SQL parser result: {sql_parsing_result}")
+
+        # Store the sql_parsing_result in run_facets for later retrieval by the DataHub listener
+        # If OpenLineage plugin is enabled and we got a result from the original parser,
+        # use OpenLineage's result but add DataHub's parsing to the facets
+        if ol_result is not None:
+            logger.debug(
+                "Using OpenLineage parser result for OperatorLineage, "
+                "adding DataHub parsing to run_facets"
+            )
+            # Add DataHub's SQL parsing result to the existing run_facets
+            # OperatorLineage is frozen (uses @define), so we need to create a new dict
+            updated_run_facets = dict(ol_result.run_facets or {})
+            updated_run_facets[DATAHUB_SQL_PARSING_RESULT_KEY] = sql_parsing_result
+
+            # Create new OperatorLineage with OpenLineage's inputs/outputs but DataHub's facet
+            operator_lineage = OperatorLineage(  # type: ignore[misc]
+                inputs=ol_result.inputs,
+                outputs=ol_result.outputs,
+                job_facets=ol_result.job_facets,
+                run_facets=updated_run_facets,
+            )
+            return operator_lineage
+
+        # OpenLineage is disabled or original parser failed - use DataHub's parsing for everything
+        logger.debug(
+            "OpenLineage plugin disabled or parser unavailable - "
+            "using DataHub parser result for OperatorLineage"
+        )
+
+        # Convert DataHub URNs to OpenLineage Dataset objects
+        def _urn_to_ol_dataset(urn: str) -> "OpenLineageDataset":
+            """Convert DataHub URN to OpenLineage Dataset format."""
+            # Parse URN to extract database, schema, table
+            # URN format: urn:li:dataset:(urn:li:dataPlatform:{platform},{database}.{schema}.{table},{env})
+            try:
+                parts = urn.split(",")
+                if len(parts) >= 2:
+                    # Extract table path from URN
+                    table_path = parts[1]  # e.g., "database.schema.table"
+
+                    # Create OL namespace and name
+                    # For now, use platform as namespace and full path as name
+                    namespace = f"{platform}://{default_database or 'default'}"
+                    name = table_path
+
+                    return OpenLineageDataset(namespace=namespace, name=name)
+            except Exception as e:
+                logger.debug(f"Error converting URN {urn} to OL Dataset: {e}")
+
+            # Fallback: use URN as name
+            return OpenLineageDataset(namespace=f"{platform}://default", name=urn)
+
+        inputs = [_urn_to_ol_dataset(urn) for urn in sql_parsing_result.in_tables]
+        outputs = [_urn_to_ol_dataset(urn) for urn in sql_parsing_result.out_tables]
+
+        run_facets = {DATAHUB_SQL_PARSING_RESULT_KEY: sql_parsing_result}
+
+        # Create OperatorLineage with DataHub's results
+        operator_lineage = OperatorLineage(  # type: ignore[misc]
+            inputs=inputs,
+            outputs=outputs,
+            job_facets={"sql": SqlJobFacet(query=sql)},
+            run_facets=run_facets,
+        )
+        return operator_lineage
+
+    except Exception as e:
+        logger.warning(
+            f"Error in DataHub SQL parser, falling back to default OpenLineage parser: {e}",
+            exc_info=True,
+        )
+        # Fall back to original implementation
+        if _original_sql_parser_method is None:
+            raise RuntimeError(
+                "Original SQLParser method not stored. patch_sqlparser() may not have been called."
+            ) from None
+        return _original_sql_parser_method(
+            self, sql, hook, database_info, database, sqlalchemy_engine, use_connection
+        )
+
+
+class SQLParserPatch:
+    """
+    Context manager for patching Airflow's SQLParser with DataHub's SQL parser.
+
+    This class encapsulates the patching logic and manages the global state properly.
+    It can be used as a context manager for automatic cleanup, or with explicit
+    patch/unpatch methods for manual control.
+
+    Usage:
+        # As a context manager (recommended for testing)
+        with SQLParserPatch():
+            # Code runs with patched SQLParser
+            pass
+        # Automatically unpatched on exit
+
+        # Or with explicit control
+        patcher = SQLParserPatch()
+        patcher.patch()
+        try:
+            # ... plugin lifetime ...
+        finally:
+            patcher.unpatch()
+
+    The patch stores the original SQLParser method and replaces it with DataHub's
+    enhanced implementation that provides column-level lineage support.
+    """
+
+    def patch(self) -> "SQLParserPatch":
+        """
+        Apply the SQLParser patch.
+
+        Stores the original SQLParser.generate_openlineage_metadata_from_sql method
+        and replaces it with DataHub's enhanced implementation.
+
+        Returns:
+            self for method chaining
+        """
+        global _original_sql_parser_method
+
+        try:
+            from airflow.providers.openlineage.sqlparser import SQLParser
+
+            # Store original method for fallback (only if not already patched)
+            if _original_sql_parser_method is None:
+                _original_sql_parser_method = (
+                    SQLParser.generate_openlineage_metadata_from_sql
+                )
+
+            SQLParser.generate_openlineage_metadata_from_sql = (  # type: ignore[assignment,method-assign]
+                _datahub_generate_openlineage_metadata_from_sql  # type: ignore[assignment,method-assign]
+            )
+            logger.debug(
+                "Patched SQLParser.generate_openlineage_metadata_from_sql with DataHub SQL parser"
+            )
+
+        except ImportError:
+            # SQLParser not available (provider package not installed or Airflow < 2.10)
+            logger.debug(
+                "SQLParser not available, skipping patch (likely Airflow < 2.10 or provider package not installed)"
+            )
+
+        return self
+
+    def unpatch(self) -> "SQLParserPatch":
+        """
+        Remove the SQLParser patch and restore the original method.
+
+        This is primarily useful for testing to ensure clean state between tests.
+        In production, the patch typically stays active for the process lifetime.
+
+        Returns:
+            self for method chaining
+        """
+        global _original_sql_parser_method
+
+        if _original_sql_parser_method is None:
+            logger.debug("SQLParser not patched, nothing to unpatch")
+            return self
+
+        try:
+            from airflow.providers.openlineage.sqlparser import SQLParser
+
+            # Restore original method
+            SQLParser.generate_openlineage_metadata_from_sql = (  # type: ignore[method-assign]
+                _original_sql_parser_method
+            )
+            logger.debug("Unpatched SQLParser, restored original method")
+
+        except ImportError:
+            logger.debug("SQLParser not available, nothing to unpatch")
+        finally:
+            # Clear the stored reference to allow re-patching
+            _original_sql_parser_method = None
+
+        return self
+
+    def __enter__(self) -> "SQLParserPatch":
+        """Context manager entry: apply the patch."""
+        return self.patch()
+
+    def __exit__(
+        self,
+        exc_type: Optional[type[BaseException]],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional[TracebackType],
+    ) -> None:
+        """Context manager exit: remove the patch."""
+        self.unpatch()
+
+
+# Global patcher instance for backward compatibility
+_global_patcher = SQLParserPatch()
+
+
+def patch_sqlparser() -> None:
+    """
+    Patch SQLParser.generate_openlineage_metadata_from_sql to use DataHub's SQL parser.
+
+    This is a convenience function that wraps SQLParserPatch.patch() for backward
+    compatibility with existing code.
+
+    This should be called early in the plugin initialization, before any SQL operators are used.
+
+    When both DataHub and OpenLineage plugins are enabled (disable_openlineage_plugin=False),
+    the patch calls BOTH parsers:
+    - OpenLineage's original parser provides inputs/outputs for OpenLineage plugin
+    - DataHub's enhanced parser (with column-level lineage) is stored in run_facets
+      for DataHub listener to extract
+
+    When only DataHub is enabled (disable_openlineage_plugin=True), only DataHub's
+    parser runs and provides both the OperatorLineage structure and the enhanced parsing.
+    """
+    _global_patcher.patch()
+
+
+def unpatch_sqlparser() -> None:
+    """
+    Remove the SQLParser patch and restore the original method.
+
+    This is a convenience function that wraps SQLParserPatch.unpatch() for consistency.
+
+    This is primarily useful for testing to ensure clean state between tests.
+    In production, the patch typically stays active for the process lifetime.
+    """
+    _global_patcher.unpatch()

datahub_airflow_plugin/airflow2/_airflow_compat.py

@@ -0,0 +1,95 @@
+# Airflow 2.x compatibility module
+# This module must be imported before any Airflow imports in any of our files.
+
+from datahub.utilities._markupsafe_compat import MARKUPSAFE_PATCHED
+
+# Critical safety check: Ensure MarkupSafe compatibility patch is applied
+# This must happen before importing Airflow to prevent MarkupSafe version conflicts
+# Using explicit exception instead of assert to ensure it runs even with python -O
+if not MARKUPSAFE_PATCHED:
+    raise RuntimeError(
+        "MarkupSafe compatibility patch must be applied before importing Airflow modules. "
+        "This is a critical safety check that cannot be disabled. "
+        "The patch ensures compatibility between different MarkupSafe versions used by "
+        "Airflow and DataHub dependencies."
+    )
+
+# Apply SQLParser patch for Airflow 2.10+ with apache-airflow-providers-openlineage
+# When using the provider package, SQL operators call SQLParser.generate_openlineage_metadata_from_sql()
+# directly (similar to Airflow 3.x), so we need to patch that method to use DataHub's SQL parser.
+#
+# For legacy openlineage-airflow package (Airflow 2.5-2.9), we use the extractor-based approach
+# in _extractors.py instead.
+import importlib.util
+import logging
+
+logger = logging.getLogger(__name__)
+
+# Check if OpenLineage provider package is available
+# Use try-except because find_spec can raise ModuleNotFoundError if parent module doesn't exist
+try:
+    has_openlineage_provider = (
+        importlib.util.find_spec("airflow.providers.openlineage.sqlparser") is not None
+    )
+except (ModuleNotFoundError, ImportError, ValueError):
+    # Parent module doesn't exist or other import error
+    has_openlineage_provider = False
+
+if has_openlineage_provider:
+    # Provider package detected - apply SQL parser patch
+    from datahub_airflow_plugin.airflow2._airflow2_sql_parser_patch import (
+        patch_sqlparser,
+    )
+
+    patch_sqlparser()
+else:
+    # Provider package not available - using legacy openlineage-airflow package
+    # No patching needed, extractors will handle SQL parsing
+    pass
+
+# Apply operator-specific patches for provider mode
+# These patches work for both Airflow 2.x and 3.x when using OpenLineage provider
+try:
+    from datahub_airflow_plugin._config import get_lineage_config
+
+    config = get_lineage_config()
+    enable_extractors = config.enable_extractors
+    extract_teradata_operator = config.extract_teradata_operator
+except Exception:
+    # If config loading fails, apply patches by default (backward compatibility)
+    enable_extractors = True
+    extract_teradata_operator = True
+
+if enable_extractors and extract_teradata_operator:
+    # TeradataOperator patch - works for both Airflow 2.x provider mode and Airflow 3.x
+    # The patch checks for method existence, so it's safe to import from airflow3 module
+    # Note: We defer the import to avoid potential issues with Airflow 3.x specific imports
+    # in Airflow 2.x environments. The patch function itself handles version compatibility.
+    import logging
+
+    logger = logging.getLogger(__name__)
+    try:
+        logger.debug("Attempting to import and apply TeradataOperator patch")
+        # Use importlib to safely import the patch module
+        import importlib.util
+
+        patch_module_path = (
+            "datahub_airflow_plugin.airflow3._teradata_openlineage_patch"
+        )
+        patch_module = importlib.import_module(patch_module_path)
+        patch_teradata_operator = patch_module.patch_teradata_operator
+
+        patch_teradata_operator()
+        logger.debug("TeradataOperator patch import and call completed")
+    except ImportError as e:
+        # Teradata provider not installed or patch not available
+        logger.debug(f"Could not import TeradataOperator patch: {e}")
+    except Exception as e:
+        # Log error but don't fail - this is optional functionality
+        logger.warning(f"Error applying TeradataOperator patch: {e}", exc_info=True)
+
+AIRFLOW_PATCHED = True
+
+__all__ = [
+    "AIRFLOW_PATCHED",
+]