Add example dag and system test for S3ToGCSOperator (#10951)

apache · Sep 16, 2020 · 76545bb · 76545bb
1 parent 2aec99c
commit 76545bb
Show file tree

Hide file tree

Showing 6 changed files with 182 additions and 5 deletions.
diff --git a/airflow/providers/google/cloud/example_dags/example_s3_to_gcs.py b/airflow/providers/google/cloud/example_dags/example_s3_to_gcs.py
@@ -0,0 +1,77 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import os
+
+from airflow import models
+from airflow.operators.python import PythonOperator
+from airflow.providers.amazon.aws.hooks.s3 import S3Hook
+from airflow.providers.amazon.aws.operators.s3_bucket import S3CreateBucketOperator, S3DeleteBucketOperator
+from airflow.providers.google.cloud.operators.gcs import GCSCreateBucketOperator, GCSDeleteBucketOperator
+from airflow.providers.google.cloud.transfers.s3_to_gcs import S3ToGCSOperator
+from airflow.utils.dates import days_ago
+
+GCP_PROJECT_ID = os.environ.get('GCP_PROJECT_ID', 'gcp-project-id')
+S3BUCKET_NAME = os.environ.get('S3BUCKET_NAME', 'example-s3bucket-name')
+GCS_BUCKET = os.environ.get('GCP_GCS_BUCKET', 'example-gcsbucket-name')
+UPLOAD_FILE = '/tmp/example-file.txt'
+PREFIX = 'TESTS'
+
+
+def upload_file():
+ """A callable to upload file to AWS bucket"""
+ s3_hook = S3Hook()
+ s3_hook.load_file(filename=UPLOAD_FILE, key=PREFIX, bucket_name=S3BUCKET_NAME)
+
+
+with models.DAG(
+ 'example_s3_to_gcs',
+ schedule_interval=None,
+ start_date=days_ago(2),
+ tags=['example'],
+) as dag:
+ create_s3_bucket = S3CreateBucketOperator(
+ task_id="create_s3_bucket", bucket_name=S3BUCKET_NAME, region_name='us-east-1'
+ )
+
+ upload_to_s3 = PythonOperator(task_id='upload_file_to_s3', python_callable=upload_file)
+
+ create_gcs_bucket = GCSCreateBucketOperator(
+ task_id="create_bucket",
+ bucket_name=GCS_BUCKET,
+ project_id=GCP_PROJECT_ID,
+ )
+ # [START howto_transfer_s3togcs_operator]
+ transfer_to_gcs = S3ToGCSOperator(
+ task_id='s3_to_gcs_task', bucket=S3BUCKET_NAME, prefix=PREFIX, dest_gcs="gs://" + GCS_BUCKET
+ )
+ # [END howto_transfer_s3togcs_operator]
+
+ delete_s3_bucket = S3DeleteBucketOperator(
+ task_id='delete_s3_bucket', bucket_name=S3BUCKET_NAME, force_delete=True
+ )
+
+ delete_gcs_bucket = GCSDeleteBucketOperator(task_id='delete_gcs_bucket', bucket_name=GCS_BUCKET)
+
+ (
+ create_s3_bucket
+ >> upload_to_s3
+ >> create_gcs_bucket
+ >> transfer_to_gcs
+ >> delete_s3_bucket
+ >> delete_gcs_bucket
+ )
diff --git a/docs/build_docs.py b/docs/build_docs.py
@@ -450,7 +450,6 @@ def check_enforce_code_block():
  'mssql_to_gcs',
  'mysql_to_gcs',
  'postgres_to_gcs',
- 's3_to_gcs',
  'sql_to_gcs',
  'tasks',
 }

