Implement wrapper job

This commit is contained in:
Stuart Axelbrooke 2025-04-18 12:16:52 -07:00
parent 8ca2b4a00b
commit 5403d2ebfc
No known key found for this signature in database
GPG key ID: 1B0A848C29D46A35

101
rules.bzl
View file

@ -2,9 +2,7 @@ def databuild_job(
name,
configure,
execute,
deps = [],
visibility = None,
**kwargs):
visibility = None):
"""Creates a DataBuild job target with configuration and execution capabilities.
Args:
@ -25,9 +23,15 @@ def databuild_job(
# Create the main rule that serves as a provider for other targets
_databuild_job_exec_rule(
name = name + ".exec",
configure = ":%s.cfg" % name,
execute = execute,
deps = deps,
visibility = visibility,
)
# Create a job target that configures then executes
_databuild_job_rule(
name = name,
configure = ":%s.cfg" % name,
execute = ":%s.exec" % name,
visibility = visibility,
)
@ -80,11 +84,6 @@ _databuild_job_cfg_rule = rule(
)
def _databuild_job_exec_impl(ctx):
deps_providers = []
for dep in ctx.attr.deps:
if DataBuildJobInfo in dep:
deps_providers.append(dep[DataBuildJobInfo])
execute_file = ctx.executable.execute
jq_file = ctx.executable._jq
@ -115,11 +114,6 @@ def _databuild_job_exec_impl(ctx):
executable = script,
runfiles = runfiles,
),
DataBuildJobInfo(
configure = ctx.attr.configure,
execute = script,
deps = deps_providers,
),
]
# Define the provider
@ -135,20 +129,12 @@ DataBuildJobInfo = provider(
_databuild_job_exec_rule = rule(
implementation = _databuild_job_exec_impl,
attrs = {
"configure": attr.label(
doc = "Target that implements the configuration logic",
mandatory = True,
),
"execute": attr.label(
doc = "Target that implements the execution logic",
mandatory = True,
executable = True,
cfg = "exec",
),
"deps": attr.label_list(
doc = "Dependencies (other job targets)",
default = [],
),
"_template": attr.label(
default = "@databuild//job:execute_wrapper.sh.tpl",
allow_single_file = True,
@ -166,6 +152,75 @@ _databuild_job_exec_rule = rule(
executable = True,
)
def _databuild_job_impl(ctx):
"""Wraps the configure and execute targets in a shell script."""
script = ctx.actions.declare_file(ctx.label.name)
ctx.actions.write(
output = script,
is_executable = True,
content = """#!/bin/bash
# TODO should this be extracted to shared init script
# Get the directory where the script is located
if [[ -z "${{RUNFILES_DIR:-}}" ]]; then
SCRIPT_DIR="$(readlink -f "${{BASH_SOURCE[0]}}")"
# Set RUNFILES_DIR relative to the script location
export RUNFILES_DIR="${{SCRIPT_DIR}}.runfiles"
fi
# --- begin runfiles.bash initialization v3 ---
# Copy-pasted from the Bazel Bash runfiles library v3.
set -uo pipefail; set +e; f=bazel_tools/tools/bash/runfiles/runfiles.bash
source "${{RUNFILES_DIR:-/dev/null}}/$f" 2>/dev/null || \
source "$(grep -sm1 "^$f " "${{RUNFILES_MANIFEST_FILE:-/dev/null}}" | cut -f2- -d' ')" 2>/dev/null || \
source "$0.runfiles/$f" 2>/dev/null || \
source "$(grep -sm1 "^$f " "$0.runfiles_manifest" | cut -f2- -d' ')" 2>/dev/null || \
source "$(grep -sm1 "^$f " "$0.exe.runfiles_manifest" | cut -f2- -d' ')" 2>/dev/null || \
{{ echo>&2 "ERROR: cannot find $f"; exit 1; }}; f=; set -e
# --- end runfiles.bash initialization v3 ---
$(rlocation _main/{configure_path}) $@ | $(rlocation _main/{execute_path})
""".format(
configure_path = ctx.attr.configure.files_to_run.executable.short_path,
execute_path = ctx.attr.execute.files_to_run.executable.short_path,
),
)
runfiles = ctx.runfiles(
files = [ctx.executable.execute, ctx.executable.configure],
).merge(ctx.attr.execute.default_runfiles).merge(ctx.attr.configure.default_runfiles)
return [
DefaultInfo(
executable = script,
runfiles = runfiles,
),
DataBuildJobInfo(
configure = ctx.attr.configure,
execute = script,
),
]
_databuild_job_rule = rule(
implementation = _databuild_job_impl,
attrs = {
"configure": attr.label(
doc = "Target that implements the configuration logic",
mandatory = True,
executable = True,
# TODO all these cdf=execs are probably a problem for deployment
cfg = "exec",
),
"execute": attr.label(
doc = "Target that implements the execution logic",
mandatory = True,
executable = True,
cfg = "exec",
),
},
executable = True,
)
def databuild_graph(name, jobs, plan, visibility = None):
"""Creates a databuild graph target."""