123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124 |
- import argparse
- import os
- import subprocess
- import sys
- import time
- from collections import OrderedDict
- from typing import List
- import requests
- from aws_requests_auth.boto_utils import BotoAWSRequestsAuth
- def retry(f):
- def inner():
- resp = None
- for _ in range(5):
- resp = f()
- print("Getting Presigned URL, status_code", resp.status_code)
- if resp.status_code >= 500:
- print("errored, retrying...")
- print(resp.text)
- time.sleep(5)
- else:
- return resp
- if resp is None or resp.status_code >= 500:
- print("still errorred after many retries")
- sys.exit(1)
- return inner
- @retry
- def perform_auth():
- auth = BotoAWSRequestsAuth(
- aws_host="vop4ss7n22.execute-api.us-west-2.amazonaws.com",
- aws_region="us-west-2",
- aws_service="execute-api",
- )
- resp = requests.get(
- "https://vop4ss7n22.execute-api.us-west-2.amazonaws.com/endpoint/",
- auth=auth,
- params={"job_id": os.environ["BUILDKITE_JOB_ID"]},
- )
- return resp
- def handle_docker_login(resp):
- pwd = resp.json()["docker_password"]
- subprocess.call(
- ["docker", "login", "--username", "raytravisbot", "--password", pwd]
- )
- def gather_paths(dir_path) -> List[str]:
- dir_path = dir_path.replace("/", os.path.sep)
- assert os.path.exists(dir_path)
- if os.path.isdir(dir_path):
- paths = [os.path.join(dir_path, f) for f in os.listdir(dir_path)]
- else:
- paths = [dir_path]
- return paths
- dest_resp_mapping = {
- "wheels": "presigned_resp_prod_wheels",
- "branch_wheels": "presigned_resp_prod_wheels",
- "jars": "presigned_resp_prod_wheels",
- "branch_jars": "presigned_resp_prod_wheels",
- "logs": "presigned_logs",
- }
- def upload_paths(paths, resp, destination):
- dest_key = dest_resp_mapping[destination]
- c = resp.json()[dest_key]
- of = OrderedDict(c["fields"])
- sha = os.environ["BUILDKITE_COMMIT"]
- branch = os.environ["BUILDKITE_BRANCH"]
- bk_job_id = os.environ["BUILDKITE_JOB_ID"]
- current_os = sys.platform
- for path in paths:
- fn = os.path.split(path)[-1]
- of["key"] = {
- "wheels": f"latest/{fn}",
- "branch_wheels": f"{branch}/{sha}/{fn}",
- "jars": f"jars/latest/{current_os}/{fn}",
- "branch_jars": f"jars/{branch}/{sha}/{current_os}/{fn}",
- "logs": f"bazel_events/{branch}/{sha}/{bk_job_id}/{fn}",
- }[destination]
- of["file"] = open(path, "rb")
- r = requests.post(c["url"], files=of)
- print(f"Uploaded {path} to {of['key']}", r.status_code)
- if __name__ == "__main__":
- parser = argparse.ArgumentParser(
- description="Helper script to upload files to S3 bucket"
- )
- parser.add_argument("--path", type=str, required=False)
- parser.add_argument("--destination", type=str)
- args = parser.parse_args()
- assert args.destination in {
- "branch_jars",
- "branch_wheels",
- "jars",
- "logs",
- "wheels",
- "docker_login",
- }
- assert "BUILDKITE_JOB_ID" in os.environ
- assert "BUILDKITE_COMMIT" in os.environ
- resp = perform_auth()
- if args.destination == "docker_login":
- handle_docker_login(resp)
- else:
- paths = gather_paths(args.path)
- print("Planning to upload", paths)
- upload_paths(paths, resp, args.destination)
|