recipes/workflows.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.
"""Run './pw' tools and builds."""

from __future__ import annotations

import contextlib
import re
import shlex
from collections.abc import Iterator, Mapping, Sequence
from typing import Any

from recipe_engine import (
    config_types,
    post_process,
    recipe_api,
    recipe_test_api,
)

from PB.pigweed.pw_build.proto import build_driver as build_driver_pb
from PB.pigweed.pw_build.proto import workflows as workflows_pb
from PB.recipe_engine import result as result_pb
from PB.recipes.pigweed.workflows import InputProperties

from RECIPE_MODULES.pigweed.checkout import api as checkout_api
from RECIPE_MODULES.pigweed.gcs_upload import api as gcs_upload_api
from RECIPE_MODULES.pigweed.raw_result import api as raw_result_api

DEPS = [
    'fuchsia/buildbucket_util',
    'fuchsia/macos_sdk',
    'fuchsia/sso',
    'pigweed/bazel',
    'pigweed/checkout',
    'pigweed/default_timeout',
    'pigweed/gcs_upload',
    'pigweed/initialization',
    'pigweed/raw_result',
    'pigweed/trigger',
    'recipe_engine/buildbucket',
    'recipe_engine/cipd',
    'recipe_engine/context',
    'recipe_engine/defer',
    'recipe_engine/file',
    'recipe_engine/path',
    'recipe_engine/platform',
    'recipe_engine/properties',
    'recipe_engine/proto',
    'recipe_engine/raw_io',
    'recipe_engine/step',
]

PROPERTIES = InputProperties


def _test_request() -> build_driver_pb.BuildDriverRequest:
    result = build_driver_pb.BuildDriverRequest()

    build = result.jobs.add().build
    build.name = 'test_build'
    build.description = 'test_build description'
    build.build_config.name = 'test_build config'
    build.build_config.description = 'test_build config description'
    build.build_config.build_type = 'test_build config build_type'
    build.build_config.args.extend(('--arg-1', '--arg-2'))
    build.rerun_shortcut = 'build test_build'

    tool = result.jobs.add().tool
    tool.name = 'test_tool'
    tool.description = 'test_tool description'
    tool.build_config.name = 'test_tool config'
    tool.build_config.description = 'test_tool config description'
    tool.build_config.build_type = 'test_tool config build_type'
    tool.build_config.args.extend(('--arg-a', '--arg-b'))
    tool.analyzer_friendly_args.append('--check')
    tool.rerun_shortcut = 'check test_tool'

    return result


