Ejemplo n.º 1
0
    def test_should_copy_single_file(self, mock_named_temporary_file,
                                     mock_gdrive, mock_gcs_hook):
        type(mock_named_temporary_file.return_value.__enter__.return_value
             ).name = mock.PropertyMock(side_effect=["TMP1"])
        task = GCSToGoogleDriveOperator(
            task_id="copy_single_file",
            source_bucket="data",
            source_object="sales/sales-2017/january.avro",
            destination_object="copied_sales/2017/january-backup.avro",
        )

        task.execute(mock.MagicMock())

        mock_gcs_hook.assert_has_calls([
            mock.call(
                delegate_to=None,
                google_cloud_storage_conn_id="google_cloud_default",
                impersonation_chain=None,
            ),
            mock.call().download(bucket_name="data",
                                 filename="TMP1",
                                 object_name="sales/sales-2017/january.avro"),
        ])

        mock_gdrive.assert_has_calls([
            mock.call(
                delegate_to=None,
                gcp_conn_id="google_cloud_default",
                impersonation_chain=None,
            ),
            mock.call().upload_file(
                local_location="TMP1",
                remote_location="copied_sales/2017/january-backup.avro"),
        ])
Ejemplo n.º 2
0
 def test_should_raise_exception_on_multiple_wildcard(
         self, mock_named_temporary_file, mock_gdrive, mock_gcs_hook):
     task = GCSToGoogleDriveOperator(task_id="move_files",
                                     source_bucket="data",
                                     source_object="sales/*/*.avro",
                                     move_object=True)
     with pytest.raises(AirflowException, match="Only one wildcard"):
         task.execute(mock.MagicMock())
Ejemplo n.º 3
0
    def test_should_move_files(self, mock_named_temporary_file, mock_gdrive,
                               mock_gcs_hook):
        type(mock_named_temporary_file.return_value.__enter__.return_value
             ).name = mock.PropertyMock(side_effect=["TMP1", "TMP2", "TMP3"])
        mock_gcs_hook.return_value.list.return_value = [
            "sales/A.avro", "sales/B.avro", "sales/C.avro"
        ]
        task = GCSToGoogleDriveOperator(
            task_id="move_files",
            source_bucket="data",
            source_object="sales/sales-2017/*.avro",
            move_object=True,
            impersonation_chain=IMPERSONATION_CHAIN,
        )

        task.execute(mock.MagicMock())
        mock_gcs_hook.assert_has_calls([
            mock.call(
                delegate_to=None,
                google_cloud_storage_conn_id="google_cloud_default",
                impersonation_chain=IMPERSONATION_CHAIN,
            ),
            mock.call().list("data",
                             delimiter=".avro",
                             prefix="sales/sales-2017/"),
            mock.call().download(bucket_name="data",
                                 filename="TMP1",
                                 object_name="sales/A.avro"),
            mock.call().delete("data", "sales/A.avro"),
            mock.call().download(bucket_name="data",
                                 filename="TMP2",
                                 object_name="sales/B.avro"),
            mock.call().delete("data", "sales/B.avro"),
            mock.call().download(bucket_name="data",
                                 filename="TMP3",
                                 object_name="sales/C.avro"),
            mock.call().delete("data", "sales/C.avro"),
        ])

        mock_gdrive.assert_has_calls([
            mock.call(
                delegate_to=None,
                gcp_conn_id="google_cloud_default",
                impersonation_chain=IMPERSONATION_CHAIN,
            ),
            mock.call().upload_file(local_location="TMP1",
                                    remote_location="sales/A.avro"),
            mock.call().upload_file(local_location="TMP2",
                                    remote_location="sales/B.avro"),
            mock.call().upload_file(local_location="TMP3",
                                    remote_location="sales/C.avro"),
        ])
Ejemplo n.º 4
0
from airflow import models
from airflow.providers.google.suite.transfers.gcs_to_gdrive import GCSToGoogleDriveOperator

GCS_TO_GDRIVE_BUCKET = os.environ.get("GCS_TO_DRIVE_BUCKET", "example-object")

with models.DAG(
    "example_gcs_to_gdrive",
    schedule_interval=None,  # Override to match your needs,
    start_date=datetime(2021, 1, 1),
    catchup=False,
    tags=['example'],
) as dag:
    # [START howto_operator_gcs_to_gdrive_copy_single_file]
    copy_single_file = GCSToGoogleDriveOperator(
        task_id="copy_single_file",
        source_bucket=GCS_TO_GDRIVE_BUCKET,
        source_object="sales/january.avro",
        destination_object="copied_sales/january-backup.avro",
    )
    # [END howto_operator_gcs_to_gdrive_copy_single_file]
    # [START howto_operator_gcs_to_gdrive_copy_files]
    copy_files = GCSToGoogleDriveOperator(
        task_id="copy_files",
        source_bucket=GCS_TO_GDRIVE_BUCKET,
        source_object="sales/*",
        destination_object="copied_sales/",
    )
    # [END howto_operator_gcs_to_gdrive_copy_files]
    # [START howto_operator_gcs_to_gdrive_move_files]
    move_files = GCSToGoogleDriveOperator(
        task_id="move_files",
        source_bucket=GCS_TO_GDRIVE_BUCKET,