Add missing template fields to DatabricksWorkflowTaskGroup

providers/databricks/src/airflow/providers/databricks/operators/databricks_workflow.py

@@ -84,15 +84,31 @@ class _CreateDatabricksWorkflowOperator(BaseOperator):
     :param existing_clusters: A list of existing clusters to use for the workflow.
     :param extra_job_params: A dictionary of extra properties which will override the default Databricks
         Workflow Job definitions.
+    :param jar_params: A list of jar parameters to pass to the workflow. These parameters will be passed to
+        all jar tasks in the workflow.
     :param job_clusters: A list of job clusters to use for the workflow.
     :param max_concurrent_runs: The maximum number of concurrent runs for the workflow.
     :param notebook_params: A dictionary of notebook parameters to pass to the workflow. These parameters
         will be passed to all notebooks in the workflow.
+    :param python_params: A list of python parameters to pass to the workflow. These parameters will be
+        passed to all python tasks in the workflow.
+    :param spark_submit_params: A list of spark submit parameters to pass to the workflow. These parameters
+        will be passed to all spark submit tasks.
     :param tasks_to_convert: A dict of tasks to convert to a Databricks workflow. This list can also be
         populated after instantiation using the `add_task` method.
     """
 
-    template_fields = ("notebook_params", "job_clusters")
+    template_fields = (
+        "databricks_conn_id",
+        "existing_clusters",
+        "extra_job_params",
+        "jar_params",
+        "job_clusters",
+        "max_concurrent_runs",
+        "notebook_params",
+        "python_params",
+        "spark_submit_params",
+    )
     caller = "_CreateDatabricksWorkflowOperator"
     # Conditionally set operator_extra_links based on Airflow version
     if AIRFLOW_V_3_0_PLUS:
@@ -111,18 +127,24 @@ def __init__(
         databricks_conn_id: str,
         existing_clusters: list[str] | None = None,
         extra_job_params: dict[str, Any] | None = None,
+        jar_params: list[str] | None = None,
         job_clusters: list[dict[str, object]] | None = None,
         max_concurrent_runs: int = 1,
         notebook_params: dict | None = None,
+        python_params: list | None = None,
+        spark_submit_params: list | None = None,
         tasks_to_convert: dict[str, BaseOperator] | None = None,
         **kwargs,
     ):
         self.databricks_conn_id = databricks_conn_id
         self.existing_clusters = existing_clusters or []
         self.extra_job_params = extra_job_params or {}
+        self.jar_params = jar_params or []
         self.job_clusters = job_clusters or []
         self.max_concurrent_runs = max_concurrent_runs
         self.notebook_params = notebook_params or {}
+        self.python_params = python_params or []
+        self.spark_submit_params = spark_submit_params or []
         self.tasks_to_convert = tasks_to_convert or {}
         self.relevant_upstreams = [task_id]
         self.workflow_run_metadata: WorkflowRunMetadata | None = None
@@ -213,10 +235,10 @@ def execute(self, context: Context) -> Any:
         run_id = self._hook.run_now(
             {
                 "job_id": job_id,
-                "jar_params": self.task_group.jar_params,
+                "jar_params": self.jar_params,
                 "notebook_params": self.notebook_params,
-                "python_params": self.task_group.python_params,
-                "spark_submit_params": self.task_group.spark_submit_params,
+                "python_params": self.python_params,
+                "spark_submit_params": self.spark_submit_params,
             }
         )
 
@@ -336,9 +358,12 @@ def __exit__(
             databricks_conn_id=self.databricks_conn_id,
             existing_clusters=self.existing_clusters,
             extra_job_params=self.extra_job_params,
+            jar_params=self.jar_params,
             job_clusters=self.job_clusters,
             max_concurrent_runs=self.max_concurrent_runs,
             notebook_params=self.notebook_params,
+            python_params=self.python_params,
+            spark_submit_params=self.spark_submit_params,
         )
 
         for task in tasks:

providers/databricks/tests/unit/databricks/operators/test_databricks_workflow.py

@@ -140,9 +140,6 @@ def test_execute(mock_databricks_hook, context, mock_task_group):
     """Test that _CreateDatabricksWorkflowOperator.execute runs the task group."""
     operator = _CreateDatabricksWorkflowOperator(task_id="test_task", databricks_conn_id="databricks_default")
     operator.task_group = mock_task_group
-    mock_task_group.jar_params = {}
-    mock_task_group.python_params = {}
-    mock_task_group.spark_submit_params = {}
 
     mock_hook_instance = mock_databricks_hook.return_value
     mock_hook_instance.run_now.return_value = 789
@@ -219,9 +216,12 @@ def test_task_group_exit_creates_operator(mock_databricks_workflow_operator):
         databricks_conn_id="databricks_conn",
         existing_clusters=[],
         extra_job_params={},
+        jar_params=[],
         job_clusters=[],
         max_concurrent_runs=1,
         notebook_params={},
+        python_params=[],
+        spark_submit_params=[],
     )