Add overrides to template fields of Google Cloud Run Jobs Execute Ope…

…rator (#36133) * feat: add overrides in the template fields * feat: add test on overrides in the template fields
apache · Dec 11, 2023 · 3dddfb4 · 3dddfb4
1 parent df23df5
commit 3dddfb4
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 2 deletions.
diff --git a/airflow/providers/google/cloud/operators/cloud_run.py b/airflow/providers/google/cloud/operators/cloud_run.py
@@ -264,7 +264,7 @@ class CloudRunExecuteJobOperator(GoogleCloudBaseOperator):
     :param deferrable: Run operator in the deferrable mode
     """
 
-    template_fields = ("project_id", "region", "gcp_conn_id", "impersonation_chain", "job_name")
+    template_fields = ("project_id", "region", "gcp_conn_id", "impersonation_chain", "job_name", "overrides")
 
     def __init__(
         self,

diff --git a/tests/providers/google/cloud/operators/test_cloud_run.py b/tests/providers/google/cloud/operators/test_cloud_run.py
@@ -40,6 +40,12 @@
 PROJECT_ID = "testproject"
 REGION = "us-central1"
 JOB_NAME = "jobname"
+OVERRIDES = {
+    "container_overrides": [{"args": ["python", "main.py"]}],
+    "task_count": 1,
+    "timeout": "60s",
+}
+
 JOB = Job()
 JOB.name = JOB_NAME
 
@@ -78,11 +84,12 @@ def test_create(self, hook_mock):
 class TestCloudRunExecuteJobOperator:
     def test_template_fields(self):
         operator = CloudRunExecuteJobOperator(
-            task_id=TASK_ID, project_id=PROJECT_ID, region=REGION, job_name=JOB_NAME
+            task_id=TASK_ID, project_id=PROJECT_ID, region=REGION, job_name=JOB_NAME, overrides=OVERRIDES
         )
 
         _assert_common_template_fields(operator.template_fields)
         assert "job_name" in operator.template_fields
+        assert "overrides" in operator.template_fields
 
     @mock.patch(CLOUD_RUN_HOOK_PATH)
     def test_execute_success(self, hook_mock):