SQL: Implement binary format support for SQL clear cursor

docs/changelog/84230.yaml

@@ -0,0 +1,6 @@
+pr: 84230
+summary: Implement binary format support for SQL clear cursor
+area: SQL
+type: bug
+issues:
+ - 53359

x-pack/plugin/sql/qa/jdbc/single-node/src/test/java/org/elasticsearch/xpack/sql/qa/jdbc/single_node/JdbcCloseCursorIT.java

@@ -0,0 +1,11 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+package org.elasticsearch.xpack.sql.qa.jdbc.single_node;
+
+import org.elasticsearch.xpack.sql.qa.jdbc.CloseCursorTestCase;
+
+public class JdbcCloseCursorIT extends CloseCursorTestCase {}

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

@@ -0,0 +1,57 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+package org.elasticsearch.xpack.sql.qa.jdbc;
+
+import org.junit.Before;
+
+import java.io.IOException;
+import java.sql.Connection;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.sql.Statement;
+
+public abstract class CloseCursorTestCase extends JdbcIntegrationTestCase {
+
+    @Before
+    public void initIndex() throws IOException {
+        index("library", "1", builder -> { builder.field("name", "foo"); });
+        index("library", "2", builder -> { builder.field("name", "bar"); });
+        index("library", "3", builder -> { builder.field("name", "baz"); });
+    }
+
+    public void testCloseCursor() throws SQLException {
+        try (Connection connection = createConnection(connectionProperties()); Statement statement = connection.createStatement()) {
+            statement.setFetchSize(1);
+            ResultSet results = statement.executeQuery(" SELECT name FROM library");
+            assertTrue(results.next());
+            results.close(); // force sending a cursor close since more pages are available
+            assertTrue(results.isClosed());
+        }
+    }
+
+    public void testCloseConsumedCursor() throws SQLException {
+        try (Connection connection = createConnection(connectionProperties()); Statement statement = connection.createStatement()) {
+            statement.setFetchSize(1);
+            ResultSet results = statement.executeQuery(" SELECT name FROM library");
+            for (int i = 0; i < 3; i++) {
+                assertTrue(results.next());
+            }
+            assertFalse(results.next());
+            results.close();
+            assertTrue(results.isClosed());
+        }
+    }
+
+    public void testCloseNoCursor() throws SQLException {
+        try (Connection connection = createConnection(connectionProperties()); Statement statement = connection.createStatement()) {
+            statement.setFetchSize(1);
+            ResultSet results = statement.executeQuery(" SELECT name FROM library where name = 'zzz'");
+            results.close();
+            assertTrue(results.isClosed());
+        }
+    }
+}

x-pack/plugin/sql/sql-action/src/main/java/org/elasticsearch/xpack/sql/action/SqlClearCursorRequest.java

@@ -10,6 +10,7 @@
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.xcontent.ConstructingObjectParser;
+import org.elasticsearch.xcontent.ParseField;
 import org.elasticsearch.xcontent.XContentParser;
 import org.elasticsearch.xpack.sql.proto.Mode;
 import org.elasticsearch.xpack.sql.proto.RequestInfo;
@@ -24,12 +25,16 @@
 import static org.elasticsearch.xpack.sql.action.AbstractSqlQueryRequest.CURSOR;
 import static org.elasticsearch.xpack.sql.action.AbstractSqlQueryRequest.MODE;
 import static org.elasticsearch.xpack.sql.action.AbstractSqlQueryRequest.VERSION;
