docs(samples): Add dataflow_batch_write_to_storage snippet (GoogleCloudPlatform#10658)

VeronicaWasson · web-flow · commit b61f8bc8f4fc · 2023-10-19T12:46:23.000+02:00
* docs(samples): Add dataflow_batch_write_to_storage snippet * Bump Beam SDK version * Fix linter error * Remove a stray debugging line * Remove __init__.py files * Revert "Remove __init__.py files" This reverts commit 21ed33e. * Update copyright headers Update year Add to __init__ files * Incorporate PR feedback - Add a test condition - Add comments re: pipeline options --------- Co-authored-by: Veronica Wasson <veronicawasson@google.com>
diff --git a/dataflow/snippets/__init__.py b/dataflow/snippets/__init__.py
@@ -0,0 +1,13 @@
+#  Copyright 2023 Google LLC
+#
+#  Licensed 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.
diff --git a/dataflow/snippets/batch_write_storage.py b/dataflow/snippets/batch_write_storage.py
@@ -0,0 +1,43 @@
+#!/usr/bin/env python
+#  Copyright 2023 Google LLC
+#
+#  Licensed 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.
+
+# [START dataflow_batch_write_to_storage]
+import apache_beam as beam
+from apache_beam.io.textio import WriteToText
+from apache_beam.options.pipeline_options import PipelineOptions
+
+
+def write_to_cloud_storage(argv=None):
+    # Parse the pipeline options passed into the application.
+    class MyOptions(PipelineOptions):
+        @classmethod
+        # Define a custom pipeline option that specfies the Cloud Storage bucket.
+        def _add_argparse_args(cls, parser):
+            parser.add_argument("--output", required=True)
+
+    wordsList = ["1", "2", "3", "4"]
+    options = MyOptions()
+
+    with beam.Pipeline(options=options) as pipeline:
+        (
+            pipeline
+            | "Create elements" >> beam.Create(wordsList)
+            | "Write Files" >> WriteToText(options.output, file_name_suffix=".txt")
+        )
+# [END dataflow_batch_write_to_storage]
+
+
+if __name__ == "__main__":
+    write_to_cloud_storage()
diff --git a/dataflow/snippets/requirements-test.txt b/dataflow/snippets/requirements-test.txt
@@ -0,0 +1 @@
+pytest==7.0.1
diff --git a/dataflow/snippets/requirements.txt b/dataflow/snippets/requirements.txt
@@ -0,0 +1 @@
+apache-beam[gcp]==2.50.0
diff --git a/dataflow/snippets/tests/__init__.py b/dataflow/snippets/tests/__init__.py
@@ -0,0 +1,13 @@
+#  Copyright 2023 Google LLC
+#
+#  Licensed 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.
diff --git a/dataflow/snippets/tests/test_batch_write_storage.py b/dataflow/snippets/tests/test_batch_write_storage.py
@@ -0,0 +1,45 @@
+#  Copyright 2023 Google LLC
+#
+#  Licensed 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 sys
+import uuid
+
+from google.cloud import storage
+
+import pytest
+
+from ..batch_write_storage import write_to_cloud_storage
+
+
+bucket_name = f'test-bucket-{uuid.uuid4()}'
+storage_client = storage.Client()
+
+
+@pytest.fixture(scope="function")
+def setup_and_teardown():
+    try:
+        bucket = storage_client.create_bucket(bucket_name)
+        yield
+    finally:
+        bucket.delete(force=True)
+
+
+def test_write_to_cloud_storage(setup_and_teardown):
+    sys.argv = ['', f'--output=gs://{bucket_name}/output/out-']
+    write_to_cloud_storage()
+
+    blobs = list(storage_client.list_blobs(bucket_name))
+    # Ensure the pipeline wrote files to Cloud Storage
+    assert blobs
+    for blob in blobs:
+        assert blob.name.endswith(".txt")
