Change return type of query and handling of promises for json object

surrealdb · Aug 23, 2023 · 403607f · 403607f
1 parent 0b3c59a
commit 403607f
Show file tree

Hide file tree

Showing 7 changed files with 71 additions and 20 deletions.
diff --git a/src/intTest/java/com/surrealdb/refactor/DemoScenarioTest.java b/src/intTest/java/com/surrealdb/refactor/DemoScenarioTest.java
@@ -9,6 +9,7 @@
 import com.surrealdb.refactor.driver.*;
 import com.surrealdb.refactor.types.Credentials;
 import com.surrealdb.refactor.types.Param;
+import com.surrealdb.refactor.types.QueryBlockResult;
 import com.surrealdb.refactor.types.surrealdb.Value;
 import java.net.URI;
 import java.util.List;
@@ -46,17 +47,17 @@ public void testDemoScenario() throws Exception {
                         new Param("year", Value.fromJson(new JsonPrimitive(2013))));
 
         // Execute the query
-        List<Value> results = surrealDB.query(query.toString(), params);
+        QueryBlockResult results = surrealDB.query(query.toString(), params);
 
         // Validate the results of the multi-statement query
-        assertEquals(results.size(), 3, results.toString());
+        assertEquals(results.getResult().size(), 3, results.toString());
         assertEquals(
-                results.get(0).intoJson(),
+                results.getResult().get(0).getResult().get(0).intoJson(),
                 asJson(
                         Tuple.of("name", new JsonPrimitive("leslie")),
                         Tuple.of("id", new JsonPrimitive("person:lamport"))));
         assertEquals(
-                results.get(1).intoJson(),
+                results.getResult().get(1).getResult().get(0).intoJson(),
                 asJson(
                         Tuple.of("name", new JsonPrimitive("leslie")),
                         Tuple.of("id", new JsonPrimitive("person:lamport"))));

diff --git a/src/main/java/com/surrealdb/refactor/driver/HttpConnection.java b/src/main/java/com/surrealdb/refactor/driver/HttpConnection.java
@@ -5,6 +5,7 @@
 import com.surrealdb.refactor.exception.SurrealDBUnimplementedException;
 import com.surrealdb.refactor.types.Credentials;
 import com.surrealdb.refactor.types.Param;
+import com.surrealdb.refactor.types.QueryBlockResult;
 import com.surrealdb.refactor.types.surrealdb.Value;
 import java.net.URI;
 import java.util.Arrays;