@raw_result_api.wrap_run_steps
def RunSteps(
    api: recipe_api.RecipeScriptApi, props: InputProperties
) -> result_pb.RawResult | None:
    exit_early, checkout = api.initialization(props.checkout_options)
    if exit_early:
        return exit_early  # pragma: no cover
    assert checkout

    # b/445709816 - Old Bazel servers can persist on macOS bots across
    # different runs, which causes remote cache authentication/connection
    # problems.
    if api.platform.is_mac:
        build = api.buildbucket.build
        if 'pigweed.skip_mac_bazel_shutdown' not in build.input.experiments:
            with api.step.nest('shut down old Bazel server'):
                bazel = api.bazel.ensure_bazelisk()
                with api.context(cwd=checkout.root):
                    api.step('shut down old Bazel server', [bazel, 'shutdown'])

    with contextlib.ExitStack() as stack:
        stack.enter_context(api.step.nest('describe steps'))
        stack.enter_context(_bazel_cache_context(api))

        experiments = api.buildbucket.build.input.experiments
        assert not (
            props.bazel_version and 'pigweed.latest_bazel' in experiments
        )

        if props.bazel_version:
            stack.enter_context(
                api.context(env={'USE_BAZEL_VERSION': props.bazel_version}),
            )
        if 'pigweed.latest_bazel' in experiments:
            stack.enter_context(
                api.context(env={'USE_BAZEL_VERSION': 'latest'}),
            )

        bazel = api.bazel.ensure_bazelisk()

        workflow_output_base = api.path.cleanup_dir / 'workflows_launcher'

        extra_bazel_arguments = _process_extra_bazel_arguments(
            props.extra_bazel_arguments,
            {'CHECKOUT_ROOT': checkout.root},
        )

        bazel_args = [
            *api.bazel.rbe_arguments(remote=False, remote_cache=True),
            *api.bazel.override_arguments(checkout),
            *api.bazel.airlock_arguments(),
            *extra_bazel_arguments,
        ]
        wrapped_bazel_args = [f'--extra-arg=bazel={x}' for x in bazel_args]

        build_args = [
            bazel,
            f'--output_base={workflow_output_base}',
            'build',
            f'--symlink_prefix={workflow_output_base}/bazel-',
            *bazel_args,
            # Don't upload BES results whenever we launch the workflows entry
            # point. This prevents noisy BES results urls at the top of each
            # build page that are unrelated to the launched steps.
            '--bes_results_url=',
            '//:pw',
        ]

        cmd_prefix = [
            *['run' if arg == 'build' else arg for arg in build_args],
            '--',
            *wrapped_bazel_args,
        ]

        describe_cmd = [
            *cmd_prefix,
            '--directory',
            checkout.root,
            '--no-banner',
            'describe',
            '--infra-metadata',
            *props.steps,
        ]

        # TODO: b/441786225 - Remove bazel.parent from PATH.
        with api.context(
            cwd=checkout.root,
            env_prefixes={'PATH': [bazel.parent]},
        ):
            with api.macos_sdk():
                api.step('build pw', build_args)
                steps_raw = api.step(
                    'describe steps',
                    describe_cmd,
                    stdout=api.raw_io.output_text(),
                    step_test_data=(
                        lambda: api.raw_io.test_api.stream_output_text(
                            api.proto.encode(_test_request(), api.proto.TEXTPB),
                        )
                    ),
                ).stdout

        with api.step.nest('steps') as pres:
            pres.logs['steps'] = steps_raw

            steps_proto = api.proto.decode(
                steps_raw,
                build_driver_pb.BuildDriverRequest,
                api.proto.TEXTPB,
            )

            for step in steps_proto.jobs:
                if step.HasField('build'):
                    api.step.empty(step.build.rerun_shortcut)
                elif step.HasField('tool'):
                    api.step.empty(step.tool.rerun_shortcut)
                else:
                    raise ValueError('empty oneof')  # pragma: no cover

    keep_going_base = ()
    for change in checkout.changes:
        if 'build-errors: continue' in change.commit_message.lower():
            keep_going_base = ('--keep_going',)

    gcs_raw_result: result_pb.RawResult | None = None

    with contextlib.ExitStack() as outer_stack:
        defer = outer_stack.enter_context(api.defer.context())

        gcs: gcs_upload_api.GcsUploadContext | None = None
        if props.gcs_bucket:
            gcs = outer_stack.enter_context(
                api.gcs_upload(gcs_bucket=props.gcs_bucket),
            )
            with api.step.nest('upload metadata'):
                checkout_snapshot = api.path.mkdtemp() / '_checkout'
                checkout.snapshot_to_dir(checkout_snapshot)
                gcs.upload(checkout_snapshot, '_checkout')

        outer_stack.enter_context(_bazel_cache_context(api))

        for job in steps_proto.jobs:
            if job.HasField('build'):
                rerun_shortcut = job.build.rerun_shortcut
                keep_going = keep_going_base
            elif job.HasField('tool'):
                rerun_shortcut = job.tool.rerun_shortcut
                keep_going = ()
            else:
                raise ValueError('empty oneof')  # pragma: no cover

            with contextlib.ExitStack() as stack:
                # TODO: b/441786225 - Remove bazel.parent from PATH.
                stack.enter_context(
                    api.context(
                        cwd=checkout.root,
                        env={'LUCI_BUILD_ID': str(api.buildbucket.build.id)},
                        env_prefixes={'PATH': [bazel.parent]},
                    ),
                )
                pres = stack.enter_context(api.step.nest(rerun_shortcut))
                stack.enter_context(api.default_timeout())
                stack.enter_context(api.macos_sdk())

                artifacts_manifest: str | None = None
                artifacts_manifest_args: tuple[str, ...] = ()

                if gcs or props.cipd_yaml_file:
                    artifacts_manifest = (
                        api.path.mkdtemp() / 'artifacts-manifest.textproto'
                    )
                    artifacts_manifest_args = (
                        f'--artifacts-manifest={artifacts_manifest}',
                    )

                step_cmd = [
                    *cmd_prefix,
                    *artifacts_manifest_args,
                    *shlex.split(rerun_shortcut),
                    *keep_going,
                ]

                defer(api.step, 'run', step_cmd)

                if gcs or props.cipd_yaml_file:
                    artifacts = read_artifacts(api, artifacts_manifest)

                if gcs:
                    num_uploaded = upload_artifacts_to_gcs(
                        api,
                        gcs,
                        artifacts,
                    )
                    if num_uploaded:
                        entries = 'entries' if num_uploaded > 1 else 'entry'
                        pres.step_summary_text = (
                            f'uploaded {num_uploaded} {entries}'
                        )

                if props.cipd_yaml_file:
                    if not api.buildbucket_util.is_dev_or_try:
                        pin = upload_artifacts_to_cipd(
                            api,
                            checkout=checkout,
                            artifacts=artifacts,
                            yaml_file=props.cipd_yaml_file,
                            refs=props.cipd_refs,
                        )
                        pres.links['@'.join(pin)] = api.cipd.make_link(*pin)

        if gcs:
            gcs_raw_result = gcs.raw_result()

    api.trigger(props.trigger_options, checkout, verbose=True)

    return gcs_raw_result


