databuild/rules.bzl

def databuild_job(
        name,
        configure,
        execute,
        deps = [],
        visibility = None,
        **kwargs):
    """Creates a DataBuild job target with configuration and execution capabilities.

    Args:
        name: Name of the job target
        configure: Target that implements the configuration logic
        execute: Target that implements the execution logic
        deps: List of other job_targets this job depends on
        visibility: Visibility specification
        **kwargs: Additional attributes to pass to the underlying rule
    """

    _databuild_job_cfg_rule(
        name = name + ".cfg",
        configure = configure,
        visibility = visibility,
    )

    # 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,
    )

def _databuild_job_cfg_impl(ctx):
    configure_file = ctx.executable.configure
    configure_path = ctx.attr.configure.files_to_run.executable.path
    script = ctx.actions.declare_file(ctx.label.name)

    ctx.actions.expand_template(
        template = ctx.file._template,
        output = script,
        substitutions = {
            "%{CONFIGURE_PATH}": configure_path,
        },
        is_executable = True,
    )

    runfiles = ctx.runfiles(
        files = [configure_file],
    ).merge(ctx.attr.configure.default_runfiles).merge(
        ctx.attr._bash_runfiles.default_runfiles,
    )

    return [
        DefaultInfo(
            executable = script,
            runfiles = runfiles,
        ),
    ]

_databuild_job_cfg_rule = rule(
    implementation = _databuild_job_cfg_impl,
    attrs = {
        "configure": attr.label(
            doc = "Target that implements the configuration logic",
            executable = True,
            cfg = "exec",
            mandatory = True,
        ),
        "_template": attr.label(
            default = "@databuild//job:configure_wrapper.sh.tpl",
            allow_single_file = True,
        ),
        "_bash_runfiles": attr.label(
            default = Label("@bazel_tools//tools/bash/runfiles"),
            allow_files = True,
        ),
    },
    executable = True,
)

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

    script = ctx.actions.declare_file(ctx.label.name)

    # Get the correct runfiles paths
    jq_path = ctx.attr._jq.files_to_run.executable.path
    execute_path = ctx.attr.execute.files_to_run.executable.path

    ctx.actions.expand_template(
        template = ctx.file._template,
        output = script,
        substitutions = {
            "%{JQ_PATH}": jq_path,
            "%{EXECUTE_PATH}": execute_path,
        },
        is_executable = True,
    )

    runfiles = ctx.runfiles(
        files = [jq_file, execute_file],
    ).merge(ctx.attr.execute.default_runfiles).merge(ctx.attr._jq.default_runfiles)

    return [
        DefaultInfo(
            executable = script,
            runfiles = runfiles,
        ),
        DataBuildJobInfo(
            configure = ctx.attr.configure,
            execute = script,
            deps = deps_providers,
        ),
    ]

# Define the provider
DataBuildJobInfo = provider(
    doc = "Information about a DataBuild job",
    fields = {
        "configure": "Target that implements the configuration logic",
        "execute": "Target that implements the execution logic",
        "deps": "List of dependencies (other DataBuildJobInfo providers)",
    },
)

_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,
        ),
        "_jq": attr.label(
            default = "@databuild//runtime:jq",
            executable = True,
            cfg = "exec",
        ),
    },
    executable = True,
)

def databuild_graph(name, jobs, plan, visibility = None):
    """Creates a databuild graph target."""

#def _graph_impl(name):
#    """
#
#    """
#
#    # Lets do this
#    pass
#
#databuild_graph = rule(
#    implementation = _graph_impl,
#    attrs = {
#        "jobs": attr.label_list(
#            doc = "The list of jobs that are candidates for building partitions in this databuild graph",
#            allow_empty = False,
#        ),
#        "plan": attr.label(
#            doc = "The binary that is run to produce a `JobGraph` that builds the requested partition refs",
#            executable = True,
#            cfg = "exec",
#        ),
#    },
#)