diff --git a/providers/google/src/airflow/providers/google/cloud/operators/gcs.py b/providers/google/src/airflow/providers/google/cloud/operators/gcs.py index b55709c309eea..51afcc631f1e4 100644 --- a/providers/google/src/airflow/providers/google/cloud/operators/gcs.py +++ b/providers/google/src/airflow/providers/google/cloud/operators/gcs.py @@ -603,8 +603,8 @@ def __init__( super().__init__(**kwargs) self.source_bucket = source_bucket self.source_object = source_object - self.destination_bucket = destination_bucket or self.source_bucket - self.destination_object = destination_object or self.source_object + self.destination_bucket = destination_bucket + self.destination_object = destination_object self.gcp_conn_id = gcp_conn_id self.transform_script = transform_script @@ -612,6 +612,8 @@ def __init__( self.impersonation_chain = impersonation_chain def execute(self, context: Context) -> None: + destination_bucket = self.destination_bucket or self.source_bucket + destination_object = self.destination_object or self.source_object hook = GCSHook(gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain) with NamedTemporaryFile() as source_file, NamedTemporaryFile() as destination_file: @@ -637,15 +639,15 @@ def execute(self, context: Context) -> None: self.log.info("Transformation succeeded. Output temporarily located at %s", destination_file.name) - self.log.info("Uploading file to %s as %s", self.destination_bucket, self.destination_object) + self.log.info("Uploading file to %s as %s", destination_bucket, destination_object) FileDetailsLink.persist( context=context, - uri=f"{self.destination_bucket}/{self.destination_object}", + uri=f"{destination_bucket}/{destination_object}", project_id=hook.project_id, ) hook.upload( - bucket_name=self.destination_bucket, - object_name=self.destination_object, + bucket_name=destination_bucket, + object_name=destination_object, filename=destination_file.name, ) @@ -653,13 +655,15 @@ def get_openlineage_facets_on_start(self): from airflow.providers.common.compat.openlineage.facet import Dataset from airflow.providers.openlineage.extractors import OperatorLineage + destination_bucket = self.destination_bucket or self.source_bucket + destination_object = self.destination_object or self.source_object input_dataset = Dataset( namespace=f"gs://{self.source_bucket}", name=self.source_object, ) output_dataset = Dataset( - namespace=f"gs://{self.destination_bucket}", - name=self.destination_object, + namespace=f"gs://{destination_bucket}", + name=destination_object, ) return OperatorLineage(inputs=[input_dataset], outputs=[output_dataset]) diff --git a/providers/google/tests/unit/google/cloud/operators/test_gcs.py b/providers/google/tests/unit/google/cloud/operators/test_gcs.py index 664308cd29188..e1647224920d9 100644 --- a/providers/google/tests/unit/google/cloud/operators/test_gcs.py +++ b/providers/google/tests/unit/google/cloud/operators/test_gcs.py @@ -336,6 +336,49 @@ def test_execute(self, mock_hook, mock_subprocess, mock_tempfile): filename=destination, ) + @mock.patch("airflow.providers.google.cloud.operators.gcs.NamedTemporaryFile") + @mock.patch("airflow.providers.google.cloud.operators.gcs.subprocess") + @mock.patch("airflow.providers.google.cloud.operators.gcs.GCSHook") + def test_execute_defaults_destination_to_source(self, mock_hook, mock_subprocess, mock_tempfile): + source_bucket = TEST_BUCKET + source_object = "test.txt" + transform_script = "script.py" + + source = "source" + destination = "destination" + + mock1 = mock.Mock() + mock2 = mock.Mock() + mock1.name = source + mock2.name = destination + + mock_tempfile.return_value.__enter__.side_effect = [mock1, mock2] + + mock_proc = mock.MagicMock() + mock_proc.returncode = 0 + mock_proc.stdout.readline = lambda: b"" + mock_proc.wait.return_value = None + mock_popen = mock.MagicMock() + mock_popen.return_value.__enter__.return_value = mock_proc + + mock_subprocess.Popen = mock_popen + mock_subprocess.PIPE = "pipe" + mock_subprocess.STDOUT = "stdout" + + op = GCSFileTransformOperator( + task_id=TASK_ID, + source_bucket=source_bucket, + source_object=source_object, + transform_script=transform_script, + ) + op.execute(context=mock.MagicMock()) + + mock_hook.return_value.upload.assert_called_with( + bucket_name=source_bucket, + object_name=source_object, + filename=destination, + ) + def test_get_openlineage_facets_on_start(self): expected_input = Dataset( namespace=f"gs://{TEST_BUCKET}", diff --git a/scripts/ci/prek/validate_operators_init_exemptions.txt b/scripts/ci/prek/validate_operators_init_exemptions.txt index 345248323f529..ff1f79c194e78 100644 --- a/scripts/ci/prek/validate_operators_init_exemptions.txt +++ b/scripts/ci/prek/validate_operators_init_exemptions.txt @@ -17,7 +17,6 @@ providers/google/src/airflow/providers/google/cloud/operators/cloud_storage_tran providers/google/src/airflow/providers/google/cloud/operators/dataproc.py::DataprocCreateClusterOperator providers/google/src/airflow/providers/google/cloud/operators/dataproc.py::DataprocSubmitJobOperator providers/google/src/airflow/providers/google/cloud/operators/functions.py::CloudFunctionDeployFunctionOperator -providers/google/src/airflow/providers/google/cloud/operators/gcs.py::GCSFileTransformOperator providers/google/src/airflow/providers/google/cloud/sensors/bigquery_dts.py::BigQueryDataTransferServiceTransferRunSensor providers/google/src/airflow/providers/google/cloud/sensors/cloud_composer.py::CloudComposerExternalTaskSensor providers/google/src/airflow/providers/google/cloud/transfers/azure_fileshare_to_gcs.py::AzureFileShareToGCSOperator