This is an automated email from the ASF dual-hosted git repository.
potiuk pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new ed68491d8be Prepare Cloud Batch jobs before template rendering (#70835)
ed68491d8be is described below
commit ed68491d8be71d18e1100db9708a852b7eea360c
Author: Shivam Rastogi <[email protected]>
AuthorDate: Tue Aug 18 12:41:44 2026 -0700
Prepare Cloud Batch jobs before template rendering (#70835)
---
.../google/cloud/operators/cloud_batch.py | 10 ++--
.../google/cloud/operators/test_cloud_batch.py | 63 ++++++++++++++++++++--
.../ci/prek/validate_operators_init_exemptions.txt | 1 -
3 files changed, 66 insertions(+), 8 deletions(-)
diff --git
a/providers/google/src/airflow/providers/google/cloud/operators/cloud_batch.py
b/providers/google/src/airflow/providers/google/cloud/operators/cloud_batch.py
index 914559070b7..0e6d8a9491a 100644
---
a/providers/google/src/airflow/providers/google/cloud/operators/cloud_batch.py
+++
b/providers/google/src/airflow/providers/google/cloud/operators/cloud_batch.py
@@ -78,10 +78,6 @@ class CloudBatchSubmitJobOperator(GoogleCloudBaseOperator):
self.region = region
self.job_name = job_name
self.job = job
- # Normalize Job protobuf to dict so Airflow's template renderer can
descend
- # into nested fields (e.g. runnable.container.commands). See #37217.
- if isinstance(job, Job):
- self.job = Job.to_dict(job)
self.polling_period_seconds = polling_period_seconds
self.timeout_seconds = timeout_seconds
self.gcp_conn_id = gcp_conn_id
@@ -89,6 +85,12 @@ class CloudBatchSubmitJobOperator(GoogleCloudBaseOperator):
self.deferrable = deferrable
self.polling_period_seconds = polling_period_seconds
+ def prepare_template(self) -> None:
+ # Normalize Job protobuf to dict so Airflow's template renderer can
descend
+ # into nested fields (e.g. runnable.container.commands). See #37217.
+ if isinstance(self.job, Job):
+ self.job = Job.to_dict(self.job)
+
def execute(self, context: Context):
hook: CloudBatchHook = CloudBatchHook(self.gcp_conn_id,
self.impersonation_chain)
job = hook.submit_batch_job(
diff --git
a/providers/google/tests/unit/google/cloud/operators/test_cloud_batch.py
b/providers/google/tests/unit/google/cloud/operators/test_cloud_batch.py
index 1b688bb65df..b231907ab59 100644
--- a/providers/google/tests/unit/google/cloud/operators/test_cloud_batch.py
+++ b/providers/google/tests/unit/google/cloud/operators/test_cloud_batch.py
@@ -32,6 +32,8 @@ from airflow.providers.google.cloud.operators.cloud_batch
import (
CloudBatchSubmitJobOperator,
)
+from tests_common.test_utils.compat import DagSerialization
+
CLOUD_BATCH_HOOK_PATH =
"airflow.providers.google.cloud.operators.cloud_batch.CloudBatchHook"
TASK_ID = "test"
PROJECT_ID = "testproject"
@@ -54,7 +56,7 @@ class TestCloudBatchSubmitJobOperator:
assert completed_job["name"] == JOB_NAME
mock.return_value.submit_batch_job.assert_called_with(
- job_name=JOB_NAME, job=batch_v1.Job.to_dict(JOB), region=REGION,
project_id=PROJECT_ID
+ job_name=JOB_NAME, job=JOB, region=REGION, project_id=PROJECT_ID
)
mock.return_value.wait_for_job.assert_called()
@@ -119,6 +121,32 @@ class TestCloudBatchSubmitJobOperatorTemplating:
def test_template_fields_includes_job(self):
assert "job" in CloudBatchSubmitJobOperator.template_fields
+ @pytest.mark.parametrize(
+ ("job_input_factory", "job_input_type"),
+ [
+ pytest.param(lambda d: d, dict, id="dict"),
+ pytest.param(lambda d: batch_v1.Job.from_json(json.dumps(d)),
batch_v1.Job, id="protobuf-Job"),
+ ],
+ )
+ def test_job_is_unchanged_until_template_preparation(self,
job_input_factory, job_input_type):
+ job = job_input_factory(_job_dict_with_template())
+ operator = CloudBatchSubmitJobOperator(
+ task_id=TASK_ID,
+ project_id=PROJECT_ID,
+ region=REGION,
+ job_name=JOB_NAME,
+ job=job,
+ )
+
+ assert operator.job is job
+ assert isinstance(operator.job, job_input_type)
+
+ operator.resolve_template_files()
+ assert isinstance(operator.job, dict)
+ prepared_job = operator.job
+ operator.resolve_template_files()
+ assert operator.job is prepared_job
+
@pytest.mark.db_test
@pytest.mark.parametrize(
"job_input_factory",
@@ -127,7 +155,10 @@ class TestCloudBatchSubmitJobOperatorTemplating:
pytest.param(lambda d: batch_v1.Job.from_json(json.dumps(d)),
id="protobuf-Job"),
],
)
- def test_jinja_in_job_commands_is_rendered(self,
create_task_instance_of_operator, job_input_factory):
+ def test_jinja_in_job_commands_is_rendered(
+ self, dag_maker, create_task_instance_of_operator, job_input_factory
+ ):
+ job = job_input_factory(_job_dict_with_template())
ti = create_task_instance_of_operator(
CloudBatchSubmitJobOperator,
dag_id="test_cloud_batch_render",
@@ -135,9 +166,12 @@ class TestCloudBatchSubmitJobOperatorTemplating:
project_id=PROJECT_ID,
region=REGION,
job_name=JOB_NAME,
- job=job_input_factory(_job_dict_with_template()),
+ job=job,
logical_date=datetime(2026, 1, 15),
)
+ operator = dag_maker.dag.get_task(TASK_ID)
+ assert isinstance(operator.job, dict)
+
task = ti.render_templates()
assert isinstance(task.job, dict)
@@ -146,6 +180,29 @@ class TestCloudBatchSubmitJobOperatorTemplating:
# dag_maker's default run_id is "test"; the point is {{ run_id }} got
substituted at all.
assert task.job["labels"]["run_id"] == "test"
+ @pytest.mark.db_test
+ @pytest.mark.need_serialized_dag
+ def test_protobuf_job_is_prepared_before_serialization(self, dag_maker):
+ job = batch_v1.Job.from_json(json.dumps(_job_dict_with_template()))
+
+ with dag_maker(dag_id="test_cloud_batch_serialization"):
+ operator = CloudBatchSubmitJobOperator(
+ task_id=TASK_ID,
+ project_id=PROJECT_ID,
+ region=REGION,
+ job_name=JOB_NAME,
+ job=job,
+ )
+ assert operator.job is job
+ dag_maker.dag.resolve_template_files()
+
+ serialized_dag =
DagSerialization.deserialize_dag(dag_maker.get_serialized_data()["dag"])
+ serialized_job = serialized_dag.get_task(TASK_ID).job
+
+ assert isinstance(serialized_job, dict)
+ command =
serialized_job["task_groups"][0]["task_spec"]["runnables"][0]["container"]["commands"][1]
+ assert command == "echo {{ ds }}"
+
class TestCloudBatchDeleteJobOperator:
@mock.patch(CLOUD_BATCH_HOOK_PATH)
diff --git a/scripts/ci/prek/validate_operators_init_exemptions.txt
b/scripts/ci/prek/validate_operators_init_exemptions.txt
index b7182ebdda3..73dd45c9731 100644
--- a/scripts/ci/prek/validate_operators_init_exemptions.txt
+++ b/scripts/ci/prek/validate_operators_init_exemptions.txt
@@ -9,7 +9,6 @@
providers/amazon/src/airflow/providers/amazon/aws/operators/neptune.py::NeptuneS
providers/amazon/src/airflow/providers/amazon/aws/operators/neptune.py::NeptuneStopDbClusterOperator
providers/amazon/src/airflow/providers/amazon/aws/transfers/gcs_to_s3.py::GCSToS3Operator
providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod.py::KubernetesPodOperator
-providers/google/src/airflow/providers/google/cloud/operators/cloud_batch.py::CloudBatchSubmitJobOperator
providers/google/src/airflow/providers/google/cloud/operators/cloud_build.py::CloudBuildCreateBuildOperator
providers/google/src/airflow/providers/google/cloud/operators/cloud_storage_transfer_service.py::CloudDataTransferServiceCreateJobOperator
providers/google/src/airflow/providers/google/cloud/operators/dataproc.py::DataprocCreateClusterOperator
