Refactor read file function and remove gcsio dependency

benEng · benEng · commit 08e88a81a066 · 2022-10-21T16:06:11.000-05:00
diff --git a/sdks/python/apache_beam/examples/complete/autocomplete_it_test.py b/sdks/python/apache_beam/examples/complete/autocomplete_it_test.py
@@ -28,7 +28,7 @@
 from apache_beam.examples.complete import autocomplete
 from apache_beam.testing.test_pipeline import TestPipeline
 from apache_beam.testing.test_utils import create_file
-from apache_beam.testing.test_utils import read_gcs_output_file
+from apache_beam.testing.test_utils import read_files_from_pattern
 
 
 def format_output_file(output_string):
@@ -85,7 +85,7 @@ def test_autocomplete_output_files_on_small_input(self):
     autocomplete.run(test_pipeline.get_full_options_as_args(**extra_opts))
 
     # Load result file and compare.
-    result = read_gcs_output_file(output).strip()
+    result = read_files_from_pattern('%s*' % output).strip()
 
     self.assertEqual(
         sorted(self.EXPECTED_PREFIXES), sorted(format_output_file(result)))
diff --git a/sdks/python/apache_beam/examples/complete/distribopt_test.py b/sdks/python/apache_beam/examples/complete/distribopt_test.py
@@ -31,7 +31,7 @@
 
 from apache_beam.testing.test_pipeline import TestPipeline
 from apache_beam.testing.test_utils import create_file
-from apache_beam.testing.test_utils import read_gcs_output_file
+from apache_beam.testing.test_utils import read_files_from_pattern
 
 FILE_CONTENTS = 'OP01,8,12,0,12\n' \
                 'OP02,30,14,3,12\n' \
@@ -72,7 +72,7 @@ def test_basics(self):
           save_main_session=False)
 
     # Load result file and compare.
-    lines = read_gcs_output_file(output).splitlines()
+    lines = read_files_from_pattern('%s*' % output).splitlines()
 
     # Only 1 result
     self.assertEqual(len(lines), 1)
diff --git a/sdks/python/apache_beam/examples/complete/estimate_pi_it_test.py b/sdks/python/apache_beam/examples/complete/estimate_pi_it_test.py
@@ -27,7 +27,7 @@
 
 from apache_beam.examples.complete import estimate_pi
 from apache_beam.testing.test_pipeline import TestPipeline
-from apache_beam.testing.test_utils import read_gcs_output_file
+from apache_beam.testing.test_utils import read_files_from_pattern
 
 
 class EstimatePiIT(unittest.TestCase):
@@ -41,7 +41,7 @@ def test_estimate_pi_output_file(self):
     extra_opts = {'output': output}
     estimate_pi.run(test_pipeline.get_full_options_as_args(**extra_opts))
     # Load result file and compare.
-    result = read_gcs_output_file(output)
+    result = read_files_from_pattern('%s*' % output)
     [_, _, estimated_pi] = json.loads(result.strip())
     # Note: Probabilistically speaking this test can fail with a probability
     # that is very small (VERY) given that we run at least 100 thousand
diff --git a/sdks/python/apache_beam/examples/complete/tfidf.py b/sdks/python/apache_beam/examples/complete/tfidf.py
@@ -30,7 +30,7 @@
 import apache_beam as beam
 from apache_beam.io import ReadFromText
 from apache_beam.io import WriteToText
-from apache_beam.io.gcp import gcsio
+from apache_beam.io.filesystems import FileSystems
 from apache_beam.options.pipeline_options import PipelineOptions
 from apache_beam.options.pipeline_options import SetupOptions
 from apache_beam.pvalue import AsSingleton
@@ -197,11 +197,12 @@ def run(argv=None, save_main_session=True):
   # workflow rely on global context (e.g., a module imported at module level).
   pipeline_options = PipelineOptions(pipeline_args)
   pipeline_options.view_as(SetupOptions).save_main_session = save_main_session
-  gcs = gcsio.GcsIO()
   with beam.Pipeline(options=pipeline_options) as p:
 
     # Read documents specified by the uris command line option.
-    pcoll = read_documents(p, gcs.list_prefix(known_args.uris).keys())
+    metadata_list = FileSystems.match([known_args.uris])[0].metadata_list
+    uris = [metadata.path for metadata in metadata_list]
+    pcoll = read_documents(p, uris)
     # Compute TF-IDF information for each word.
     output = pcoll | TfIdf()
     # Write the output using a "Write" transform that has side effects.
diff --git a/sdks/python/apache_beam/examples/complete/tfidf_it_test.py b/sdks/python/apache_beam/examples/complete/tfidf_it_test.py
@@ -29,13 +29,7 @@
 from apache_beam.examples.complete import tfidf
 from apache_beam.testing.test_pipeline import TestPipeline
 from apache_beam.testing.test_utils import create_file
