Support v2 "async" transactions returning multipart Arrow responses (#17)

* Support v2 "async" transactions returning multipart Arrow responses

Part 1: WIP basic GET/PUT calls for the v2 API.

* Parse the MultiPartForms + Arrow results into Arrow.Tables[]

* Add tests for v2 async/sync responses

* Update src/api.jl

* Add support for v1_inputs to v2-style exec :) (#23)

* Add support for v1_inputs to v2-style exec :)

* Adjust to latest /problems endpoint implementation from frontend

* Better comment + version patch on parse_multipart_form

* clarify test

* Rename exec to exec_async

* Fixups:

- support empty responses
- export the new names
- add an "is finished" function

* Delete the struct ResultPhysicalRelation, and just return a pair

* prefer haskey

* Switch back to returning a Dict, since i think it's more normal.

* Add examples/exec_async.jl:

```bash
$ julia --proj examples/exec_async.jl "nhd-test-1" "nhd-s" "def x = 1 def x = x+1 def output = x"
Transaction is created...
JSON3.Object{Vector{UInt8}, Vector{UInt64}} with 2 entries:
  :id => "c616a8ae-3be2-0d51-fb0e-b7f862bc367d"
  :state => "CREATED"
```

* Fix bug in transaction_is_done after switching to dicts

Author: NHDaly

Project.toml

@@ -4,10 +4,12 @@ version = "0.0.1"
 
 [deps]
 ArgParse = "c7e460c6-2fb9-53a9-8c5b-16f535851c63"
+Arrow = "69666777-d1a9-59fb-9406-91d4454c9d45"
 ConfParser = "88353bc9-fd38-507d-a820-d3b43837d6b9"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 HTTP = "cd3eb016-35fb-5094-929b-558a96fad6f3"
 JSON3 = "0f8b85d8-7281-11e9-16c2-39a750bddbf1"
+Mocking = "78c3b35d-d492-501b-9361-3d52fe80e533"
 Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c"
 
 [compat]

examples/exec_async.jl

@@ -0,0 +1,59 @@
+# 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
+
+# Execute the given query string.
+# Example:
+# $ julia --proj=. examples/exec_async.jl "nhd-test-1" "nhd-s" "2+2"
+# Transaction is done...
+# JSON3.Object{Vector{UInt8}, Vector{UInt64}} with 2 entries:
+#   :id    => "261f5d59-f1b4-f778-4c13-a7993871c972"
+#   :state => "CREATED"
+
+import RAI
+using RAI: Context, HTTPError, exec_async, load_config, show_result, get_transaction
+
+include("parseargs.jl")
+
+function run(database, engine, source; profile)
+    conf = load_config(; profile = profile)
+    ctx = Context(conf)
+    txn = exec_async(ctx, database, engine, source)
+    println("Transaction is created...")
+    display(txn)
+    println()
+end
+
+function main()
+    args = parseargs(
+        "database", Dict(:help => "database name", :required => true),
+        "engine", Dict(:help => "engine name", :required => true),
+        "command", Dict(:help => "rel source string"),
+        ["--file", "-f"], Dict(:help => "rel source file"),
+        "--readonly", Dict(:help => "readonly query (default: false)", :action => "store_true"),
+        "--profile", Dict(:help => "config profile (default: default)"))
+    try
+        source = nothing
+        if !isnothing(args.command)
+            source = args.command
+        elseif !isnothing(args.file)
+            source = open(args.file, "r")
+        end
+        isnothing(source) && return  # nothing to execute
+        run(args.database, args.engine, source; profile = args.profile)
+    catch e
+        e isa HTTPError ? show(e) : rethrow()
+    end
+end
+
+main()

src/RAI.jl

@@ -69,10 +69,17 @@ export
 
 export
     exec,
+    exec_async,
     list_edbs,
     load_csv,
     load_json
 
+export
+    get_transaction,
+    get_transaction_metadata,
+    get_transaction_problems,
+    get_transaction_results
+
 export
     Relation,
     TransactionResult,

src/api.jl

@@ -18,12 +18,17 @@
 # functionality as direclty as possible, but in a way that is natural for the
 # Julia language.
 
+#using Infiltrator
 import JSON3
