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

@@ -705,25 +705,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.gotCompletedResult)
+                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(
@@ -775,12 +772,7 @@ 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,
-                Arrays.asList(metadataResponse),
-                problemsResult
-        );
+        return new TransactionAsyncResult(transactionResponse, results, Arrays.asList(metadataResponse), problemsResult, true);
     }
 
     public TransactionAsyncSingleResponse getTransaction(String id) throws HttpError, IOException, InterruptedException {

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

@@ -4,15 +4,32 @@
 
 public class TransactionAsyncResult extends Entity {
 
+    public Boolean gotCompletedResult;
     public TransactionAsyncCompactResponse transaction;
     public List<ArrowRelation> results;
     public List<TransactionAsyncMetadataResponse> metadata;
     public List<Object> problems;
 
-    public TransactionAsyncResult(TransactionAsyncCompactResponse transaction, List<ArrowRelation> results, List<TransactionAsyncMetadataResponse> metadata, List<Object> problems) {
+    public TransactionAsyncResult(
+            TransactionAsyncCompactResponse transaction,
+            List<ArrowRelation> results,
+            List<TransactionAsyncMetadataResponse> metadata,
+            List<Object> problems
+    ) {
         this.transaction = transaction;
         this.results = results;
         this.metadata = metadata;
         this.problems = problems;
+        this.gotCompletedResult = false;
+    }
+    public TransactionAsyncResult(
+            TransactionAsyncCompactResponse transaction,
+            List<ArrowRelation> results,
+            List<TransactionAsyncMetadataResponse> metadata,
+            List<Object> problems,
+            Boolean gotCompletedResult
+    ) {
+        this(transaction, results, metadata, problems);
+        this.gotCompletedResult = false;
     }
 }