sdks/python/apache_beam/io/filebasedsink_test.py

# -*- coding: utf-8 -*-
#
# 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.
#

"""Unit tests for file sinks."""

# pytype: skip-file

import glob
import logging
import os
import shutil
import tempfile
import unittest

import hamcrest as hc
import mock

import apache_beam as beam
from apache_beam.coders import coders
from apache_beam.io import filebasedsink
from apache_beam.io.filesystem import BeamIOError
from apache_beam.options.value_provider import StaticValueProvider
from apache_beam.testing.test_pipeline import TestPipeline
from apache_beam.transforms.display import DisplayData
from apache_beam.transforms.display_test import DisplayDataItemMatcher

_LOGGER = logging.getLogger(__name__)


# TODO: Refactor code so all io tests are using same library
# TestCaseWithTempDirCleanup class.
class _TestCaseWithTempDirCleanUp(unittest.TestCase):
  """Base class for TestCases that deals with TempDir clean-up.

  Inherited test cases will call self._new_tempdir() to start a temporary dir
  which will be deleted at the end of the tests (when tearDown() is called).
  """
  def setUp(self):
    self._tempdirs = []

  def tearDown(self):
    for path in self._tempdirs:
      if os.path.exists(path):
        shutil.rmtree(path)
    self._tempdirs = []

  def _new_tempdir(self):
    result = tempfile.mkdtemp()
    self._tempdirs.append(result)
    return result

  def _create_temp_file(self, name='', suffix='', dir=None, content=None):
    if not name:
      name = tempfile.template
    if not dir:
      dir = self._new_tempdir()
    file_name = tempfile.NamedTemporaryFile(
        delete=False, prefix=name, dir=dir, suffix=suffix).name

    if content:
      with open(file_name, 'w') as f:
        f.write(content)
    return file_name


class MyFileBasedSink(filebasedsink.FileBasedSink):
  def open(self, temp_path):
    # TODO: Fix main session pickling.
    # file_handle = super().open(temp_path)
    file_handle = filebasedsink.FileBasedSink.open(self, temp_path)
    file_handle.write(b'[start]')
    return file_handle

  def write_encoded_record(self, file_handle, encoded_value):
    file_handle.write(b'[')
    file_handle.write(encoded_value)
    file_handle.write(b']')

  def close(self, file_handle):
    file_handle.write(b'[end]')
    # TODO: Fix main session pickling.
    # file_handle = super().close(file_handle)
    file_handle = filebasedsink.FileBasedSink.close(self, file_handle)


