recipe_modules/gcs_upload/api.py

# Copyright 2025 The Pigweed Authors
#
# Licensed under the Apache License, Version 2.0 (the "License"); you may not
# use this file except in compliance with the License. You may obtain a copy of
# the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations under
# the License.
"""This module provides an API for uploading artifacts to Google Cloud Storage.

It supports namespacing uploads (e.g., by build ID or custom prefix) and
running uploads in parallel using the futures module. It handles both files
and directories, and can generate links to the uploaded artifacts.
"""

from __future__ import annotations

import contextlib
import dataclasses
from collections.abc import Generator

from recipe_engine import config_types, recipe_api, util

from PB.go.chromium.org.luci.buildbucket.proto import common as common_pb
from PB.recipe_engine import result as result_pb

from RECIPE_MODULES.recipe_engine.futures import api as futures_api


@dataclasses.dataclass
class GcsUploadContext:
    _api: recipe_api.RecipeApi
    gcs_bucket: str
    namespace: str | None
    _futures: list[futures_api.Future] = dataclasses.field(default_factory=list)

    def __post_init__(self):
        if not self.namespace:
            self.namespace = self._api.buildbucket.build.id

    def subcontext(self, subdirectory: str) -> GcsUploadContext:
        return GcsUploadContext(
            _api=self._api,
            gcs_bucket=self.gcs_bucket,
            namespace=f'{self.namespace}/{subdirectory.strip("/")}',
            _futures=self._futures,
        )

    def upload(
        self,
        source: config_types.Path | util.InputPlaceholder,
        dest: str,
    ) -> None:
        with self._api.step.nest(f'upload {dest}'):
            if not self._api.path.exists(source):
                self._api.step.empty('does not exist')  # pragma: no cover

            elif self._api.path.isfile(source):
                self._api.step.empty('file')

                # Copy any files to another location immediately. Future steps
                # might reuse the same output directory.
                temp = self._api.path.mkdtemp() / source.name
                self._api.file.copy('copy', source, temp)

                self._futures.append(
                    self._api.futures.spawn(
                        self._api.gsutil.upload_namespaced_file,
                        bucket=self.gcs_bucket,
                        namespace=self.namespace,
                        source=temp,
                        subpath=dest,
                    ),
                )

            elif self._api.path.isdir(source):
                self._api.step.empty('directory')

                # Copy any files to another location immediately. Future steps
                # might reuse the same output directory.
                temp = self._api.path.mkdtemp() / source.name
                self._api.file.copytree('copytree', source, temp)

                self._futures.append(
                    self._api.futures.spawn(
                        self._api.file.listdir,
                        'ls',
                        temp,
                        recursive=True,
                    )
                )
                self._futures.append(
                    self._api.futures.spawn(
                        self._api.gsutil.upload_namespaced_directory,
                        bucket=self.gcs_bucket,
                        namespace=self.namespace,
                        source=temp,
                        subpath=dest,
                    ),
                )

    def wait(self) -> None:
        self._api.futures.wait(self._futures)

    @property
    def link(self) -> str:
        return self._api.gsutil.namespaced_directory_url(self.gcs_bucket)

    def set_output_properties(self) -> None:
        pres = self._api.step.empty('set output properties').presentation
        pres.properties['artifacts'] = self.link

    def raw_result(self) -> result_pb.RawResult:
        return result_pb.RawResult(
            summary_markdown=f'[artifacts]({self.link})',
            status=common_pb.SUCCESS,
        )


class GcsUploadApi(recipe_api.RecipeApi):
    """API for uploading files and directories to Google Cloud Storage."""

    @contextlib.contextmanager
    def __call__(
        self,
        *,
        gcs_bucket: str | None = None,
        namespace: str | None = None,
        parent: GcsUploadContext | None = None,
        subdirectory: str | None = None,
    ) -> Generator[GcsUploadContext, None, None]:
        """Enter a GcsUploadContext. See also init()."""
        ctx: GcsUploadContext | None = None
        try:
            ctx = self.init(
                gcs_bucket=gcs_bucket,
                namespace=namespace,
                parent=parent,
                subdirectory=subdirectory,
            )
            yield ctx
        finally:
            if ctx:
                ctx.wait()
                ctx.set_output_properties()

    def init(
        self,
        *,
        gcs_bucket: str | None = None,
        namespace: str | None = None,
        parent: GcsUploadContext | None = None,
        subdirectory: str | None = None,
    ) -> GcsUploadContext:
        """Create a GcsUploadContext object.

        Callers can pass gcs_bucket (with an optional namespace), or they can
        pass parent and subdirectory. Other combinations are not allowed.

        Args:
            gcs_bucket: Destination bucket.
            namespace: Prefix in destination bucket.
            parent: Parent context object.
            subdirectory: Path relative to parent.

        Returns:
            Context object.
        """
        assert parent or gcs_bucket, f'parent={parent} gcs_bucket={gcs_bucket}'

        if parent:
            assert subdirectory
            assert not gcs_bucket
            assert not namespace
            return parent.subcontext(subdirectory)

        if gcs_bucket:
            assert not parent
            assert not subdirectory
            return GcsUploadContext(
                _api=self.m,
                gcs_bucket=gcs_bucket,
                namespace=namespace,
            )