SQL: Use request flavored methods in tests (#30345)

Modifies the SQL tests to use the new `Request` object flavored methods
introduced onto the `RestClient` in #29623. We'd like to remove the old
methods eventually so we should stop using them.

x-pack/qa/sql/multinode/src/test/java/org/elasticsearch/xpack/qa/sql/multinode/RestSqlMultinodeIT.java

@@ -8,6 +8,7 @@
 import org.apache.http.HttpHost;
 import org.apache.http.entity.ContentType;
 import org.apache.http.entity.StringEntity;
+import org.elasticsearch.client.Request;
 import org.elasticsearch.client.Response;
 import org.elasticsearch.client.RestClient;
 import org.elasticsearch.common.Strings;
@@ -53,7 +54,7 @@ public void testIndexOnWrongNode() throws IOException {
         String firstHostName = null;
 
         String match = firstHost.getHostName() + ":" + firstHost.getPort();
-        Map<String, Object> nodesInfo = responseToMap(client().performRequest("GET", "/_nodes"));
+        Map<String, Object> nodesInfo = responseToMap(client().performRequest(new Request("GET", "/_nodes")));
         @SuppressWarnings("unchecked")
         Map<String, Object> nodes = (Map<String, Object>) nodesInfo.get("nodes");
         for (Map.Entry<String, Object> node : nodes.entrySet()) {
@@ -74,7 +75,9 @@ public void testIndexOnWrongNode() throws IOException {
         }
         index.endObject();
         index.endObject();
-        client().performRequest("PUT", "/test", emptyMap(), new StringEntity(Strings.toString(index), ContentType.APPLICATION_JSON));
+        Request request = new Request("PUT", "/test");
+        request.setJsonEntity(Strings.toString(index));
+        client().performRequest(request);
         int documents = between(10, 100);
         createTestData(documents);
 
@@ -84,6 +87,9 @@ public void testIndexOnWrongNode() throws IOException {
     }
 
     private void createTestData(int documents) throws UnsupportedCharsetException, IOException {
+        Request request = new Request("PUT", "/test/test/_bulk");
+        request.addParameter("refresh", "true");
+
         StringBuilder bulk = new StringBuilder();
         for (int i = 0; i < documents; i++) {
             int a = 3 * i;
@@ -92,8 +98,9 @@ private void createTestData(int documents) throws UnsupportedCharsetException, I
             bulk.append("{\"index\":{\"_id\":\"" + i + "\"}\n");
             bulk.append("{\"a\": " + a + ", \"b\": " + b + ", \"c\": " + c + "}\n");
         }
-        client().performRequest("PUT", "/test/test/_bulk", singletonMap("refresh", "true"),
-                new StringEntity(bulk.toString(), ContentType.APPLICATION_JSON));
+        request.setJsonEntity(bulk.toString());
+
+        client().performRequest(request);
     }
 
     private Map<String, Object> responseToMap(Response response) throws IOException {
@@ -108,14 +115,12 @@ private void assertCount(RestClient client, int count) throws IOException {
         expected.put("columns", singletonList(columnInfo(mode, "COUNT(1)", "long", JDBCType.BIGINT, 20)));
         expected.put("rows", singletonList(singletonList(count)));
 
-        Map<String, String> params = new TreeMap<>();
-        params.put("format", "json");        // JSON is easier to parse then a table
-        if (Strings.hasText(mode)) {
-            params.put("mode", mode);        // JDBC or PLAIN mode
+        Request request = new Request("POST", "/_xpack/sql");
+        if (false == mode.isEmpty()) {
+            request.addParameter("mode", mode);
         }
-
-        Map<String, Object> actual = responseToMap(client.performRequest("POST", "/_xpack/sql", params,
-                new StringEntity("{\"query\": \"SELECT COUNT(*) FROM test\"}", ContentType.APPLICATION_JSON)));
+        request.setJsonEntity("{\"query\": \"SELECT COUNT(*) FROM test\"}");
+        Map<String, Object> actual = responseToMap(client.performRequest(request));
 
         if (false == expected.equals(actual)) {
             NotEqualMessageBuilder message = new NotEqualMessageBuilder();

x-pack/qa/sql/security/src/test/java/org/elasticsearch/xpack/qa/sql/security/RestSqlSecurityIT.java

@@ -10,6 +10,7 @@
 import org.apache.http.entity.ContentType;
 import org.apache.http.entity.StringEntity;
 import org.apache.http.message.BasicHeader;
+import org.elasticsearch.client.Request;
 import org.elasticsearch.client.Response;
 import org.elasticsearch.client.ResponseException;
 import org.elasticsearch.common.Nullable;
@@ -176,14 +177,15 @@ private static Map<String, Object> runSql(@Nullable String asUser, String mode,
         }
 
         private static Map<String, Object> runSql(@Nullable String asUser, String mode, HttpEntity entity) throws IOException {
-            Map<String, String> params = new TreeMap<>();
-            params.put("format", "json");        // JSON is easier to parse then a table
-            if (Strings.hasText(mode)) {
-                params.put("mode", mode);        // JDBC or PLAIN mode
+            Request request = new Request("POST", "/_xpack/sql");
+            if (false == mode.isEmpty()) {
+                request.addParameter("mode", mode);
             }
-            Header[] headers = asUser == null ? new Header[0] : new Header[] {new BasicHeader("es-security-runas-user", asUser)};
-            Response response = client().performRequest("POST", "/_xpack/sql", params, entity, headers);
-            return toMap(response);
+            if (asUser != null) {
+                request.setHeaders(new BasicHeader("es-security-runas-user", asUser));
+            }
+            request.setEntity(entity);
+            return toMap(client().performRequest(request));
         }
 
         private static void assertResponse(Map<String, Object> expected, Map<String, Object> actual) {

x-pack/qa/sql/security/src/test/java/org/elasticsearch/xpack/qa/sql/security/SqlSecurityTestCase.java

@@ -11,6 +11,7 @@
 import org.elasticsearch.SpecialPermission;
 import org.elasticsearch.action.admin.indices.get.GetIndexAction;
 import org.elasticsearch.action.admin.indices.get.GetIndexRequest;
+import org.elasticsearch.client.Request;
 import org.elasticsearch.client.ResponseException;
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.settings.Settings;
@@ -41,7 +42,6 @@
 import java.util.function.Function;
 import java.util.regex.Pattern;
 
-import static java.util.Collections.emptyMap;
 import static java.util.Collections.singletonMap;
 import static org.hamcrest.Matchers.contains;
 import static org.hamcrest.Matchers.empty;
@@ -135,15 +135,18 @@ public void oneTimeSetup() throws Exception {
              * write the test data once. */
             return;
         }
+        Request request = new Request("PUT", "/_bulk");
+        request.addParameter("refresh", "true");
+
         StringBuilder bulk = new StringBuilder();
         bulk.append("{\"index\":{\"_index\": \"test\", \"_type\": \"doc\", \"_id\":\"1\"}\n");
         bulk.append("{\"a\": 1, \"b\": 2, \"c\": 3}\n");
         bulk.append("{\"index\":{\"_index\": \"test\", \"_type\": \"doc\", \"_id\":\"2\"}\n");
         bulk.append("{\"a\": 4, \"b\": 5, \"c\": 6}\n");
         bulk.append("{\"index\":{\"_index\": \"bort\", \"_type\": \"doc\", \"_id\":\"1\"}\n");
         bulk.append("{\"a\": \"test\"}\n");
-        client().performRequest("PUT", "/_bulk", singletonMap("refresh", "true"),
-                new StringEntity(bulk.toString(), ContentType.APPLICATION_JSON));
+        request.setJsonEntity(bulk.toString());
+        client().performRequest(request);
         oneTimeSetup = true;
     }
 
@@ -173,7 +176,7 @@ public void setInitialAuditLogOffset() {
     @AfterClass
     public static void wipeIndicesAfterTests() throws IOException {
         try {
-            adminClient().performRequest("DELETE", "*");
+            adminClient().performRequest(new Request("DELETE", "*"));
         } catch (ResponseException e) {
             // 404 here just means we had no indexes
             if (e.getResponse().getStatusLine().getStatusCode() != 404) {
@@ -472,13 +475,15 @@ public void testNoGetIndex() throws Exception {
     }
 
     protected static void createUser(String name, String role) throws IOException {
-        XContentBuilder user = JsonXContent.contentBuilder().prettyPrint().startObject(); {
+        Request request = new Request("PUT", "/_xpack/security/user/" + name);
+        XContentBuilder user = JsonXContent.contentBuilder().prettyPrint();
+        user.startObject(); {
             user.field("password", "testpass");
             user.field("roles", role);
         }
         user.endObject();
-        client().performRequest("PUT", "/_xpack/security/user/" + name, emptyMap(),
-                new StringEntity(Strings.toString(user), ContentType.APPLICATION_JSON));
+        request.setJsonEntity(Strings.toString(user));
+        client().performRequest(request);
     }
 
     protected AuditLogAsserter createAuditLogAsserter() {

x-pack/qa/sql/src/main/java/org/elasticsearch/xpack/qa/sql/cli/CliIntegrationTestCase.java

@@ -5,9 +5,9 @@
  */
 package org.elasticsearch.xpack.qa.sql.cli;
 
-import org.apache.http.HttpEntity;
 import org.apache.http.entity.ContentType;
 import org.apache.http.entity.StringEntity;
+import org.elasticsearch.client.Request;
 import org.elasticsearch.common.CheckedConsumer;
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.xcontent.XContentBuilder;
@@ -19,7 +19,6 @@
 
 import java.io.IOException;
 
-import static java.util.Collections.singletonMap;
 import static org.elasticsearch.xpack.qa.sql.rest.RestSqlTestCase.assertNoSearchContexts;
 
 public abstract class CliIntegrationTestCase extends ESRestTestCase {
@@ -60,11 +59,13 @@ protected SecurityConfig securityConfig() {
     }
 
     protected void index(String index, CheckedConsumer<XContentBuilder, IOException> body) throws IOException {
+        Request request = new Request("PUT", "/" + index + "/doc/1");
+        request.addParameter("refresh", "true");
         XContentBuilder builder = JsonXContent.contentBuilder().startObject();
         body.accept(builder);
         builder.endObject();
-        HttpEntity doc = new StringEntity(Strings.toString(builder), ContentType.APPLICATION_JSON);
-        client().performRequest("PUT", "/" + index + "/doc/1", singletonMap("refresh", "true"), doc);
+        request.setJsonEntity(Strings.toString(builder));
+        client().performRequest(request);
     }
 
     public String command(String command) throws IOException {

x-pack/qa/sql/src/main/java/org/elasticsearch/xpack/qa/sql/cli/ErrorsTestCase.java

@@ -8,8 +8,7 @@
 import java.io.IOException;
 import org.apache.http.entity.ContentType;
 import org.apache.http.entity.StringEntity;
-
-import static java.util.Collections.emptyMap;
+import org.elasticsearch.client.Request;
 
 import static org.hamcrest.Matchers.startsWith;
 
@@ -41,7 +40,9 @@ public void testSelectFromMissingIndex() throws IOException {
     @Override
     public void testSelectFromIndexWithoutTypes() throws Exception {
         // Create an index without any types
-        client().performRequest("PUT", "/test", emptyMap(), new StringEntity("{}", ContentType.APPLICATION_JSON));
+        Request request = new Request("PUT", "/test");
+        request.setJsonEntity("{}");
+        client().performRequest(request);
 
         assertFoundOneProblem(command("SELECT * FROM test"));
         assertEquals("line 1:15: [test] doesn't have any types so it is incompatible with sql" + END, readLine());

x-pack/qa/sql/src/main/java/org/elasticsearch/xpack/qa/sql/cli/FetchSizeTestCase.java

@@ -7,24 +7,27 @@
 
 import org.apache.http.entity.ContentType;
 import org.apache.http.entity.StringEntity;
+import org.elasticsearch.client.Request;
 
 import java.io.IOException;
 
-import static java.util.Collections.singletonMap;
 import static org.hamcrest.Matchers.containsString;
 
 /**
  * Test for setting the fetch size.
  */
 public abstract class FetchSizeTestCase extends CliIntegrationTestCase {
     public void testSelect() throws IOException {
+        Request request = new Request("PUT", "/test/doc/_bulk");
+        request.addParameter("refresh", "true");
         StringBuilder bulk = new StringBuilder();
         for (int i = 0; i < 20; i++) {
             bulk.append("{\"index\":{}}\n");
             bulk.append("{\"test_field\":" + i + "}\n");
         }
-        client().performRequest("PUT", "/test/doc/_bulk", singletonMap("refresh", "true"),
-                new StringEntity(bulk.toString(), ContentType.APPLICATION_JSON));
+        request.setJsonEntity(bulk.toString());
+        client().performRequest(request);
+
         assertEquals("[?1l>[?1000l[?2004lfetch size set to [90m4[0m", command("fetch size = 4"));
         assertEquals("[?1l>[?1000l[?2004lfetch separator set to \"[90m -- fetch sep -- [0m\"",
             command("fetch separator = \" -- fetch sep -- \""));

x-pack/qa/sql/src/main/java/org/elasticsearch/xpack/qa/sql/jdbc/DataLoader.java

@@ -8,6 +8,7 @@
 import org.apache.http.HttpHost;
 import org.apache.http.entity.ContentType;
 import org.apache.http.entity.StringEntity;
+import org.elasticsearch.client.Request;
 import org.elasticsearch.client.RestClient;
 import org.elasticsearch.common.CheckedBiConsumer;
 import org.elasticsearch.common.Strings;
@@ -55,6 +56,7 @@ private static void createString(String name, XContentBuilder builder) throws Ex
        .endObject();
     }
     protected static void loadDatasetIntoEs(RestClient client, String index) throws Exception {
+        Request request = new Request("PUT", "/" + index);
         XContentBuilder createIndex = JsonXContent.contentBuilder().startObject();
         createIndex.startObject("settings");
         {
@@ -91,11 +93,9 @@ protected static void loadDatasetIntoEs(RestClient client, String index) throws
             createIndex.endObject();
         }
         createIndex.endObject().endObject();
-
-        client.performRequest("PUT", "/" + index, emptyMap(), new StringEntity(Strings.toString(createIndex),
-                        ContentType.APPLICATION_JSON));
+        request.setJsonEntity(Strings.toString(createIndex));
+        client.performRequest(request);
 
-
         Map<String, String> deps = new LinkedHashMap<>();
         csvToLines("departments", (titles, fields) -> deps.put(fields.get(0), fields.get(1)));
 
@@ -119,6 +119,8 @@ protected static void loadDatasetIntoEs(RestClient client, String index) throws
             list.add(dep);
         });
 
+        request = new Request("POST", "/" + index + "/emp/_bulk");
+        request.addParameter("refresh", "true");
         StringBuilder bulk = new StringBuilder();
         csvToLines("employees", (titles, fields) -> {
             bulk.append("{\"index\":{}}\n");
@@ -146,17 +148,16 @@ protected static void loadDatasetIntoEs(RestClient client, String index) throws
                 bulk.setLength(bulk.length() - 1);
                 bulk.append("]");
             }
-            
+
             bulk.append("}\n");
         });
-
-        client.performRequest("POST", "/" + index + "/emp/_bulk", singletonMap("refresh", "true"),
-                new StringEntity(bulk.toString(), ContentType.APPLICATION_JSON));
+        request.setJsonEntity(bulk.toString());
+        client.performRequest(request);
     }
 
     protected static void makeAlias(RestClient client, String aliasName, String... indices) throws Exception {
         for (String index : indices) {
-            client.performRequest("POST", "/" + index + "/_alias/" + aliasName);
+            client.performRequest(new Request("POST", "/" + index + "/_alias/" + aliasName));
         }
     }
 

x-pack/qa/sql/src/main/java/org/elasticsearch/xpack/qa/sql/jdbc/ErrorsTestCase.java

@@ -9,8 +9,7 @@
 import java.sql.SQLException;
 import org.apache.http.entity.ContentType;
 import org.apache.http.entity.StringEntity;
-
-import static java.util.Collections.emptyMap;
+import org.elasticsearch.client.Request;
 
 import static org.hamcrest.Matchers.startsWith;
 
@@ -37,7 +36,9 @@ public void testSelectFromMissingIndex() throws SQLException {
     @Override
     public void testSelectFromIndexWithoutTypes() throws Exception {
         // Create an index without any types
-        client().performRequest("PUT", "/test", emptyMap(), new StringEntity("{}", ContentType.APPLICATION_JSON));
+        Request request = new Request("PUT", "/test");
+        request.setJsonEntity("{}");
+        client().performRequest(request);
 
         try (Connection c = esJdbc()) {
             SQLException e = expectThrows(SQLException.class, () -> c.prepareStatement("SELECT * FROM test").executeQuery());

x-pack/qa/sql/src/main/java/org/elasticsearch/xpack/qa/sql/jdbc/FetchSizeTestCase.java

@@ -7,6 +7,7 @@
 
 import org.apache.http.entity.ContentType;
 import org.apache.http.entity.StringEntity;
+import org.elasticsearch.client.Request;
 import org.junit.Before;
 
 import java.io.IOException;
@@ -15,7 +16,6 @@
 import java.sql.SQLException;
 import java.sql.Statement;
 
-import static java.util.Collections.singletonMap;
 import static org.elasticsearch.xpack.qa.sql.rest.RestSqlTestCase.assertNoSearchContexts;
 
 /**
@@ -25,13 +25,15 @@
 public class FetchSizeTestCase extends JdbcIntegrationTestCase {
     @Before
     public void createTestIndex() throws IOException {
+        Request request = new Request("PUT", "/test/doc/_bulk");
+        request.addParameter("refresh", "true");
         StringBuilder bulk = new StringBuilder();
         for (int i = 0; i < 20; i++) {
             bulk.append("{\"index\":{}}\n");
             bulk.append("{\"test_field\":" + i + "}\n");
         }
-        client().performRequest("PUT", "/test/doc/_bulk", singletonMap("refresh", "true"),
-                new StringEntity(bulk.toString(), ContentType.APPLICATION_JSON));
+        request.setJsonEntity(bulk.toString());
+        client().performRequest(request);
     }
 
     /**

x-pack/qa/sql/src/main/java/org/elasticsearch/xpack/qa/sql/jdbc/JdbcIntegrationTestCase.java

@@ -9,6 +9,7 @@
 import org.apache.http.entity.ContentType;
 import org.apache.http.entity.StringEntity;
 import org.apache.http.util.EntityUtils;
+import org.elasticsearch.client.Request;
 import org.elasticsearch.common.CheckedConsumer;
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.xcontent.XContentBuilder;
@@ -85,16 +86,18 @@ protected Connection useDataSource() throws SQLException {
     }
 
     public static void index(String index, CheckedConsumer<XContentBuilder, IOException> body) throws IOException {
+        Request request = new Request("PUT", "/" + index + "/doc/1");
+        request.addParameter("refresh", "true");
         XContentBuilder builder = JsonXContent.contentBuilder().startObject();
         body.accept(builder);
         builder.endObject();
-        HttpEntity doc = new StringEntity(Strings.toString(builder), ContentType.APPLICATION_JSON);
-        client().performRequest("PUT", "/" + index + "/doc/1", singletonMap("refresh", "true"), doc);
+        request.setJsonEntity(Strings.toString(builder));
+        client().performRequest(request);
     }
 
     protected String clusterName() {
         try {
-            String response = EntityUtils.toString(client().performRequest("GET", "/").getEntity());
+            String response = EntityUtils.toString(client().performRequest(new Request("GET", "/")).getEntity());
             return XContentHelper.convertToMap(JsonXContent.jsonXContent, response, false).get("cluster_name").toString();
         } catch (IOException e) {
             throw new RuntimeException(e);