feat(ingest/airflow): fix materialize_iolets bug (#10613)

metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/client/airflow_generator.py

@@ -129,10 +129,8 @@ def _get_dependencies(
 
  @staticmethod
  def generate_dataflow(
- cluster: str,
+ config: DatahubLineageConfig,
  dag: "DAG",
- capture_owner: bool = True,
- capture_tags: bool = True,
  ) -> DataFlow:
  """
  Generates a Dataflow object from an Airflow DAG
@@ -146,7 +144,10 @@ def generate_dataflow(
  orchestrator = "airflow"
  description = "\n\n".join(filter(None, [dag.description, dag.doc_md])) or None
  data_flow = DataFlow(
- env=cluster, id=id, orchestrator=orchestrator, description=description
+ env=config.cluster,
+ id=id,
+ orchestrator=orchestrator,
+ description=description,
  )
 
  flow_property_bag: Dict[str, str] = {}
@@ -173,10 +174,10 @@ def generate_dataflow(
  base_url = conf.get("webserver", "base_url")
  data_flow.url = f"{base_url}/tree?dag_id={dag.dag_id}"
 
- if capture_owner and dag.owner:
+ if config.capture_ownership_info and dag.owner:
  data_flow.owners.update(owner.strip() for owner in dag.owner.split(","))
 
- if capture_tags and dag.tags:
+ if config.capture_tags_info and dag.tags:
  data_flow.tags.update(dag.tags)
 
  return data_flow
@@ -311,14 +312,14 @@ def create_datajob_instance(
  @staticmethod
  def run_dataflow(
  emitter: Emitter,
- cluster: str,
+ config: DatahubLineageConfig,
  dag_run: "DagRun",
  start_timestamp_millis: Optional[int] = None,
  dataflow: Optional[DataFlow] = None,
  ) -> None:
  if dataflow is None:
  assert dag_run.dag
- dataflow = AirflowGenerator.generate_dataflow(cluster, dag_run.dag)
+ dataflow = AirflowGenerator.generate_dataflow(config, dag_run.dag)
 
  if start_timestamp_millis is None:
  assert dag_run.execution_date
@@ -357,13 +358,15 @@ def run_dataflow(
  dpi.properties.update(property_bag)
 
  dpi.emit_process_start(
- emitter=emitter, start_timestamp_millis=start_timestamp_millis
+ emitter=emitter,
+ start_timestamp_millis=start_timestamp_millis,
+ materialize_iolets=config.materialize_iolets,
  )
 
  @staticmethod
  def complete_dataflow(
  emitter: Emitter,
- cluster: str,
+ config: DatahubLineageConfig,
  dag_run: "DagRun",
  end_timestamp_millis: Optional[int] = None,
  dataflow: Optional[DataFlow] = None,
@@ -378,7 +381,7 @@ def complete_dataflow(
  """
  if dataflow is None:
  assert dag_run.dag
- dataflow = AirflowGenerator.generate_dataflow(cluster, dag_run.dag)
+ dataflow = AirflowGenerator.generate_dataflow(config, dag_run.dag)
 
  assert dag_run.run_id
  dpi = DataProcessInstance.from_dataflow(dataflow=dataflow, id=dag_run.run_id)
@@ -409,28 +412,27 @@ def complete_dataflow(
  @staticmethod
  def run_datajob(
  emitter: Emitter,
- cluster: str,
  ti: "TaskInstance",
  dag: "DAG",
  dag_run: "DagRun",
+ config: DatahubLineageConfig,
  start_timestamp_millis: Optional[int] = None,
  datajob: Optional[DataJob] = None,
  attempt: Optional[int] = None,
  emit_templates: bool = True,
- config: Optional[DatahubLineageConfig] = None,
  ) -> DataProcessInstance:
  if datajob is None:
  assert ti.task is not None
  datajob = AirflowGenerator.generate_datajob(
- cluster, ti.task, dag, config=config
+ config.cluster, ti.task, dag, config=config
  )
 
  assert dag_run.run_id
  dpi = DataProcessInstance.from_datajob(
  datajob=datajob,
  id=f"{dag.dag_id}_{ti.task_id}_{dag_run.run_id}",
- clone_inlets=config is None or config.materialize_iolets,
- clone_outlets=config is None or config.materialize_iolets,
+ clone_inlets=True,
+ clone_outlets=True,
  )
  job_property_bag: Dict[str, str] = {}
  job_property_bag["run_id"] = str(dag_run.run_id)
@@ -481,6 +483,7 @@ def run_datajob(
  start_timestamp_millis=start_timestamp_millis,
  attempt=attempt,
  emit_template=emit_templates,
+ materialize_iolets=config.materialize_iolets,
  )
  return dpi
 

metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/datahub_listener.py

@@ -408,13 +408,12 @@ def on_task_instance_running(
  if self.config.capture_executions:
  dpi = AirflowGenerator.run_datajob(
  emitter=self.emitter,
- cluster=self.config.cluster,
+ config=self.config,
  ti=task_instance,
  dag=dag,
  dag_run=dagrun,
  datajob=datajob,
  emit_templates=False,
- config=self.config,
  )
  logger.debug(f"Emitted DataHub DataProcess Instance start: {dpi}")
 
@@ -530,10 +529,8 @@ def on_dag_start(self, dag_run: "DagRun") -> None:
  return
 
  dataflow = AirflowGenerator.generate_dataflow(
- cluster=self.config.cluster,
+ config=self.config,
  dag=dag,
- capture_tags=self.config.capture_tags_info,
- capture_owner=self.config.capture_ownership_info,
  )
  dataflow.emit(self.emitter, callback=self._make_emit_callback())
 

metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/datahub_plugin_v22.py

@@ -106,10 +106,8 @@ def datahub_task_status_callback(context, status):
  )
 
  dataflow = AirflowGenerator.generate_dataflow(
- cluster=config.cluster,
+ config=config,
  dag=dag,
- capture_tags=config.capture_tags_info,
- capture_owner=config.capture_ownership_info,
  )
  task.log.info(f"Emitting Datahub Dataflow: {dataflow}")
  dataflow.emit(emitter, callback=_make_emit_callback(task.log))
@@ -139,13 +137,12 @@ def datahub_task_status_callback(context, status):
  if config.capture_executions:
  dpi = AirflowGenerator.run_datajob(
  emitter=emitter,
- cluster=config.cluster,
+ config=config,
  ti=ti,
  dag=dag,
  dag_run=context["dag_run"],
  datajob=datajob,
  start_timestamp_millis=int(ti.start_date.timestamp() * 1000),
- config=config,
  )
 
  task.log.info(f"Emitted Start Datahub Dataprocess Instance: {dpi}")
@@ -207,13 +204,12 @@ def datahub_pre_execution(context):
  if config.capture_executions:
  dpi = AirflowGenerator.run_datajob(
  emitter=emitter,
- cluster=config.cluster,
+ config=config,
  ti=ti,
  dag=dag,
  dag_run=context["dag_run"],
  datajob=datajob,
  start_timestamp_millis=int(ti.start_date.timestamp() * 1000),
- config=config,
  )
 
  task.log.info(f"Emitting Datahub Dataprocess Instance: {dpi}")

metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/lineage/_lineage_core.py

@@ -37,10 +37,8 @@ def send_lineage_to_datahub(
  emitter = hook.make_emitter()
 
  dataflow = AirflowGenerator.generate_dataflow(
- cluster=config.cluster,
+ config=config,
  dag=dag,
- capture_tags=config.capture_tags_info,
- capture_owner=config.capture_ownership_info,
  )
  dataflow.emit(emitter)
  operator.log.info(f"Emitted from Lineage: {dataflow}")
@@ -68,7 +66,7 @@ def send_lineage_to_datahub(
 
  dpi = AirflowGenerator.run_datajob(
  emitter=emitter,
- cluster=config.cluster,
+ config=config,
  ti=ti,
  dag=dag,
  dag_run=dag_run,