Call query_async and poll for state in run_query.py (#78)

* Return higher-level structures for txn v2 results

* run_query to call query_async and poll for state

* add get_transaction_results_and_problems

* nits

* nits

* nits

* changes after code review

* bring names in line with other SDKs

* no need for show.results(multipart)

* show.results for query is physical

* file renamed

* adding execute.py

* query rename to exec_v1

Co-authored-by: George Kollias <286415+geokollias@users.noreply.github.com>
Co-authored-by: Sameera Salameh <sameera@sameeras-mbp.relational.ai.beta.tailscale.net>
Co-authored-by: Sameera Salameh <sameera@sameeras-mbp.mynet>
Co-authored-by: Pete Vilter <pete.vilter@gmail.com>

Author: 5 people

examples/execute.py

@@ -0,0 +1,41 @@
+# 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
+
+from argparse import ArgumentParser
+from urllib.request import HTTPError
+from railib import api, config, show
+
+
+def run(database: str, engine: str, command: str, readonly: bool, profile: str):
+    cfg = config.read(profile=profile)
+    ctx = api.Context(**cfg)
+    rsp = api.exec(ctx, database, engine, command, readonly=readonly)
+    show.results(rsp, "wire")
+
+
+if __name__ == "__main__":
+    p = ArgumentParser()
+    p.add_argument("database", type=str, help="database name")
+    p.add_argument("engine", type=str, help="engine name")
+    p.add_argument("command", type=str, help="rel source string")
+    p.add_argument("--readonly", action="store_true", default=False,
+                   help="readonly query (default: false)")
+    p.add_argument("-p", "--profile", type=str, default="default",
+                   help="profile name")
+    args = p.parse_args()
+    try:
+        run(args.database, args.engine, args.command, args.readonly,
+            args.profile)
+    except HTTPError as e:
+        show.http_error(e)

examples/run_query_async.py examples/execute_async.pyexamples/run_query_async.py renamed to examples/execute_async.py

@@ -21,7 +21,7 @@
 def run(database: str, engine: str, command: str, readonly: bool, profile: str):
     cfg = config.read(profile=profile)
     ctx = api.Context(**cfg)
-    rsp = api.query_async(ctx, database, engine, command, readonly=readonly)
+    rsp = api.exec_async(ctx, database, engine, command, readonly=readonly)
     print(json.dumps(rsp, indent=2))
 
 

examples/run_query.py examples/execute_v1.pyexamples/run_query.py renamed to examples/execute_v1.py

@@ -20,7 +20,7 @@
 def run(database: str, engine: str, command: str, readonly: bool, profile: str):
     cfg = config.read(profile=profile)
     ctx = api.Context(**cfg)
-    rsp = api.query(ctx, database, engine, command, readonly=readonly)
+    rsp = api.exec_v1(ctx, database, engine, command, readonly=readonly)
     show.results(rsp)
 
 

railib/api.py

@@ -17,6 +17,7 @@
 import io
 import json
 import pyarrow as pa
+import time
 from enum import Enum, unique
 from typing import List, Union
 from . import rest
@@ -655,7 +656,7 @@ def load_csv(ctx: Context, database: str, engine: str, relation: str,
     command = _gen_syntax_config(syntax)
     command += ("def config:data = data\n"
                 "def insert:%s = load_csv[config]" % relation)
-    return query(ctx, database, engine, command, inputs=inputs, readonly=False)
+    return exec_v1(ctx, database, engine, command, inputs=inputs, readonly=False)
 
 
 def load_json(ctx: Context, database: str, engine: str, relation: str,
@@ -669,16 +670,38 @@ def load_json(ctx: Context, database: str, engine: str, relation: str,
     inputs = {'data': data}
     command = ("def config:data = data\n"
                "def insert:%s = load_json[config]" % relation)
-    return query(ctx, database, engine, command, inputs=inputs, readonly=False)
+    return exec_v1(ctx, database, engine, command, inputs=inputs, readonly=False)
 
-
-def query(ctx: Context, database: str, engine: str, command: str,
+def exec_v1(ctx: Context, database: str, engine: str, command: str,
           inputs: dict = None, readonly: bool = True) -> dict:
     tx = Transaction(database, engine, readonly=readonly)
     return tx.run(ctx, _query_action(command, inputs=inputs))
 
+# Answers if the given transaction state is a terminal state.
+def is_txn_term_state(state: str) -> bool:
+    return state == "COMPLETED" or state == "ABORTED"
+
+def exec(ctx: Context, database: str, engine: str, command: str,
+          inputs: dict = None, readonly: bool = True) -> list:
+    async_result = exec_async(ctx, database, engine, command, readonly=readonly)
+    if isinstance(async_result, list):  # in case of if short-path, return results directly, no need to poll for state
+        return async_result
+
+    rsp = []
+    while True:
+        time.sleep(3)
+        txn = get_transaction(ctx, async_result["id"])
+        if is_txn_term_state(txn["state"]):
+            rsp.append(txn)
+            rsp.append(get_transaction_metadata(ctx, txn["id"]))
+            rsp.append(get_transaction_problems(ctx, txn["id"]))
+            rsp.append(get_transaction_results(ctx, txn["id"]))
+            break
+
+    return rsp
+
 
-def query_async(ctx: Context, database: str, engine: str, command: str,
+def exec_async(ctx: Context, database: str, engine: str, command: str,
                 readonly: bool = True, inputs: dict = None) -> Union[dict, list]:
     tx = TransactionAsync(database, engine, readonly=readonly)
     return tx.run(ctx, command, inputs=inputs)