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

datahub_airflow_plugin/_extractors.py

@@ -0,0 +1,365 @@
+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.task_to_extractor.extractors["TeradataOperator"] = (
+            TeradataOperatorExtractor
+        )
+
+        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();")
+
+
+class TeradataOperatorExtractor(BaseExtractor):
+    """Extractor for Teradata SQL operations.
+
+    Extracts lineage from TeradataOperator tasks by parsing the SQL queries
+    and understanding Teradata's two-tier database.table naming convention.
+    """
+
+    def extract(self) -> Optional[TaskMetadata]:
+        from airflow.providers.teradata.operators.teradata import TeradataOperator
+
+        operator: "TeradataOperator" = self.operator
+        sql = operator.sql
+        if not sql:
+            self.log.warning("No query found in TeradataOperator")
+            return None
+
+        return _parse_sql_into_task_metadata(
+            self,
+            sql,
+            platform="teradata",
+            default_database=None,
+            default_schema=None,
+        )

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.5"
+__version__ = "1.3.1.5rc1"

datahub_airflow_plugin/client/airflow_generator.py

@@ -1,6 +1,5 @@
-import json
-from datetime import datetime, tzinfo
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Tuple, Union, cast
+from datetime import datetime
+from typing import TYPE_CHECKING, Dict, List, Optional, Set, Tuple, Union, cast
 
 from airflow.configuration import conf
 
@@ -13,60 +12,12 @@ from datahub.emitter.generic_emitter import Emitter
 from datahub.metadata.schema_classes import DataProcessTypeClass
 from datahub.utilities.urns.data_flow_urn import DataFlowUrn
 from datahub.utilities.urns.data_job_urn import DataJobUrn
-from datahub_airflow_plugin._airflow_version_specific import (
-    get_task_instance_attributes,
-)
 from datahub_airflow_plugin._config import DatahubLineageConfig, DatajobUrl
 
 if TYPE_CHECKING:
     from airflow import DAG
     from airflow.models import DagRun, TaskInstance
-
-    from datahub_airflow_plugin._airflow_shims import Operator
-
-    try:
-        from airflow.serialization.serialized_objects import (
-            SerializedBaseOperator,
-            SerializedDAG,
-        )
-
-        DagType = Union[DAG, SerializedDAG]
-        OperatorType = Union[Operator, SerializedBaseOperator]
-    except ImportError:
-        DagType = DAG  # type: ignore[misc]
-        OperatorType = Operator  # type: ignore[misc]
-
-    # Add type ignore for ti.task which can be MappedOperator from different modules
-    # airflow.models.mappedoperator.MappedOperator (2.x) vs airflow.sdk.definitions.mappedoperator.MappedOperator (3.x)
-    TaskType = Union[OperatorType, Any]  # type: ignore[misc]
-
-
-def _get_base_url() -> str:
-    """
-    Get the Airflow base URL for constructing web UI links.
-
-    Tries multiple configuration sources for backward compatibility:
-    1. webserver.base_url (Airflow 2.x and 3.x with computed default)
-    2. api.base_url (Airflow 3.x alternative configuration)
-    3. Fallback to http://localhost:8080 (safe default)
-
-    Returns:
-        str: The base URL for the Airflow web UI
-    """
-    # Try webserver.base_url first (works in both Airflow 2.x and 3.x)
-    # In Airflow 3.x, this is computed from web_server_host + web_server_port
-    base_url = conf.get("webserver", "base_url", fallback=None)
-    if base_url:
-        return base_url
-
-    # Fallback to api.base_url for environments that use it
-    # Some Airflow 3.x deployments may set this explicitly
-    api_base_url = conf.get("api", "base_url", fallback=None)
-    if api_base_url:
-        return api_base_url
-
-    # Final fallback to localhost (safe default for development/testing)
-    return "http://localhost:8080"
+    from airflow.models.operator import Operator
 
 
 def _task_downstream_task_ids(operator: "Operator") -> Set[str]:
@@ -78,8 +29,8 @@ def _task_downstream_task_ids(operator: "Operator") -> Set[str]:
 class AirflowGenerator:
     @staticmethod
     def _get_dependencies(
-        task: "OperatorType",
-        dag: "DagType",
+        task: "Operator",
+        dag: "DAG",
         flow_urn: DataFlowUrn,
         config: Optional[DatahubLineageConfig] = None,
     ) -> List[DataJobUrn]:
@@ -116,18 +67,14 @@ class AirflowGenerator:
 
         # subdags are always named with 'parent.child' style or Airflow won't run them
         # add connection from subdag trigger(s) if subdag task has no upstreams
