Extract to helper function

datahub-project · Jul 17, 2024 · 6d24978 · 6d24978
1 parent eb567cc
commit 6d24978
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/...a-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/client/airflow_generator.py b/...a-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/client/airflow_generator.py
@@ -127,6 +127,10 @@ def _get_dependencies(
         )
         return upstream_tasks
 
+    @staticmethod
+    def _extract_owners(dag: "DAG") -> list[str]:
+        return [owner.strip() for owner in dag.owner.split(",")]
+
     @staticmethod
     def generate_dataflow(
         config: DatahubLineageConfig,
@@ -175,7 +179,7 @@ def generate_dataflow(
         data_flow.url = f"{base_url}/tree?dag_id={dag.dag_id}"
 
         if config.capture_ownership_info and dag.owner:
-            owners = [owner.strip() for owner in dag.owner.split(",")]
+            owners = AirflowGenerator._extract_owners(dag)
             if config.capture_ownership_as_group:
                 data_flow.group_owners.update(owners)
             else:
@@ -283,7 +287,7 @@ def generate_datajob(
 
         if capture_owner and dag.owner:
             if config and config.capture_ownership_info:
-                owners = [owner.strip() for owner in dag.owner.split(",")]
+                owners = AirflowGenerator._extract_owners(dag)
                 if config.capture_ownership_as_group:
                     datajob.group_owners.update(owners)
                 else: