Revert of [Telemetry][Wpr] Remove old version of archive info.

telemetry/telemetry/wpr/archive_info2.py

Index: telemetry/telemetry/wpr/archive_info2.py
diff --git a/telemetry/telemetry/wpr/archive_info2.py b/telemetry/telemetry/wpr/archive_info2.py
new file mode 100644
index 0000000000000000000000000000000000000000..56ad36ecdbe83654b06103048dd9ff399a2d94f1
--- /dev/null
+++ b/telemetry/telemetry/wpr/archive_info2.py
@@ -0,0 +1,219 @@
+# Copyright 2013 The Chromium Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+import json
+import logging
+import os
+import re
+import shutil
+import tempfile
+
+from py_utils import cloud_storage  # pylint: disable=import-error
+
+
+_DEFAULT_PLATFORM = 'DEFAULT'
+_ALL_PLATFORMS = ['mac', 'linux', 'android', 'win', _DEFAULT_PLATFORM]
+
+
+# TODO(rnephew): Move this file to archive_info.py when it is the default.
+# These changes were chosen to not be merged with archive_info.py directly
+# because all code in archive_info.py will be deleted once everything is
+# switched to using this version.
+def AssertValidCloudStorageBucket(bucket):
+  is_valid = bucket in (None,
+                        cloud_storage.PUBLIC_BUCKET,
+                        cloud_storage.PARTNER_BUCKET,
+                        cloud_storage.INTERNAL_BUCKET)
+  if not is_valid:
+    raise ValueError("Cloud storage privacy bucket %s is invalid" % bucket)
+
+
+class WprArchiveInfo(object):
+  def __init__(self, file_path, data, bucket):
+    AssertValidCloudStorageBucket(bucket)
+    self._file_path = file_path
+    self._base_dir = os.path.dirname(file_path)
+    self._data = data
+    self._bucket = bucket
+    self.temp_target_wpr_file_path = None
+    # Ensure directory exists.
+    if not os.path.exists(self._base_dir):
+      os.makedirs(self._base_dir)
+
+    assert data.get('platform_specific', False), (
+        'archive_info2.WprArchiveInfo is exclusively for platform specific WPR '
+        'archives. Platform agnostic archives should use '
+        'archive_info.WprArchiveInfo archive_info2 will become archive_info '
+        'once all data files are switched.')
+
+    self._story_name_to_wpr_file = data['archives']
+
+  @classmethod
+  def FromFile(cls, file_path, bucket):
+    """ Generates an archive_info2 instance with the given json file. """
+    if os.path.exists(file_path):
+      with open(file_path, 'r') as f:
+        data = json.load(f)
+        return cls(file_path, data, bucket)
+    return cls(file_path, {'archives': {}}, bucket)
+
+  def DownloadArchivesIfNeeded(self, target_platforms=None):
+    """Downloads archives iff the Archive has a bucket parameter and the user
+    has permission to access the bucket.
+
+    Raises cloud storage Permissions or Credentials error when there is no
+    local copy of the archive and the user doesn't have permission to access
+    the archive's bucket.
+
+    Warns when a bucket is not specified or when the user doesn't have
+    permission to access the archive's bucket but a local copy of the archive
+    exists.
+    """
+    # If no target platform is set, download all platforms.
+    if target_platforms is None:
+      target_platforms = _ALL_PLATFORMS
+    else:
+      assert isinstance(target_platforms, list), 'Must pass platforms as a list'
+      target_platforms = target_platforms + [_DEFAULT_PLATFORM]
+    # Download all .wpr files.
+    if not self._bucket:
+      logging.warning('Story set in %s has no bucket specified, and '
+                      'cannot be downloaded from cloud_storage.', )
+      return
+    assert 'archives' in self._data, ("Invalid data format in %s. 'archives' "
+                                      "field is needed" % self._file_path)
+
+    def download_if_needed(path):
+      try:
+        cloud_storage.GetIfChanged(path, self._bucket)
+      except (cloud_storage.CredentialsError, cloud_storage.PermissionError):
+        if os.path.exists(path):
+          # If the archive exists, assume the user recorded their own and warn
+          # them that they do not have the proper credentials to download.
+          logging.warning('Need credentials to update WPR archive: %s', path)
+        else:
+          logging.error("You either aren't authenticated or don't have "
+                        "permission to use the archives for this page set."
+                        "\nYou may need to run gsutil config."
+                        "\nYou can find instructions for gsutil config at: "
+                        "http://www.chromium.org/developers/telemetry/"
+                        "upload_to_cloud_storage")
+          raise
+
+    story_archives = self._data['archives']
+    for story in story_archives:
+      for target_platform in target_platforms:
+        if story_archives[story].get(target_platform):
+          archive_path = self._WprFileNameToPath(
+              story_archives[story][target_platform])
+          download_if_needed(archive_path)
+
+  def WprFilePathForStory(self, story, target_platform=_DEFAULT_PLATFORM):
+    if self.temp_target_wpr_file_path:
+      return self.temp_target_wpr_file_path
+
+    wpr_file = self._story_name_to_wpr_file.get(story.display_name, None)
+    if wpr_file is None and hasattr(story, 'url'):
+      # Some old pages always use the URL to identify a page rather than the
+      # display_name, so try to look for that.
+      wpr_file = self._story_name_to_wpr_file.get(story.url, None)
+    if wpr_file:
+      if target_platform in wpr_file:
+        return self._WprFileNameToPath(wpr_file[target_platform])
+      return self._WprFileNameToPath(wpr_file[_DEFAULT_PLATFORM])
+    return None
+
+  def AddNewTemporaryRecording(self, temp_wpr_file_path=None):
+    if temp_wpr_file_path is None:
+      temp_wpr_file_handle, temp_wpr_file_path = tempfile.mkstemp()
+      os.close(temp_wpr_file_handle)
+    self.temp_target_wpr_file_path = temp_wpr_file_path
+
+  def AddRecordedStories(self, stories, upload_to_cloud_storage=False,
+                         target_platform=_DEFAULT_PLATFORM):
+    if not stories:
+      os.remove(self.temp_target_wpr_file_path)
+      return
+
+    (target_wpr_file, target_wpr_file_path) = self._NextWprFileName()
+    for story in stories:
+      # Check to see if the platform has been manually overrided.
+      if not story.platform_specific:
+        current_target_platform = _DEFAULT_PLATFORM
+      else:
+        current_target_platform = target_platform
+      self._SetWprFileForStory(
+          story.display_name, target_wpr_file, current_target_platform)
+    shutil.move(self.temp_target_wpr_file_path, target_wpr_file_path)
+
+    # Update the hash file.
+    target_wpr_file_hash = cloud_storage.CalculateHash(target_wpr_file_path)
+    with open(target_wpr_file_path + '.sha1', 'wb') as f:
+      f.write(target_wpr_file_hash)
+      f.flush()
+
+    self._WriteToFile()
+
+    # Upload to cloud storage
+    if upload_to_cloud_storage:
+      if not self._bucket:
+        logging.warning('StorySet must have bucket specified to upload '
+                        'stories to cloud storage.')
+        return
+      try:
+        cloud_storage.Insert(self._bucket, target_wpr_file_hash,
+                             target_wpr_file_path)
+      except cloud_storage.CloudStorageError, e:
+        logging.warning('Failed to upload wpr file %s to cloud storage. '
+                        'Error:%s' % target_wpr_file_path, e)
+
+  def _WriteToFile(self):
+    """Writes the metadata into the file passed as constructor parameter."""
+    metadata = dict()
+    metadata['description'] = (
+        'Describes the Web Page Replay archives for a story set. '
+        'Don\'t edit by hand! Use record_wpr for updating.')
+    metadata['archives'] = self._story_name_to_wpr_file.copy()
+    metadata['platform_specific'] = True
+
+    with open(self._file_path, 'w') as f:
+      json.dump(metadata, f, indent=4, sort_keys=True, separators=(',', ': '))
+      f.flush()
+
+  def _WprFileNameToPath(self, wpr_file):
+    return os.path.abspath(os.path.join(self._base_dir, wpr_file))
+
+  def _NextWprFileName(self):
+    """Creates a new file name for a wpr archive file."""
+    # The names are of the format "some_thing_number.wpr". Read the numbers.
+    highest_number = -1
+    base = None
+    wpr_files = []
+    for story in self._data['archives']:
+      for p in self._data['archives'][story]:
+        wpr_files.append(self._data['archives'][story][p])
+
+    for wpr_file in wpr_files:
+      match = re.match(r'(?P<BASE>.*)_(?P<NUMBER>[0-9]+)\.wpr', wpr_file)
+      if not match:
+        raise Exception('Illegal wpr file name ' + wpr_file)
+      highest_number = max(int(match.groupdict()['NUMBER']), highest_number)
+      if base and match.groupdict()['BASE'] != base:
+        raise Exception('Illegal wpr file name ' + wpr_file +
+                        ', doesn\'t begin with ' + base)
+      base = match.groupdict()['BASE']
+    if not base:
+      # If we're creating a completely new info file, use the base name of the
+      # story set file.
+      base = os.path.splitext(os.path.basename(self._file_path))[0]
+    new_filename = '%s_%03d.wpr' % (base, highest_number + 1)
+    return new_filename, self._WprFileNameToPath(new_filename)
+
+  def _SetWprFileForStory(self, story_name, wpr_file, target_platform):
+    """For modifying the metadata when we're going to record a new archive."""
+    if story_name not in self._data['archives']:
+      # If there is no other recording we want the first to be the default
+      # until a new default is recorded.
+      self._data['archives'][story_name] = {_DEFAULT_PLATFORM: wpr_file}
+    self._data['archives'][story_name][target_platform] = wpr_file