-from apache_beam.testing.test_utils import read_gcs_output_file
-
-# Protect against environments where gcsio library is not available.
-try:
-  from apache_beam.io.gcp import gcsio
-except ImportError:
-  gcsio = None
+from apache_beam.testing.test_utils import read_files_from_pattern
 
 EXPECTED_RESULTS = set([
     ('ghi', '1.txt', 0.3662040962227032), ('abc', '1.txt', 0.0),
@@ -46,7 +40,6 @@
 EXPECTED_LINE_RE = r'\(u?\'([a-z]*)\', \(\'.*([0-9]\.txt)\', (.*)\)\)'
 
 
-@unittest.skipIf(gcsio is None, 'GCP dependencies are not installed')
 class TfIdfIT(unittest.TestCase):
   @pytest.mark.examples_postcommit
   def test_basics(self):
@@ -60,14 +53,14 @@ def test_basics(self):
     create_file('/'.join([input_folder, '3.txt']), 'abc')
     output = '/'.join([temp_location, str(uuid.uuid4()), 'result'])
 
-    extra_opts = {'uris': input_folder, 'output': output}
+    extra_opts = {'uris': '%s/**' % input_folder, 'output': output}
     tfidf.run(
         test_pipeline.get_full_options_as_args(**extra_opts),
         save_main_session=False)
 
     # Parse result file and compare.
     results = []
-    lines = read_gcs_output_file(output).splitlines()
+    lines = read_files_from_pattern('%s*' % output).splitlines()
     for line in lines:
       match = re.search(EXPECTED_LINE_RE, line)
       logging.info('Result line: %s', line)
diff --git a/sdks/python/apache_beam/examples/complete/tfidf_test.py b/sdks/python/apache_beam/examples/complete/tfidf_test.py
@@ -28,20 +28,13 @@
 from apache_beam.testing.util import assert_that
 from apache_beam.testing.util import equal_to
 
-# Protect against environments where gcsio library is not available.
-try:
-  from apache_beam.io.gcp import gcsio
-except ImportError:
-  gcsio = None
-
 EXPECTED_RESULTS = set([
     ('ghi', '1.txt', 0.3662040962227032), ('abc', '1.txt', 0.0),
     ('abc', '3.txt', 0.0), ('abc', '2.txt', 0.0),
     ('def', '1.txt', 0.13515503603605478), ('def', '2.txt', 0.2027325540540822)
 ])
 
 
-@unittest.skipIf(gcsio is None, 'GCP dependencies are not installed')
 class TfIdfTest(unittest.TestCase):
   def test_tfidf_transform(self):
     with TestPipeline() as p:
diff --git a/sdks/python/apache_beam/examples/complete/top_wikipedia_sessions_it_test.py b/sdks/python/apache_beam/examples/complete/top_wikipedia_sessions_it_test.py
@@ -28,7 +28,7 @@
 from apache_beam.examples.complete import top_wikipedia_sessions
 from apache_beam.testing.test_pipeline import TestPipeline
 from apache_beam.testing.test_utils import create_file
-from apache_beam.testing.test_utils import read_gcs_output_file
+from apache_beam.testing.test_utils import read_files_from_pattern
 
 
 class ComputeTopSessionsIT(unittest.TestCase):
@@ -88,7 +88,7 @@ def test_top_wikipedia_sessions_output_files_on_small_input(self):
         test_pipeline.get_full_options_as_args(**extra_opts))
 
     # Load result file and compare.
-    result = read_gcs_output_file(output).strip().splitlines()
+    result = read_files_from_pattern('%s*' % output).strip().splitlines()
 
     self.assertEqual(self.EXPECTED, sorted(result, key=lambda x: x.split()[0]))
 
diff --git a/sdks/python/apache_beam/examples/cookbook/coders_it_test.py b/sdks/python/apache_beam/examples/cookbook/coders_it_test.py
@@ -28,7 +28,7 @@
 from apache_beam.examples.cookbook import coders
 from apache_beam.testing.test_pipeline import TestPipeline
 from apache_beam.testing.test_utils import create_file
-from apache_beam.testing.test_utils import read_gcs_output_file
+from apache_beam.testing.test_utils import read_files_from_pattern
 
 
 def format_result(result_string):
@@ -72,7 +72,7 @@ def test_coders_output_files_on_small_input(self):
     coders.run(test_pipeline.get_full_options_as_args(**extra_opts))
 
     # Load result file and compare.
-    result = read_gcs_output_file(output).strip()
+    result = read_files_from_pattern('%s*' % output).strip()
 
     self.assertEqual(
         sorted(self.EXPECTED_RESULT), sorted(format_result(result)))
diff --git a/sdks/python/apache_beam/examples/cookbook/custom_ptransform_it_test.py b/sdks/python/apache_beam/examples/cookbook/custom_ptransform_it_test.py
@@ -27,7 +27,7 @@
 from apache_beam.examples.cookbook import custom_ptransform
 from apache_beam.testing.test_pipeline import TestPipeline
 from apache_beam.testing.test_utils import create_file
-from apache_beam.testing.test_utils import read_gcs_output_file
+from apache_beam.testing.test_utils import read_files_from_pattern
 
 
 def format_result(result_string):
@@ -61,7 +61,7 @@ def test_custom_ptransform_output_files_on_small_input(self):
     custom_ptransform.run(test_pipeline.get_full_options_as_args(**extra_opts))
 
     # Load result file and compare.
-    result = read_gcs_output_file(output).strip()
+    result = read_files_from_pattern('%s*' % output).strip()
     self.assertEqual(result, self.EXPECTED_RESULT)
 
 
diff --git a/sdks/python/apache_beam/examples/cookbook/group_with_coder_test.py b/sdks/python/apache_beam/examples/cookbook/group_with_coder_test.py
@@ -27,7 +27,7 @@
 
 from apache_beam.examples.cookbook import group_with_coder
 from apache_beam.testing.test_pipeline import TestPipeline
-from apache_beam.testing.test_utils import read_gcs_output_file
+from apache_beam.testing.test_utils import read_files_from_pattern
 
 # Protect against environments where gcsio library is not available.
 try:
@@ -88,7 +88,7 @@ def test_basics_with_type_check(self):
         save_main_session=False)
     # Parse result file and compare.
     results = []
