mark isPartial at the end

Author: dnhatn

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

@@ -175,7 +175,14 @@ public void markEndQuery() {
         assert relativeStartNanos != null : "Relative start time must be set when markEndQuery is called";
         overallTook = new TimeValue(System.nanoTime() - relativeStartNanos, TimeUnit.NANOSECONDS);
         if (isPartial == false) {
-            isPartial = clusterInfo.values().stream().anyMatch(c -> c.failedShards != null && c.failedShards > 0);
+            // TODO: Mark individual clusters as partial if failed shards exist
+            isPartial = clusterInfo.values()
+                .stream()
+                .anyMatch(
+                    c -> c.status == Cluster.Status.PARTIAL
+                        || c.status == Cluster.Status.SKIPPED
+                        || c.failedShards != null && c.failedShards > 0
+                );
         }
     }
 

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

@@ -234,16 +234,6 @@ private void innerExecute(Task task, EsqlQueryRequest request, ActionListener<Es
             planRunner,
             services,
             ActionListener.wrap(result -> {
-                // If we had any skipped or partial clusters, the result is partial
-                if (executionInfo.getClusters()
-                    .values()
-                    .stream()
-                    .anyMatch(
-                        c -> c.getStatus() == EsqlExecutionInfo.Cluster.Status.SKIPPED
-                            || c.getStatus() == EsqlExecutionInfo.Cluster.Status.PARTIAL
-                    )) {
-                    executionInfo.markAsPartial();
-                }
                 recordCCSTelemetry(task, executionInfo, request, null);
                 listener.onResponse(toResponse(task, request, configuration, result));
             }, ex -> {