SQL: Add project routing support to JDBC (elastic#138756)

Author: luigidellaquila

docs/changelog/138756.yaml

@@ -0,0 +1,5 @@
+pr: 138756
+summary: Add project routing support to JDBC
+area: SQL
+type: enhancement
+issues: []

x-pack/plugin/sql/jdbc/src/main/java/org/elasticsearch/xpack/sql/jdbc/JdbcHttpClient.java

@@ -75,7 +75,8 @@ Cursor query(String sql, List<SqlTypedParamValue> params, RequestMeta meta) thro
             conCfg.fieldMultiValueLeniency(),
             conCfg.indexIncludeFrozen(),
             conCfg.binaryCommunication(),
-            conCfg.allowPartialSearchResults()
+            conCfg.allowPartialSearchResults(),
+            conCfg.projectRouting()
         );
         ResponseWithWarnings<SqlQueryResponse> response = httpClient.query(sqlRequest);
         return new DefaultCursor(
@@ -99,7 +100,8 @@ Tuple<String, List<List<Object>>> nextPage(String cursor, RequestMeta meta) thro
             TimeValue.timeValueMillis(meta.pageTimeoutInMs()),
             new RequestInfo(Mode.JDBC),
             conCfg.binaryCommunication(),
-            conCfg.allowPartialSearchResults()
+            conCfg.allowPartialSearchResults(),
+            conCfg.projectRouting()
         );
         SqlQueryResponse response = httpClient.query(sqlRequest).response();
         return new Tuple<>(response.cursor(), response.rows());

x-pack/plugin/sql/jdbc/src/test/java/org/elasticsearch/xpack/sql/jdbc/JdbcConfigurationTests.java

@@ -116,6 +116,12 @@ public void testDebugFlushAlways() throws Exception {
         assertThat(ci.flushAlways(), is(false));
     }
 
+    public void testProjectRouting() throws Exception {
+        JdbcConfiguration ci = ci(jdbcPrefix() + "a:1/?project.routing=foo");
+        assertThat(ci.baseUri().toString(), is("http://a:1/"));
+        assertThat(ci.projectRouting(), is("foo"));
+    }
+
     public void testTypeInParam() throws Exception {
         Exception e = expectThrows(JdbcSQLException.class, () -> ci(jdbcPrefix() + "a:1/foo/bar/tar?debug=true&debug.out=jdbc.out"));
         assertEquals("Unknown parameter [debug.out]; did you mean [debug.output]", e.getMessage());

x-pack/plugin/sql/sql-action/src/test/java/org/elasticsearch/xpack/sql/action/TestSqlQueryRequest.java

@@ -87,7 +87,8 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
             ProtoShim.toProto(this.waitForCompletionTimeout()),
             this.keepOnCompletion(),
             ProtoShim.toProto(this.keepAlive()),
-            this.allowPartialSearchResults()
+            this.allowPartialSearchResults(),
+            null
         );
         return SqlTestUtils.toXContentBuilder(builder, this, protoInstance);
     }

x-pack/plugin/sql/sql-action/src/test/java/org/elasticsearch/xpack/sql/action/TestSqlTranslateRequest.java

@@ -62,7 +62,8 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
             null,
             false,
             null,
-            false
+            false,
+            null
         );
         return SqlTestUtils.toXContentBuilder(builder, this, protoInstance);
     }

x-pack/plugin/sql/sql-client/src/main/java/org/elasticsearch/xpack/sql/client/ConnectionConfiguration.java

