diff options
author | Daniel Baumann <daniel.baumann@progress-linux.org> | 2024-04-19 00:47:55 +0000 |
---|---|---|
committer | Daniel Baumann <daniel.baumann@progress-linux.org> | 2024-04-19 00:47:55 +0000 |
commit | 26a029d407be480d791972afb5975cf62c9360a6 (patch) | |
tree | f435a8308119effd964b339f76abb83a57c29483 /tools/moztreedocs/upload.py | |
parent | Initial commit. (diff) | |
download | firefox-26a029d407be480d791972afb5975cf62c9360a6.tar.xz firefox-26a029d407be480d791972afb5975cf62c9360a6.zip |
Adding upstream version 124.0.1.upstream/124.0.1
Signed-off-by: Daniel Baumann <daniel.baumann@progress-linux.org>
Diffstat (limited to 'tools/moztreedocs/upload.py')
-rw-r--r-- | tools/moztreedocs/upload.py | 171 |
1 files changed, 171 insertions, 0 deletions
diff --git a/tools/moztreedocs/upload.py b/tools/moztreedocs/upload.py new file mode 100644 index 0000000000..a607135174 --- /dev/null +++ b/tools/moztreedocs/upload.py @@ -0,0 +1,171 @@ +# This Source Code Form is subject to the terms of the Mozilla Public +# License, v. 2.0. If a copy of the MPL was not distributed with this +# file, # You can obtain one at http://mozilla.org/MPL/2.0/. + +import io +import mimetypes +import os +import sys +from concurrent import futures +from pprint import pprint + +import boto3 +import botocore +import requests +from mozbuild.util import memoize + + +@memoize +def create_aws_session(): + """ + This function creates an aws session that is + shared between upload and delete both. + """ + region = "us-west-2" + level = os.environ.get("MOZ_SCM_LEVEL", "1") + bucket = { + "1": "gecko-docs.mozilla.org-l1", + "2": "gecko-docs.mozilla.org-l2", + "3": "gecko-docs.mozilla.org", + }[level] + secrets_url = "http://taskcluster/secrets/v1/secret/" + secrets_url += "project/releng/gecko/build/level-{}/gecko-docs-upload".format(level) + + # Get the credentials from the TC secrets service. Note that these + # differ per SCM level + if "TASK_ID" in os.environ: + print("Using AWS credentials from the secrets service") + session = requests.Session() + res = session.get(secrets_url) + res.raise_for_status() + secret = res.json()["secret"] + session = boto3.session.Session( + aws_access_key_id=secret["AWS_ACCESS_KEY_ID"], + aws_secret_access_key=secret["AWS_SECRET_ACCESS_KEY"], + region_name=region, + ) + else: + print("Trying to use your AWS credentials..") + session = boto3.session.Session(region_name=region) + + s3 = session.client("s3", config=botocore.client.Config(max_pool_connections=20)) + + return s3, bucket + + +@memoize +def get_s3_keys(s3, bucket): + kwargs = {"Bucket": bucket} + all_keys = [] + while True: + response = s3.list_objects_v2(**kwargs) + for obj in response["Contents"]: + all_keys.append(obj["Key"]) + + try: + kwargs["ContinuationToken"] = response["NextContinuationToken"] + except KeyError: + break + + return all_keys + + +def s3_set_redirects(redirects): + s3, bucket = create_aws_session() + + configuration = {"IndexDocument": {"Suffix": "index.html"}, "RoutingRules": []} + + for path, redirect in redirects.items(): + rule = { + "Condition": {"KeyPrefixEquals": path}, + "Redirect": {"ReplaceKeyPrefixWith": redirect}, + } + if os.environ.get("MOZ_SCM_LEVEL") == "3": + rule["Redirect"]["HostName"] = "firefox-source-docs.mozilla.org" + + configuration["RoutingRules"].append(rule) + + s3.put_bucket_website( + Bucket=bucket, + WebsiteConfiguration=configuration, + ) + + +def s3_delete_missing(files, key_prefix=None): + """Delete files in the S3 bucket. + + Delete files on the S3 bucket that doesn't match the files + given as the param. If the key_prefix is not specified, missing + files that has main/ as a prefix will be removed. Otherwise, it + will remove files with the same prefix as key_prefix. + """ + s3, bucket = create_aws_session() + files_on_server = get_s3_keys(s3, bucket) + if key_prefix: + files_on_server = [ + path for path in files_on_server if path.startswith(key_prefix) + ] + else: + files_on_server = [ + path for path in files_on_server if not path.startswith("main/") + ] + files = [key_prefix + "/" + path if key_prefix else path for path, f in files] + files_to_delete = [path for path in files_on_server if path not in files] + + query_size = 1000 + while files_to_delete: + keys_to_remove = [{"Key": key} for key in files_to_delete[:query_size]] + response = s3.delete_objects( + Bucket=bucket, + Delete={ + "Objects": keys_to_remove, + }, # NOQA + ) + pprint(response, indent=2) + files_to_delete = files_to_delete[query_size:] + + +def s3_upload(files, key_prefix=None): + """Upload files to an S3 bucket. + + ``files`` is an iterable of ``(path, BaseFile)`` (typically from a + mozpack Finder). + + Keys in the bucket correspond to source filenames. If ``key_prefix`` is + defined, key names will be ``<key_prefix>/<path>``. + """ + s3, bucket = create_aws_session() + + def upload(f, path, bucket, key, extra_args): + # Need to flush to avoid buffering/interleaving from multiple threads. + sys.stdout.write("uploading %s to %s\n" % (path, key)) + sys.stdout.flush() + s3.upload_fileobj(f, bucket, key, ExtraArgs=extra_args) + + fs = [] + with futures.ThreadPoolExecutor(20) as e: + for path, f in files: + content_type, content_encoding = mimetypes.guess_type(path) + extra_args = {} + if content_type: + if content_type.startswith("text/"): + content_type += '; charset="utf-8"' + extra_args["ContentType"] = content_type + if content_encoding: + extra_args["ContentEncoding"] = content_encoding + + if key_prefix: + key = "%s/%s" % (key_prefix, path) + else: + key = path + + # The file types returned by mozpack behave like file objects. But + # they don't accept an argument to read(). So we wrap in a BytesIO. + fs.append( + e.submit(upload, io.BytesIO(f.read()), path, bucket, key, extra_args) + ) + + s3_delete_missing(files, key_prefix) + # Need to do this to catch any exceptions. + for f in fs: + f.result() |