From 5403d2ebfcdd549ba2e4fc46205f8fbd8bc9022b Mon Sep 17 00:00:00 2001
From: Stuart Axelbrooke <stuart@axelbrooke.com>
Date: Fri, 18 Apr 2025 12:16:52 -0700
Subject: [PATCH] Implement wrapper job

---
 rules.bzl | 101 +++++++++++++++++++++++++++++++++++++++++-------------
 1 file changed, 78 insertions(+), 23 deletions(-)

diff --git a/rules.bzl b/rules.bzl
index 352d363..5847d18 100644
--- a/rules.bzl
+++ b/rules.bzl
@@ -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."""