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

toolchains/extensions.bzl

@@ -0,0 +1,53 @@
+"""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:
+        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.label_list(default = []),
+})
+
+java_repository = module_extension(
+    _java_repository_impl,
+    tag_classes = {
+        "local": _local,
+        "remote": _remote,
+    },
+)