Upload script for first class gcs deps

This script allows users to upload one file or one or more compressed
directories to GS. It will also output the DEPS blob with the required
sha256sum and size_bytes.

Verified script locally.

Bug: b/328065301
Change-Id: Ic8e894917eb4efde4f287ffb6869a22be1e89bb3
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/tools/depot_tools/+/5381863
Reviewed-by: Joanna Wang <jojwang@chromium.org>
Commit-Queue: Stephanie Kim <kimstephanie@google.com>

gclient.py

@@ -114,6 +114,7 @@ import scm as scm_git
 import setup_color
 import subcommand
 import subprocess2
+import upload_to_google_storage_first_class
 from third_party.repo.progress import Progress
 
 # TODO: Should fix these warnings.
@@ -2574,18 +2575,6 @@ class GcsDependency(Dependency):
             return True
         return False
 
-    def GetSha256Sum(self, filename):
-        sha = hashlib.sha256()
-        with open(filename, 'rb') as f:
-            while True:
-                # Read in 1mb chunks, so it doesn't all have to be loaded into
-                # memory.
-                chunk = f.read(1024 * 1024)
-                if not chunk:
-                    break
-                sha.update(chunk)
-        return sha.hexdigest()
-
     def ValidateTarFile(self, tar, prefixes):
 
         def _validate(tarinfo):
@@ -2662,7 +2651,8 @@ class GcsDependency(Dependency):
             calculated_sha256sum = 'abcd123'
             calculated_size_bytes = 10000
         else:
-            calculated_sha256sum = self.GetSha256Sum(output_file)
+            calculated_sha256sum = (
+                upload_to_google_storage_first_class.get_sha256sum(output_file))
             calculated_size_bytes = os.path.getsize(output_file)
 
         if calculated_sha256sum != self.sha256sum:

tests/download_from_google_storage_unittest.py

@@ -223,14 +223,16 @@ class DownloadTests(unittest.TestCase):
         for item in download_from_google_storage.enumerate_input(
                 self.base_path, True, True, False, None, False, False):
             self.queue.put(item)
-        expected_queue = [('e6c4fbd4fe7607f3e6ebf68b2ea4ef694da7b4fe',
+        expected_queue = [
-                           os.path.join(self.base_path, 'rootfolder_text.txt')),
+            ('e6c4fbd4fe7607f3e6ebf68b2ea4ef694da7b4fe',
-                          ('7871c8e24da15bad8b0be2c36edc9dc77e37727f',
+             os.path.join(self.base_path, 'rootfolder_text.txt')),
-                           os.path.join(self.base_path,
+            ('7871c8e24da15bad8b0be2c36edc9dc77e37727f',
-                                        'uploaded_lorem_ipsum.txt')),
+             os.path.join(self.base_path, 'uploaded_lorem_ipsum.txt')),
-                          ('b5415aa0b64006a95c0c409182e628881d6d6463',
+            ('b5415aa0b64006a95c0c409182e628881d6d6463',
-                           os.path.join(self.base_path, 'subfolder',
+             os.path.join(self.base_path, 'subfolder', 'subfolder_text.txt')),
-                                        'subfolder_text.txt'))]
+            ('b5415aa0b64006a95c0c409182e628881d6d6463',
+             os.path.join(self.base_path, 'subfolder2', 'subfolder_text.txt')),
+        ]
         self.assertEqual(sorted(expected_queue), sorted(self.queue.queue))
 
     def test_download_worker_single_file(self):

tests/gstools/download_test_data/subfolder2/subfolder_text.txt

@@ -0,0 +1,2 @@
+This is a test file.
+This file exists in a subfolder

tests/gstools/download_test_data/subfolder2/subfolder_text.txt.sha1

@@ -0,0 +1 @@
+b5415aa0b64006a95c0c409182e628881d6d6463

tests/upload_to_google_storage_first_class_unittest.py