+import Arrow
+
+using Mocking: Mocking, @mock  # For unit testing, by mocking API server responses
 
 const PATH_DATABASE = "/database"
 const PATH_ENGINE = "/compute"
 const PATH_OAUTH_CLIENTS = "/oauth-clients"
 const PATH_TRANSACTION = "/transaction"
+const PATH_ASYNC_TRANSACTIONS = "/transactions"
 const PATH_USERS = "/users"
 
 struct HTTPError <: Exception
@@ -57,7 +62,7 @@ function _mkurl(ctx::Context, path)
 end
 
 function _print_request(method, path, query, body)
-    println("$method $path") 
+    println("$method $path")
     !isnothing(query) && for (k, v) in query
         println("$k: $v")
     end
@@ -338,7 +343,7 @@ end
 
 # Execute the given query string, using any given optioanl query inputs.
 # todo: consider create_transaction
-# todo: consider create_transaction to better align with future transaciton
+# todo: consider create_transaction to better align with future transaction
 #   resource model
 function exec(ctx::Context, database::AbstractString, engine::AbstractString, source; inputs = nothing, readonly = false, kw...)
     source isa IO && (source = read(source, String))
@@ -349,6 +354,113 @@ end
 # todo: when we have async transactions, add a variation that dispatches and
 #   waits .. consider creating two entry points for readonly and readwrite.
 
+function exec_async(ctx::Context, database::AbstractString, engine::AbstractString, source; inputs = nothing, readonly = false, kw...)
+    source isa IO && (source = read(source, String))
+    tx_body = Dict(
+        "dbname" => database,
+        "engine_name" => engine,
+        "query" => source,
+        #"nowait_durable" => self.nowait_durable, # TODO: currently unsupported
+        "readonly" => readonly,
+        # "sync_mode" => "async"
+    )
+    if inputs !== nothing
+        tx_body["v1_inputs"] = [_make_query_action_input(k, v) for (k, v) in inputs]
+    end
+    body = JSON3.write(tx_body)
+    path = _mkurl(ctx, PATH_ASYNC_TRANSACTIONS)
+    rsp = @mock request(ctx, "POST", path; body = body, kw...)
+    return _parse_response(rsp)
+end
+
+function _parse_response(rsp)
+    content_type = HTTP.header(rsp, "Content-Type")
+    if lowercase(content_type) == "application/json"
+        content = HTTP.body(rsp)
+        # async mode
+        return JSON3.read(content)
+    elseif occursin("multipart/form-data", lowercase(content_type))
+        # sync mode
+        return _parse_multipart_fastpath_sync_response(rsp)
+    else
+        error("Unknown response content-type, for response:\n$(rsp)")
+    end
+end
+
+function get_transaction(ctx::Context, id::AbstractString; kw...)
+    path = PATH_ASYNC_TRANSACTIONS * "/$id"
+    rsp = _get(ctx, path; kw...)
+    return rsp.transaction
+end
+
+function transaction_is_done(txn)
+    if haskey(txn, "transaction")
+        txn = txn["transaction"]
+    end
+    return txn["state"] ∈ ("COMPLETED", "ABORTED")
+end
+
+function get_transaction_metadata(ctx::Context, id::AbstractString; kw...)
+    path = PATH_ASYNC_TRANSACTIONS * "/$id/metadata"
+    rsp = _get(ctx, path; kw...)
+    return rsp
+end
+
+function get_transaction_problems(ctx::Context, id::AbstractString; kw...)
+    path = PATH_ASYNC_TRANSACTIONS * "/$id/problems"
+    rsp = _get(ctx, path; kw...)
+    return rsp
+end
+
+function get_transaction_results(ctx::Context, id::AbstractString; kw...)
+    path = PATH_ASYNC_TRANSACTIONS * "/$id/results"
+    path = _mkurl(ctx, path)
+    rsp = request(ctx, "GET", path; kw...)
+    content_type = HTTP.header(rsp, "Content-Type")
+    if !occursin("multipart/form-data", content_type)
+        throw(HTTPError(400, "Unexpected response content-type for rsp:\n$rsp"))
+    end
+    return _parse_multipart_results_response(rsp)
+end
+
+function _parse_multipart_fastpath_sync_response(msg)
+    # TODO: in-place conversion to Arrow without copying the bytes.
+    #   ... HTTP.parse_multipart_form() copies the bytes into IOBuffers.
+    parts = _parse_multipart_form(msg)
+    @assert parts[1].name == "transaction"
+    @assert parts[2].name == "metadata"
+
+    problems_idx = findfirst(p->p.name == "problems", parts)
+    problems = JSON3.read(parts[problems_idx])
+
+    results_start_idx = findfirst(p->startswith(p.name, '/'), parts)
+    if results_start_idx === nothing
+        results = []
+    else
+        results = _extract_multipart_results_response(@view(parts[results_start_idx:end]))
+    end
+
+    return Dict(
+        "transaction" => JSON3.read(parts[1]),
+        "metadata" => JSON3.read(parts[2]),
+        "problems" => problems,
+        "results" => results,
+    )
+end
+
+function _parse_multipart_results_response(msg)
+    # TODO: in-place conversion to Arrow without copying the bytes.
+    #   ... HTTP.parse_multipart_form() copies the bytes into IOBuffers.
+    parts = _parse_multipart_form(msg)
+    return _extract_multipart_results_response(parts)
+end
+function _extract_multipart_results_response(parts)
+    return [
+        (part.name => Arrow.Table(part.data)) for part in parts
+    ]
+end
+
+
 function list_edbs(ctx::Context, database::AbstractString, engine::AbstractString; kw...)
     tx = Transaction(ctx.region, database, engine, "OPEN"; readonly = true)
     data = body(tx, _make_list_edb_action())
