Add a module extension for registering local/remote jdks

Author: hvadehra

java/extensions.bzl

@@ -23,6 +23,7 @@ load(
     "remote_jdk21_repos",
     "remote_jdk8_repos",
 )
+load("//toolchains:extensions.bzl", _java_repository = "java_repository")
 
 def _toolchains_impl(module_ctx):
     java_tools_repos()
@@ -38,3 +39,5 @@ def _toolchains_impl(module_ctx):
         return None
 
 toolchains = module_extension(_toolchains_impl)
+
+java_repository = _java_repository

test/repo/BUILD.bazel

@@ -53,4 +53,6 @@ default_java_toolchain(
     name = "my_funky_toolchain",
     bootclasspath = ["@bazel_tools//tools/jdk:platformclasspath"],
     configuration = NONPREBUILT_TOOLCHAIN_CONFIGURATION,
+    java_runtime = "@my_funky_jdk//:jdk",
+    exec_compatible_with = ["@platforms//os:linux", "@platforms//cpu:x86_64"],
 )

test/repo/MODULE.bazel

@@ -49,6 +49,21 @@ use_repo(
     "remotejdk21_win",
 )
 
+custom_jdk = use_extension("@rules_java//java:extensions.bzl", "java_repository")
+custom_jdk.remote(
+  name = "my_funky_jdk",
+  version = "24",
+  urls = [
+    "https://cdn.azul.com/zulu/bin/zulu24.32.13-ca-jdk24.0.2-linux_x64.tar.gz",
+  ],
+  strip_prefix = "zulu24.32.13-ca-jdk24.0.2-linux_x64",
+  prefix = "funky",
+  target_compatible_with = ["@platforms//os:linux", "@platforms//cpu:x86_64"],
+)
+use_repo(custom_jdk, "my_funky_jdk", "my_funky_jdk_toolchain_config_repo")
+register_toolchains("@my_funky_jdk_toolchain_config_repo//:all")
+
 register_toolchains("//:all")
 
 bazel_dep(name = "rules_shell", version = "0.4.0", dev_dependency = True)
+bazel_dep(name = "platforms", version = "0.0.11", dev_dependency = True)

toolchains/extensions.bzl

@@ -0,0 +1,59 @@
+"""Module extensions for local and remote java repositories"""
+
+load(":local_java_repository.bzl", "local_java_repository")
+load(":remote_java_repository.bzl", "remote_java_repository")
+
+visibility(["//java"])
+
+def _java_repository_impl(mctx):
+    for mod in mctx.modules:
+        if not mod.is_root:
+            fail(
+                """This module extension may only be used in the root module. {name}
+                must set `dev_dependency` = True on it's usage of this extension,
+                if {name} can be dependency of other modules.""".format(name = mod.name),
+            )
+        for local in mod.tags.local:
+            local_java_repository(
+                local.name,
+                java_home = local.java_home,
+                version = local.version,
+                build_file = local.build_file,
+                build_file_content = local.build_file_content,
+            )
+        for remote in mod.tags.remote:
+            remote_java_repository(
+                remote.name,
+                remote.version,
+                target_compatible_with = remote.target_compatible_with,
+                prefix = remote.prefix,
+                sha256 = remote.sha256,
+                strip_prefix = remote.strip_prefix,
+                urls = remote.urls,
+            )
+
+_local = tag_class(attrs = {
+    "name": attr.string(mandatory = True),
+    "build_file": attr.label(default = None),
+    "build_file_content": attr.string(default = ""),
+    "java_home": attr.string(default = ""),
+    "version": attr.string(default = ""),
+})
+
+_remote = tag_class(attrs = {
+    "name": attr.string(mandatory = True),
+    "version": attr.string(mandatory = True),
+    "urls": attr.string_list(mandatory = True),
+    "prefix": attr.string(default = ""),
+    "sha256": attr.string(default = ""),
+    "strip_prefix": attr.string(default = ""),
+    "target_compatible_with": attr.string_list(default = []),
+})
+
+java_repository = module_extension(
+    _java_repository_impl,
+    tag_classes = {
+        "local": _local,
+        "remote": _remote,
+    },
+)