From 75df6c36b8faef0b24b205eb008c5bfa5e29a7f6 Mon Sep 17 00:00:00 2001
From: Jingwen Chen <jingwen@google.com>
Date: Tue, 5 Nov 2024 13:08:52 +0000
Subject: [PATCH] Allow root module's override tags to take precedence over the
 overrides from the transitive deps.

---
 MODULE.bazel                                  | 20 ++++++++++++
 private/extensions/maven.bzl                  | 17 +++++++---
 tests/integration/override_targets/BUILD      | 17 ++++++++++
 .../override_targets/module/MODULE.bazel      | 18 +++++++++++
 ...odule_can_override_protobuf_to_javapoet.sh | 31 +++++++++++++++++++
 5 files changed, 99 insertions(+), 4 deletions(-)
 create mode 100644 tests/integration/override_targets/module/MODULE.bazel
 create mode 100755 tests/integration/override_targets/root_module_can_override_protobuf_to_javapoet.sh

diff --git a/MODULE.bazel b/MODULE.bazel
index 0a304db4d..13045dc99 100644
--- a/MODULE.bazel
+++ b/MODULE.bazel
@@ -732,6 +732,25 @@ dev_maven.install(
     ],
 )
 
+
+dev_maven.install(
+    name = "root_module_can_override",
+    artifacts = ["com.squareup:javapoet:1.11.1"],
+)
+
+bazel_dep(name = "transitive_module_can_override", version = "0.0.0")
+local_path_override(
+    module_name = "transitive_module_can_override",
+    path = "tests/integration/override_targets/module",
+)
+
+dev_maven.override(
+    # This override demonstrates that this root module's override takes precedence over that transitive override definition.
+    # Use something absurd for testing, like overriding protobuf-java to javapoet.
+    coordinates = "com.google.protobuf:protobuf-java",
+    target = "@root_module_can_override//:com_squareup_javapoet",
+)
+
 # Where there are file locks, the pinned and unpinned repos are listed
 # next to each other. Where compat repositories are created, they are
 # listed next to the repo that created them. The list is otherwise kept
@@ -815,6 +834,7 @@ use_repo(
     "starlark_aar_import_test",
     "starlark_aar_import_with_sources_test",
     "strict_visibility_testing",
+    "root_module_can_override",
 
     # Repo with compat repos
     "com_google_http_client_google_http_client_gson",
diff --git a/private/extensions/maven.bzl b/private/extensions/maven.bzl
index cbe406479..93f3c5639 100644
--- a/private/extensions/maven.bzl
+++ b/private/extensions/maven.bzl
@@ -227,13 +227,22 @@ def maven_impl(mctx):
     # module attempts to update a maven repo (which is normally undesired behaviour)
     repo_name_2_module_name = {}
 
-    for mod in mctx.modules:
+    # First compute the overrides. The order of the transitive overrides do not matter, but the root
+    # overrides take precedence over all transitive ones.
+    for idx, mod in enumerate(reversed(mctx.modules)):
+        # Rotate the root module to the last to be visited.
+        is_root_module = idx == (len(mctx.modules) - 1)
         for override in mod.tags.override:
             value = str(override.target)
-            current = overrides.get(override.coordinates, None)
-            to_use = _fail_if_different("Target of override for %s" % override.coordinates, current, value, [None])
-            overrides.update({override.coordinates: to_use})
+            if is_root_module:
+                # Allow the root module's overrides to take precedence over any transitive overrides.
+                to_use = value
+            else:
+                current = overrides.get(override.coordinates, None)
+                to_use = _fail_if_different("Target of override for %s" % override.coordinates, current, value, [None])
+            overrides.update({override.coordinates: value})
 
+    for mod in mctx.modules:
         for artifact in mod.tags.artifact:
             _check_repo_name(repo_name_2_module_name, artifact.name, mod.name)
 
diff --git a/tests/integration/override_targets/BUILD b/tests/integration/override_targets/BUILD
index 3702d9ad4..3baf98d92 100644
--- a/tests/integration/override_targets/BUILD
+++ b/tests/integration/override_targets/BUILD
@@ -69,3 +69,20 @@ sh_test(
         "@bazel_tools//tools/bash/runfiles",
     ],
 )
+
+genquery(
+    name = "root_module_can_override",
+    expression = "deps(@root_module_can_override//:com_google_protobuf_protobuf_java)",
+    opts = [
+        "--nohost_deps",
+        "--noimplicit_deps",
+    ],
+    scope = ["@root_module_can_override//:com_google_protobuf_protobuf_java"],
+)
+
+sh_test(
+    name = "root_module_can_override_protobuf_to_javapoet",
+    srcs = ["root_module_can_override_protobuf_to_javapoet.sh"],
+    data = [":root_module_can_override"],
+    deps = ["@bazel_tools//tools/bash/runfiles"],
+)
diff --git a/tests/integration/override_targets/module/MODULE.bazel b/tests/integration/override_targets/module/MODULE.bazel
new file mode 100644
index 000000000..d839d57d2
--- /dev/null
+++ b/tests/integration/override_targets/module/MODULE.bazel
@@ -0,0 +1,18 @@
+module(name = "transitive_module_can_override", version = "0.0.0")
+
+bazel_dep(name = "rules_jvm_external", version = "0.0")
+local_path_override(
+    module_name = "rules_jvm_external",
+    path = "../../../..",
+)
+
+maven = use_extension("@rules_jvm_external//:extensions.bzl", "maven")
+maven.install(
+    name = "root_module_can_override",
+    artifacts = ["com.google.protobuf:protobuf-java:3.25.5"],
+)
+
+maven.override(
+    coordinates = "com.google.protobuf:protobuf-java",
+    target = "//:poison_pill_non_existent_target",
+)
diff --git a/tests/integration/override_targets/root_module_can_override_protobuf_to_javapoet.sh b/tests/integration/override_targets/root_module_can_override_protobuf_to_javapoet.sh
new file mode 100755
index 000000000..8faf367ec
--- /dev/null
+++ b/tests/integration/override_targets/root_module_can_override_protobuf_to_javapoet.sh
@@ -0,0 +1,31 @@
+# --- begin runfiles.bash initialization v2 ---
+# Copy-pasted from the Bazel Bash runfiles library v2.
+set -uo pipefail; 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 v2 ---
+
+set -euox pipefail
+
+deps_file=$(rlocation rules_jvm_external/tests/integration/override_targets/root_module_can_override)
+
+function clean_up_workspace_names() {
+  local file_name="$1"
+  local target="$2"
+  # The first `sed` command replaces `@@` with `@`. The second extracts the visible name
+  # from the bzlmod mangled workspace name
+  cat "$file_name" | sed -e 's|^@@|@|g; s|\r||g' | sed -e 's|^@[^/]*[+~]|@|g; s|\r||g' | grep "$target"
+  cat "$file_name" | sed -e 's|^@@|@|g; s|\r||g' | sed -e 's|^@[^/]*[+~]|@|g; s|\r||g' | grep -q "$target"
+}
+
+if ! clean_up_workspace_names "$deps_file" "@root_module_can_override//:com_google_protobuf_protobuf_java"; then
+  exit 1
+fi
+
+if ! clean_up_workspace_names "$deps_file" "@root_module_can_override//:com_squareup_javapoet"; then
+  exit 1
+fi