diff --git a/ci/deploy_mz-debug/pipeline.template.yml b/ci/deploy_mz-debug/pipeline.template.yml index ef8df122bc9df..12e2da721d8de 100644 --- a/ci/deploy_mz-debug/pipeline.template.yml +++ b/ci/deploy_mz-debug/pipeline.template.yml @@ -34,6 +34,8 @@ steps: command: bin/pyactivate -m ci.deploy_mz-debug.macos agents: queue: mac-aarch64 + env: + MZ_BINARIES_UPLOAD_ROLE_ARN: arn:aws:iam::834237029485:role/buildkite-macos-deploy timeout_in_minutes: 30 concurrency: 1 concurrency_group: deploy-mz-debug/macos/aarch64 diff --git a/ci/deploy_mz-deploy/pipeline.template.yml b/ci/deploy_mz-deploy/pipeline.template.yml index 5ab0a5b2e2c34..8e6ae81cc3558 100644 --- a/ci/deploy_mz-deploy/pipeline.template.yml +++ b/ci/deploy_mz-deploy/pipeline.template.yml @@ -34,6 +34,8 @@ steps: command: bin/pyactivate -m ci.deploy_mz-deploy.macos agents: queue: mac-aarch64 + env: + MZ_BINARIES_UPLOAD_ROLE_ARN: arn:aws:iam::834237029485:role/buildkite-macos-deploy timeout_in_minutes: 30 concurrency: 1 concurrency_group: deploy-mz-deploy/macos/aarch64 diff --git a/ci/deploy_mz/pipeline.template.yml b/ci/deploy_mz/pipeline.template.yml index 010b3d41f6eeb..1417be476ec0d 100644 --- a/ci/deploy_mz/pipeline.template.yml +++ b/ci/deploy_mz/pipeline.template.yml @@ -42,9 +42,12 @@ steps: concurrency: 1 concurrency_group: deploy-mz/docker - - command: bin/pyactivate -m ci.deploy_mz.macos + - id: deploy_mz-macos-aarch64 + command: bin/pyactivate -m ci.deploy_mz.macos agents: queue: mac-aarch64 + env: + MZ_BINARIES_UPLOAD_ROLE_ARN: arn:aws:iam::834237029485:role/buildkite-macos-deploy timeout_in_minutes: 30 concurrency: 1 concurrency_group: deploy-mz/macos/aarch64 diff --git a/ci/tarball_uploader.py b/ci/tarball_uploader.py index 9c51eb5237568..072610e06979f 100644 --- a/ci/tarball_uploader.py +++ b/ci/tarball_uploader.py @@ -10,10 +10,12 @@ # This file contains a class that can be used to upload a tarball to the materialize-binaries S3 bucket. import os +import subprocess import tarfile import tempfile import time from pathlib import Path +from typing import Any import boto3 import humanize @@ -23,12 +25,55 @@ BINARIES_BUCKET = "materialize-binaries" +# When set, uploads use short-lived credentials for this role, obtained by +# exchanging the job's Buildkite OIDC token, rather than the ambient AWS +# credentials. +UPLOAD_ROLE_ARN_ENV = "MZ_BINARIES_UPLOAD_ROLE_ARN" + + +def _s3_client() -> Any: + role_arn = os.environ.get(UPLOAD_ROLE_ARN_ENV) + if not role_arn: + return boto3.client("s3") + + # Requested immediately before use: OIDC tokens are short-lived and the + # build that precedes the upload can outlast one. + token = subprocess.check_output( + [ + "buildkite-agent", + "oidc", + "request-token", + "--audience", + "sts.amazonaws.com", + ], + text=True, + ).strip() + session = boto3.Session() + creds = session.client("sts").assume_role_with_web_identity( + RoleArn=role_arn, + RoleSessionName=f"buildkite-{os.environ['BUILDKITE_JOB_ID']}"[:64], + WebIdentityToken=token, + DurationSeconds=3600, + )["Credentials"] + return boto3.client( + "s3", + aws_access_key_id=creds["AccessKeyId"], + aws_secret_access_key=creds["SecretAccessKey"], + aws_session_token=creds["SessionToken"], + ) + class TarballUploader: def __init__(self, package_name: str, version: TypedVersionBase): self.package_name = package_name self.version = version self.version_str = f"v{version.str_without_prefix()}" + self._s3: Any = None + + def _client(self) -> Any: + if self._s3 is None: + self._s3 = _s3_client() + return self._s3 def _create_tardir(self, name: str) -> tarfile.TarInfo: """Takes a dir path and creates a TarInfo. Sets the @@ -46,7 +91,7 @@ def _sanitize_tarinfo(self, tarinfo: tarfile.TarInfo) -> tarfile.TarInfo: def _upload_redirect(self, key: str, to: str) -> None: with tempfile.NamedTemporaryFile() as empty: - boto3.client("s3").upload_fileobj( + self._client().upload_fileobj( Fileobj=empty, Bucket=BINARIES_BUCKET, Key=key, @@ -56,7 +101,7 @@ def _upload_redirect(self, key: str, to: str) -> None: def _upload_tarball(self, tarball: Path, platform: str) -> None: """Uploads the tarball (.tar) to the S3 bucket.""" s3_object = f"{self.package_name}-{self.version_str}-{platform}.tar.gz" - boto3.client("s3").upload_file( + self._client().upload_file( Filename=str(tarball), Bucket=BINARIES_BUCKET, Key=s3_object,