ProtoBuf metadata (#60)

* Update CHANGELOG

* Bump version

* Add EnumX and ProtocolBuffers packages

* ProtoBuf codegen

* Add methods to handle and show new ProtoBuf metadata

* Ensure Accept: application/x-protobuf header for requests

* Deserialize metadata using ProtoBuf decoder

* Add tests

* Add v2_async_response to test

* Move encode workaround out of generated file

* Add comment about strings and symbol constant types

* Fix typo

* JuliaFormatter

* Migrate to ProtoBuf.jl 1.0

* Fix merge conflict: support show for empty tuples.

* Fix merge conflict

* Regenerate protos from new schema

* Remove SYMBOL reference

* Change show representation of RelationId

* Adapt tests

Co-authored-by: Nathan Daly <[email protected]>
Co-authored-by: Tomáš Drvoštěp <[email protected]>

Author: 3 people

CHANGELOG.md

@@ -1,13 +1,27 @@
 # Changelog
 
-## main
+## v0.2.0
+* Deprecate metadata JSON format.
+* Add support to deserialize ProtoBuf metadata.
+* `get_transaction_metadata` returns a `MetadataInfo`, see `src/proto`
+  and `src/gen/relationalai/protocol/` for more information.
 
-* Add find_user to api.jl
-* Rename examples/get-userid.jl to examples/find-user.jl
-* Fixed bug in support for custom extra `headers` in SDK. For example:
-```
-create_engine(ctx, engine; size = size, headers=["my-custom-header" => "custom header value"])
-```
+## v0.1.5
+* Add support for optional audience field to Config
+
+## v0.1.4
+* Retry retryable HTTP errors
+
+## v0.1.0
+* Implement V2 `show_result` method
+
+## v0.0.4
+* Properly filter transaction results based on multi-part content type
+* Set `HTTP.jl` compat to `1.0`
+
+## v0.0.3
+* New access token for each request
+* Anticipate access token expiration
 
 ## v0.0.2
 
@@ -16,3 +30,12 @@ create_engine(ctx, engine; size = size, headers=["my-custom-header" => "custom h
     - If you cancel the polling via `ctrl-C`, the error log will print the transaction ID, so you can still
       recover the transaction or cancel it.
 * Consistent return format (`Dict`) from `exec_async()` and `exec()`, regardless of whether you get synchronous results ([#24](https://github.com/RelationalAI/rai-sdk-julia/pull/24)).
+
+## main
+
+* Add find_user to api.jl
+* Rename examples/get-userid.jl to examples/find-user.jl
+* Fixed bug in support for custom extra `headers` in SDK. For example:
+```
+create_engine(ctx, engine; size = size, headers=["my-custom-header" => "custom header value"])
+```

Project.toml

@@ -1,25 +1,29 @@
 name = "RAI"
 uuid = "9c30249a-7e08-11ec-0e99-a323e937e79f"
-version = "0.1.8"
+version = "0.2.0"
 
 [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"
+EnumX = "4e289a0a-7415-4d19-859d-a7e5c4648b56"
 HTTP = "cd3eb016-35fb-5094-929b-558a96fad6f3"
 JSON3 = "0f8b85d8-7281-11e9-16c2-39a750bddbf1"
 Mocking = "78c3b35d-d492-501b-9361-3d52fe80e533"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
+ProtoBuf = "3349acd9-ac6a-5e09-bcdb-63829b23a429"
 Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c"
 
 [compat]
 ArgParse = "1"
 Arrow = "2"
 ConfParser = "0.1"
+EnumX = "1"
 HTTP = "1.0"
 JSON3 = "1"
 Mocking = "0.7"
+ProtoBuf = "1"
 Tables = "1"
 julia = "1.6"
 

src/RAI.jl

@@ -96,6 +96,11 @@ export
 include("creds.jl")
 include("config.jl")
 include("rest.jl")
+# Generated with:
+# ProtoBuf.protojl(readdir("src/proto"), "src/proto", "src/gen", parametrize_oneofs=true)
+include("gen/relationalai/relationalai.jl")
+import .relationalai.protocol
+include("metadata.jl")
 include("response.jl")
 include("api.jl")
 include("results.jl")

src/api.jl

@@ -19,6 +19,7 @@
 # Julia language.
 
 import Arrow
+import ProtoBuf
 using Base.Threads: @spawn
 import Dates
 import JSON3
@@ -564,10 +565,15 @@ function exec_async(ctx::Context, database::AbstractString, engine::AbstractStri
     end
     body = JSON3.write(tx_body)
     path = _mkurl(ctx, PATH_ASYNC_TRANSACTIONS)
+    headers = _ensure_proto_accept_header(get(kw, :headers, []))
     rsp = @mock request(ctx, "POST", path; body = body, kw...)
     return _parse_response(rsp)
 end
 
+# We **only** support ProtoBuf metadata, so we overwrite the `Accept` header.
+_ensure_proto_accept_header(headers) =
+    collect(merge(Dict(headers), Dict("Accept" => "application/x-protobuf")))
+
 function _parse_response(rsp)
     content_type = HTTP.header(rsp, "Content-Type")
     if lowercase(content_type) == "application/json"
@@ -605,8 +611,12 @@ end
 
 function get_transaction_metadata(ctx::Context, id::AbstractString; kw...)
     path = PATH_ASYNC_TRANSACTIONS * "/$id/metadata"
-    rsp = _get(ctx, path; kw...)
-    return rsp
+    path = _mkurl(ctx, path)
+    headers = _ensure_proto_accept_header(get(kw, :headers, []))
+    rsp = request(ctx, "GET", path; kw..., headers)
+    d = ProtoBuf.ProtoDecoder(IOBuffer(rsp.body));
+    metadata = ProtoBuf.decode(d, protocol.MetadataInfo)
+    return metadata
 end
 
 function get_transaction_problems(ctx::Context, id::AbstractString; kw...)
@@ -631,12 +641,16 @@ function _parse_multipart_fastpath_sync_response(msg)
     #   ... HTTP.parse_multipart_form() copies the bytes into IOBuffers.
     parts = _parse_multipart_form(msg)
     @assert parts[1].name == "transaction"
-    @assert parts[2].name == "metadata"
 
     transaction = JSON3.read(parts[1])
-    metadata = JSON3.read(parts[2])
+
+    metadata_idx = findfirst(p->p.name == "metadata.proto", parts)
+    d = ProtoBuf.ProtoDecoder(parts[metadata_idx].data);
+    metadata = ProtoBuf.decode(d, protocol.MetadataInfo)
+
     problems_idx = findfirst(p->p.name == "problems", parts)
     problems = JSON3.read(parts[problems_idx])
+
     results = _extract_multipart_results_response(parts)
 
     return TransactionResponse(transaction, metadata, problems, results)

src/gen/relationalai/protocol/message_pb.jl

@@ -0,0 +1,74 @@
+# Autogenerated using ProtoBuf.jl v1.0.5 on 2022-08-31T16:58:14.621
+# original file: /Users/david/Projects/rai-sdk-julia/src/proto/message.proto (proto3 syntax)
+
+import ProtoBuf as PB
+using ProtoBuf: OneOf
+using EnumX: @enumx
+
+export RelationMetadata, MetadataInfo
+
+struct RelationMetadata
+    relation_id::Union{Nothing,RelationId}
+    file_name::String
+end
+PB.default_values(::Type{RelationMetadata}) = (;relation_id = nothing, file_name = "")
+PB.field_numbers(::Type{RelationMetadata}) = (;relation_id = 1, file_name = 2)
+
+function PB.decode(d::PB.AbstractProtoDecoder, ::Type{<:RelationMetadata})
+    relation_id = Ref{Union{Nothing,RelationId}}(nothing)
+    file_name = ""
+    while !PB.message_done(d)
+        field_number, wire_type = PB.decode_tag(d)
+        if field_number == 1
+            PB.decode!(d, relation_id)
+        elseif field_number == 2
+            file_name = PB.decode(d, String)
+        else
+            PB.skip(d, wire_type)
+        end
+    end
+    return RelationMetadata(relation_id[], file_name)
+end
+
+function PB.encode(e::PB.AbstractProtoEncoder, x::RelationMetadata)
+    initpos = position(e.io)
+    !isnothing(x.relation_id) && PB.encode(e, 1, x.relation_id)
+    !isempty(x.file_name) && PB.encode(e, 2, x.file_name)
+    return position(e.io) - initpos
+end
+function PB._encoded_size(x::RelationMetadata)
+    encoded_size = 0
+    !isnothing(x.relation_id) && (encoded_size += PB._encoded_size(x.relation_id, 1))
+    !isempty(x.file_name) && (encoded_size += PB._encoded_size(x.file_name, 2))
+    return encoded_size
+end
+
+struct MetadataInfo
+    relations::Vector{RelationMetadata}
+end
+PB.default_values(::Type{MetadataInfo}) = (;relations = Vector{RelationMetadata}())
+PB.field_numbers(::Type{MetadataInfo}) = (;relations = 1)
+
+function PB.decode(d::PB.AbstractProtoDecoder, ::Type{<:MetadataInfo})
+    relations = PB.BufferedVector{RelationMetadata}()
+    while !PB.message_done(d)
+        field_number, wire_type = PB.decode_tag(d)
+        if field_number == 1
+            PB.decode!(d, relations)
+        else
+            PB.skip(d, wire_type)
+        end
+    end
+    return MetadataInfo(relations[])
+end
+
+function PB.encode(e::PB.AbstractProtoEncoder, x::MetadataInfo)
+    initpos = position(e.io)
+    !isempty(x.relations) && PB.encode(e, 1, x.relations)
+    return position(e.io) - initpos
+end
+function PB._encoded_size(x::MetadataInfo)
+    encoded_size = 0
+    !isempty(x.relations) && (encoded_size += PB._encoded_size(x.relations, 1))
+    return encoded_size
+end

src/gen/relationalai/protocol/protocol.jl

@@ -0,0 +1,6 @@
+module protocol
+
+include("schema_pb.jl")
+include("message_pb.jl")
+
+end # module protocol