-    lines = read_gcs_output_file(output).splitlines()
+    lines = read_files_from_pattern('%s*' % output).splitlines()
     for line in lines:
       name, points = line.split(',')
       results.append((name, int(points)))
diff --git a/sdks/python/apache_beam/examples/cookbook/mergecontacts_test.py b/sdks/python/apache_beam/examples/cookbook/mergecontacts_test.py
@@ -28,7 +28,7 @@
 from apache_beam.examples.cookbook import mergecontacts
 from apache_beam.testing.test_pipeline import TestPipeline
 from apache_beam.testing.test_utils import create_file
-from apache_beam.testing.test_utils import read_gcs_output_file
+from apache_beam.testing.test_utils import read_files_from_pattern
 
 
 class MergeContactsTest(unittest.TestCase):
@@ -167,7 +167,7 @@ def test_mergecontacts(self):
     mergecontacts.run(
         pipeline_opts, assert_results=(2, 1, 3), save_main_session=False)
 
-    contents = read_gcs_output_file('%s.tsv' % result_prefix)
+    contents = read_files_from_pattern('%s*' % result_prefix)
     self.assertEqual(self.EXPECTED_TSV, self.normalize_tsv_results(contents))
 
 
diff --git a/sdks/python/apache_beam/examples/cookbook/multiple_output_pardo_test.py b/sdks/python/apache_beam/examples/cookbook/multiple_output_pardo_test.py
@@ -29,7 +29,7 @@
 from apache_beam.examples.cookbook import multiple_output_pardo
 from apache_beam.testing.test_pipeline import TestPipeline
 from apache_beam.testing.test_utils import create_file
-from apache_beam.testing.test_utils import read_gcs_output_file
+from apache_beam.testing.test_utils import read_files_from_pattern
 
 
 class MultipleOutputParDo(unittest.TestCase):
@@ -41,7 +41,7 @@ class MultipleOutputParDo(unittest.TestCase):
 
   def get_wordcount_results(self, result_path):
     results = []
-    lines = read_gcs_output_file(result_path).splitlines()
+    lines = read_files_from_pattern(result_path).splitlines()
     for line in lines:
       match = re.search(r'([A-Za-z]+): ([0-9]+)', line)
       if match is not None:
@@ -64,13 +64,13 @@ def test_multiple_output_pardo(self):
         save_main_session=False)
 
     expected_char_count = len(''.join(self.SAMPLE_TEXT.split('\n')))
-    contents = read_gcs_output_file(result_prefix + '-chars-')
+    contents = read_files_from_pattern(result_prefix + '-chars*')
     self.assertEqual(expected_char_count, int(contents))
 
-    short_words = self.get_wordcount_results(result_prefix + '-short-words-')
+    short_words = self.get_wordcount_results(result_prefix + '-short-words*')
     self.assertEqual(sorted(short_words), sorted(self.EXPECTED_SHORT_WORDS))
 
-    words = self.get_wordcount_results(result_prefix + '-words-')
+    words = self.get_wordcount_results(result_prefix + '-words*')
     self.assertEqual(sorted(words), sorted(self.EXPECTED_WORDS))
 
 
