Optimize exec async

rai-sdk-examples/src/main/java/com/relationalai/examples/Execute.java

@@ -0,0 +1,71 @@
+/*
+ * Copyright 2022 RelationalAI, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"): you may
+ * not use this file except in compliance with the License. You may obtain
+ * a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ */
+
+package com.relationalai.examples;
+
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+
+import com.relationalai.Client;
+import com.relationalai.Config;
+import com.relationalai.HttpError;
+import com.relationalai.Json;
+
+public class Execute implements Runnable {
+    boolean readonly;
+    String database, engine, command, filename, profile;
+
+    // Returns the name of the file, without extension.
+    static String readFile(String fname) throws IOException {
+        return Files.readAllBytes(Path.of(fname)).toString();
+    }
+
+    String getCommand() throws IOException {
+        if (command != null)
+            return command; // prefer command line
+        if (filename != null)
+            return readFile(filename);
+        return null;
+    }
+
+    public void parseArgs(String[] args) {
+        var c = Command.create("Execute")
+                .addArgument("database")
+                .addArgument("engine")
+                .addOption("c", "rel source string")
+                .addOption("f", "rel source file")
+                .addFlag("readonly", "readonly query (default: false)")
+                .parseArgs(args);
+        this.database = c.getValue("database");
+        this.engine = c.getValue("engine");
+        this.command = c.getValue("c");
+        this.filename = c.getValue("f");
+        this.readonly = c.getValue("readonly", Boolean.class);
+        this.profile = c.getValue("profile");
+    }
+
+    public void run(String[] args) throws HttpError, InterruptedException, IOException {
+        parseArgs(args);
+        var cfg = Config.loadConfig("~/.rai/config", profile);
+        var client = new Client(cfg);
+        String source = getCommand();
+        if (source == null)
+            return; // nothing to execute
+        var rsp = client.execute(database, engine, source, readonly);
+        Json.print(rsp);
+    }
+}

rai-sdk/src/main/java/com/relationalai/Client.java

@@ -730,20 +730,22 @@ public TransactionAsyncResult execute(
             String source, boolean readonly,
             Map<String, String> inputs) throws HttpError, IOException, InterruptedException {
 
-        var id = executeAsync(database, engine, source, readonly, inputs).transaction.id;
+        var rsp = executeAsync(database, engine, source, readonly, inputs);
 
-        var transaction = getTransaction(id).transaction;
+        if (rsp.gotCompleteResult)
+                return rsp;
 
+        var transaction = getTransaction(rsp.transaction.id).transaction;
         while ( !("COMPLETED".equals(transaction.state) || "ABORTED".equals(transaction.state)) ) {
             Thread.sleep(2000);
-            transaction = getTransaction(id).transaction;
+            transaction = getTransaction(transaction.id).transaction;
         }
 
-        var results = getTransactionResults(id);
-        var metadata = getTransactionMetadata(id);
-        var problems = getTransactionProblems(id);
+        var results = getTransactionResults(transaction.id);
+        var metadata = getTransactionMetadata(transaction.id);
+        var problems = getTransactionProblems(transaction.id);
 
-        return new TransactionAsyncResult(transaction, results, metadata, problems);
+        return new TransactionAsyncResult(transaction, results, metadata, problems, true);
     }
 
     public TransactionAsyncResult executeAsync(
@@ -797,12 +799,8 @@ private TransactionAsyncResult readTransactionAsyncResults(List<TransactionAsync
         var problemsResult = parseProblemsResult(new String(problems.get(0).data, StandardCharsets.UTF_8));
 
         var results = readArrowFiles(files);
-        return new TransactionAsyncResult(
-                transactionResponse,
-                results,
-                metadataInfoResult,
-                problemsResult
-        );
+
+        return new TransactionAsyncResult(transactionResponse, results, metadataInfoResult, problemsResult, true);
     }
 
     public TransactionAsyncSingleResponse getTransaction(String id) throws HttpError, IOException, InterruptedException {

rai-sdk/src/main/java/com/relationalai/TransactionAsyncResult.java

@@ -6,6 +6,7 @@
 
 public class TransactionAsyncResult extends Entity {
 
+    public Boolean gotCompleteResult;
     public TransactionAsyncCompactResponse transaction;
     public List<ArrowRelation> results;
     public Message.MetadataInfo metadata;
@@ -20,6 +21,17 @@ public TransactionAsyncResult(
         this.results = results;
         this.metadata = metadata;
         this.problems = problems;
+        this.gotCompleteResult = false;
+    }
+    public TransactionAsyncResult(
+            TransactionAsyncCompactResponse transaction,
+            List<ArrowRelation> results,
+            Message.MetadataInfo metadata,
+            List<Object> problems,
+            Boolean gotCompleteResult
+    ) {
+        this(transaction, results, metadata, problems);
+        this.gotCompleteResult = gotCompleteResult;
     }
 
     @Override

rai-sdk/src/test/resources/metadata.pb

@@ -2,8 +2,8 @@
 >
 3
 "
-

-�
output
+

+�
output
 

 

 


rai-sdk/src/test/resources/metadata.pb.txt

@@ -5,11 +5,11 @@ relations {
       constant_type {
         rel_type {
           tag: PRIMITIVE_TYPE
-          primitive_type: SYMBOL
+          primitive_type: STRING
         }
         value {
           arguments {
-            tag: SYMBOL
+            tag: STRING
             string_val: "output"
           }
         }