@@ -26,7 +27,7 @@ public UnusedSurrealDB<StatelessSurrealDB> authenticate(Credentials credentials)
                 StatelessSurrealDB surrealdb =
                         new StatelessSurrealDB() {
                             @Override
-                            public List<Value> query(String query, List<Param> params) {
+                            public QueryBlockResult query(String query, List<Param> params) {
                                 throw new SurrealDBUnimplementedException(
                                         "https://github.com/surrealdb/surrealdb.java/issues/61",
                                         "HTTP connections are not yet implemented");

diff --git a/src/main/java/com/surrealdb/refactor/driver/StatelessSurrealDB.java b/src/main/java/com/surrealdb/refactor/driver/StatelessSurrealDB.java
@@ -1,6 +1,7 @@
 package com.surrealdb.refactor.driver;
 
 import com.surrealdb.refactor.types.Param;
+import com.surrealdb.refactor.types.QueryBlockResult;
 import com.surrealdb.refactor.types.surrealdb.Value;
 import java.util.List;
 
@@ -12,5 +13,5 @@ public interface StatelessSurrealDB {
      * @param query
      * @return
      */
-    List<Value> query(String query, List<Param> params);
+    QueryBlockResult query(String query, List<Param> params);
 }
diff --git a/src/main/java/com/surrealdb/refactor/driver/SurrealDBWebsocketClientProtocolHandler.java b/src/main/java/com/surrealdb/refactor/driver/SurrealDBWebsocketClientProtocolHandler.java
@@ -9,6 +9,7 @@
 import com.surrealdb.refactor.exception.UnknownResponseToRequest;
 import com.surrealdb.refactor.types.Credentials;
 import com.surrealdb.refactor.types.Param;
+import com.surrealdb.refactor.types.QueryBlockResult;
 import io.netty.channel.*;
 import io.netty.handler.codec.http.websocketx.*;
 import io.netty.util.concurrent.Promise;
@@ -27,7 +28,7 @@ public class SurrealDBWebsocketClientProtocolHandler
             Logger.getLogger(SurrealDBWebsocketClientProtocolHandler.class.toString());
     private static final String PROPERTY_REQUEST_ID = "id";
     private ChannelPromise handshakeFuture;
-    private final ConcurrentMap<String, Promise<Object>> requestMap = new ConcurrentHashMap();
+    private final ConcurrentMap<String, Promise<JsonObject>> requestMap = new ConcurrentHashMap();
 
     private Channel channel;
 
@@ -61,7 +62,7 @@ protected void channelRead0(ChannelHandlerContext ctx, TextWebSocketFrame msg) {
                         "Received a message presumed to be a response without a request id");
             }
             String requestID = obj.getAsJsonPrimitive(PROPERTY_REQUEST_ID).getAsString();
-            Promise<Object> promise = requestMap.remove(requestID);
+            Promise<JsonObject> promise = requestMap.remove(requestID);
             if (promise == null) {
                 promise.setFailure(
                         new UnknownResponseToRequest(
@@ -86,7 +87,7 @@ public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) {
         ctx.close();
     }
 
-    public Future<Object> query(String requestID, String query, List<Param> params) {
+    public Future<JsonObject> query(String requestID, String query, List<Param> params) {
         String method = "query";
         checkChannelAndThrow(method);
         QueryMessage queryMessage = new QueryMessage(requestID, query, params);
@@ -95,11 +96,11 @@ public Future<Object> query(String requestID, String query, List<Param> params)
         return sendAndPromise(method, requestID, serialised);
     }
 
-    public Future<Object> signin(Credentials credentials) {
+    public Future<JsonObject> signin(Credentials credentials) {
         return signin(UUID.randomUUID().toString(), credentials);
     }
 
-    public Future<Object> signin(String requestID, Credentials credentials) {
+    public Future<JsonObject> signin(String requestID, Credentials credentials) {
         String method = "signin";
         checkChannelAndThrow(method);
         // Construct message to be sent
@@ -109,11 +110,11 @@ public Future<Object> signin(String requestID, Credentials credentials) {
         return sendAndPromise(method, requestID, new Gson().toJson(signinMessage));
     }
 
-    public Future<Object> use(String namespace, String database) {
+    public Future<JsonObject> use(String namespace, String database) {
         return use(UUID.randomUUID().toString(), namespace, database);
     }
 
-    public Future<Object> use(String requestID, String namespace, String database) {
+    public Future<JsonObject> use(String requestID, String namespace, String database) {
         String method = "use";
         checkChannelAndThrow(method);
         UseMessage useMessage = new UseMessage(requestID, namespace, database);
@@ -130,9 +131,9 @@ private void checkChannelAndThrow(String method) {
         }
     }
 
-    private Promise<Object> sendAndPromise(
+    private Promise<JsonObject> sendAndPromise(
             String method, String requestID, String textFrameContent) {
-        Promise<Object> promise = channel.eventLoop().newPromise();
+        Promise<JsonObject> promise = channel.eventLoop().newPromise();
         registerRequest(requestID, promise);
         try {
             channel.writeAndFlush(new TextWebSocketFrame(textFrameContent)).sync();
@@ -144,8 +145,8 @@ private Promise<Object> sendAndPromise(
         return promise;
     }
 
-    private void registerRequest(String requestID, Promise<Object> promise) {
-        Promise<Object> popped = requestMap.putIfAbsent(requestID, promise);
+    private void registerRequest(String requestID, Promise<JsonObject> promise) {
+        Promise<JsonObject> popped = requestMap.putIfAbsent(requestID, promise);
         if (popped != null) {
             // Reinsert whatever we removed; This is actually quite problematic, and we should do a
             // contains check before in case

diff --git a/src/main/java/com/surrealdb/refactor/driver/WsPlaintextConnection.java b/src/main/java/com/surrealdb/refactor/driver/WsPlaintextConnection.java
@@ -2,9 +2,12 @@
 
 import com.google.gson.Gson;
 import com.google.gson.GsonBuilder;
+import com.google.gson.JsonObject;
 import com.surrealdb.refactor.exception.SurrealDBUnimplementedException;
 import com.surrealdb.refactor.types.Credentials;
 import com.surrealdb.refactor.types.Param;
+import com.surrealdb.refactor.types.QueryBlockResult;
+import com.surrealdb.refactor.types.QueryResult;
 import com.surrealdb.refactor.types.surrealdb.Value;
 import io.netty.bootstrap.Bootstrap;
 import io.netty.channel.*;
@@ -64,15 +67,15 @@ public UnusedSurrealDB<BidirectionalSurrealDB> authenticate(Credentials credenti
                         new BidirectionalSurrealDB() {
 
                             @Override
-                            public List<Value> query(String query, List<Param> params) {
-                                Object resp = null;
+                            public QueryBlockResult query(String query, List<Param> params) {
+                                JsonObject resp = null;
                                 try {
                                     resp = srdbHandler.query(UUID.randomUUID().toString(), query, params).get(2, TimeUnit.SECONDS);
                                 } catch (InterruptedException | ExecutionException | TimeoutException e) {
                                     throw new RuntimeException(e);
                                 }
                                 List<Value> casted = Arrays.asList(new Value(resp.toString()));
-                                return casted;
+                                return new QueryBlockResult(List.of(new QueryResult(casted, "change this status", "change this time")));
                             }
                         };
                 return new UnusedSurrealDB<>() {

diff --git a/src/main/java/com/surrealdb/refactor/types/QueryBlockResult.java b/src/main/java/com/surrealdb/refactor/types/QueryBlockResult.java
@@ -0,0 +1,23 @@
+package com.surrealdb.refactor.types;
+
+import lombok.Getter;
+import lombok.ToString;
+
+import java.util.List;
+
+@ToString
+public class QueryBlockResult {
+    private final List<QueryResult> result;
+
+    public QueryBlockResult(List<QueryResult> result) {
+        this.result = result;
+    }
+
+    /**
+     * Get all the query results from the query
+     * @return the result of each individual statement in a query block
+     */
+    public List<QueryResult> getResult() {
+        return result;
+    }
+}
diff --git a/src/main/java/com/surrealdb/refactor/types/QueryResult.java b/src/main/java/com/surrealdb/refactor/types/QueryResult.java
@@ -0,0 +1,21 @@
+package com.surrealdb.refactor.types;
+
+import com.surrealdb.refactor.types.surrealdb.Value;
+import lombok.Getter;
+import lombok.ToString;
+
+import java.util.List;
+
+@Getter
+@ToString
+public class QueryResult {
+    private final List<Value> result;
+    private final String status;
+    private final String time;
+
+    public QueryResult(List<Value> result, String status, String time) {
+        this.result = result;
+        this.status = status;
+        this.time = time;
+    }
+}