+import static org.elasticsearch.xpack.sql.proto.CoreProtocol.BINARY_COMMUNICATION;
+import static org.elasticsearch.xpack.sql.proto.CoreProtocol.BINARY_FORMAT_NAME;
 
 /**
  * Request to clean all SQL resources associated with the cursor
  */
 public class SqlClearCursorRequest extends AbstractSqlRequest {
 
+    static final ParseField BINARY_FORMAT = new ParseField(BINARY_FORMAT_NAME);
+
     private static final ConstructingObjectParser<SqlClearCursorRequest, Void> PARSER =
         // here the position in "objects" is the same as the fields parser declarations below
         new ConstructingObjectParser<>(SqlClearCursorAction.NAME, objects -> {
@@ -43,9 +48,11 @@ public class SqlClearCursorRequest extends AbstractSqlRequest {
         PARSER.declareString(optionalConstructorArg(), MODE);
         PARSER.declareString(optionalConstructorArg(), CLIENT_ID);
         PARSER.declareString(optionalConstructorArg(), VERSION);
+        PARSER.declareBoolean(SqlClearCursorRequest::binaryCommunication, BINARY_FORMAT);
     }
 
     private String cursor;
+    private Boolean binaryCommunication = BINARY_COMMUNICATION;
 
     public SqlClearCursorRequest() {}
 
@@ -82,6 +89,15 @@ public SqlClearCursorRequest(StreamInput in) throws IOException {
         cursor = in.readString();
     }
 
+    public SqlClearCursorRequest binaryCommunication(Boolean binaryCommunication) {
+        this.binaryCommunication = binaryCommunication;
+        return this;
+    }
+
+    public Boolean binaryCommunication() {
+        return this.binaryCommunication;
+    }
+
     @Override
     public void writeTo(StreamOutput out) throws IOException {
         super.writeTo(out);

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

@@ -14,6 +14,7 @@
 import org.elasticsearch.xcontent.XContentParser;
 import org.elasticsearch.xcontent.XContentParserConfiguration;
 import org.elasticsearch.xcontent.XContentType;
+import org.elasticsearch.xpack.sql.proto.CoreProtocol;
 import org.elasticsearch.xpack.sql.proto.Mode;
 import org.elasticsearch.xpack.sql.proto.SqlTypedParamValue;
 
@@ -61,23 +62,27 @@ public void testClearCursorRequestParser() throws IOException {
               "cursor": "whatever",
               "mode": "%s",
               "client_id": "bla",
-              "version": "1.2.3"
+              "version": "1.2.3",
+              "binary_format": true
             }""".formatted(randomMode), SqlClearCursorRequest::fromXContent);
         assertNull(request.clientId());
         assertNull(request.version());
         assertEquals(randomMode, request.mode());
         assertEquals("whatever", request.getCursor());
+        assertTrue(request.binaryCommunication());
 
         randomMode = randomFrom(Mode.values());
         request = generateRequest("""
             {
               "cursor": "whatever",
               "mode": "%s",
-              "client_id": "bla"
+              "client_id": "bla",
+              "binary_format": false
             }""".formatted(randomMode.toString()), SqlClearCursorRequest::fromXContent);
         assertNull(request.clientId());
         assertEquals(randomMode, request.mode());
         assertEquals("whatever", request.getCursor());
+        assertFalse(request.binaryCommunication());
 
         request = generateRequest("{\"cursor\" : \"whatever\"}", SqlClearCursorRequest::fromXContent);
         assertNull(request.clientId());
@@ -94,6 +99,7 @@ public void testClearCursorRequestParser() throws IOException {
         assertNull(request.version());
         assertEquals(Mode.PLAIN, request.mode());
         assertEquals("whatever", request.getCursor());
+        assertEquals(CoreProtocol.BINARY_COMMUNICATION, request.binaryCommunication());
 
         request = generateRequest("""
             {"cursor" : "whatever", "client_id" : "cANVAs"}""", SqlClearCursorRequest::fromXContent);

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

@@ -33,7 +33,8 @@ public TestSqlClearCursorRequest(RequestInfo requestInfo, String cursor) {
     public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
         org.elasticsearch.xpack.sql.proto.SqlClearCursorRequest protoInstance = new org.elasticsearch.xpack.sql.proto.SqlClearCursorRequest(
             this.getCursor(),
-            this.requestInfo()
+            this.requestInfo(),
+            false
         );
         return SqlTestUtils.toXContentBuilder(builder, g -> Payloads.generate(g, protoInstance));
     }

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

@@ -104,7 +104,7 @@ public SqlQueryResponse nextPage(String cursor) throws SQLException {
     public boolean queryClose(String cursor, Mode mode) throws SQLException {
         SqlClearCursorResponse response = post(
             CoreProtocol.CLEAR_CURSOR_REST_ENDPOINT,
-            new SqlClearCursorRequest(cursor, new RequestInfo(mode)),
+            new SqlClearCursorRequest(cursor, new RequestInfo(mode), cfg.binaryCommunication()),
             Payloads::parseClearCursorResponse
         );
         return response.isSucceeded();

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

@@ -170,6 +170,7 @@ public static void generate(JsonGenerator generator, SqlClearCursorRequest reque
         generator.writeStringField("mode", request.mode().toString());
         writeIfValid(generator, "client_id", request.clientId());
         writeIfValidAsString(generator, "version", request.version());
+        generator.writeBooleanField("binary_format", request.binaryCommunication());
 
         generator.writeEndObject();
     }

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

@@ -14,16 +14,22 @@
 public class SqlClearCursorRequest extends AbstractSqlRequest {
 
     private final String cursor;
+    private final boolean binaryCommunication;
 
-    public SqlClearCursorRequest(String cursor, RequestInfo requestInfo) {
+    public SqlClearCursorRequest(String cursor, RequestInfo requestInfo, boolean binaryCommunication) {
         super(requestInfo);
         this.cursor = cursor;
+        this.binaryCommunication = binaryCommunication;
     }
 
     public String getCursor() {
         return cursor;
     }
 
+    public boolean binaryCommunication() {
+        return binaryCommunication;
+    }
+
     @Override
     public boolean equals(Object o) {
         if (this == o) return true;

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

@@ -10,12 +10,19 @@
 import org.elasticsearch.client.internal.node.NodeClient;
 import org.elasticsearch.core.RestApiVersion;
 import org.elasticsearch.rest.BaseRestHandler;
+import org.elasticsearch.rest.BytesRestResponse;
 import org.elasticsearch.rest.RestRequest;
-import org.elasticsearch.rest.action.RestToXContentListener;
+import org.elasticsearch.rest.RestResponse;
+import org.elasticsearch.rest.RestStatus;
+import org.elasticsearch.rest.action.RestResponseListener;
+import org.elasticsearch.xcontent.XContentBuilder;
 import org.elasticsearch.xcontent.XContentParser;
+import org.elasticsearch.xcontent.XContentType;
 import org.elasticsearch.xpack.sql.action.Protocol;
 import org.elasticsearch.xpack.sql.action.SqlClearCursorAction;
 import org.elasticsearch.xpack.sql.action.SqlClearCursorRequest;
+import org.elasticsearch.xpack.sql.action.SqlClearCursorResponse;
+import org.elasticsearch.xpack.sql.proto.Mode;
 
 import java.io.IOException;
 import java.util.List;
@@ -40,7 +47,18 @@ protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient cli
             sqlRequest = SqlClearCursorRequest.fromXContent(parser);
         }
 
-        return channel -> client.executeLocally(SqlClearCursorAction.INSTANCE, sqlRequest, new RestToXContentListener<>(channel));
+        return channel -> client.executeLocally(SqlClearCursorAction.INSTANCE, sqlRequest, new RestResponseListener<>(channel) {
+            @Override
+            public RestResponse buildResponse(SqlClearCursorResponse response) throws Exception {
+                Boolean binaryRequest = sqlRequest.binaryCommunication();
+                XContentType type = Boolean.TRUE.equals(binaryRequest) || (binaryRequest == null && Mode.isDriver(sqlRequest.mode()))
+                    ? XContentType.CBOR
+                    : XContentType.JSON;
+                XContentBuilder builder = channel.newBuilder(request.getXContentType(), type, false);
+                response.toXContent(builder, request);
+                return new BytesRestResponse(RestStatus.OK, builder);
+            }
+        });
     }
 
     @Override