Models actions from v1 to v2 protocol (#76)

* moving models actions from v1 to v2

* models integration tests

* refactoring && cleanup

* remove unused _list_models

* Update test/integration.jl

Co-authored-by: Nathan Daly <nathan.daly@relational.ai>

* updates

* fix

* changelog updates

* addressing PR comments

Co-authored-by: Nathan Daly <nathan.daly@relational.ai>

Author: NRHelmi

CHANGELOG.md

@@ -1,5 +1,8 @@
 # Changelog
 
+## main
+* Update models actions to use v2 protocol
+* Update `load_model` to `load_models`
 ## v0.2.1
 * Increased `connection_limit` to 4096
 

examples/delete-model.jl renamed to examples/delete-models.jl

@@ -21,7 +21,7 @@ include("parseargs.jl")
 function run(database, engine, model; profile)
     cfg = load_config(; profile = profile)
     ctx = Context(cfg)
-    rsp = delete_model(ctx, database, engine, model)
+    rsp = delete_models(ctx, database, engine, model)
     println(rsp)
 end
 

examples/load-model.jl renamed to examples/load-models.jl

@@ -14,7 +14,7 @@
 
 # Load the given Rel model into the given database.
 
-using RAI: Context, HTTPError, load_config, load_model
+using RAI: Context, HTTPError, load_config, load_models
 
 include("parseargs.jl")
 
@@ -25,7 +25,7 @@ function run(database, engine, fullname; profile)
     models = Dict(_sansext(fullname) => read(fullname, String))
     cfg = load_config(; profile = profile)
     ctx = Context(cfg)
-    rsp = load_model(ctx, database, engine, models)
+    rsp = load_models(ctx, database, engine, models)
     println(rsp)
 end
 

src/RAI.jl

@@ -46,10 +46,10 @@ export
     list_engines
 
 export
-    delete_model,
+    delete_models,
     get_model,
     list_models,
-    load_model
+    load_models
 
 export
     create_oauth_client,
@@ -66,7 +66,7 @@ export
     update_user
 
 export
-    delete_model,
+    delete_models,
     get_model,
     list_models
 

src/api.jl

@@ -234,11 +234,8 @@ function delete_engine(ctx::Context, engine::AbstractString; kw...)
     return _delete(ctx, PATH_ENGINE; body = JSON3.write(data), kw...)
 end
 
-function delete_model(ctx::Context, database::AbstractString, engine::AbstractString, model::AbstractString; kw...)
-    tx = Transaction(ctx.region, database, engine, "OPEN"; readonly = false)
-    actions = [_make_delete_models_action([model])]
-    return _post(ctx, PATH_TRANSACTION; query = query(tx), body = body(tx, actions...), kw...)
-end
+# escape rel special string
+_escape_string_for_rel(str) = replace(repr(str), '%' => "\\%")
 
 function delete_oauth_client(ctx::Context, id::AbstractString; kw...)
     return _delete(ctx, joinpath(PATH_OAUTH_CLIENTS, id); kw...)
@@ -270,15 +267,6 @@ function get_database(ctx::Context, database::AbstractString; kw...)
     return rsp[1]
 end
 
-# todo: move to rel query
-function get_model(ctx::Context, database::AbstractString, engine::AbstractString, name::AbstractString; kw...)
-    models = _list_models(ctx, database, engine; kw...)
-    for model in models
-        model["name"] == name && return model["value"]
-    end
-    throw(HTTPError(404))
-end
-
 function get_oauth_client(ctx::Context, id::AbstractString; kw...)
     return _get(ctx, joinpath(PATH_OAUTH_CLIENTS, id); kw...).client
 end
@@ -387,22 +375,6 @@ function _make_actions(actions...)
     return result
 end
 
-function _make_delete_models_action(models::Vector)
-    return Dict(
-        "type" => "ModifyWorkspaceAction",
-        "delete_source" => models)
-end
-
-function _make_load_model_action(name, model)
-    return Dict(
-        "type" => "InstallAction",
-        "sources" => [_make_query_source(name, model)])
-end
-
-function _make_list_models_action()
-    return Dict("type" => "ListSourceAction")
-end
-
 function _make_list_edb_action()
     return Dict("type" => "ListEdbAction")
 end
@@ -695,19 +667,6 @@ function list_edbs(ctx::Context, database::AbstractString, engine::AbstractStrin
     return rsp.actions[1].result.rels
 end
 
-function _list_models(ctx::Context, database::AbstractString, engine::AbstractString; kw...)
-    tx = Transaction(ctx.region, database, engine, "OPEN"; readonly = true)
-    data = body(tx, _make_list_models_action())
-    rsp = _post(ctx, PATH_TRANSACTION; query = query(tx), body = data, kw...).actions
-    length(rsp) == 0 && return []
-    return rsp[1].result.sources
-end
-
-function list_models(ctx::Context, database::AbstractString, engine::AbstractString; kw...)
-    models = _list_models(ctx, database, engine; kw...)
-    return [model["name"] for model in models]
-end
-
 function _gen_literal(value)
     return "$value"
 end
@@ -777,13 +736,82 @@ function load_json(ctx::Context, database::AbstractString, engine::AbstractStrin
     return exec(ctx, database, engine, source; inputs = inputs, readonly = false, kw...)
 end
 
-function load_model(ctx::Context, database::AbstractString, engine::AbstractString, models::Dict; kw...)
-    tx = Transaction(ctx.region, database, engine, "OPEN"; readonly = false)
-    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...)
+function load_models(ctx::Context, database::AbstractString, engine::AbstractString, models::Dict; kw...)
+    queries = []
+    queries_inputs = Dict()
+    rand_uint = rand(UInt64)
+
+    index = 0
+    for (name, value) in models
+        input_name = string("input_", rand_uint, "_", index)
+        push!(queries, """
+            def delete:rel:catalog:model["$name"] = rel:catalog:model["$name"]
+            def insert:rel:catalog:model["$name"] = $input_name
+        """)
+
+        queries_inputs[input_name] = value
+        index+=1
+    end
+
+    return exec(ctx, database, engine, join(queries, "\n"); inputs = queries_inputs, readonly = false, kw...)
 end
 
+function load_models_async(ctx::Context, database::AbstractString, engine::AbstractString, models::Dict; kw...)
+    queries = []
+    queries_inputs = Dict()
+    rand_uint = rand(UInt64)
+
+    index = 0
+    for (name, value) in models
+        input_name = string("input_", rand_uint, "_", index)
+        push!(queries, """
+            def delete:rel:catalog:model["$name"] = rel:catalog:model["$name"]
+            def insert:rel:catalog:model["$name"] = $input_name
+        """)
 
+        queries_inputs[input_name] = value
+        index+=1
+    end
+
+    return exec_async(ctx, database, engine, join(queries, "\n"); inputs = queries_inputs, readonly = false, kw...)
+end
+
+function list_models(ctx::Context, database::AbstractString, engine::AbstractString; kw...)
+    out_name = "model$(rand(UInt64))"
+    query = """ def output:$out_name[name] = rel:catalog:model(name, _) """
+    resp = exec(ctx, database, engine, query)
+    for result in resp.results
+        if occursin("/:output/:$out_name", result.first)
+            return [name for name in result.second.v1]
+        end
+    end
+end
+
+function get_model(ctx::Context, database::AbstractString, engine::AbstractString, name::AbstractString; kw...)
+    out_name = "model$(rand(UInt64))"
+    query = """def output:$out_name = rel:catalog:model[$(_escape_string_for_rel(name))]"""
+    resp = exec(ctx, database, engine, query)
+    for result in resp.results
+        if occursin("/:output/:$out_name", result.first)
+            return first(result.second.v1)
+        end
+    end
+    throw(HTTPError(404))
+end
+
+function delete_models(ctx::Context, database::AbstractString, engine::AbstractString, models::Vector{String}; kw...)
+    queries = ["""
+            def delete:rel:catalog:model[$(_escape_string_for_rel(model))] = rel:catalog:model[$(_escape_string_for_rel(model))]
+        """ for model in models]
+    return exec(ctx, database, engine, join(queries, "\n"); readonly=false, kw...)
+end
+
+function delete_models_async(ctx::Context, database::AbstractString, engine::AbstractString, model::AbstractString; kw...)
+    queries = ["""
+            def delete:rel:catalog:model[$(_escape_string_for_rel(model))] = rel:catalog:model[$(_escape_string_for_rel(model))
+        """ for model in models]
+    return exec_async(ctx, database, engine, join(queries, "\n"); readonly=false, kw...)
+end
 
 # --- utils -------------------------
 # Patch for older versions of HTTP package that don't support parsing multipart responses:

test/integration.jl

@@ -278,7 +278,72 @@ with_engine(CTX) do engine_name
 
         # -----------------------------------
         # models
-        @testset "models" begin end
+        @testset "models" begin
+            models = list_models(CTX, database_name, engine_name)
+            @test length(models) > 0
+
+            models = Dict("test_model" => "def foo = :bar")
+            resp = load_models(CTX, database_name, engine_name, models)
+            @test resp.transaction.state == "COMPLETED"
+
+            value = get_model(CTX, database_name, engine_name, "test_model")
+            @test models["test_model"] == value
+
+            models = list_models(CTX, database_name, engine_name)
+            @test "test_model" in models
+
+            resp = delete_models(CTX, database_name, engine_name, ["test_model"])
+            @test resp.transaction.state == "COMPLETED"
+            @test length(resp.problems) == 0
+
+            models = list_models(CTX, database_name, engine_name)
+            @test !("test_model" in models)
+
+            # test escape special rel character
+            models = Dict("percent" => "def foo = \"98%\"")
+            resp = load_models(CTX, database_name, engine_name, models)
+            @test resp.transaction.state == "COMPLETED"
+            @test length(resp.problems) > 0
+            resp = delete_models(CTX, database_name, engine_name, ["percent"])
+            @test resp.transaction.state == "COMPLETED"
+            @test length(resp.problems) == 0
+
+            models = Dict("percent" => "def foo = \"98\\%\"")
+            resp = load_models(CTX, database_name, engine_name, models)
+            @test resp.transaction.state == "COMPLETED"
+            @test length(resp.problems) == 0
+            value = get_model(CTX, database_name, engine_name, "percent")
+            @test models["percent"] == value
+
+            models = list_models(CTX, database_name, engine_name)
+            @test "percent" in models
+
+            resp = delete_models(CTX, database_name, engine_name, ["percent"])
+            @test resp.transaction.state == "COMPLETED"
+            @test length(resp.problems) == 0
+
+            models = list_models(CTX, database_name, engine_name)
+            @test !("percent" in models)
+
+            # test escape """
+            models = Dict("triple_quoted" => "def foo = \"\"\"98\\%\"\"\" ")
+            resp = load_models(CTX, database_name, engine_name, models)
+            @test resp.transaction.state == "COMPLETED"
+            @test length(resp.problems) == 0
+            value = get_model(CTX, database_name, engine_name, "triple_quoted")
+            @test models["triple_quoted"] == value
+
+            models = list_models(CTX, database_name, engine_name)
+            @test "triple_quoted" in models
+            @test length(resp.problems) == 0
+
+            resp = delete_models(CTX, database_name, engine_name, ["triple_quoted"])
+            @test resp.transaction.state == "COMPLETED"
+            @test length(resp.problems) == 0
+
+            models = list_models(CTX, database_name, engine_name)
+            @test !("triple_quoted" in models)
+        end
     end
 end