@@ -0,0 +1,178 @@
+#!/usr/bin/env vpython3
+# Copyright (c) 2024 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.
+"""Unit tests for upload_to_google_storage_first_class.py."""
+
+from io import StringIO
+import optparse
+import os
+import posixpath
+
+import shutil
+import sys
+import tarfile
+import tempfile
+import unittest
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+import upload_to_google_storage_first_class
+from download_from_google_storage_unittest import GsutilMock
+from download_from_google_storage_unittest import ChangedWorkingDirectory
+
+# ../third_party/gsutil/gsutil
+GSUTIL_DEFAULT_PATH = os.path.join(
+    os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'third_party',
+    'gsutil', 'gsutil')
+TEST_DIR = os.path.dirname(os.path.abspath(__file__))
+
+
+class UploadTests(unittest.TestCase):
+
+    def setUp(self):
+        self.gsutil = GsutilMock(GSUTIL_DEFAULT_PATH, None)
+        self.temp_dir = tempfile.mkdtemp(prefix='gstools_test')
+        self.base_path = os.path.join(self.temp_dir, 'gstools')
+        shutil.copytree(os.path.join(TEST_DIR, 'gstools'), self.base_path)
+        self.base_url = 'gs://sometesturl'
+        self.parser = optparse.OptionParser()
+        self.lorem_ipsum = os.path.join(self.base_path, 'lorem_ipsum.txt')
+
+    def tearDown(self):
+        shutil.rmtree(self.temp_dir)
+        sys.stdin = sys.__stdin__
+
+    def test_upload_single_file(self):
+        file = self.lorem_ipsum
+        object_name = 'gs_object_name/version123'
+        upload_to_google_storage_first_class.upload_to_google_storage(
+            file=file,
+            base_url=self.base_url,
+            object_name=object_name,
+            gsutil=self.gsutil,
+            force=True,
+            gzip=False,
+            dry_run=False)
+        self.assertEqual(self.gsutil.history, [
+            ('check_call', ('ls', '%s/%s' % (self.base_url, object_name))),
+            ('check_call', ('-h', 'Cache-Control:public, max-age=31536000',
+                            'cp', file, '%s/%s' % (self.base_url, object_name)))
+        ])
+
+    def test_upload_single_file_remote_exists(self):
+        file = self.lorem_ipsum
+        object_name = 'gs_object_name/version123'
+        etag_string = 'ETag: 634d7c1ed3545383837428f031840a1e'
+        self.gsutil.add_expected(0, b'', b'')
+        self.gsutil.add_expected(0, etag_string, b'')
+
+        with self.assertRaises(Exception):
+            upload_to_google_storage_first_class.upload_to_google_storage(
+                file=file,
+                base_url=self.base_url,
+                object_name=object_name,
+                gsutil=self.gsutil,
+                force=False,
+                gzip=False,
+                dry_run=False)
+
+        self.assertEqual(self.gsutil.history,
+                         [('check_call',
+                           ('ls', '%s/%s' % (self.base_url, object_name))),
+                          ('check_call', ('ls', '-L', '%s/%s' %
+                                          (self.base_url, object_name)))])
+
+    def test_create_archive(self):
+        work_dir = os.path.join(self.base_path, 'download_test_data')
+        with ChangedWorkingDirectory(work_dir):
+            dirname = 'subfolder'
+            dirs = [dirname]
+            self.assertTrue(
+                upload_to_google_storage_first_class.validate_archive_dirs(
+                    dirs))
+            tar_filename = upload_to_google_storage_first_class.create_archive(
+                dirs)
+            with tarfile.open(tar_filename, 'r:gz') as tar:
+                content = map(lambda x: x.name, tar.getmembers())
+                self.assertIn(dirname, content)
+
+    def test_create_archive_multiple_dirs(self):
+        work_dir = os.path.join(self.base_path, 'download_test_data')
+        with ChangedWorkingDirectory(work_dir):
+            dirs = ['subfolder', 'subfolder2']
+            self.assertTrue(
+                upload_to_google_storage_first_class.validate_archive_dirs(
+                    dirs))
+            tar_filename = upload_to_google_storage_first_class.create_archive(
+                dirs)
+            with tarfile.open(tar_filename, 'r:gz') as tar:
+                content = map(lambda x: x.name, tar.getmembers())
+                for dirname in dirs:
+                    self.assertIn(dirname, content)
+
+    @unittest.skipIf(sys.platform == 'win32',
+                     'os.symlink does not exist on win')
+    def test_validate_archive_dirs_fails(self):
+        work_dir = os.path.join(self.base_path, 'download_test_data')
+        with ChangedWorkingDirectory(work_dir):
+            symlink = 'link'
+            os.symlink(os.path.join(self.base_path, 'subfolder'), symlink)
+        self.assertFalse(
+            upload_to_google_storage_first_class.validate_archive_dirs(
+                [symlink]))
+        self.assertFalse(
+            upload_to_google_storage_first_class.validate_archive_dirs(
+                ['foobar']))
+
+    def test_dry_run(self):
+        file = self.lorem_ipsum
+        object_name = 'gs_object_name/version123'
+
+        upload_to_google_storage_first_class.upload_to_google_storage(
+            file=file,
+            base_url=self.base_url,
+            object_name=object_name,
+            gsutil=self.gsutil,
+            force=False,
+            gzip=False,
+            dry_run=True)
+
+        self.assertEqual(self.gsutil.history,
+                         [('check_call',
+                           ('ls', '%s/%s' % (self.base_url, object_name))),
+                          ('check_call', ('ls', '-L', '%s/%s' %
+                                          (self.base_url, object_name)))])
+
+    def test_get_targets_no_args(self):
+        try:
+            upload_to_google_storage_first_class.get_targets([], self.parser,
+                                                             False)
+            self.fail()
+        except SystemExit as e:
+            self.assertEqual(e.code, 2)
+
+    def test_get_targets_passthrough(self):
+        result = upload_to_google_storage_first_class.get_targets(
+            ['a', 'b', 'c', 'd', 'e'], self.parser, False)
+        self.assertEqual(result, ['a', 'b', 'c', 'd', 'e'])
+
+    def test_get_targets_multiple_stdin(self):
+        inputs = ['a', 'b', 'c', 'd', 'e']
+        sys.stdin = StringIO(os.linesep.join(inputs))
+        result = upload_to_google_storage_first_class.get_targets(['-'],
+                                                                  self.parser,
+                                                                  False)
+        self.assertEqual(result, inputs)
+
+    def test_get_targets_multiple_stdin_null(self):
+        inputs = ['a', 'b', 'c', 'd', 'e']
+        sys.stdin = StringIO('\0'.join(inputs))
+        result = upload_to_google_storage_first_class.get_targets(['-'],
+                                                                  self.parser,
+                                                                  True)
+        self.assertEqual(result, inputs)
+
+
+if __name__ == '__main__':
+    unittest.main()