diff --git a/sdks/python/apache_beam/examples/dataframe/wordcount_test.py b/sdks/python/apache_beam/examples/dataframe/wordcount_test.py
@@ -31,7 +31,7 @@
 from apache_beam.examples.dataframe import wordcount
 from apache_beam.testing.test_pipeline import TestPipeline
 from apache_beam.testing.test_utils import create_file
-from apache_beam.testing.test_utils import read_gcs_output_file
+from apache_beam.testing.test_utils import read_files_from_pattern
 
 
 class WordCountTest(unittest.TestCase):
@@ -57,7 +57,7 @@ def test_basics(self):
     wordcount.run(test_pipeline.get_full_options_as_args(**extra_opts))
     # Parse result file and compare.
     results = []
-    lines = read_gcs_output_file(temp_path + '.result-').splitlines()
+    lines = read_files_from_pattern(temp_path + '.result*').splitlines()
     for line in lines:
       match = re.search(r'(\S+),([0-9]+)', line)
       if match is not None:
diff --git a/sdks/python/apache_beam/examples/wordcount_debugging_test.py b/sdks/python/apache_beam/examples/wordcount_debugging_test.py
@@ -29,7 +29,7 @@
 from apache_beam.examples import wordcount_debugging
 from apache_beam.testing.test_pipeline import TestPipeline
 from apache_beam.testing.test_utils import create_file
-from apache_beam.testing.test_utils import read_gcs_output_file
+from apache_beam.testing.test_utils import read_files_from_pattern
 
 
 @pytest.mark.examples_postcommit
@@ -39,7 +39,7 @@ class WordCountDebuggingTest(unittest.TestCase):
 
   def get_results(self, temp_path):
     results = []
-    lines = read_gcs_output_file(temp_path + '.result-').splitlines()
+    lines = read_files_from_pattern(temp_path + '.result*').splitlines()
     for line in lines:
       match = re.search(r'([A-Za-z]+): ([0-9]+)', line)
       if match is not None:
diff --git a/sdks/python/apache_beam/examples/wordcount_minimal_test.py b/sdks/python/apache_beam/examples/wordcount_minimal_test.py
@@ -43,7 +43,7 @@
 from apache_beam.examples import wordcount_minimal
 from apache_beam.testing.test_pipeline import TestPipeline
 from apache_beam.testing.test_utils import create_file
-from apache_beam.testing.test_utils import read_gcs_output_file
+from apache_beam.testing.test_utils import read_files_from_pattern
 
 
 @pytest.mark.examples_postcommit
@@ -70,7 +70,7 @@ def test_basics(self):
 
     # Parse result file and compare.
     results = []
-    lines = read_gcs_output_file(temp_path + '.result-').splitlines()
+    lines = read_files_from_pattern(temp_path + '.result*').splitlines()
     for line in lines:
       match = re.search(r'([a-z]+): ([0-9]+)', line)
       if match is not None:
diff --git a/sdks/python/apache_beam/examples/wordcount_test.py b/sdks/python/apache_beam/examples/wordcount_test.py
@@ -44,7 +44,7 @@
 from apache_beam.examples import wordcount
 from apache_beam.testing.test_pipeline import TestPipeline
 from apache_beam.testing.test_utils import create_file
-from apache_beam.testing.test_utils import read_gcs_output_file
+from apache_beam.testing.test_utils import read_files_from_pattern
 
 
 @pytest.mark.examples_postcommit
@@ -68,7 +68,7 @@ def test_basics(self):
         save_main_session=False)
     # Parse result file and compare.
     results = []
-    lines = read_gcs_output_file(temp_path + '.result-').splitlines()
+    lines = read_files_from_pattern(temp_path + '.result*').splitlines()
     for line in lines:
       match = re.search(r'(\S+): ([0-9]+)', line)
       if match is not None:
diff --git a/sdks/python/apache_beam/testing/test_utils.py b/sdks/python/apache_beam/testing/test_utils.py
@@ -215,13 +215,11 @@ def create_file(path, contents):
   return path
 
 
-def read_gcs_output_file(file_pattern):
-  """Reads the gcs output file of a pipeline by prefix"""
-  from apache_beam.io.gcp import gcsio
-  gcs = gcsio.GcsIO()
-  file_names = gcs.list_prefix(file_pattern).keys()
+def read_files_from_pattern(file_pattern):
+  """Reads the files that match a pattern"""
+  metadata_list = FileSystems.match([file_pattern])[0].metadata_list
   output = []
-  for file_name in file_names:
-    with FileSystems.open(file_name) as f:
+  for metadata in metadata_list:
+    with FileSystems.open(metadata.path) as f:
       output.append(f.read().decode('utf-8').strip())
   return '\n'.join(output)