@@ -73,6 +73,9 @@ public class ConnectionConfiguration {
     public static final String ALLOW_PARTIAL_SEARCH_RESULTS = "allow.partial.search.results";
     public static final String ALLOW_PARTIAL_SEARCH_RESULTS_DEFAULT = "false";
 
+    // CPS project routing
+    public static final String PROJECT_ROUTING = "project.routing";
+
     protected static final Set<String> OPTION_NAMES = new LinkedHashSet<>(
         Arrays.asList(
             PROPERTIES_VALIDATION,
@@ -85,7 +88,8 @@ public class ConnectionConfiguration {
             AUTH_USER,
             AUTH_PASS,
             CATALOG,
-            ALLOW_PARTIAL_SEARCH_RESULTS
+            ALLOW_PARTIAL_SEARCH_RESULTS,
+            PROJECT_ROUTING
         )
     );
 
@@ -116,6 +120,8 @@ public class ConnectionConfiguration {
 
     private final boolean allowPartialSearchResults;
 
+    private final String projectRouting;
+
     @SuppressWarnings("this-escape")
     public ConnectionConfiguration(URI baseURI, String connectionString, Properties props) throws ClientException {
         this.connectionString = connectionString;
@@ -157,6 +163,8 @@ public ConnectionConfiguration(URI baseURI, String connectionString, Properties
             settings.getProperty(ALLOW_PARTIAL_SEARCH_RESULTS, ALLOW_PARTIAL_SEARCH_RESULTS_DEFAULT),
             Boolean::parseBoolean
         );
+
+        projectRouting = settings.getProperty(PROJECT_ROUTING);
     }
 
     public ConnectionConfiguration(
@@ -173,7 +181,8 @@ public ConnectionConfiguration(
         String pass,
         SslConfig sslConfig,
         ProxyConfig proxyConfig,
-        boolean allowPartialSearchResults
+        boolean allowPartialSearchResults,
+        String projectRouting
     ) throws ClientException {
         this.validateProperties = validateProperties;
         this.binaryCommunication = binaryCommunication;
@@ -194,6 +203,7 @@ public ConnectionConfiguration(
         this.baseURI = baseURI;
 
         this.allowPartialSearchResults = allowPartialSearchResults;
+        this.projectRouting = projectRouting;
     }
 
     private static URI normalizeSchema(URI uri, String connectionString, boolean isSSLEnabled) {
@@ -311,4 +321,8 @@ public String connectionString() {
     public boolean allowPartialSearchResults() {
         return allowPartialSearchResults;
     }
+
+    public String projectRouting() {
+        return projectRouting;
+    }
 }

x-pack/plugin/sql/sql-client/src/main/java/org/elasticsearch/xpack/sql/client/HttpClient.java

@@ -85,11 +85,21 @@ public SqlQueryResponse basicQuery(String query, int fetchSize) throws SQLExcept
     }
 
     public SqlQueryResponse basicQuery(String query, int fetchSize, boolean fieldMultiValueLeniency) throws SQLException {
-        return basicQuery(query, fetchSize, fieldMultiValueLeniency, cfg.allowPartialSearchResults());
+        return basicQuery(query, fetchSize, fieldMultiValueLeniency, cfg.allowPartialSearchResults(), cfg.projectRouting());
     }
 
     public SqlQueryResponse basicQuery(String query, int fetchSize, boolean fieldMultiValueLeniency, boolean allowPartialSearchResults)
         throws SQLException {
+        return basicQuery(query, fetchSize, fieldMultiValueLeniency, allowPartialSearchResults, cfg.projectRouting());
+    }
+
+    public SqlQueryResponse basicQuery(
+        String query,
+        int fetchSize,
+        boolean fieldMultiValueLeniency,
+        boolean allowPartialSearchResults,
+        String projectRouting
+    ) throws SQLException {
         // TODO allow customizing the time zone - this is what session set/reset/get should be about
         // method called only from CLI
         SqlQueryRequest sqlRequest = new SqlQueryRequest(
@@ -106,7 +116,8 @@ public SqlQueryResponse basicQuery(String query, int fetchSize, boolean fieldMul
             fieldMultiValueLeniency,
             false,
             cfg.binaryCommunication(),
-            allowPartialSearchResults
+            allowPartialSearchResults,
+            projectRouting
         );
         return query(sqlRequest).response();
     }
@@ -123,7 +134,8 @@ public SqlQueryResponse nextPage(String cursor) throws SQLException {
             TimeValue.timeValueMillis(cfg.pageTimeout()),
             new RequestInfo(Mode.CLI),
             cfg.binaryCommunication(),
-            cfg.allowPartialSearchResults()
+            cfg.allowPartialSearchResults(),
+            cfg.projectRouting()
         );
         return post(CoreProtocol.SQL_QUERY_REST_ENDPOINT, sqlRequest, Payloads::parseQueryResponse).response();
     }
@@ -179,7 +191,8 @@ private boolean head(String path, long timeoutInMs) throws SQLException {
             cfg.authPass(),
             cfg.sslConfig(),
             cfg.proxyConfig(),
-            CoreProtocol.ALLOW_PARTIAL_SEARCH_RESULTS
+            CoreProtocol.ALLOW_PARTIAL_SEARCH_RESULTS,
+            cfg.projectRouting()
         );
         try {
             return java.security.AccessController.doPrivileged(

x-pack/plugin/sql/sql-client/src/test/java/org/elasticsearch/xpack/sql/client/HttpClientRequestTests.java

@@ -175,7 +175,8 @@ private void assertBinaryRequestForDrivers(XContentType xContentType) throws URI
             randomBoolean(),
             randomBoolean(),
             isBinary,
-            randomBoolean()
+            randomBoolean(),
+            null
         );
 
         prepareMockResponse();

x-pack/plugin/sql/sql-proto/src/main/java/org/elasticsearch/xpack/sql/proto/Payloads.java

@@ -41,6 +41,7 @@
 import static org.elasticsearch.xpack.sql.proto.CoreProtocol.MODE_NAME;
 import static org.elasticsearch.xpack.sql.proto.CoreProtocol.PAGE_TIMEOUT_NAME;
 import static org.elasticsearch.xpack.sql.proto.CoreProtocol.PARAMS_NAME;
+import static org.elasticsearch.xpack.sql.proto.CoreProtocol.PROJECT_ROUTING_NAME;
 import static org.elasticsearch.xpack.sql.proto.CoreProtocol.QUERY_NAME;
 import static org.elasticsearch.xpack.sql.proto.CoreProtocol.REQUEST_TIMEOUT_NAME;
 import static org.elasticsearch.xpack.sql.proto.CoreProtocol.ROWS_NAME;
@@ -242,6 +243,9 @@ public static void generate(
         if (request.allowPartialSearchResults()) {
             generator.writeBooleanField(ALLOW_PARTIAL_SEARCH_RESULTS_NAME, request.allowPartialSearchResults());
         }
+        if (request.projectRouting() != null) {
+            generator.writeStringField(PROJECT_ROUTING_NAME, request.projectRouting());
+        }
 
         if (extraFields != null) {
             extraFields.accept(generator);

x-pack/plugin/sql/sql-proto/src/main/java/org/elasticsearch/xpack/sql/proto/SqlQueryRequest.java

@@ -39,6 +39,7 @@ public class SqlQueryRequest extends AbstractSqlRequest {
     private final TimeValue keepAlive;
 
     private final boolean allowPartialSearchResults;
+    private final String projectRouting;
 
     public SqlQueryRequest(
         String query,
@@ -57,7 +58,8 @@ public SqlQueryRequest(
         TimeValue waitForCompletionTimeout,
         boolean keepOnCompletion,
         TimeValue keepAlive,
-        boolean allowPartialSearchResults
+        boolean allowPartialSearchResults,
+        String projectRouting
     ) {
         super(requestInfo);
         this.query = query;
@@ -76,6 +78,7 @@ public SqlQueryRequest(
         this.keepOnCompletion = keepOnCompletion;
         this.keepAlive = keepAlive;
         this.allowPartialSearchResults = allowPartialSearchResults;
+        this.projectRouting = projectRouting;
     }
 
     public SqlQueryRequest(
@@ -92,7 +95,8 @@ public SqlQueryRequest(
         boolean fieldMultiValueLeniency,
         boolean indexIncludeFrozen,
         Boolean binaryCommunication,
-        boolean allowPartialSearchResults
+        boolean allowPartialSearchResults,
+        String projectRouting
     ) {
         this(
             query,
@@ -111,7 +115,8 @@ public SqlQueryRequest(
             CoreProtocol.DEFAULT_WAIT_FOR_COMPLETION_TIMEOUT,
             CoreProtocol.DEFAULT_KEEP_ON_COMPLETION,
             CoreProtocol.DEFAULT_KEEP_ALIVE,
-            allowPartialSearchResults
+            allowPartialSearchResults,
+            projectRouting
         );
     }
 
@@ -121,7 +126,8 @@ public SqlQueryRequest(
         TimeValue pageTimeout,
         RequestInfo requestInfo,
         boolean binaryCommunication,
-        boolean allowPartialSearchResults
+        boolean allowPartialSearchResults,
+        String projectRouting
     ) {
         this(
             "",
@@ -137,7 +143,8 @@ public SqlQueryRequest(
             CoreProtocol.FIELD_MULTI_VALUE_LENIENCY,
             CoreProtocol.INDEX_INCLUDE_FROZEN,
             binaryCommunication,
-            allowPartialSearchResults
+            allowPartialSearchResults,
+            projectRouting
         );
     }
 
@@ -227,6 +234,10 @@ public TimeValue keepAlive() {
         return keepAlive;
     }
 
+    public String projectRouting() {
+        return projectRouting;
+    }
+
     public boolean allowPartialSearchResults() {
         return allowPartialSearchResults;
     }