@@ -444,3 +556,27 @@ function load_model(ctx::Context, database::AbstractString, engine::AbstractStri
     actions = [_make_load_model_action(name, model) for (name, model) in models]
     return _post(ctx, PATH_TRANSACTION; query = query(tx), body = body(tx, actions...), kw...)
 end
+
+
+
+# --- utils -------------------------
+# Patch for older versions of HTTP package that don't support parsing multipart responses:
+if hasmethod(HTTP.MultiPartParsing.parse_multipart_form, (HTTP.Response,))
+    # Available as of HTTP v0.9.18:
+    _parse_multipart_form = HTTP.MultiPartParsing.parse_multipart_form
+else
+    # This function is copied directly from this PR: https://github.com/JuliaWeb/HTTP.jl/pull/817
+    function _parse_multipart_form(msg::HTTP.Message)
+        # parse boundary from Content-Type
+        m = match(r"multipart/form-data; boundary=(.*)$", msg["Content-Type"])
+        m === nothing && return nothing
+
+        boundary_delimiter = m[1]
+
+        # [RFC2046 5.1.1](https://tools.ietf.org/html/rfc2046#section-5.1.1)
+        length(boundary_delimiter) > 70 && error("boundary delimiter must not be greater than 70 characters")
+
+        return HTTP.MultiPartParsing.parse_multipart_body(HTTP.payload(msg), boundary_delimiter)
+    end
+end
+# -----------------------------------

test/api.jl

@@ -0,0 +1,98 @@
+using RAI
+using Test
+import HTTP, Arrow
+using Mocking
+
+Mocking.activate()
+
+# -----------------------------------
+# v2 transactions
+
+make_patch(response) = @patch RAI.request(ctx::Context, args...; kw...) = response
+
+const v2_async_response = HTTP.Response(200, [
+        "Content-Type" => "application/json",
+    ],
+    body = """{"id":"1fc9001b-1b88-8685-452e-c01bc6812429","state":"CREATED"}""")
+
+const v2_get_results_response() = join([
+        "--8a89e52be8efe57f0b68ea75388314a3",
+        "Content-Disposition: form-data; name=\"/:output/Int64\"; filename=\"/:output/Int64\"",
+        "Content-Type: application/vnd.apache.arrow.stream",
+        "",
+        "\xff\xff\xff\xffx\0\0\0\x10\0\0\0\0\0\n\0\f\0\n\0\b\0\x04\0\n\0\0\0\x10\0\0\0\x01\0\x04\0\b\0\b\0\0\0\x04\0\b\0\0\0\x04\0\0\0\x01\0\0\0\x14\0\0\0\x10\0\x14\0\x10\0\0\0\x0e\0\b\0\0\0\x04\0\x10\0\0\0\x10\0\0\0\x18\0\0\0\0\0\x02\0\x1c\0\0\0\0\0\0\0\b\0\f\0\b\0\a\0\b\0\0\0\0\0\0\x01@\0\0\0\x02\0\0\0v1\0\0\xff\xff\xff\xff\x88\0\0\0\x14\0\0\0\0\0\0\0\f\0\x16\0\x14\0\x12\0\f\0\x04\0\f\0\0\0\b\0\0\0\0\0\0\0\x14\0\0\0\0\0\x03\0\x04\0\n\0\x18\0\f\0\b\0\x04\0\n\0\0\0\x14\0\0\08\0\0\0\x01\0\0\0\0\0\0\0\0\0\0\0\x02\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\b\0\0\0\0\0\0\0\0\0\0\0\x01\0\0\0\x01\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\x04\0\0\0\0\0\0\0\xff\xff\xff\xff\0\0\0\0",
+    ], "\r\n")
+
+const v2_fastpath_response = HTTP.Response(200, [
+        "Content-Type" => "Content-Type: multipart/form-data; boundary=8a89e52be8efe57f0b68ea75388314a3",
+        "Transfer-Encoding" => "chunked",
+    ],
+    body = join([
+    "",
+    "--8a89e52be8efe57f0b68ea75388314a3",
+    "Content-Disposition: form-data; name=\"transaction\"; filename=\"\"",
+    "Content-Type: application/json",
+    "",
+    """{"id":"a3e3bc91-0a98-50ba-733c-0987e160eb7d","results_format_version":"2.0.1","state":"COMPLETED"}""",
+    "--8a89e52be8efe57f0b68ea75388314a3",
+    "Content-Disposition: form-data; name=\"metadata\"; filename=\"\"",
+    "Content-Type: application/json",
+    "",
+    """[{"relationId":"/:output/Int64","types":[":output","Int64"]}]""",
+    "--8a89e52be8efe57f0b68ea75388314a3",
+    "Content-Disposition: form-data; name=\"problems\"; filename=\"\"",
+    "Content-Type: application/json",
+    "",
+    """[]""",
+    v2_get_results_response(),
+    "--8a89e52be8efe57f0b68ea75388314a3--",
+    "",
+], "\r\n"))
+
+function make_arrow_table(vals)
+    io = IOBuffer()
+    Arrow.write(io, (v1=vals,))
+    seekstart(io)
+    return Arrow.Table(io)
+end
+
+@testset "exec_async" begin
+    ctx = Context("region", "scheme", "host", "2342", nothing)
+
+    @testset "async response" begin
+        patch = make_patch(v2_async_response)
+
+        apply(patch) do
+            rsp = RAI.exec_async(ctx, "engine", "database", "2+2")
+            @test rsp == JSON3.read("""{"id":"1fc9001b-1b88-8685-452e-c01bc6812429","state":"CREATED"}""")
+        end
+    end
+
+    @testset "sync response" begin
+        patch = make_patch(v2_fastpath_response)
+
+        apply(patch) do
+            rsp = RAI.exec_async(ctx, "engine", "database", "2+2")
+            @test rsp["transaction"] == JSON3.read("""{
+                    "id": "a3e3bc91-0a98-50ba-733c-0987e160eb7d",
+                    "results_format_version": "2.0.1",
+                    "state": "COMPLETED"
+                }""")
+            @test rsp["metadata"] == [JSON3.read("""{
+                "relationId": "/:output/Int64",
+                    "types": [
+                                ":output",
+                                "Int64"
+                            ]
+            }""")]
+            @test rsp["problems"] == Union{}[]
+
+            # Test for the expected arrow data:
+            expected_data = make_arrow_table([4])
+            # Arrow.Tables can't be compared via == (https://github.com/apache/arrow-julia/issues/310)
+            @test length(rsp["results"]) == 1
+            @test rsp["results"][1][1] == "/:output/Int64"
+            @test collect(rsp["results"][1][2]) == collect(expected_data)
+        end
+    end
+end

test/runtests.jl

@@ -18,6 +18,10 @@ import Tables
 using RAI
 using Test
 
+@testset "api.jl" begin
+    include("api.jl")
+end
+
 # def output =
 #     1, "foo", 3.4, :foo;
 #     2, "bar", 5.6, :foo