gslib/parallel_tracker_file.py

# -*- coding: utf-8 -*-
# Copyright 2015 Google Inc. All Rights Reserved.
#
# 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.
"""Helper functions for composite upload tracker file functionality."""

from __future__ import absolute_import
from __future__ import print_function
from __future__ import division
from __future__ import unicode_literals

from collections import namedtuple
import errno
import json
import random

import six

import gslib
from gslib.exception import CommandException
from gslib.tracker_file import (WriteJsonDataToTrackerFile,
                                RaiseUnwritableTrackerFileException)
from gslib.utils.constants import UTF8

ObjectFromTracker = namedtuple('ObjectFromTracker', 'object_name generation')


class _CompositeUploadTrackerEntry(object):
  """Enum class for composite upload tracker file JSON keys."""
  COMPONENTS_LIST = 'components'
  COMPONENT_NAME = 'component_name'
  COMPONENT_GENERATION = 'component_generation'
  ENC_SHA256 = 'encryption_key_sha256'
  PREFIX = 'prefix'


def ReadParallelUploadTrackerFile(tracker_file_name, logger):
  """Read the tracker file from the last parallel composite upload attempt.

  If it exists, the tracker file is of the format described in
  WriteParallelUploadTrackerFile or a legacy format. If the file doesn't exist
  or is formatted incorrectly, then the upload will start from the beginning.

  This function is not thread-safe and must be protected by a lock if
  called within Command.Apply.

  Args:
    tracker_file_name: The name of the tracker file to read parse.
    logger: logging.Logger for outputting log messages.

  Returns:
    enc_key_sha256: Encryption key SHA256 used to encrypt the existing
        components, or None if an encryption key was not used.
    component_prefix: String prefix used in naming the existing components, or
        None if no prefix was found.
    existing_components: A list of ObjectFromTracker objects representing
        the set of files that have already been uploaded.
  """
  enc_key_sha256 = None
  prefix = None
  existing_components = []
  tracker_file = None

  # If we already have a matching tracker file, get the serialization data
  # so that we can resume the upload.
  try:
    tracker_file = open(tracker_file_name, 'r')
    tracker_data = tracker_file.read()
    tracker_json = json.loads(tracker_data)
    enc_key_sha256 = tracker_json[_CompositeUploadTrackerEntry.ENC_SHA256]
    prefix = tracker_json[_CompositeUploadTrackerEntry.PREFIX]
    for component in tracker_json[_CompositeUploadTrackerEntry.COMPONENTS_LIST]:
      existing_components.append(
          ObjectFromTracker(
              component[_CompositeUploadTrackerEntry.COMPONENT_NAME],
              component[_CompositeUploadTrackerEntry.COMPONENT_GENERATION]))
  except IOError as e:
    # Ignore non-existent file (happens first time a upload is attempted on an
    # object, or when re-starting an upload after a
    # ResumableUploadStartOverException), but warn user for other errors.
    if e.errno != errno.ENOENT:
      logger.warn(
          'Couldn\'t read upload tracker file (%s): %s. Restarting '
          'parallel composite upload from scratch.', tracker_file_name,
          e.strerror)
  except (KeyError, ValueError) as e:
    # Legacy format did not support user-supplied encryption.
    enc_key_sha256 = None
    (prefix, existing_components) = _ParseLegacyTrackerData(tracker_data)
  finally:
    if tracker_file:
      tracker_file.close()

  return (enc_key_sha256, prefix, existing_components)


def _ParseLegacyTrackerData(tracker_data):
  """Parses a legacy parallel composite upload tracker file.

  Args:
    tracker_data: Legacy tracker file contents.

  Returns:
    component_prefix: The prefix used in naming the existing components, or
        None if no prefix was found.
    existing_components: A list of ObjectFromTracker objects representing
        the set of files that have already been uploaded.
  """
  # Old tracker files used a non-JSON format.
  # The first line represents the prefix, followed by line pairs of object_name
  # and generation. Discard the last blank line.
  old_tracker_data = tracker_data.split('\n')[:-1]
  prefix = None
  existing_components = []
  if old_tracker_data:
    prefix = old_tracker_data[0]
    i = 1
    while i < len(old_tracker_data) - 1:
      (name, generation) = (old_tracker_data[i], old_tracker_data[i + 1])
      if not generation:
        # Cover the '' case.
        generation = None
      existing_components.append(ObjectFromTracker(name, generation))
      i += 2
  return (prefix, existing_components)