def _process_extra_bazel_arguments(
    extra_bazel_arguments: Sequence[str],
    variables: Mapping[str, str],
):
    result: list[str] = []
    for arg in extra_bazel_arguments:

        def repl(match: re.Match) -> str:
            return str(variables.get(match.group(1), match.group(0)))

        final_arg = arg
        final_arg = re.sub(r'\$([\w]+)', repl, final_arg)
        final_arg = re.sub(r'\$\{([\w]+)\}', repl, final_arg)
        result.append(final_arg)

    return result


def _bazel_cache_context(api: recipe_api.RecipeScriptApi):
    # See b/441999854. Don't allow builds that have github/RBE disabled and
    # builds that don't to share Bazel caches. For now, do this by disabling
    # the local cache for builds that disable github.
    test_tmpdir = api.path.cache_dir / 'bazel'

    if (
        'pigweed.disable_github' in api.buildbucket.build.input.experiments
        or 'pigweed.disable_rbe' in api.buildbucket.build.input.experiments
    ):
        test_tmpdir = api.path.mkdtemp()

    return api.context(env={'TEST_TMPDIR': test_tmpdir})


def _test_artifacts(
    api: recipe_api.RecipeScriptApi,
) -> workflows_pb.BuildArtifacts:
    result = workflows_pb.BuildArtifacts()
    output_root = api.path.start_dir / 'output'
    result.output_root = str(output_root)
    group_one = result.output_groups.add()
    group_one.add_prefix = 'one/'
    group_one.matching_files.append('foo/bar.txt')
    group_two = result.output_groups.add()
    group_two.strip_prefix = 'prefix/'
    group_two.matching_files.append('prefix/ham.txt')
    group_two.matching_files.append('prefix/spam.txt')
    group_two.matching_files.append('prefix/eggs.txt')

    for group in result.output_groups:
        for file in group.matching_files:
            api.path.mock_add_file(output_root / file)

    return result


def read_artifacts(
    api: recipe_api.RecipeScriptApi,
    manifest: config_types.Path,
) -> workflows_pb.BuildArtifacts:
    api.path.mock_add_file(manifest)
    assert api.path.isfile(manifest), manifest

    raw_artifacts = api.file.read_text(
        'read artifacts manifest',
        manifest,
        test_data=api.proto.encode(_test_artifacts(api), api.proto.TEXTPB),
    )

    return api.proto.decode(
        raw_artifacts,
        workflows_pb.BuildArtifacts,
        api.proto.TEXTPB,
    )


def upload_artifacts_to_gcs(
    api: recipe_api.RecipeScriptApi,
    gcs: gcs_upload_api.GcsUploadContext,
    artifacts: workflows_pb.BuildArtifacts,
) -> int:
    num_uploaded = 0

    output_root = api.context.cwd / artifacts.output_root

    for group in artifacts.output_groups:
        for file in group.matching_files:
            source = output_root / file
            api.path.mock_add_file(source)
            destination = file.removeprefix(group.strip_prefix or '').lstrip(
                '/'
            )
            if group.add_prefix:
                prefix = group.add_prefix.strip('/')
                destination = f'{prefix}/{destination}'
            gcs.upload(source, destination)
            num_uploaded += 1

    return num_uploaded


