ESQL: Load csv data as ClassRule in single/multi-node

x-pack/plugin/esql/qa/server/build.gradle

@@ -10,4 +10,5 @@ dependencies {
   // Requirement for some ESQL-specific utilities
   implementation project(':x-pack:plugin:esql')
   api project(xpackModule('esql:qa:testFixtures'))
+  api project(':test:test-clusters')
 }

x-pack/plugin/esql/qa/server/mixed-cluster/src/javaRestTest/java/org/elasticsearch/xpack/esql/qa/mixed/MixedClusterEsqlSpecIT.java

@@ -8,6 +8,8 @@
 package org.elasticsearch.xpack.esql.qa.mixed;
 
 import org.elasticsearch.Version;
+import org.elasticsearch.client.RestClient;
+import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.features.NodeFeature;
 import org.elasticsearch.test.cluster.ElasticsearchCluster;
 import org.elasticsearch.test.rest.TestFeatureService;
@@ -16,15 +18,26 @@
 import org.junit.AfterClass;
 import org.junit.Before;
 import org.junit.ClassRule;
+import org.junit.rules.RuleChain;
+import org.junit.rules.TestRule;
 
 import java.io.IOException;
 
 import static org.elasticsearch.xpack.esql.CsvTestUtils.isEnabled;
 import static org.elasticsearch.xpack.esql.qa.rest.EsqlSpecTestCase.Mode.ASYNC;
 
 public class MixedClusterEsqlSpecIT extends EsqlSpecTestCase {
-    @ClassRule
     public static ElasticsearchCluster cluster = Clusters.mixedVersionCluster();
+    public static ClosingTestRule<RestClient> client = new ClosingTestRule<>() {
+        @Override
+        protected RestClient provideObject() throws IOException {
+            return startClient(cluster, Settings.builder().build());
+        }
+    };
+    public static CsvLoader loader = new CsvLoader(client);
+
+    @ClassRule
+    public static TestRule clusterRule = RuleChain.outerRule(cluster).around(client).around(loader);
 
     @Override
     protected String getTestRestCluster() {

x-pack/plugin/esql/qa/server/multi-clusters/src/javaRestTest/java/org/elasticsearch/xpack/esql/ccq/MultiClusterSpecIT.java

@@ -62,9 +62,17 @@ public class MultiClusterSpecIT extends EsqlSpecTestCase {
 
     static ElasticsearchCluster remoteCluster = Clusters.remoteCluster();
     static ElasticsearchCluster localCluster = Clusters.localCluster(remoteCluster);
+    public static ClosingTestRule<RestClient> client = new ClosingTestRule<>() {
+        @Override
+        protected RestClient provideObject() throws IOException {
+            HttpHost[] localHosts = parseClusterHostsStatic(localCluster.getHttpAddresses()).toArray(HttpHost[]::new);
+            return doBuildClient(Settings.builder().build(), localHosts);
+        }
+    };
+    public static CsvLoader loader = new CsvLoader(client);
 
     @ClassRule
-    public static TestRule clusterRule = RuleChain.outerRule(remoteCluster).around(localCluster);
+    public static TestRule clusterRule = RuleChain.outerRule(remoteCluster).around(localCluster).around(client).around(loader);
 
     private static TestFeatureService remoteFeaturesService;
     private static RestClient remoteClusterClient;
@@ -133,9 +141,13 @@ protected String getTestRestCluster() {
 
     @Override
     protected RestClient buildClient(Settings settings, HttpHost[] localHosts) throws IOException {
-        RestClient localClient = super.buildClient(settings, localHosts);
-        HttpHost[] remoteHosts = parseClusterHosts(remoteCluster.getHttpAddresses()).toArray(HttpHost[]::new);
-        RestClient remoteClient = super.buildClient(settings, remoteHosts);
+        return doBuildClient(settings, localHosts);
+    }
+
+    private static RestClient doBuildClient(Settings settings, HttpHost[] localHosts) throws IOException {
+        RestClient localClient = buildClientStatic(settings, localHosts);
+        HttpHost[] remoteHosts = parseClusterHostsStatic(remoteCluster.getHttpAddresses()).toArray(HttpHost[]::new);
+        RestClient remoteClient = buildClientStatic(settings, remoteHosts);
         return twoClients(localClient, remoteClient);
     }
 

x-pack/plugin/esql/qa/server/multi-node/src/javaRestTest/java/org/elasticsearch/xpack/esql/qa/multi_node/EsqlSpecIT.java

@@ -7,14 +7,29 @@
 
 package org.elasticsearch.xpack.esql.qa.multi_node;
 
+import org.elasticsearch.client.RestClient;
+import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.test.cluster.ElasticsearchCluster;
 import org.elasticsearch.xpack.esql.core.CsvSpecReader.CsvTestCase;
 import org.elasticsearch.xpack.esql.qa.rest.EsqlSpecTestCase;
 import org.junit.ClassRule;
+import org.junit.rules.RuleChain;
+import org.junit.rules.TestRule;
+
+import java.io.IOException;
 
 public class EsqlSpecIT extends EsqlSpecTestCase {
-    @ClassRule
     public static ElasticsearchCluster cluster = Clusters.testCluster();
+    public static ClosingTestRule<RestClient> client = new ClosingTestRule<>() {
+        @Override
+        protected RestClient provideObject() throws IOException {
+            return startClient(cluster, Settings.builder().build());
+        }
+    };
+    public static CsvLoader loader = new CsvLoader(client);
+
+    @ClassRule
+    public static TestRule clusterRule = RuleChain.outerRule(cluster).around(client).around(loader);
 
     @Override
     protected String getTestRestCluster() {

x-pack/plugin/esql/qa/server/single-node/src/javaRestTest/java/org/elasticsearch/xpack/esql/qa/single_node/EsqlSpecIT.java

@@ -9,16 +9,31 @@
 
 import com.carrotsearch.randomizedtesting.annotations.ThreadLeakFilters;
 
+import org.elasticsearch.client.RestClient;
+import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.test.TestClustersThreadFilter;
 import org.elasticsearch.test.cluster.ElasticsearchCluster;
 import org.elasticsearch.xpack.esql.core.CsvSpecReader.CsvTestCase;
 import org.elasticsearch.xpack.esql.qa.rest.EsqlSpecTestCase;
 import org.junit.ClassRule;
+import org.junit.rules.RuleChain;
+import org.junit.rules.TestRule;
+
+import java.io.IOException;
 
 @ThreadLeakFilters(filters = TestClustersThreadFilter.class)
 public class EsqlSpecIT extends EsqlSpecTestCase {
-    @ClassRule
     public static ElasticsearchCluster cluster = Clusters.testCluster();
+    public static ClosingTestRule<RestClient> client = new ClosingTestRule<>() {
+        @Override
+        protected RestClient provideObject() throws IOException {
+            return startClient(cluster, Settings.builder().build());
+        }
+    };
+    public static CsvLoader loader = new CsvLoader(client);
+
+    @ClassRule
+    public static TestRule clusterRule = RuleChain.outerRule(cluster).around(client).around(loader);
 
     @Override
     protected String getTestRestCluster() {

x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/EsqlSpecTestCase.java

@@ -10,20 +10,25 @@
 import com.carrotsearch.randomizedtesting.annotations.TimeoutSuite;
 
 import org.apache.http.HttpEntity;
+import org.apache.http.HttpHost;
 import org.apache.lucene.tests.util.TimeUnits;
 import org.elasticsearch.Build;
 import org.elasticsearch.Version;
 import org.elasticsearch.client.Request;
 import org.elasticsearch.client.ResponseException;
 import org.elasticsearch.client.RestClient;
+import org.elasticsearch.client.RestClientBuilder;
+import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.xcontent.XContentHelper;
+import org.elasticsearch.core.IOUtils;
 import org.elasticsearch.features.NodeFeature;
 import org.elasticsearch.geometry.Geometry;
 import org.elasticsearch.geometry.Point;
 import org.elasticsearch.geometry.utils.GeometryValidator;
 import org.elasticsearch.geometry.utils.WellKnownText;
 import org.elasticsearch.logging.LogManager;
 import org.elasticsearch.logging.Logger;
+import org.elasticsearch.test.cluster.ElasticsearchCluster;
 import org.elasticsearch.test.rest.ESRestTestCase;
 import org.elasticsearch.test.rest.TestFeatureService;
 import org.elasticsearch.xcontent.XContentType;
@@ -36,7 +41,11 @@
 import org.junit.After;
 import org.junit.AfterClass;
 import org.junit.Before;
+import org.junit.rules.TestRule;
+import org.junit.runner.Description;
+import org.junit.runners.model.Statement;
 
+import java.io.Closeable;
 import java.io.IOException;
 import java.net.URL;
 import java.util.ArrayList;
@@ -48,6 +57,7 @@
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
+import static java.util.Collections.unmodifiableList;
 import static org.apache.lucene.geo.GeoEncodingUtils.decodeLatitude;
 import static org.apache.lucene.geo.GeoEncodingUtils.decodeLongitude;
 import static org.apache.lucene.geo.GeoEncodingUtils.encodeLatitude;
@@ -59,7 +69,6 @@
 import static org.elasticsearch.xpack.esql.CsvTestUtils.ExpectedResults;
 import static org.elasticsearch.xpack.esql.CsvTestUtils.isEnabled;
 import static org.elasticsearch.xpack.esql.CsvTestUtils.loadCsvSpecValues;
-import static org.elasticsearch.xpack.esql.CsvTestsDataLoader.CSV_DATASET_MAP;
 import static org.elasticsearch.xpack.esql.CsvTestsDataLoader.loadDataSetIntoEs;
 import static org.elasticsearch.xpack.esql.core.CsvSpecReader.specParser;
 import static org.elasticsearch.xpack.esql.core.TestUtils.classpathResources;
@@ -119,18 +128,12 @@ protected EsqlSpecTestCase(String fileName, String groupName, String testName, I
         this.mode = mode;
     }
 
-    @Before
-    public void setup() throws IOException {
-        if (indexExists(CSV_DATASET_MAP.keySet().iterator().next()) == false) {
-            loadDataSetIntoEs(client());
-        }
-    }
-
     protected boolean supportsAsync() {
         return clusterHasFeature(ASYNC_QUERY_FEATURE_ID); // the Async API was introduced in 8.13.0
     }
 
     @AfterClass
+    // TODO: This should become part of a classrule for consistency.
     public static void wipeTestData() throws IOException {
         try {
             adminClient().performRequest(new Request("DELETE", "/*"));
@@ -323,4 +326,89 @@ public static void assertRequestBreakerEmpty() throws Exception {
             }
         });
     }
+
+    public static RestClient startClient(ElasticsearchCluster cluster, Settings clientSettings) throws IOException {
+        String address = cluster.getHttpAddress(0);
+        int portSeparator = address.lastIndexOf(':');
+        if (portSeparator < 0) {
+            throw new IllegalArgumentException("Illegal cluster url [" + address + "]");
+        }
+        String host = address.substring(0, portSeparator);
+        int port = Integer.parseInt(address.substring(portSeparator + 1));
+        HttpHost[] httpHosts = new HttpHost[] { new HttpHost(host, port) };
+
+        return buildClientStatic(clientSettings, httpHosts);
+    }
+
+    /**
+     * Like {@link ESRestTestCase#buildClient(Settings, HttpHost[])} but static.
+     */
+    public static RestClient buildClientStatic(Settings settings, HttpHost[] hosts) throws IOException {
+        RestClientBuilder builder = RestClient.builder(hosts);
+        doConfigureClient(builder, settings);
+        builder.setStrictDeprecationMode(true);
+        return builder.build();
+    }
+
+    /**
+     * Like {@link ESRestTestCase#parseClusterHosts(String)} but static.
+     */
+    protected static List<HttpHost> parseClusterHostsStatic(String hostsString) {
+        String[] stringUrls = hostsString.split(",");
+        List<HttpHost> hosts = new ArrayList<>(stringUrls.length);
+        for (String stringUrl : stringUrls) {
+            int portSeparator = stringUrl.lastIndexOf(':');
+            if (portSeparator < 0) {
+                throw new IllegalArgumentException("Illegal cluster url [" + stringUrl + "]");
+            }
+            String host = stringUrl.substring(0, portSeparator);
+            int port = Integer.valueOf(stringUrl.substring(portSeparator + 1));
+            hosts.add(new HttpHost(host, port, "http"));
+        }
+        return unmodifiableList(hosts);
+    }
+
+    public static class CsvLoader implements TestRule {
+        ClosingTestRule<RestClient> client;
+
+        public CsvLoader(ClosingTestRule<RestClient> client) {
+            this.client = client;
+        }
+
+        @Override
+        public Statement apply(Statement base, Description description) {
+            return new Statement() {
+                @Override
+                public void evaluate() throws Throwable {
+                    loadDataSetIntoEs(client.get());
+                    base.evaluate();
+                }
+            };
+        }
+    }
+
+    public abstract static class ClosingTestRule<T extends Closeable> implements TestRule {
+        private T providedObject;
+
+        public T get() {
+            return providedObject;
+        }
+
+        protected abstract T provideObject() throws IOException;
+
+        @Override
+        public Statement apply(Statement base, Description description) {
+            return new Statement() {
+                @Override
+                public void evaluate() throws Throwable {
+                    try {
+                        providedObject = provideObject();
+                        base.evaluate();
+                    } finally {
+                        IOUtils.close(providedObject);
+                    }
+                }
+            };
+        }
+    }
 }