class TestFileBasedSink(_TestCaseWithTempDirCleanUp):
  def _common_init(self, sink):
    # Manually invoke the generic Sink API.
    init_token = sink.initialize_write()

    writer1 = sink.open_writer(init_token, '1')
    writer1.write('a')
    writer1.write('b')
    res1 = writer1.close()

    writer2 = sink.open_writer(init_token, '2')
    writer2.write('x')
    writer2.write('y')
    writer2.write('z')
    res2 = writer2.close()

    return init_token, [res1, res2]

  def test_file_sink_writing(self):
    temp_path = os.path.join(self._new_tempdir(), 'FileBasedSink')
    sink = MyFileBasedSink(
        temp_path, file_name_suffix='.output', coder=coders.ToBytesCoder())

    init_token, writer_results = self._common_init(sink)

    pre_finalize_results = sink.pre_finalize(init_token, writer_results)
    finalize_res1 = list(
        sink.finalize_write(init_token, writer_results, pre_finalize_results))
    # Retry the finalize operation (as if the first attempt was lost).
    finalize_res2 = list(
        sink.finalize_write(init_token, writer_results, pre_finalize_results))

    # Check the results.
    shard1 = temp_path + '-00000-of-00002.output'
    shard2 = temp_path + '-00001-of-00002.output'
    self.assertEqual(finalize_res1, [shard1, shard2])
    self.assertEqual(finalize_res2, [])
    self.assertEqual(open(shard1).read(), '[start][a][b][end]')
    self.assertEqual(open(shard2).read(), '[start][x][y][z][end]')

    # Check that any temp files are deleted.
    self.assertCountEqual([shard1, shard2], glob.glob(temp_path + '*'))

  def test_file_sink_display_data(self):
    temp_path = os.path.join(self._new_tempdir(), 'display')
    sink = MyFileBasedSink(
        temp_path, file_name_suffix='.output', coder=coders.ToBytesCoder())
    dd = DisplayData.create_from(sink)
    expected_items = [
        DisplayDataItemMatcher('compression', 'auto'),
        DisplayDataItemMatcher(
            'file_pattern',
            '{}{}'.format(
                temp_path, '-%(shard_num)05d-of-%(num_shards)05d.output'))
    ]
    hc.assert_that(dd.items, hc.contains_inanyorder(*expected_items))

  def test_empty_write(self):
    temp_path = tempfile.NamedTemporaryFile().name
    sink = MyFileBasedSink(
        temp_path, file_name_suffix='.output', coder=coders.ToBytesCoder())
    with TestPipeline() as p:
      p | beam.Create([]) | beam.io.Write(sink)  # pylint: disable=expression-not-assigned
    self.assertEqual(
        open(temp_path + '-00000-of-00001.output').read(), '[start][end]')

  def test_static_value_provider_empty_write(self):
    temp_path = StaticValueProvider(
        value_type=str, value=tempfile.NamedTemporaryFile().name)
    sink = MyFileBasedSink(
        temp_path,
        file_name_suffix=StaticValueProvider(value_type=str, value='.output'),
        coder=coders.ToBytesCoder())
    with TestPipeline() as p:
      p | beam.Create([]) | beam.io.Write(sink)  # pylint: disable=expression-not-assigned
    self.assertEqual(
        open(temp_path.get() + '-00000-of-00001.output').read(), '[start][end]')

  def test_fixed_shard_write(self):
    temp_path = os.path.join(self._new_tempdir(), 'empty')
    sink = MyFileBasedSink(
        temp_path,
        file_name_suffix='.output',
        num_shards=3,
        shard_name_template='_NN_SSS_',
        coder=coders.ToBytesCoder())
    with TestPipeline() as p:
      p | beam.Create(['a', 'b']) | beam.io.Write(sink)  # pylint: disable=expression-not-assigned

    concat = ''.join(
        open(temp_path + '_03_%03d_.output' % shard_num).read()
        for shard_num in range(3))
    self.assertTrue('][a][' in concat, concat)
    self.assertTrue('][b][' in concat, concat)

  # Not using 'test' in name so that 'nose' doesn't pick this as a test.
  def run_temp_dir_check(
      self,
      no_dir_path,
      dir_path,
      no_dir_root_path,
      dir_root_path,
      prefix,
      separator):
    def _get_temp_dir(file_path_prefix):
      sink = MyFileBasedSink(
          file_path_prefix,
          file_name_suffix='.output',
          coder=coders.ToBytesCoder())
      return sink.initialize_write()

    temp_dir = _get_temp_dir(no_dir_path)
    self.assertTrue(temp_dir.startswith(prefix))
    last_sep = temp_dir.rfind(separator)
    self.assertTrue(temp_dir[last_sep + 1:].startswith('beam-temp'))

    temp_dir = _get_temp_dir(dir_path)
    self.assertTrue(temp_dir.startswith(prefix))
    last_sep = temp_dir.rfind(separator)
    self.assertTrue(temp_dir[last_sep + 1:].startswith('beam-temp'))

    with self.assertRaises(ValueError):
      _get_temp_dir(no_dir_root_path)

    with self.assertRaises(ValueError):
      _get_temp_dir(dir_root_path)

  def test_temp_dir_uniqueness(self):
    temp_path = os.path.join(self._new_tempdir(), 'unique')
    sink = MyFileBasedSink(temp_path, coder=coders.ToBytesCoder())
    init_list = [''] * 1000
    temp_dir_list = [sink._create_temp_dir(temp_path) for _ in init_list]
    temp_dir_set = set(temp_dir_list)
    self.assertEqual(len(temp_dir_list), len(temp_dir_set))

  def test_temp_dir_gcs(self):
    try:
      self.run_temp_dir_check(
          'gs://aaa/bbb',
          'gs://aaa/bbb/',
          'gs://aaa',
          'gs://aaa/',
          'gs://',
          '/')
    except ValueError:
      _LOGGER.debug('Ignoring test since GCP module is not installed')

  @mock.patch('apache_beam.io.localfilesystem.os')
  def test_temp_dir_local(self, filesystem_os_mock):
    # Here we test a unix-like mock file-system
    # (not really testing Unix or Windows since we mock the function of 'os'
    # module).

    def _fake_unix_split(path):
      sep = path.rfind('/')
      if sep < 0:
        raise ValueError('Path must contain a separator')
      return (path[:sep], path[sep + 1:])

    def _fake_unix_join(base, path):
      return base + '/' + path

    filesystem_os_mock.path.abspath = lambda a: a
    filesystem_os_mock.path.split.side_effect = _fake_unix_split
    filesystem_os_mock.path.join.side_effect = _fake_unix_join
    self.run_temp_dir_check('/aaa/bbb', '/aaa/bbb/', '/', '/', '/', '/')

  def test_file_sink_multi_shards(self):
    temp_path = os.path.join(self._new_tempdir(), 'multishard')
    sink = MyFileBasedSink(
        temp_path, file_name_suffix='.output', coder=coders.ToBytesCoder())

    # Manually invoke the generic Sink API.
    init_token = sink.initialize_write()

    num_shards = 1000
    writer_results = []
    for i in range(num_shards):
      uuid = 'uuid-%05d' % i
      writer = sink.open_writer(init_token, uuid)
      writer.write('a')
      writer.write('b')
      writer.write(uuid)
      writer_results.append(writer.close())

    pre_finalize_results = sink.pre_finalize(init_token, writer_results)
    res = sorted(
        sink.finalize_write(init_token, writer_results, pre_finalize_results))

    for i in range(num_shards):
      shard_name = '%s-%05d-of-%05d.output' % (temp_path, i, num_shards)
      uuid = 'uuid-%05d' % i
      self.assertEqual(res[i], shard_name)
      self.assertEqual(
          open(shard_name).read(), ('[start][a][b][%s][end]' % uuid))

    # Check that any temp files are deleted.
    self.assertCountEqual(res, glob.glob(temp_path + '*'))

  @mock.patch.object(filebasedsink.FileSystems, 'rename')
  def test_file_sink_rename_error(self, rename_mock):
    temp_path = os.path.join(self._new_tempdir(), 'rename_error')
    sink = MyFileBasedSink(
        temp_path, file_name_suffix='.output', coder=coders.ToBytesCoder())
    init_token, writer_results = self._common_init(sink)
    pre_finalize_results = sink.pre_finalize(init_token, writer_results)

    error_str = 'mock rename error description'
    rename_mock.side_effect = BeamIOError(
        'mock rename error', {('src', 'dst'): error_str})
    with self.assertRaisesRegex(Exception, error_str):
      list(
          sink.finalize_write(init_token, writer_results, pre_finalize_results))

  def test_file_sink_src_missing(self):
    temp_path = os.path.join(self._new_tempdir(), 'src_missing')
    sink = MyFileBasedSink(
        temp_path, file_name_suffix='.output', coder=coders.ToBytesCoder())
    init_token, writer_results = self._common_init(sink)
    pre_finalize_results = sink.pre_finalize(init_token, writer_results)

    os.remove(writer_results[0])
    with self.assertRaisesRegex(Exception, r'not exist'):
      list(
          sink.finalize_write(init_token, writer_results, pre_finalize_results))

  def test_file_sink_dst_matches_src(self):
    temp_path = os.path.join(self._new_tempdir(), 'dst_matches_src')
    sink = MyFileBasedSink(
        temp_path, file_name_suffix='.output', coder=coders.ToBytesCoder())
    init_token, [res1, res2] = self._common_init(sink)

    pre_finalize_results = sink.pre_finalize(init_token, [res1, res2])
    list(sink.finalize_write(init_token, [res1, res2], pre_finalize_results))

    self.assertFalse(os.path.exists(res1))
    self.assertFalse(os.path.exists(res2))
    shard1 = temp_path + '-00000-of-00002.output'
    shard2 = temp_path + '-00001-of-00002.output'
    self.assertEqual(open(shard1).read(), '[start][a][b][end]')
    self.assertEqual(open(shard2).read(), '[start][x][y][z][end]')

    os.makedirs(os.path.dirname(res1))
    shutil.copyfile(shard1, res1)
    shutil.copyfile(shard2, res2)
    list(sink.finalize_write(init_token, [res1, res2], pre_finalize_results))

  def test_pre_finalize(self):
    temp_path = os.path.join(self._new_tempdir(), 'pre_finalize')
    sink = MyFileBasedSink(
        temp_path, file_name_suffix='.output', coder=coders.ToBytesCoder())
    init_token, [res1, res2] = self._common_init(sink)

    # no-op
    sink.pre_finalize(init_token, [res1, res2])

    # Create finalized outputs from a previous run, which pre_finalize should
    # delete.
    shard1 = temp_path + '-00000-of-00002.output'
    shard2 = temp_path + '-00001-of-00002.output'
    with open(shard1, 'w') as f:
      f.write('foo')
    with open(shard2, 'w') as f:
      f.write('foo')
    self.assertTrue(os.path.exists(res1))
    self.assertTrue(os.path.exists(res2))
    self.assertTrue(os.path.exists(shard1))
    self.assertTrue(os.path.exists(shard2))

    sink.pre_finalize(init_token, [res1, res2])
    self.assertTrue(os.path.exists(res1))
    self.assertTrue(os.path.exists(res2))
    self.assertFalse(os.path.exists(shard1))
    self.assertFalse(os.path.exists(shard2))

  @mock.patch.object(filebasedsink.FileSystems, 'delete')
  def test_pre_finalize_error(self, delete_mock):
    temp_path = os.path.join(self._new_tempdir(), 'pre_finalize')
    sink = MyFileBasedSink(
        temp_path, file_name_suffix='.output', coder=coders.ToBytesCoder())
    init_token, [res1, res2] = self._common_init(sink)

    # no-op
    sink.pre_finalize(init_token, [res1, res2])

    # Create finalized outputs from a previous run, which pre_finalize should
    # delete.
    shard1 = temp_path + '-00000-of-00002.output'
    shard2 = temp_path + '-00001-of-00002.output'
    with open(shard1, 'w') as f:
      f.write('foo')
    with open(shard2, 'w') as f:
      f.write('foo')

    error_str = 'mock rename error description'
    delete_mock.side_effect = BeamIOError(
        'mock rename error', {shard2: error_str})
    with self.assertRaisesRegex(Exception, error_str):
      sink.pre_finalize(init_token, [res1, res2])


if __name__ == '__main__':
  logging.getLogger().setLevel(logging.INFO)
  unittest.main()