diff --git a/docs/howto/operator/google/transfer/s3_to_gcs.rst b/docs/howto/operator/google/transfer/s3_to_gcs.rst
@@ -0,0 +1,53 @@
+ .. Licensed to the Apache Software Foundation (ASF) under one
+ or more contributor license agreements. See the NOTICE file
+ distributed with this work for additional information
+ regarding copyright ownership. The ASF licenses this file
+ to you under the Apache License, Version 2.0 (the
+ "License"); you may not use this file except in compliance
+ with the License. You may obtain a copy of the License at
+
+ .. http://www.apache.org/licenses/LICENSE-2.0
+
+ .. Unless required by applicable law or agreed to in writing,
+ software distributed under the License is distributed on an
+ "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ KIND, either express or implied. See the License for the
+ specific language governing permissions and limitations
+ under the License.
+
+
+
+Transfer Data from Amazon S3 to Google Cloud Storage
+====================================================
+The `Google Cloud Storage <https://cloud.google.com/storage/>`__ (GCS) is used to store large
+data from various applications. This is also the same with `Amazon Simple Storage Service <https://docs.aws.amazon.com/AmazonS3/latest/dev/Introduction.html>`__.
+This page shows how to transfer data from Amazon S3 to GCS.
+
+.. contents::
+ :depth: 1
+ :local:
+
+Prerequisite Tasks
+^^^^^^^^^^^^^^^^^^
+
+.. include::/howto/operator/google/_partials/prerequisite_tasks.rst
+
+.. _howto/operator:S3ToGCSOperator:
+
+Use the :class:`~airflow.providers.google.cloud.transfers.s3_to_gcs.S3ToGCSOperator`
+to transfer data from Amazon S3 to Google Cloud Storage.
+
+.. exampleinclude::/../airflow/providers/google/cloud/example_dags/example_s3_to_gcs.py
+ :language: python
+ :start-after: [START howto_transfer_s3togcs_operator]
+ :end-before: [END howto_transfer_s3togcs_operator]
+
+Reference
+^^^^^^^^^
+
+For further information, look at:
+
+* `GCS Client Library Documentation <https://googleapis.dev/python/storage/latest/index.html>`__
+* `GCS Product Documentation <https://cloud.google.com/storage/docs/>`__
+* `S3 Client Library Documentation <https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/s3.html>`__
+* `S3 Product Documentation <https://docs.aws.amazon.com/AmazonS3/latest/dev/Introduction.html>`__
diff --git a/docs/operators-and-hooks-ref.rst b/docs/operators-and-hooks-ref.rst
@@ -936,9 +936,8 @@ These integrations allow you to copy data from/to Google Cloud.
 
  * - `Amazon Simple Storage Service (S3) <https://aws.amazon.com/s3/>`__
  - `Google Cloud Storage (GCS) <https://cloud.google.com/gcs/>`__
- - :doc:`How to use <howto/operator/google/cloud/cloud_storage_transfer_service>`
- - :mod:`airflow.providers.google.cloud.transfers.s3_to_gcs`,
- :mod:`airflow.providers.google.cloud.operators.cloud_storage_transfer_service`
+ - :doc:`How to use <howto/operator/google/transfer/s3_to_gcs>`
+ - :mod:`airflow.providers.google.cloud.transfers.s3_to_gcs`
 
  * - `Apache Cassandra <http://cassandra.apache.org/>`__
  - `Google Cloud Storage (GCS) <https://cloud.google.com/gcs/>`__

diff --git a/tests/providers/google/cloud/transfers/test_s3_to_gcs_system.py b/tests/providers/google/cloud/transfers/test_s3_to_gcs_system.py
@@ -0,0 +1,50 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import pytest
+from airflow.providers.google.cloud.example_dags.example_s3_to_gcs import UPLOAD_FILE
+from tests.providers.google.cloud.utils.gcp_authenticator import GCP_GCS_KEY
+from tests.test_utils.gcp_system_helpers import GoogleSystemTest, provide_gcp_context, CLOUD_DAG_FOLDER
+
+FILENAME = UPLOAD_FILE.split('/')[-1]
+
+
+@pytest.mark.backend("mysql", "postgres")
+@pytest.mark.credential_file(GCP_GCS_KEY)
+class S3ToGCSSystemTest(GoogleSystemTest):
+ """System test for S3 to GCS transfer operator.
+
+ This test requires the following environment variables:
+
+ GCP_PROJECT_ID=your-gcp-project-id
+ GCP_GCS_BUCKET=unique-bucket-name-to-create
+ S3BUCKET_NAME=unique-s3-bucket-name
+ AWS_ACCESS_KEY_ID=your-aws-access-key
+ AWS_SECRET_ACCESS_KEY=your-aws-secret-access-key
+ """
+
+ def setUp(self) -> None:
+ super().setUp()
+ self.create_dummy_file(FILENAME)
+
+ def tearDown(self) -> None:
+ self.delete_dummy_file(FILENAME, dir_path='/tmp')
+ super().tearDown()
+
+ @provide_gcp_context(GCP_GCS_KEY)
+ def test_run_example_dag_s3_to_gcs(self):
+ self.run_dag('example_s3_to_gcs', CLOUD_DAG_FOLDER)
diff --git a/tests/test_project_structure.py b/tests/test_project_structure.py
@@ -92,7 +92,6 @@ class TestGoogleProviderProjectStructure(unittest.TestCase):
  MISSING_EXAMPLE_DAGS = {
  ('cloud', 'adls_to_gcs'),
  ('cloud', 'sql_to_gcs'),
- ('cloud', 's3_to_gcs'),
  ('cloud', 'bigquery_to_mysql'),
  ('cloud', 'cassandra_to_gcs'),
  ('cloud', 'mysql_to_gcs'),