feat: add common_task_parameters support

Nike-Inc · Sep 15, 2024 · c92cab1 · c92cab1
1 parent b8e700b
commit c92cab1
Show file tree

Hide file tree

Showing 6 changed files with 18 additions and 5 deletions.
diff --git a/brickflow/codegen/databricks_bundle.py b/brickflow/codegen/databricks_bundle.py
@@ -555,7 +555,7 @@ def _build_native_spark_python_task(
  task_libraries: List[JobsTasksLibraries],
  task_settings: TaskSettings,
  depends_on: List[JobsTasksDependsOn],
- **_kwargs: Any,
+ **kwargs: Any,
  ) -> JobsTasks:
  try:
  spark_python_task: JobsTasksSparkPythonTask = task.task_func()
@@ -569,6 +569,15 @@ def _build_native_spark_python_task(
  spark_python_task.python_file = self.adjust_file_path(
  file_path=spark_python_task.python_file
  )
+ workflow: Optional[Workflow] = kwargs.get("workflow")
+ commom_task_parameters = workflow.common_task_parameters if workflow else None
+
+ if commom_task_parameters:
+ spark_python_task.parameters = spark_python_task.parameters or []
+ for k, v in commom_task_parameters.items():
+ if k not in spark_python_task.parameters:
+ spark_python_task.parameters.append(k)
+ spark_python_task.parameters.append(v)
 
  return JobsTasks(
  **task_settings.to_tf_dict(),

diff --git a/tests/codegen/expected_bundles/dev_bundle_monorepo.yml b/tests/codegen/expected_bundles/dev_bundle_monorepo.yml
@@ -124,7 +124,7 @@ targets:
  min_retry_interval_millis: null
  spark_python_task:
  python_file: ./products/test-project/spark/python/src/run_task.py
- parameters: ["--param1", "World!"]
+ parameters: ["--param1", "World!", "all_tasks1", "test", "all_tasks3", "123"]
  source: GIT
  retry_on_timeout: null
  task_key: spark_python_task_a

diff --git a/tests/codegen/expected_bundles/dev_bundle_polyrepo.yml b/tests/codegen/expected_bundles/dev_bundle_polyrepo.yml
@@ -124,7 +124,7 @@ targets:
  min_retry_interval_millis: null
  spark_python_task:
  python_file: ./products/test-project/spark/python/src/run_task.py
- parameters: ["--param1", "World!"]
+ parameters: ["--param1", "World!", "all_tasks1", "test", "all_tasks3", "123"]
  source: GIT
  retry_on_timeout: null
  task_key: spark_python_task_a

diff --git a/tests/codegen/expected_bundles/dev_bundle_polyrepo_with_auto_libs.yml b/tests/codegen/expected_bundles/dev_bundle_polyrepo_with_auto_libs.yml
@@ -209,7 +209,7 @@ targets:
  min_retry_interval_millis: null
  spark_python_task:
  python_file: ./products/test-project/spark/python/src/run_task.py
- parameters: ["--param1", "World!"]
+ parameters: ["--param1", "World!", "all_tasks1", "test", "all_tasks3", "123"]
  source: GIT
  retry_on_timeout: null
  task_key: spark_python_task_a

diff --git a/tests/codegen/expected_bundles/local_bundle.yml b/tests/codegen/expected_bundles/local_bundle.yml
@@ -247,6 +247,10 @@
  "parameters":
  - "--param1"
  - "World!"
+ - "all_tasks1"
+ - "test"
+ - "all_tasks3"
+ - "123"
  "task_key": "spark_python_task_a"
  "timeout_seconds": null
  - "depends_on": []

diff --git a/tests/codegen/expected_bundles/local_bundle_prefix_suffix.yml b/tests/codegen/expected_bundles/local_bundle_prefix_suffix.yml
@@ -121,7 +121,7 @@ targets:
  min_retry_interval_millis: null
  spark_python_task:
  python_file: Workspace/Users/${workspace.current_user.userName}/.brickflow_bundles/test-project/local/files/spark/python/src/run_task.py
- parameters: ["--param1", "World!"]
+ parameters: ["--param1", "World!", "all_tasks1", "test", "all_tasks3", "123"]
  source: WORKSPACE
  retry_on_timeout: null
  task_key: spark_python_task_a