Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -78,17 +78,20 @@ def __init__(
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
self.impersonation_chain = impersonation_chain
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) before rendering runs.
# 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(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,7 @@ def test_execute(self, mock):
operator = CloudBatchSubmitJobOperator(
task_id=TASK_ID, project_id=PROJECT_ID, region=REGION, job_name=JOB_NAME, job=JOB
)
operator.prepare_template()

completed_job = operator.execute(context=mock.MagicMock())

Expand Down Expand Up @@ -119,6 +120,18 @@ class TestCloudBatchSubmitJobOperatorTemplating:
def test_template_fields_includes_job(self):
assert "job" in CloudBatchSubmitJobOperator.template_fields

def test_protobuf_job_is_normalized_by_prepare_template_not_init(self):
job = batch_v1.Job.from_json(json.dumps(_job_dict_with_template()))
operator = CloudBatchSubmitJobOperator(
task_id=TASK_ID, project_id=PROJECT_ID, region=REGION, job_name=JOB_NAME, job=job
)

assert isinstance(operator.job, batch_v1.Job)

operator.prepare_template()

assert isinstance(operator.job, dict)

@pytest.mark.db_test
@pytest.mark.parametrize(
"job_input_factory",
Expand Down
1 change: 0 additions & 1 deletion scripts/ci/prek/validate_operators_init_exemptions.txt
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down