Solves a bug

Author: ncordon

x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/TestAnalyzer.java

@@ -392,7 +392,9 @@ public TestAnalyzer addLookupIndex(String name, String mappingLocation) {
     }
 
     /**
-     * Add an error resolving enrich indices.
+     * Add an error resolving enrich indices. If {@code policyName}/{@code mode} already has a pending registration (from an
+     * earlier {@link #addEnrichPolicy}/{@link #addEnrichError} call), this one is matched to the next occurrence of that
+     * policy name/mode in the query instead of the first - see {@link #resolveEnrichResolution}.
      */
     public TestAnalyzer addEnrichError(String policyName, Enrich.Mode mode, String reason) {
         pendingEnrichResolutions.add(new PendingEnrich(policyName, mode, null, reason));
@@ -450,7 +452,10 @@ public TestAnalyzer addEnrichPolicy(String policyType, String policy, String fie
     }
 
     /**
-     * Adds an enrich policy resolution with a specific mode by loading the mapping from a resource file.
+     * Adds an enrich policy resolution with a specific mode by loading the mapping from a resource file. If
+     * {@code policy}/{@code mode} already has a pending registration (from an earlier {@link #addEnrichPolicy}/
+     * {@link #addEnrichError} call), this one is matched to the next occurrence of that policy name/mode in the query
+     * instead of the first - see {@link #resolveEnrichResolution}.
      */
     public TestAnalyzer addEnrichPolicy(Enrich.Mode mode, String policyType, String policy, String field, String index, String mapping) {
         IndexResolution indexResolution = loadMapping(mapping, index, IndexMode.STANDARD);
@@ -464,7 +469,11 @@ public TestAnalyzer addEnrichPolicy(Enrich.Mode mode, String policyType, String
     }
 
     /**
-     * Adds an enrich policy resolution with a specific mode by loading the mapping from a resource file.
+     * Adds an enrich policy resolution with a specific mode. If {@code policy}/{@code mode} already has a pending
+     * registration (from an earlier {@link #addEnrichPolicy}/{@link #addEnrichError} call), this one is matched to the next
+     * occurrence of that policy name/mode in the query instead of the first - see {@link #resolveEnrichResolution}. This
+     * lets a test give two occurrences of the same policy name/mode (e.g. one per subquery branch) different resolutions,
+     * simply by calling {@link #addEnrichPolicy}/{@link #addEnrichError} more than once for that policy name/mode.
      */
     public TestAnalyzer addEnrichPolicy(Enrich.Mode mode, String policy, ResolvedEnrichPolicy resolved) {
         pendingEnrichResolutions.add(new PendingEnrich(policy, mode, resolved, null));
@@ -475,23 +484,38 @@ public TestAnalyzer addEnrichPolicy(Enrich.Mode mode, String policy, ResolvedEnr
      * Matches pending {@link #addEnrichPolicy}/{@link #addEnrichError} registrations (queued by policy name + mode before the
      * query was known) against the actual {@link Enrich} occurrences in the now-parsed plan, and registers each match into the
      * real {@link #enrichResolution} keyed by that occurrence's {@code Source} - mirroring how {@code EnrichPolicyResolver}
-     * keys production resolutions. If several occurrences share the same policy name and mode, every one of them is
-     * registered with the same resolution.
+     * keys production resolutions.
+     * <p>
+     * Registrations for the same policy name/mode are consumed in registration order, matched 1:1 against occurrences of
+     * that policy name/mode in the order {@link Enrich} nodes are visited by {@link LogicalPlan#forEachUp} over the parsed
+     * plan (for a top-to-bottom or left-to-right query this matches the order the occurrences appear in the query text): the
+     * first registration goes to the first occurrence, the second to the second, and so on. If there are more occurrences
+     * than registrations, the extra occurrences reuse the last registration - which is also why a single registration for a
+     * given policy name/mode applies to every occurrence of it, as it always has before this per-occurrence matching
+     * existed.
      */
     private void resolveEnrichResolution(LogicalPlan plan) {
+        Map<PolicyKey, List<PendingEnrich>> byKey = pendingEnrichResolutions.stream()
+            .collect(groupingBy(p -> new PolicyKey(p.policyName(), p.mode()), LinkedHashMap::new, Collectors.toList()));
+        Map<PolicyKey, Integer> occurrencesSeen = new HashMap<>();
         plan.forEachUp(Enrich.class, enrich -> {
-            for (PendingEnrich pending : pendingEnrichResolutions) {
-                if (pending.policyName().equals(enrich.resolvedPolicyName()) && pending.mode() == enrich.mode()) {
-                    if (pending.resolved() != null) {
-                        enrichResolution.addResolvedPolicy(enrich.source(), pending.resolved());
-                    } else {
-                        enrichResolution.addError(enrich.source(), pending.error());
-                    }
-                }
+            PolicyKey key = new PolicyKey(enrich.resolvedPolicyName(), enrich.mode());
+            List<PendingEnrich> candidates = byKey.get(key);
+            if (candidates == null) {
+                return;
+            }
+            int occurrence = occurrencesSeen.merge(key, 1, Integer::sum) - 1;
+            PendingEnrich matched = candidates.get(Math.min(occurrence, candidates.size() - 1));
+            if (matched.resolved() != null) {
+                enrichResolution.addResolvedPolicy(enrich.source(), matched.resolved());
+            } else {
+                enrichResolution.addError(enrich.source(), matched.error());
             }
         });
     }
 
+    private record PolicyKey(String policyName, Enrich.Mode mode) {}
+
     /**
      * Set external source resolution.
      */

x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/CrossClusterInSubqueryIT.java

@@ -1173,6 +1173,106 @@ public void testEnrichWithRemoteModePolicyMixedScopesSucceeds() {
         }
     }
 
+    /**
+     * Two occurrences of the same ENRICH command - same policy name, same mode - sharing the exact same {@code Source}
+     */
+    public void testEnrichWithViewReferencedTwiceSharesSource() {
+        assumeTrue("Requires FROM-subquery support", EsqlCapabilities.Cap.SUBQUERY_IN_FROM_COMMAND.isEnabled());
+        setupEnrichPolicy(client(LOCAL_CLUSTER), "values_enrich", 10);
+        setSkipUnavailable(REMOTE_CLUSTER_1, false);
+        setSkipUnavailable(REMOTE_CLUSTER_2, false);
+        createViewOnCluster(
+            LOCAL_CLUSTER,
+            "enrich_view",
+            "FROM logs-* | WHERE v > 1 AND v < 7 | ENRICH values_enrich ON v WITH enrich_name | KEEP v"
+        );
+        try {
+            try (EsqlQueryResponse resp = runQuery("""
+                FROM (FROM enrich_view), (FROM cluster-a:logs-*)
+                | WHERE v IN (FROM enrich_view)
+                | KEEP v
+                | SORT v
+                """, false)) {
+                assertThat(
+                    getValuesList(resp),
+                    equalTo(List.of(List.of(2L), List.of(3L), List.of(4L), List.of(4L), List.of(5L), List.of(6L)))
+                );
+            }
+        } finally {
+            deleteEnrichPolicy(client(LOCAL_CLUSTER), "values_enrich");
+            clearSkipUnavailable(3);
+        }
+    }
+
+    /**
+     * The same policy name used with three different modes (ANY, COORDINATOR, REMOTE) in one query: each occurrence has
+     * its own mode-specific target-cluster requirement (see {@code EnrichPolicyResolver#calculateTargetClusters}), so all
+     * three must resolve independently even though they share a policy name - the REMOTE occurrence only needs
+     * {@code cluster-a} (its own scope), the ANY occurrence needs both remotes the outer {@code *:logs-*} touches (that
+     * pattern spans {@code cluster-a} and {@code remote-b}, not the local cluster) plus {@code _local} (ANY always
+     * requires it), and the COORDINATOR occurrence only ever needs {@code _local}.
+     */
+    public void testEnrichSamePolicyDifferentModesAllSucceed() {
+        setupEnrichPolicy(client(LOCAL_CLUSTER), "values_enrich", 10);
+        setupEnrichPolicy(client(REMOTE_CLUSTER_1), "values_enrich", 10);
+        setupEnrichPolicy(client(REMOTE_CLUSTER_2), "values_enrich", 10);
+        setSkipUnavailable(REMOTE_CLUSTER_1, false);
+        setSkipUnavailable(REMOTE_CLUSTER_2, false);
+        try {
+            try (EsqlQueryResponse resp = runQuery("""
+                FROM *:logs-*
+                | WHERE v IN (
+                    FROM cluster-a:logs-*
+                    | WHERE v > 1 AND v < 7
+                    | ENRICH _remote:values_enrich ON v WITH enrich_name
+                    | KEEP v
+                  )
+                | ENRICH values_enrich ON v WITH enrich_name
+                | ENRICH _coordinator:values_enrich ON v WITH enrich_name
+                | KEEP v
+                | SORT v
+                """, false)) {
+                // *:logs-* spans cluster-a and remote-b only; each contributes a single v=4 row (i=2 -> 2^2=4).
+                assertThat(getValuesList(resp), equalTo(List.of(List.of(4L), List.of(4L))));
+            }
+        } finally {
+            deleteEnrichPolicy(client(LOCAL_CLUSTER), "values_enrich");
+            deleteEnrichPolicy(client(REMOTE_CLUSTER_1), "values_enrich");
+            deleteEnrichPolicy(client(REMOTE_CLUSTER_2), "values_enrich");
+            clearSkipUnavailable(3);
+        }
+    }
+
+    /**
+     * The same policy name used with two different modes where only one occurrence's requirement is satisfied: REMOTE
+     * mode (scoped to {@code cluster-a} only) fails because the policy is missing there, while COORDINATOR mode
+     * (scoped to {@code _local} only, elsewhere in the same query) would succeed on its own since the policy exists
+     * locally. The failure must report exactly the REMOTE occurrence's own missing cluster ({@code cluster-a}), proving
+     * the two occurrences are resolved independently rather than one's requirement leaking into the other's error.
+     */
+    public void testEnrichSamePolicyDifferentModesFailureIsolatedPerOccurrence() {
+        setupEnrichPolicy(client(LOCAL_CLUSTER), "values_enrich", 10);
+        setSkipUnavailable(REMOTE_CLUSTER_1, false);
+        setSkipUnavailable(REMOTE_CLUSTER_2, false);
+        try {
+            VerificationException ex = expectThrows(VerificationException.class, () -> runQuery("""
+                FROM logs-*
+                | WHERE v IN (
+                    FROM cluster-a:logs-*
+                    | WHERE v > 1 AND v < 7
+                    | ENRICH _remote:values_enrich ON v WITH enrich_name
+                    | KEEP v
+                  )
+                | ENRICH _coordinator:values_enrich ON v WITH enrich_name
+                | KEEP v
+                """, false));
+            assertThat(ex.getMessage(), containsString("cannot find enrich policy [values_enrich] on clusters [cluster-a]"));
+        } finally {
+            deleteEnrichPolicy(client(LOCAL_CLUSTER), "values_enrich");
+            clearSkipUnavailable(3);
+        }
+    }
+
     // ---- helpers ----
 
     /**

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlCapabilities.java

@@ -3381,12 +3381,6 @@ public enum Cap {
          */
         SPATIAL_BBOX_VALIDATION_FIX,
 
-        /**
-         * ENRICH policy resolution is scoped to the clusters that actually feed rows into that specific ENRICH (its own
-         * FROM-subquery branch or WHERE-IN subquery), instead of every cluster touched anywhere in the whole query.
-         */
-        ENRICH_SCOPED_TO_SUBQUERY,
-
         // Last capability should still have a comma for fewer merge conflicts when adding new ones :)
         // This comment prevents the semicolon from being on the previous capability when Spotless formats the file.
         ;

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/session/EsqlSession.java

@@ -1514,14 +1514,17 @@ static Set<String> computeEnrichScope(Enrich enrich, Map<IndexPattern, IndexReso
      * that failed to connect during main index resolution (skipped, e.g. behind {@code skip_unavailable=true}) can still show
      * up in a wildcard pattern's {@code originalIndices()}, but the policy should not be required there.
      */
-    private static Map<Source, Set<String>> computeEnrichScopes(
+    // package-private static so EsqlSessionTests can drive it directly, e.g. to exercise the same-Source union above
+    static Map<Source, Set<String>> computeEnrichScopes(
         List<Enrich> enriches,
         Map<IndexPattern, IndexResolution> indexResolution,
         EsqlExecutionInfo executionInfo
     ) {
         Map<Source, Set<String>> enrichScopes = new HashMap<>();
         for (Enrich enrich : enriches) {
-            Set<String> scope = EsqlCCSUtils.onlyRunning(executionInfo, computeEnrichScope(enrich, indexResolution));
+            Set<String> scope = new HashSet<>(EsqlCCSUtils.onlyRunning(executionInfo, computeEnrichScope(enrich, indexResolution)));
+            // onlyRunning can return an immutable Set (e.g. Set.of(...) when no cluster is tracked yet), so copy it into a
+            // mutable one before it's potentially unioned in place by a later same-Source occurrence, below.
             enrichScopes.merge(enrich.source(), scope, (existing, additional) -> {
                 existing.addAll(additional);
                 return existing;

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/analysis/AnalyzerTests.java

@@ -2079,6 +2079,52 @@ public void testEnrichFieldsIncludeMatchField() {
         assertThat(Expressions.names(limit.output()), contains("language_name", "language_code"));
     }
 
+    /**
+     * Two occurrences of the same policy name and mode - registered with different resolutions via two separate
+     * {@code addEnrichPolicy} calls - resolve independently rather than both getting whichever resolution was registered
+     * first: {@code TestAnalyzer} matches registrations to occurrences of the same policy name/mode in registration order
+     * (first call -> first occurrence, second call -> second occurrence, etc). This exercises
+     * {@code Analyzer.ResolveEnrich}'s per-{@code Source} dispatch: {@code EnrichResolution} is keyed by the originating
+     * {@code Enrich} node's {@code Source}, not by policy name/mode, precisely so that occurrences like these two - which
+     * share a policy name and mode but can be resolved differently (e.g. because they live in differently-scoped subquery
+     * branches in production) - don't collide.
+     */
+    public void testEnrichPolicyDifferentResolutionPerOccurrence() {
+        IndexResolution languageIndex = loadMapping("mapping-languages.json", "languages");
+        LogicalPlan plan = basic().addEnrichPolicy(
+            Enrich.Mode.ANY,
+            "languages",
+            new ResolvedEnrichPolicy(
+                "language_code",
+                EnrichPolicy.MATCH_TYPE,
+                List.of("language_name"),
+                Map.of("", "languages"),
+                languageIndex.get().mapping()
+            )
+        )
+            .addEnrichPolicy(
+                Enrich.Mode.ANY,
+                "languages",
+                new ResolvedEnrichPolicy(
+                    "language_code",
+                    EnrichPolicy.MATCH_TYPE,
+                    List.of("language_code"),
+                    Map.of("", "languages"),
+                    languageIndex.get().mapping()
+                )
+            )
+            .query("""
+                FROM test
+                | EVAL x = to_string(languages)
+                | ENRICH languages ON x
+                | EVAL y = to_string(languages)
+                | ENRICH languages ON y
+                | KEEP language_name, language_code
+                """);
+        var limit = as(plan, Limit.class);
+        assertThat(Expressions.names(limit.output()), contains("language_name", "language_code"));
+    }
+
     public void testChainedEvalFieldsUse() {
         var query = "from test | eval x0 = pow(salary, 1), x1 = pow(x0, 2), x2 = pow(x1, 3)";
         int additionalEvals = randomIntBetween(0, 5);

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/session/EsqlSessionTests.java

@@ -16,6 +16,7 @@
 import org.elasticsearch.test.ESTestCase;
 import org.elasticsearch.transport.RemoteClusterAware;
 import org.elasticsearch.xpack.esql.action.EsqlCapabilities;
+import org.elasticsearch.xpack.esql.action.EsqlExecutionInfo;
 import org.elasticsearch.xpack.esql.action.EsqlQueryRequest;
 import org.elasticsearch.xpack.esql.analysis.InSubqueryResolver;
 import org.elasticsearch.xpack.esql.analysis.PreAnalyzer;
@@ -57,6 +58,7 @@
 import static org.hamcrest.Matchers.containsString;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.hasSize;
+import static org.hamcrest.Matchers.not;
 import static org.hamcrest.Matchers.sameInstance;
 
 public class EsqlSessionTests extends ESTestCase {
@@ -568,6 +570,34 @@ public void testComputeEnrichScopeMixedSubqueries() {
         }
     }
 
+    /**
+     * Two distinct {@link Enrich} occurrences can share the exact same {@link Source} today: a view containing an ENRICH is
+     * re-parsed independently every time it's referenced, so referencing it from two differently-scoped subquery branches
+     * produces two structurally-identical (and therefore {@code Source.equals()}) {@link Enrich} nodes. Parsing the same
+     * literal query text twice below reproduces that same-Source collision without needing an actual view.
+     * <p>
+     * {@link EsqlSession#computeEnrichScopes} unions the two occurrences' scopes rather than letting the second overwrite the
+     * first. This also exercises the case where {@link EsqlCCSUtils#onlyRunning} hands back an immutable {@code Set.of(...)}
+     * (no cluster tracked yet) for the first occurrence - that value must not be mutated in place once the second,
+     * same-Source occurrence is merged into it.
+     */
+    public void testComputeEnrichScopesUnionsDuplicateSource() {
+        var plan1 = TEST_PARSER.parseQuery("ROW key = 1 | ENRICH policy ON key");
+        var plan2 = TEST_PARSER.parseQuery("ROW key = 1 | ENRICH policy ON key");
+        Enrich enrich1 = enrichNamed(plan1, "policy");
+        Enrich enrich2 = enrichNamed(plan2, "policy");
+        assertThat(enrich1, not(sameInstance(enrich2)));
+        assertThat(enrich1.source(), equalTo(enrich2.source()));
+
+        var resolution = createIndexResolution();
+        EsqlExecutionInfo executionInfo = new EsqlExecutionInfo(alias -> false, EsqlExecutionInfo.IncludeExecutionMetadata.NEVER);
+
+        Map<Source, Set<String>> scopes = EsqlSession.computeEnrichScopes(List.of(enrich1, enrich2), resolution, executionInfo);
+
+        assertThat(scopes.keySet(), hasSize(1));
+        assertThat(scopes.get(enrich1.source()), equalTo(Set.of(RemoteClusterAware.LOCAL_CLUSTER_GROUP_KEY)));
+    }
+
     /**
      * Returns the sole {@link Enrich} node with the given (resolved) policy name in {@code plan}, failing if there isn't
      * exactly one - test queries give each ENRICH occurrence under test a distinct policy name to disambiguate it.