-        # Note: is_subdag was removed in Airflow 3.x (subdags deprecated in Airflow 2.0)
-        parent_dag = getattr(dag, "parent_dag", None)
         if (
-            getattr(dag, "is_subdag", False)
-            and parent_dag is not None
+            dag.is_subdag
+            and dag.parent_dag is not None
             and len(task.upstream_task_ids) == 0
         ):
             # filter through the parent dag's tasks and find the subdag trigger(s)
             subdags = [
-                x
-                for x in parent_dag.task_dict.values()
-                if x.subdag is not None  # type: ignore[union-attr]
+                x for x in dag.parent_dag.task_dict.values() if x.subdag is not None
             ]
             matched_subdags = [
                 x for x in subdags if x.subdag and x.subdag.dag_id == dag.dag_id
@@ -137,14 +84,14 @@ class AirflowGenerator:
             subdag_task_id = matched_subdags[0].task_id
 
             # iterate through the parent dag's tasks and find the ones that trigger the subdag
-            for upstream_task_id in parent_dag.task_dict:  # type: ignore[union-attr]
-                upstream_task = parent_dag.task_dict[upstream_task_id]  # type: ignore[union-attr]
+            for upstream_task_id in dag.parent_dag.task_dict:
+                upstream_task = dag.parent_dag.task_dict[upstream_task_id]
                 upstream_task_urn = DataJobUrn.create_from_ids(
                     data_flow_urn=str(flow_urn), job_id=upstream_task_id
                 )
 
                 # if the task triggers the subdag, link it to this node in the subdag
-                if subdag_task_id in sorted(_task_downstream_task_ids(upstream_task)):  # type: ignore[arg-type]
+                if subdag_task_id in sorted(_task_downstream_task_ids(upstream_task)):
                     upstream_subdag_triggers.append(upstream_task_urn)
 
         # If the operator is an ExternalTaskSensor then we set the remote task as upstream.
@@ -153,16 +100,14 @@ class AirflowGenerator:
         external_task_upstreams = []
         if isinstance(task, ExternalTaskSensor):
             task = cast(ExternalTaskSensor, task)
-            external_task_id = getattr(task, "external_task_id", None)
-            external_dag_id = getattr(task, "external_dag_id", None)
-            if external_task_id is not None and external_dag_id is not None:
+            if hasattr(task, "external_task_id") and task.external_task_id is not None:
                 external_task_upstreams = [
                     DataJobUrn.create_from_ids(
-                        job_id=external_task_id,
+                        job_id=task.external_task_id,
                         data_flow_urn=str(
                             DataFlowUrn.create_from_ids(
                                 orchestrator=flow_urn.orchestrator,
-                                flow_id=external_dag_id,
+                                flow_id=task.external_dag_id,
                                 env=flow_urn.cluster,
                                 platform_instance=config.platform_instance
                                 if config
@@ -185,13 +130,13 @@ class AirflowGenerator:
         return upstream_tasks
 
     @staticmethod
-    def _extract_owners(dag: "DagType") -> List[str]:
+    def _extract_owners(dag: "DAG") -> List[str]:
         return [owner.strip() for owner in dag.owner.split(",")]
 
     @staticmethod
     def generate_dataflow(
         config: DatahubLineageConfig,
-        dag: "DagType",
+        dag: "DAG",
     ) -> DataFlow:
         """
         Generates a Dataflow object from an Airflow DAG
@@ -228,34 +173,12 @@ class AirflowGenerator:
             "timezone",
         ]
 
-        def _serialize_dag_property(value: Any) -> str:
-            """Serialize DAG property values to string format (JSON-compatible when possible)."""
-            if value is None:
-                return ""
-            elif isinstance(value, bool):
-                return "true" if value else "false"
-            elif isinstance(value, datetime):
-                return value.isoformat()
-            elif isinstance(value, (set, frozenset)):
-                # Convert set to JSON array string
-                return json.dumps(sorted(list(value)))
-            elif isinstance(value, tzinfo):
-                return str(value.tzname(None))
-            elif isinstance(value, (int, float)):
-                return str(value)
-            elif isinstance(value, str):
-                return value
-            else:
-                # For other types, convert to string but avoid repr() format
-                return str(value)
-
         for key in allowed_flow_keys:
             if hasattr(dag, key):
-                value = getattr(dag, key)
-                flow_property_bag[key] = _serialize_dag_property(value)
+                flow_property_bag[key] = repr(getattr(dag, key))
 
         data_flow.properties = flow_property_bag
-        base_url = _get_base_url()
+        base_url = conf.get("webserver", "base_url")
         data_flow.url = f"{base_url}/tree?dag_id={dag.dag_id}"
 
         if config.capture_ownership_info and dag.owner:
@@ -271,8 +194,8 @@ class AirflowGenerator:
         return data_flow
 
     @staticmethod
-    def _get_description(task: "OperatorType") -> Optional[str]:
-        from datahub_airflow_plugin._airflow_shims import BaseOperator
+    def _get_description(task: "Operator") -> Optional[str]:
+        from airflow.models.baseoperator import BaseOperator
 
         if not isinstance(task, BaseOperator):
             # TODO: Get docs for mapped operators.
@@ -293,8 +216,8 @@ class AirflowGenerator:
     @staticmethod
     def generate_datajob(
         cluster: str,
-        task: "OperatorType",
-        dag: "DagType",
+        task: "Operator",
+        dag: "DAG",
         set_dependencies: bool = True,
         capture_owner: bool = True,
         capture_tags: bool = True,
@@ -366,15 +289,11 @@ class AirflowGenerator:
                     break
 
         datajob.properties = job_property_bag
-        base_url = _get_base_url()
+        base_url = conf.get("webserver", "base_url")
 
         if config and config.datajob_url_link == DatajobUrl.GRID:
             datajob.url = f"{base_url}/dags/{dag.dag_id}/grid?task_id={task.task_id}"
-        elif config and config.datajob_url_link == DatajobUrl.TASKS:
-            # Airflow 3.x task URL format
-            datajob.url = f"{base_url}/dags/{dag.dag_id}/tasks/{task.task_id}"
         else:
-            # Airflow 2.x taskinstance list URL format
             datajob.url = f"{base_url}/taskinstance/list/?flt1_dag_id_equals={dag.dag_id}&_flt_3_task_id={task.task_id}"
 
         if capture_owner and dag.owner:
@@ -528,12 +447,8 @@ class AirflowGenerator:
     ) -> DataProcessInstance:
         if datajob is None:
             assert ti.task is not None
-            # ti.task can be MappedOperator from different modules (airflow.models vs airflow.sdk.definitions)
             datajob = AirflowGenerator.generate_datajob(
-                config.cluster,
-                ti.task,  # type: ignore[arg-type]
-                dag,
-                config=config,
+                config.cluster, ti.task, dag, config=config
             )
 
         assert dag_run.run_id
@@ -543,23 +458,26 @@ class AirflowGenerator:
             clone_inlets=True,
             clone_outlets=True,
         )
-
-        job_property_bag = get_task_instance_attributes(ti)
-
-        # Add orchestrator and DAG/task IDs
+        job_property_bag: Dict[str, str] = {}
+        job_property_bag["run_id"] = str(dag_run.run_id)
+        job_property_bag["duration"] = str(ti.duration)
+        job_property_bag["start_date"] = str(ti.start_date)
+        job_property_bag["end_date"] = str(ti.end_date)
+        job_property_bag["execution_date"] = str(ti.execution_date)
+        job_property_bag["try_number"] = str(ti.try_number - 1)
+        job_property_bag["max_tries"] = str(ti.max_tries)
+        # Not compatible with Airflow 1
+        if hasattr(ti, "external_executor_id"):
+            job_property_bag["external_executor_id"] = str(ti.external_executor_id)
+        job_property_bag["state"] = str(ti.state)
+        job_property_bag["operator"] = str(ti.operator)
+        job_property_bag["priority_weight"] = str(ti.priority_weight)
+        job_property_bag["log_url"] = ti.log_url
         job_property_bag["orchestrator"] = "airflow"
-        if "dag_id" not in job_property_bag:
-            job_property_bag["dag_id"] = str(dag.dag_id)
-        if "task_id" not in job_property_bag:
-            job_property_bag["task_id"] = str(ti.task_id)
-        if "run_id" not in job_property_bag:
-            job_property_bag["run_id"] = str(dag_run.run_id)
-
+        job_property_bag["dag_id"] = str(dag.dag_id)
+        job_property_bag["task_id"] = str(ti.task_id)
         dpi.properties.update(job_property_bag)
-
-        # Set URL if log_url is available
-        if "log_url" in job_property_bag:
-            dpi.url = job_property_bag["log_url"]
+        dpi.url = ti.log_url
 
         # This property only exists in Airflow2
         if hasattr(ti, "dag_run") and hasattr(ti.dag_run, "run_type"):
@@ -620,12 +538,8 @@ class AirflowGenerator:
         """
         if datajob is None:
             assert ti.task is not None
-            # ti.task can be MappedOperator from different modules (airflow.models vs airflow.sdk.definitions)
             datajob = AirflowGenerator.generate_datajob(
-                cluster,
-                ti.task,  # type: ignore[arg-type]
-                dag,
-                config=config,
+                cluster, ti.task, dag, config=config
             )
 
         if end_timestamp_millis is None:
@@ -652,24 +566,6 @@ class AirflowGenerator:
             clone_inlets=True,
             clone_outlets=True,
         )
-
-        job_property_bag = get_task_instance_attributes(ti)
-
-        # Add orchestrator and DAG/task IDs
-        job_property_bag["orchestrator"] = "airflow"
-        if "dag_id" not in job_property_bag:
-            job_property_bag["dag_id"] = str(dag.dag_id)
-        if "task_id" not in job_property_bag:
-            job_property_bag["task_id"] = str(ti.task_id)
-        if "run_id" not in job_property_bag:
-            job_property_bag["run_id"] = str(dag_run.run_id)
-
-        dpi.properties.update(job_property_bag)
-
-        # Set URL if log_url is available
-        if "log_url" in job_property_bag:
-            dpi.url = job_property_bag["log_url"]
-
         dpi.emit_process_end(
             emitter=emitter,
             end_timestamp_millis=end_timestamp_millis,