Return higher-level structure for txn v2 results (#76)

* Return higher-level structures for txn v2 results

* no need for show.results(multipart)

Co-authored-by: Sameera Salameh <sameera@sameeras-mbp.relational.ai.beta.tailscale.net>

Author: geokollias

examples/get_transaction_results.py

@@ -25,7 +25,7 @@ def run(id: str, profile: str):
     ctx = api.Context(**cfg)
     rsp = api.get_transaction_results(ctx, id)
     print("Results:")
-    show.results(rsp, "multipart")
+    print(json.dumps(rsp, indent=2))
 
     rsp = api.get_transaction_problems(ctx, id)
     print("\nProblems:")

examples/run_query_async.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License
 
+import json
 from argparse import ArgumentParser
 from urllib.request import HTTPError
 from railib import api, config, show
@@ -21,7 +22,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)
-    show.results(rsp)
+    print(json.dumps(rsp, indent=2))
 
 
 if __name__ == "__main__":

railib/api.py

@@ -16,6 +16,7 @@
 
 import io
 import json
+import pyarrow as pa
 from enum import Enum, unique
 from typing import List, Union
 from . import rest
@@ -172,17 +173,31 @@ def _get_collection(ctx, path: str, key=None, **kwargs):
 # Parses "multipart/form-data" responses. It returns the parts in a list.
 def _parse_multipart(content_type: str, content: bytes) -> list:
     result = []
+    content_type_json = b'application/json'
+    content_type_arrow = b'application/vnd.apache.arrow.stream'
     boundary = _extract_multipart_boundary(content_type)
+
     parts = content.split(b'\r\n' + boundary)
     for i, part in enumerate(parts):
         # fix the first part, it may contain the boundary
         if i == 0:
             part = part.split(boundary)[-1]
-        # skip the part which contains the following
+        # Skip "--\r\n"
         if b'--\r\n' in part:
             continue
-        # append the part
-        result.append(part)
+        # Part body and headers are separated with CRLFCRLF. Get the body.
+        part_value = part.split(b'\r\n\r\n')[-1]
+        if content_type_json in part:
+            result.append(json.loads(part_value))
+        # if the part has arrow stream, then decode the arrow stream
+        # the results are in a form of a tuple/table
+        elif content_type_arrow in part:
+            with pa.ipc.open_stream(part_value) as reader:
+                schema = reader.schema
+                batches = [batch for batch in reader]
+                table = pa.Table.from_batches(batches=batches, schema=schema)
+                result.append(table.to_pydict())
+
     return result
 
 

railib/show.py

@@ -76,31 +76,6 @@ def _show_rel(rsp: dict) -> None:
         print(rsp["status"])
 
 
-def _show_multipart(parts: list):
-    result = []
-    content_type_json = b'application/json'
-    content_type_arrow_stream = b'application/vnd.apache.arrow.stream'
-    for part in parts:
-        # split the part
-        # part body and headers are separated with CRLFCRLF
-        strings = part.split(b'\r\n\r\n')
-        # last part is the content/body
-        part_value = strings[len(strings) - 1]
-        # content is json, decode the part content as json
-        if content_type_json in part:
-            result.append(json.loads(part_value))
-        # if the part has arrow stream, then decode the arrow stream
-        # the results are in a form of a tuple/table
-        elif content_type_arrow_stream in part:
-            with pa.ipc.open_stream(part_value) as reader:
-                schema = reader.schema
-                batches = [batch for batch in reader]
-                table = pa.Table.from_batches(batches=batches, schema=schema)
-                result.append(table.to_pydict())
-
-    json.dump(result, sys.stdout, indent=2)
-
-
 # Print the problems in the given response dict.
 def problems(rsp: dict) -> None:
     if rsp is None:
@@ -122,17 +97,13 @@ def problems(rsp: dict) -> None:
 
 
 # Print the results contained in the given response dict.
-def results(rsp: Union[dict, list], format=None) -> None:
+def results(rsp: Union[dict, list], format="physical") -> None:
     if rsp is None:
         return
-    if format is None:
-        format = "multipart" if isinstance(rsp, list) else "wire"
     if format == "wire":
         json.dump(rsp, sys.stdout, indent=2)
     elif format == "physical":
         _show_rel(rsp)
         problems(rsp)
-    elif format == "multipart":
-        _show_multipart(rsp)
     else:
         raise Exception(f"unknown format: '{format}'")