diff --git a/providers/google/src/airflow/providers/google/cloud/transfers/azure_blob_to_gcs.py b/providers/google/src/airflow/providers/google/cloud/transfers/azure_blob_to_gcs.py index 1ace104c28297..536fa85375edc 100644 --- a/providers/google/src/airflow/providers/google/cloud/transfers/azure_blob_to_gcs.py +++ b/providers/google/src/airflow/providers/google/cloud/transfers/azure_blob_to_gcs.py @@ -59,6 +59,10 @@ class AzureBlobStorageToGCSOperator(BaseOperator): If set as a sequence, the identities from the list must grant Service Account Token Creator IAM role to the directly preceding identity, with first account from the list granting this role to the originating account. + :param unwrap_single: If True, unwrap a single-element result list into a plain string value + for backward compatibility. If False, always return a list of GCS URIs. + If not explicitly provided, defaults to True and emits a FutureWarning that + the default will change to False in a future release. """ def __init__( @@ -73,6 +77,7 @@ def __init__( filename: str, gzip: bool, impersonation_chain: str | Sequence[str] | None = None, + unwrap_single: bool | None = None, **kwargs, ) -> None: super().__init__(**kwargs) @@ -85,6 +90,18 @@ def __init__( self.filename = filename self.gzip = gzip self.impersonation_chain = impersonation_chain + if unwrap_single is None: + self.unwrap_single = True + import warnings + + warnings.warn( + "The default value of `unwrap_single` will change from True to False in a future release. " + "Please set `unwrap_single` explicitly to avoid this warning.", + FutureWarning, + stacklevel=2, + ) + else: + self.unwrap_single = unwrap_single template_fields: Sequence[str] = ( "blob_name", @@ -94,7 +111,7 @@ def __init__( "filename", ) - def execute(self, context: Context) -> str: + def execute(self, context: Context) -> str | list[str]: azure_hook = WasbHook(wasb_conn_id=self.wasb_conn_id) gcs_hook = GCSHook( gcp_conn_id=self.gcp_conn_id, @@ -122,7 +139,10 @@ def execute(self, context: Context) -> str: self.blob_name, self.bucket_name, ) - return f"gs://{self.bucket_name}/{self.object_name}" + gcs_uri = f"gs://{self.bucket_name}/{self.object_name}" + if self.unwrap_single: + return gcs_uri + return [gcs_uri] def get_openlineage_facets_on_start(self): from airflow.providers.common.compat.openlineage.facet import Dataset diff --git a/providers/google/src/airflow/providers/google/cloud/transfers/azure_fileshare_to_gcs.py b/providers/google/src/airflow/providers/google/cloud/transfers/azure_fileshare_to_gcs.py index ad362c508ba04..993ce091d5010 100644 --- a/providers/google/src/airflow/providers/google/cloud/transfers/azure_fileshare_to_gcs.py +++ b/providers/google/src/airflow/providers/google/cloud/transfers/azure_fileshare_to_gcs.py @@ -66,6 +66,8 @@ class AzureFileShareToGCSOperator(BaseOperator): If set as a sequence, the identities from the list must grant Service Account Token Creator IAM role to the directly preceding identity, with first account from the list granting this role to the originating account (templated). + :param return_gcs_uris: If True, return a list of GCS URIs. If False (default), return the legacy + list of Azure FileShare filenames and emit a deprecation warning. Note that ``share_name``, ``directory_path``, ``prefix``, and ``dest_gcs`` are templated, so you can use variables in them if you wish. @@ -92,6 +94,7 @@ def __init__( replace: bool = False, gzip: bool = False, google_impersonation_chain: str | Sequence[str] | None = None, + return_gcs_uris: bool | None = None, **kwargs, ): super().__init__(**kwargs) @@ -113,6 +116,16 @@ def __init__( self.replace = replace self.gzip = gzip self.google_impersonation_chain = google_impersonation_chain + if return_gcs_uris is None: + self.return_gcs_uris = False + warnings.warn( + "Returning a list of Azure FileShare filenames from AzureFileShareToGCSOperator is deprecated and " + "will change to list[str] of GCS URIs in a future release. Set return_gcs_uris=True to opt in.", + FutureWarning, + stacklevel=2, + ) + else: + self.return_gcs_uris = return_gcs_uris def _check_inputs(self) -> None: if self.dest_gcs and not gcs_object_is_directory(self.dest_gcs): @@ -125,7 +138,7 @@ def _check_inputs(self) -> None: 'The destination Google Cloud Storage path must end with a slash "/" or be empty.' ) - def execute(self, context: Context): + def execute(self, context: Context) -> list[str]: self._check_inputs() azure_fileshare_hook = AzureFileShareHook( share_name=self.share_name, @@ -162,6 +175,7 @@ def execute(self, context: Context): files = list(set(files) - set(existing_files)) + uploaded_gcs_uris: list[str] = [] if files: self.log.info("%s files are going to be synced.", len(files)) if self.directory_path is None: @@ -181,9 +195,12 @@ def execute(self, context: Context): # enforced at instantiation time dest_gcs_object = dest_gcs_object_prefix + file gcs_hook.upload(dest_gcs_bucket, dest_gcs_object, temp_file.name, gzip=self.gzip) + uploaded_gcs_uris.append(f"gs://{dest_gcs_bucket}/{dest_gcs_object}") self.log.info("All done, uploaded %d files to Google Cloud Storage.", len(files)) else: self.log.info("There are no new files to sync. Have a nice day!") self.log.info("In sync, no files needed to be uploaded to Google Cloud Storage") - return files + if not self.return_gcs_uris: + return files + return uploaded_gcs_uris diff --git a/providers/google/tests/system/google/cloud/azure/example_azure_fileshare_to_gcs.py b/providers/google/tests/system/google/cloud/azure/example_azure_fileshare_to_gcs.py index d03c493c8dd9c..c62db075d66ab 100644 --- a/providers/google/tests/system/google/cloud/azure/example_azure_fileshare_to_gcs.py +++ b/providers/google/tests/system/google/cloud/azure/example_azure_fileshare_to_gcs.py @@ -62,7 +62,7 @@ sync_azure_files_with_gcs = AzureFileShareToGCSOperator( task_id="sync_azure_files_with_gcs", share_name=AZURE_SHARE_NAME, - dest_gcs=BUCKET_NAME, + dest_gcs=f"gs://{BUCKET_NAME}/", directory_path=AZURE_DIRECTORY_PATH, replace=False, gzip=True, diff --git a/providers/google/tests/unit/google/cloud/transfers/test_azure_blob_to_gcs.py b/providers/google/tests/unit/google/cloud/transfers/test_azure_blob_to_gcs.py index aa575eb9df40b..ae3300af988b9 100644 --- a/providers/google/tests/unit/google/cloud/transfers/test_azure_blob_to_gcs.py +++ b/providers/google/tests/unit/google/cloud/transfers/test_azure_blob_to_gcs.py @@ -18,8 +18,12 @@ from unittest import mock +import pytest + from airflow.providers.google.cloud.transfers.azure_blob_to_gcs import AzureBlobStorageToGCSOperator +pytestmark = pytest.mark.filterwarnings("ignore::FutureWarning") + WASB_CONN_ID = "wasb_default" GCP_CONN_ID = "google_cloud_default" BLOB_NAME = "azure_blob" @@ -57,10 +61,17 @@ def test_init(self): assert operator.impersonation_chain == IMPERSONATION_CHAIN assert operator.task_id == TASK_ID + @pytest.mark.parametrize( + ("unwrap_single", "expected"), + [ + (True, f"gs://{BUCKET_NAME}/{OBJECT_NAME}"), + (False, [f"gs://{BUCKET_NAME}/{OBJECT_NAME}"]), + ], + ) @mock.patch("airflow.providers.google.cloud.transfers.azure_blob_to_gcs.WasbHook") @mock.patch("airflow.providers.google.cloud.transfers.azure_blob_to_gcs.GCSHook") @mock.patch("airflow.providers.google.cloud.transfers.azure_blob_to_gcs.tempfile") - def test_execute(self, mock_temp, mock_hook_gcs, mock_hook_wasb): + def test_execute(self, mock_temp, mock_hook_gcs, mock_hook_wasb, unwrap_single, expected): op = AzureBlobStorageToGCSOperator( wasb_conn_id=WASB_CONN_ID, gcp_conn_id=GCP_CONN_ID, @@ -71,10 +82,11 @@ def test_execute(self, mock_temp, mock_hook_gcs, mock_hook_wasb): filename=FILENAME, gzip=GZIP, impersonation_chain=IMPERSONATION_CHAIN, + unwrap_single=unwrap_single, task_id=TASK_ID, ) - op.execute(context=None) + result = op.execute(context=None) mock_hook_wasb.assert_called_once_with(wasb_conn_id=WASB_CONN_ID) mock_hook_wasb.return_value.get_file.assert_called_once_with( @@ -91,6 +103,7 @@ def test_execute(self, mock_temp, mock_hook_gcs, mock_hook_wasb): gzip=GZIP, filename=mock_temp.NamedTemporaryFile.return_value.__enter__.return_value.name, ) + assert result == expected @mock.patch("airflow.providers.google.cloud.transfers.azure_blob_to_gcs.WasbHook") def test_execute_single_file_transfer_openlineage(self, mock_hook_wasb): diff --git a/providers/google/tests/unit/google/cloud/transfers/test_azure_fileshare_to_gcs.py b/providers/google/tests/unit/google/cloud/transfers/test_azure_fileshare_to_gcs.py index 39b5bb62f3859..a70d7cc709b7d 100644 --- a/providers/google/tests/unit/google/cloud/transfers/test_azure_fileshare_to_gcs.py +++ b/providers/google/tests/unit/google/cloud/transfers/test_azure_fileshare_to_gcs.py @@ -18,8 +18,12 @@ from unittest import mock +import pytest + from airflow.providers.google.cloud.transfers.azure_fileshare_to_gcs import AzureFileShareToGCSOperator +pytestmark = pytest.mark.filterwarnings("ignore::FutureWarning") + TASK_ID = "test-azure-fileshare-to-gcs" AZURE_FILESHARE_SHARE = "test-share" AZURE_FILESHARE_DIRECTORY_PATH = "/path/to/dir" @@ -52,9 +56,10 @@ def test_init(self): assert operator.dest_gcs == GCS_PATH_PREFIX assert operator.google_impersonation_chain == IMPERSONATION_CHAIN + @pytest.mark.parametrize("return_gcs_uris", [True, False]) @mock.patch("airflow.providers.google.cloud.transfers.azure_fileshare_to_gcs.AzureFileShareHook") @mock.patch("airflow.providers.google.cloud.transfers.azure_fileshare_to_gcs.GCSHook") - def test_execute(self, gcs_mock_hook, azure_fileshare_mock_hook): + def test_execute(self, gcs_mock_hook, azure_fileshare_mock_hook, return_gcs_uris): """Test the execute function when the run is successful.""" operator = AzureFileShareToGCSOperator( @@ -65,6 +70,7 @@ def test_execute(self, gcs_mock_hook, azure_fileshare_mock_hook): gcp_conn_id=GCS_CONN_ID, dest_gcs=GCS_PATH_PREFIX, google_impersonation_chain=IMPERSONATION_CHAIN, + return_gcs_uris=return_gcs_uris, ) azure_fileshare_mock_hook.return_value.list_files.return_value = MOCK_FILES @@ -87,7 +93,12 @@ def test_execute(self, gcs_mock_hook, azure_fileshare_mock_hook): impersonation_chain=IMPERSONATION_CHAIN, ) - assert sorted(MOCK_FILES) == sorted(uploaded_files) + expected_files = ( + [f"gs://gcs-bucket/data/{file_name}" for file_name in MOCK_FILES] + if return_gcs_uris + else MOCK_FILES + ) + assert sorted(expected_files) == sorted(uploaded_files) @mock.patch("airflow.providers.google.cloud.transfers.azure_fileshare_to_gcs.AzureFileShareHook") @mock.patch("airflow.providers.google.cloud.transfers.azure_fileshare_to_gcs.GCSHook")