upload_to_google_storage_first_class.py

@@ -0,0 +1,248 @@
+#!/usr/bin/env python3
+# Copyright (c) 2024 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.
+"""Uploads files to Google Storage and output DEPS blob."""
+
+import hashlib
+import optparse
+import os
+import json
+import tempfile
+
+import re
+import sys
+import tarfile
+
+from download_from_google_storage import Gsutil
+from download_from_google_storage import GSUTIL_DEFAULT_PATH
+
+USAGE_STRING = """%prog [options] target [target2 ...].
+Target(s) is the files or directies intended to be uploaded to Google Storage.
+If a single target is a directory, it will be compressed and uploaded as a
+tar.gz file.
+If target is "-", then a list of directories will be taken from standard input.
+The list of directories will be compressed together and uploaded as one tar.gz
+file.
+
+Example usage
+------------
+./upload_to_google_storage_first_class.py --bucket gsutil-upload-playground
+--object-name my_object_name hello_world.txt
+
+./upload_to_google_storage_first_class.py --bucket gsutil-upload-playground
+--object-name my_object_name my_dir1
+
+./upload_to_google_storage_first_class.py --bucket gsutil-upload-playground
+--object-name my_object_name my_dir1 my_dir2
+
+Scan the current directory and upload all files larger than 1MB:
+find . -name .svn -prune -o -size +1000k -type f -print0 |
+./upload_to_google_storage_first_class.py --bucket gsutil-upload-playground
+--object-name my_object_name -
+"""
+
+
+def get_targets(args: list[str], parser: optparse.OptionParser,
+                use_null_terminator: bool) -> list[str]:
+    """Get target(s) to upload to GCS"""
+    if not args:
+        parser.error('Missing target.')
+
+    if len(args) == 1 and args[0] == '-':
+        # Take stdin as a newline or null separated list of files.
+        if use_null_terminator:
+            return sys.stdin.read().split('\0')
+
+        return sys.stdin.read().splitlines()
+
+    return args
+
+
+def create_archive(dirs: list[str]) -> str:
+    """Given a list of directories, compress them all into one tar file"""
+    # tarfile name cannot have a forward slash or else an error will be
+    # thrown
+    _, filename = tempfile.mkstemp(suffix='.tar.gz')
+    with tarfile.open(filename, 'w:gz') as tar:
+        for d in dirs:
+            tar.add(d)
+    return filename
+
+
+def validate_archive_dirs(dirs: list[str]) -> bool:
+    """Validate the list of directories"""
+    for d in dirs:
+        # We don't allow .. in paths in our archives.
+        if d == '..':
+            return False
+        # We only allow dirs.
+        if not os.path.isdir(d):
+            return False
+        # Symlinks must point to a target inside the dirs
+        if os.path.islink(d) and not any(
+                os.realpath(d).startswith(os.realpath(dir_prefix))
+                for dir_prefix in dirs):
+            return False
+        # We required that the subdirectories we are archiving are all just
+        # below cwd.
+        if d not in next(os.walk('.'))[1]:
+            return False
+
+    return True
+
+
+def get_sha256sum(filename: str) -> str:
+    """Get the sha256sum of the file"""
+    sha = hashlib.sha256()
+    with open(filename, 'rb') as f:
+        while True:
+            # Read in 1mb chunks, so it doesn't all have to be loaded into
+            # memory.
+            chunk = f.read(1024 * 1024)
+            if not chunk:
+                break
+            sha.update(chunk)
+    return sha.hexdigest()
+
+
+def upload_to_google_storage(file: str, base_url: str, object_name: str,
+                             gsutil: Gsutil, force: bool, gzip: str,
+                             dry_run: bool):
+    """Upload file to GCS"""
+    file_url = '%s/%s' % (base_url, object_name)
+    if gsutil.check_call('ls', file_url)[0] == 0 and not force:
+        # File exists, check MD5 hash.
+        _, out, _ = gsutil.check_call_with_retries('ls', '-L', file_url)
+        etag_match = re.search(r'ETag:\s+\S+', out)
+        if etag_match:
+            raise Exception('File with url %s already exists' % file_url)
+    if dry_run:
+        return
+    print("Uploading %s as %s" % (file, file_url))
+    gsutil_args = ['-h', 'Cache-Control:public, max-age=31536000', 'cp']
+    if gzip:
+        gsutil_args.extend(['-z', gzip])
+    gsutil_args.extend([file, file_url])
+    code, _, err = gsutil.check_call_with_retries(*gsutil_args)
+    if code != 0:
+        raise Exception(
+            code, 'Encountered error on uploading %s to %s\n%s' %
+            (file, file_url, err))
+
+
+def construct_deps_blob(bucket: str, object_name: str, file: str) -> dict:
+    """Output a blob hint that would need be added to a DEPS file"""
+    sha256sum = get_sha256sum(file)
+    size_bytes = os.path.getsize(file)
+    return {
+        '<path>': {
+            'dep_type': 'gcs',
+            'bucket': bucket,
+            'object_name': object_name,
+            'sha256sum': sha256sum,
+            'size_bytes': size_bytes,
+        }
+    }
+
+
+def main():
+    parser = optparse.OptionParser(USAGE_STRING)
+    parser.add_option('-b',
+                      '--bucket',
+                      help='Google Storage bucket to upload to.')
+    parser.add_option('-o',
+                      '--object-name',
+                      help='Optional object name of uploaded tar file. '
+                      'If empty, the sha256sum will be the object name.')
+    parser.add_option('-d',
+                      '--dry-run',
+                      action='store_true',
+                      help='Check if file already exists on GS without '
+                      'uploading it and output DEP blob.')
+    parser.add_option('-c',
+                      '--config',
+                      action='store_true',
+                      help='Alias for "gsutil config".  Run this if you want '
+                      'to initialize your saved Google Storage '
+                      'credentials.  This will create a read-only '
+                      'credentials file in ~/.boto.depot_tools.')
+    parser.add_option('-e', '--boto', help='Specify a custom boto file.')
+    parser.add_option('-f',
+                      '--force',
+                      action='store_true',
+                      help='Force upload even if remote file exists.')
+    parser.add_option('-g',
+                      '--gsutil_path',
+                      default=GSUTIL_DEFAULT_PATH,
+                      help='Path to the gsutil script.')
+    parser.add_option('-0',
+                      '--use_null_terminator',
+                      action='store_true',
+                      help='Use \\0 instead of \\n when parsing '
+                      'the file list from stdin.  This is useful if the input '
+                      'is coming from "find ... -print0".')
+    parser.add_option('-z',
+                      '--gzip',
+                      metavar='ext',
+                      help='For files which end in <ext> gzip them before '
+                      'upload. '
+                      'ext is a comma-separated list')
+    (options, args) = parser.parse_args()
+
+    # Enumerate our inputs.
+    input_filenames = get_targets(args, parser, options.use_null_terminator)
+
+    if len(input_filenames) > 1 or (len(input_filenames) == 1
+                                    and os.path.isdir(input_filenames[0])):
+        if not validate_archive_dirs(input_filenames):
+            parser.error(
+                'Only directories just below cwd are valid entries. '
+                'Entries cannot contain .. and entries can not be symlinks. '
+                'Entries was %s' % input_filenames)
+            return 1
+        file = create_archive(input_filenames)
+    else:
+        file = input_filenames[0]
+
+    object_name = options.object_name
+    if not object_name:
+        object_name = get_sha256sum(file)
+
+    # Make sure we can find a working instance of gsutil.
+    if os.path.exists(GSUTIL_DEFAULT_PATH):
+        gsutil = Gsutil(GSUTIL_DEFAULT_PATH, boto_path=options.boto)
+    else:
+        gsutil = None
+        for path in os.environ["PATH"].split(os.pathsep):
+            if os.path.exists(path) and 'gsutil' in os.listdir(path):
+                gsutil = Gsutil(os.path.join(path, 'gsutil'),
+                                boto_path=options.boto)
+        if not gsutil:
+            parser.error('gsutil not found in %s, bad depot_tools checkout?' %
+                         GSUTIL_DEFAULT_PATH)
+
+    # Passing in -g/--config will run our copy of GSUtil, then quit.
+    if options.config:
+        print('===Note from depot_tools===')
+        print('If you do not have a project ID, enter "0" when asked for one.')
+        print('===End note from depot_tools===')
+        print()
+        gsutil.check_call('version')
+        return gsutil.call('config')
+
+    base_url = 'gs://%s' % options.bucket
+
+    upload_to_google_storage(file, base_url, object_name, gsutil, options.force,
+                             options.gzip, options.dry_run)
+    print(
+        json.dumps(construct_deps_blob(options.bucket, object_name, file),
+                   indent=2))
+
+
+if __name__ == '__main__':
+    try:
+        sys.exit(main())
+    except KeyboardInterrupt:
+        sys.stderr.write('interrupted\n')
+        sys.exit(1)