def ValidateParallelCompositeTrackerData(tracker_file_name, existing_enc_sha256,
                                         existing_prefix, existing_components,
                                         current_enc_key_sha256, bucket_url,
                                         command_obj, logger, delete_func,
                                         delete_exc_handler):
  """Validates that tracker data matches the current encryption key.

  If the data does not match, makes a best-effort attempt to delete existing
  temporary component objects encrypted with the old key.

  Args:
    tracker_file_name: String file name of tracker file.
    existing_enc_sha256: Encryption key SHA256 used to encrypt the existing
        components, or None if an encryption key was not used.
    existing_prefix: String prefix used in naming the existing components, or
        None if no prefix was found.
    existing_components: A list of ObjectFromTracker objects representing
        the set of files that have already been uploaded.
    current_enc_key_sha256: Current Encryption key SHA256 that should be used
        to encrypt objects.
    bucket_url: Bucket URL in which the components exist.
    command_obj: Command class for calls to Apply.
    logger: logging.Logger for outputting log messages.
    delete_func: command.Apply-callable function for deleting objects.
    delete_exc_handler: Exception handler for delete_func.

  Returns:
    prefix: existing_prefix, or None if the encryption key did not match.
    existing_components: existing_components, or empty list if the encryption
        key did not match.
  """
  if six.PY3:
    if isinstance(existing_enc_sha256, str):
      existing_enc_sha256 = existing_enc_sha256.encode(UTF8)
    if isinstance(current_enc_key_sha256, str):
      current_enc_key_sha256 = current_enc_key_sha256.encode(UTF8)
  if existing_prefix and existing_enc_sha256 != current_enc_key_sha256:
    try:
      logger.warn(
          'Upload tracker file (%s) does not match current encryption '
          'key. Deleting old components and restarting upload from '
          'scratch with a new tracker file that uses the current '
          'encryption key.', tracker_file_name)
      components_to_delete = []
      for component in existing_components:
        url = bucket_url.Clone()
        url.object_name = component.object_name
        url.generation = component.generation

      command_obj.Apply(
          delete_func,
          components_to_delete,
          delete_exc_handler,
          arg_checker=gslib.command.DummyArgChecker,
          parallel_operations_override=command_obj.ParallelOverrideReason.SPEED)
    except:  # pylint: disable=bare-except
      # Regardless of why we can't clean up old components, need to proceed
      # with the user's original intent to upload the file, so merely warn.
      component_names = [
          component.object_name for component in existing_components
      ]
      logger.warn(
          'Failed to delete some of the following temporary objects:\n%s\n'
          '(Continuing on to re-upload components from scratch.)',
          '\n'.join(component_names))

    # Encryption keys have changed, so the old components and prefix
    # cannot be used.
    return (None, [])

  return (existing_prefix, existing_components)


def GenerateComponentObjectPrefix(encryption_key_sha256=None):
  """Generates a random prefix for component objects.

  Args:
    encryption_key_sha256: Encryption key SHA256 that will be used to encrypt
        the components. This is hashed into the prefix to avoid collision
        during resumption with a different encryption key.

  Returns:
    String prefix for use in the composite upload.
  """
  return str(
      (random.randint(1, (10**10) - 1) + hash(encryption_key_sha256)) % 10**10)


def WriteComponentToParallelUploadTrackerFile(tracker_file_name,
                                              tracker_file_lock,
                                              component,
                                              logger,
                                              encryption_key_sha256=None):
  """Rewrites an existing tracker file with info about the uploaded component.

  Follows the format described in _CreateParallelUploadTrackerFile.

  Args:
    tracker_file_name: Tracker file to append to.
    tracker_file_lock: Thread and process-safe Lock protecting the tracker file.
    component: ObjectFromTracker describing the object that was uploaded.
    logger: logging.Logger for outputting log messages.
    encryption_key_sha256: Encryption key SHA256 for use in this upload, if any.
  """
  with tracker_file_lock:
    (existing_enc_key_sha256, prefix,
     existing_components) = (ReadParallelUploadTrackerFile(
         tracker_file_name, logger))
    if existing_enc_key_sha256 != encryption_key_sha256:
      raise CommandException(
          'gsutil client error: encryption key SHA256 (%s) in tracker file '
          'does not match encryption key SHA256 (%s) of component %s' %
          (existing_enc_key_sha256, encryption_key_sha256,
           component.object_name))
    newly_completed_components = [component]
    completed_components = existing_components + newly_completed_components
    WriteParallelUploadTrackerFile(tracker_file_name,
                                   prefix,
                                   completed_components,
                                   encryption_key_sha256=encryption_key_sha256)


def WriteParallelUploadTrackerFile(tracker_file_name,
                                   prefix,
                                   components,
                                   encryption_key_sha256=None):
  """Writes information about components that were successfully uploaded.

  The tracker file is serialized JSON of the form:
  {
    "encryption_key_sha256": sha256 hash of encryption key (or null),
    "prefix": Prefix used for the component objects,
    "components": [
      {
       "component_name": Component object name,
       "component_generation": Component object generation (or null),
      }, ...
    ]
  }
  where N is the number of components that have been successfully uploaded.

  This function is not thread-safe and must be protected by a lock if
  called within Command.Apply.

  Args:
    tracker_file_name: The name of the parallel upload tracker file.
    prefix: The generated prefix that used for uploading any existing
        components.
    components: A list of ObjectFromTracker objects that were uploaded.
    encryption_key_sha256: Encryption key SHA256 for use in this upload, if any.
  """
  if six.PY3:
    if isinstance(encryption_key_sha256, bytes):
      encryption_key_sha256 = encryption_key_sha256.decode('ascii')

  tracker_components = []
  for component in components:
    tracker_components.append({
        _CompositeUploadTrackerEntry.COMPONENT_NAME: component.object_name,
        _CompositeUploadTrackerEntry.COMPONENT_GENERATION: component.generation
    })
  tracker_file_data = {
      _CompositeUploadTrackerEntry.COMPONENTS_LIST: tracker_components,
      _CompositeUploadTrackerEntry.ENC_SHA256: encryption_key_sha256,
      _CompositeUploadTrackerEntry.PREFIX: prefix
  }
  WriteJsonDataToTrackerFile(tracker_file_name, tracker_file_data)