mobuchowski commented on code in PR #72643:
URL: https://github.com/apache/airflow/pull/72643#discussion_r3959491102


##########
providers/databricks/src/airflow/providers/databricks/operators/databricks.py:
##########
@@ -932,12 +946,30 @@ def _prepare_submit_json(self, context: Context) -> 
dict[str, Any]:
                 _inject_airflow_params_into_task(json, params_dump)
 
         if self.openlineage_inject_parent_job_info or 
self.openlineage_inject_transport_info:
-            self.log.info("Automatic injection of OpenLineage information into 
Spark properties is enabled.")
+            self.log.info("Automatic injection of OpenLineage information is 
enabled.")
             json = 
self._inject_openlineage_properties_into_databricks_job(json, context)
 
         return cast("dict[str, Any]", normalise_json_content(json))
 
     def _inject_openlineage_properties_into_databricks_job(self, json: dict, 
context: Context) -> dict:
+        if self.openlineage_inject_parent_job_info:
+            try:
+                context_json = copy.deepcopy(json)
+                tasks = context_json.get("tasks")
+                if isinstance(tasks, list):
+                    for task in tasks:
+                        if isinstance(task, dict):
+                            
_inject_openlineage_context_into_task_parameters(task, context)
+                else:
+                    
_inject_openlineage_context_into_task_parameters(context_json, context)
+                json = context_json

Review Comment:
   I got: 
   
   When a runs/submit payload uses a for_each_task whose nested task is a 
notebook, wheel, SQL, or run-job task, this loop invokes the injector only on 
the outer wrapper. It never reaches for_each_task.task, so no 
OPENLINEAGE_CONTEXT is added; recurse into the nested task as the existing 
Spark-cluster traversal does.
   
   Can be done in a follow up I guess.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]

Reply via email to