def upload_artifacts_to_cipd(
    api: recipe_api.RecipeScriptApi,
    checkout: checkout_api.CheckoutContext,
    artifacts: workflows_pb.BuildArtifacts,
    yaml_file: config_types.Path,
    refs: Sequence[str],
) -> None:
    yaml_path = api.context.cwd / artifacts.output_root / yaml_file
    api.path.mock_add_file(yaml_path)
    assert api.path.isfile(yaml_path), yaml_path

    metadata = [
        api.cipd.Metadata(
            'git_repository',
            api.sso.sso_to_https(checkout.options.remote),
        ),
    ]

    tags = {
        'git_revision': checkout.revision(),
    }

    return api.cipd.create_from_yaml(
        pkg_def=yaml_path,
        refs=refs,
        tags=tags,
        metadata=metadata,
    )


def GenTests(
    api: recipe_test_api.RecipeTestApi,
) -> Iterator[recipe_test_api.TestData]:
    def properties(**kwargs: Any) -> recipe_test_api.TestData:
        props = InputProperties(**kwargs)
        props.checkout_options.CopyFrom(api.checkout.git_options())
        return api.properties(props)

    def must_upload(
        prefix: str,
        name: str,
        bucket: str = 'bucket',
    ) -> recipe_test_api.TestData:
        return api.post_process(
            post_process.MustRun,
            f'{prefix}.upload {name}.upload {name} to {bucket}',
        )

    def has_args(step, *args):
        return api.post_process(post_process.StepCommandContains, step, args)

    create_step = 'check test_tool.create cipd.yaml'

    yield api.test(
        'upload',
        properties(
            gcs_bucket='bucket',
            cipd_yaml_file='cipd.yaml',
            cipd_refs=('latest',),
        ),
        must_upload('build test_build', 'one/foo/bar.txt'),
        must_upload('build test_build', 'ham.txt'),
        must_upload('build test_build', 'spam.txt'),
        must_upload('build test_build', 'eggs.txt'),
        has_args(create_step, '-ref', 'latest'),
        has_args(create_step, '-tag', 'git_revision:HASH'),
        has_args(create_step, '-metadata', re.compile('^git_repository:.*')),
        api.post_process(post_process.DropExpectation),
    )

    yield api.test(
        'mac',
        properties(
            extra_bazel_arguments=[
                '--extra-arg=${CHECKOUT_ROOT}/with',
                '--extra-arg=$CHECKOUT_ROOT/without',
            ],
        ),
        api.buildbucket.try_build(
            experiments=['pigweed.disable_rbe', 'pigweed.latest_bazel'],
            git_repo='https://pigweed.googlesource.com/pigweed/pigweed',
        ),
        api.platform.name('mac'),
        api.post_check(  # One check that './pw' is invoked with the arg.
            post_process.StepCommandContains,
            'build test_build.run',
            [
                '--extra-arg=[START_DIR]/co/with',
                '--extra-arg=[START_DIR]/co/without',
            ],
        ),
        api.post_check(  # One check that the build is invoked with the arg.
            post_process.StepCommandContains,
            'build test_build.run',
            [
                '--extra-arg=bazel=--extra-arg=[START_DIR]/co/with',
                '--extra-arg=bazel=--extra-arg=[START_DIR]/co/without',
            ],
        ),
        api.post_process(post_process.DropExpectation),
    )

    yield api.test(
        'build-errors-continue',
        properties(bazel_version='1.2.3.4.5.6.7.8.9.0'),
        api.checkout.try_test_data(),
        api.checkout.cl_branch_parents(message='Build-Errors: continue'),
        api.post_check(
            post_process.StepCommandContains,
            'build test_build.run',
            ['--keep_going'],
        ),
        api.post_check(
            post_process.StepCommandDoesNotContain,
            'check test_tool.run',
            ['--keep_going'],
        ),
        api.post_process(post_process.DropExpectation),
    )