Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Fix assignment of template field in __init__ in CloudDataTransferServiceCreateJobOperator #36909

Merged
merged 4 commits into from
Feb 4, 2024
Merged

Fix assignment of template field in __init__ in CloudDataTransferServiceCreateJobOperator #36909

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
aac09fa
fix initialization of templated field in constructor
romsharon98 Jan 19, 2024
9a5a6c4
remove file from exclude
romsharon98 Jan 20, 2024
c6e9868
add test for templated field
romsharon98 Jan 21, 2024
2896623
change body to be realistic
romsharon98 Feb 4, 2024
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
1 change: 0 additions & 1 deletion .pre-commit-config.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -318,7 +318,6 @@ repos:
^airflow\/providers\/google\/cloud\/operators\/bigquery\.py$|
^airflow\/providers\/amazon\/aws\/transfers\/gcs_to_s3\.py$|
^airflow\/providers\/databricks\/operators\/databricks\.py$|
^airflow\/providers\/google\/cloud\/operators\/cloud_storage_transfer_service\.py$|
^airflow\/providers\/google\/cloud\/transfers\/bigquery_to_mysql\.py$|
^airflow\/providers\/amazon\/aws\/transfers\/redshift_to_s3\.py$|
^airflow\/providers\/google\/cloud\/operators\/compute\.py$|
Expand Down
4 changes: 3 additions & 1 deletion airflow/providers/google/cloud/operators/cloud_storage_transfer_service.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -236,7 +236,9 @@ def __init__(
**kwargs,
) -> None:
super().__init__(**kwargs)
self.body = deepcopy(body)
self.body = body
if isinstance(self.body, dict):
self.body = deepcopy(body)
hussein-awala marked this conversation as resolved.
Show resolved Hide resolved
self.aws_conn_id = aws_conn_id
self.gcp_conn_id = gcp_conn_id
self.api_version = api_version
Expand Down
16 changes: 12 additions & 4 deletions tests/providers/google/cloud/operators/test_cloud_storage_transfer_service.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -114,6 +114,10 @@
SCHEDULE: SCHEDULE_DICT,
TRANSFER_SPEC: {GCS_DATA_SINK: {BUCKET_NAME: GCS_BUCKET_NAME, PATH: DESTINATION_PATH}},
}
VALID_TRANSFER_JOB_JINJA = deepcopy(VALID_TRANSFER_JOB_BASE)
VALID_TRANSFER_JOB_JINJA[NAME] = "{{ dag.dag_id }}"
VALID_TRANSFER_JOB_JINJA_RENDERED = deepcopy(VALID_TRANSFER_JOB_JINJA)
VALID_TRANSFER_JOB_JINJA_RENDERED[NAME] = "TestGcpStorageTransferJobCreateOperator"
VALID_TRANSFER_JOB_GCS = deepcopy(VALID_TRANSFER_JOB_BASE)
VALID_TRANSFER_JOB_GCS[TRANSFER_SPEC].update(deepcopy(SOURCE_GCS))
VALID_TRANSFER_JOB_AWS = deepcopy(VALID_TRANSFER_JOB_BASE)
Expand Down Expand Up @@ -324,21 +328,25 @@ def test_job_create_multiple(self, aws_hook, gcp_hook):
# (could be anything else) just to test if the templating works for all
# fields
@pytest.mark.db_test
@pytest.mark.parametrize(
"body, excepted",
[(VALID_TRANSFER_JOB_JINJA, VALID_TRANSFER_JOB_JINJA_RENDERED)],
)
@mock.patch(
"airflow.providers.google.cloud.operators.cloud_storage_transfer_service.CloudDataTransferServiceHook"
)
def test_templates(self, _, create_task_instance_of_operator):
dag_id = "TestGcpStorageTransferJobCreateOperator_test_templates"
def test_templates(self, _, create_task_instance_of_operator, body, excepted):
dag_id = "TestGcpStorageTransferJobCreateOperator"
ti = create_task_instance_of_operator(
CloudDataTransferServiceCreateJobOperator,
dag_id=dag_id,
body={"description": "{{ dag.dag_id }}"},
body=body,
gcp_conn_id="{{ dag.dag_id }}",
aws_conn_id="{{ dag.dag_id }}",
task_id="task-id",
)
ti.render_templates()
assert dag_id == getattr(ti.task, "body")[DESCRIPTION]
assert excepted == getattr(ti.task, "body")
assert dag_id == getattr(ti.task, "gcp_conn_id")
assert dag_id == getattr(ti.task, "aws_conn_id")

Expand Down
Loading