Adding user authentication functionality and including authorization in CI pipeline

Author: pebeto

.github/workflows/CI.yml

@@ -34,7 +34,7 @@ jobs:
       - name: Setup mlflow locally
         run: |
           pip install -r ./requirements.txt
-          python3 /opt/hostedtoolcache/Python/3.10.13/x64/bin/mlflow server --host 0.0.0.0 --port 5000 &
+          python3 /opt/hostedtoolcache/Python/3.10.13/x64/bin/mlflow server --app-name basic-auth --host 0.0.0.0 --port 5000 &
           sleep 5
       - uses: julia-actions/setup-julia@v1
         with:

.gitignore

@@ -7,3 +7,4 @@ mlruns
 coverage
 Pipfile
 Pipfile.lock
+*.db

Project.toml

@@ -19,7 +19,8 @@ URIs = "1.0"
 julia = "1.0"
 
 [extras]
+Base64 = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test"]
+test = ["Base64", "Test"]

src/MLFlowClient.jl

@@ -81,4 +81,7 @@ export getlatestmodelversions, getmodelversion, createmodelversion, deletemodelv
     transitionmodelversionstage, setmodelversiontag, deletemodelversiontag,
     getmodelversionbyalias
 
+include("services/user.jl")
+export createuser, getuser, updateuserpassword, updateuseradmin, deleteuser
+
 end

src/services/user.jl

@@ -0,0 +1,77 @@
+"""
+    createuser(instance::MLFlow, username::String, password::String)
+
+# Arguments
+- `instance`: [`MLFlow`](@ref) configuration.
+- `username`: Username.
+- `password`: Password.
+
+# Returns
+An [`User`](@ref) object.
+"""
+function createuser(instance::MLFlow, username::String, password::String)::User
+    result = mlfpost(instance, "users/create"; username=username, password=password)
+    return result["user"] |> User
+end
+
+"""
+    getuser(instance::MLFlow, username::String)
+
+# Arguments
+- `instance`: [`MLFlow`](@ref) configuration.
+- `username`: Username.
+
+# Returns
+An [`User`](@ref) object.
+"""
+function getuser(instance::MLFlow, username::String)::User
+    result = mlfget(instance, "users/get"; username=username)
+    return result["user"] |> User
+end
+
+"""
+    deleteuser(instance::MLFlow, username::String, password::String)
+
+# Arguments
+- `instance`: [`MLFlow`](@ref) configuration.
+- `username`: Username.
+- `password`: New password.
+
+# Returns
+`true` if successful. Otherwise, raises exception.
+"""
+function updateuserpassword(instance::MLFlow, username::String, password::String)::Bool
+    mlfpatch(instance, "users/update-password"; username=username, password=password)
+    return true
+end
+
+"""
+    updateuseradmin(instance::MLFlow, username::String, is_admin::Bool)
+
+# Arguments
+- `instance`: [`MLFlow`](@ref) configuration.
+- `username`: Username.
+- `is_admin`: New admin status.
+
+# Returns
+`true` if successful. Otherwise, raises exception.
+"""
+function updateuseradmin(instance::MLFlow, username::String, is_admin::Bool)::Bool
+    mlfpatch(instance, "users/update-admin"; username=username, is_admin=is_admin)
+    return true
+end
+
+"""
+    deleteuser(instance::MLFlow, username::String)
+
+# Arguments
+- `instance`: [`MLFlow`](@ref) configuration.
+- `username`: Username.
+
+# Returns
+`true` if successful. Otherwise, raises exception.
+"""
+function deleteuser(instance::MLFlow, username::String)::Bool
+    mlfdelete(instance, "users/delete"; username=username)
+    return true
+end

src/types/user.jl

@@ -17,3 +17,7 @@ struct User
     experiment_permissions::Array{ExperimentPermission}
     registered_model_permissions::Array{RegisteredModelPermission}
 end
+User(data::Dict{String,Any}) = User(data["id"] |> string, data["username"], data["is_admin"],
+    [ExperimentPermission(permission) for permission in get(data, "experiment_permissions", [])],
+    [RegisteredModelPermission(permission) for permission in get(data, "registered_model_permissions", [])])
+Base.show(io::IO, t::User) = show(io, ShowCase(t, new_lines=true))

test/base.jl

@@ -1,7 +1,8 @@
-using MLFlowClient
 using Test
-using UUIDs
 using Dates
+using UUIDs
+using Base64
+using MLFlowClient
 
 function mlflow_server_is_running(mlf::MLFlow)
     try
@@ -16,7 +17,8 @@ end
 # skips test if mlflow is not available on default location, ENV["MLFLOW_TRACKING_URI"]
 macro ensuremlf()
     e = quote
-        mlf = MLFlow()
+        encoded_credentials = Base64.base64encode("admin:password")
+        mlf = MLFlow(headers=Dict("Authorization" => "Basic $(encoded_credentials)"))
         mlflow_server_is_running(mlf) || return nothing
     end
     eval(e)

test/runtests.jl

@@ -11,3 +11,4 @@ include("services/artifact.jl")
 include("services/experiment.jl")
 include("services/registered_model.jl")
 include("services/model_version.jl")
+include("services/user.jl")

test/services/user.jl

@@ -0,0 +1,62 @@
+@testset verbose = true "create user" begin
+    @ensuremlf
+
+    user = createuser(mlf, "missy", "gala")
+
+    @test user isa User
+    @test user.username == "missy"
+    @test user.is_admin == false
+
+    deleteuser(mlf, user.username)
+end
+
+@testset verbose = true "get user" begin
+    @ensuremlf
+
+    user = createuser(mlf, "missy", "gala")
+
+    retrieved_user = getuser(mlf, "missy")
+
+    @test retrieved_user isa User
+    @test retrieved_user.username == "missy"
+    @test retrieved_user.is_admin == false
+
+    deleteuser(mlf, retrieved_user.username)
+end
+
+@testset verbose = true "update user password" begin
+    @ensuremlf
+
+    getmlfinstance(encoded_credentials::String) =
+        MLFlow(headers=Dict("Authorization" => "Basic $(encoded_credentials)"))
+
+    user = createuser(mlf, "missy", "gala")
+    encoded_credentials = Base64.base64encode("$(user.username):gala")
+
+    updateuserpassword(getmlfinstance(encoded_credentials), "missy", "ana")
+    encoded_credentials = Base64.base64encode("$(user.username):ana")
+
+    @test_nowarn searchexperiments(getmlfinstance(encoded_credentials))
+    deleteuser(mlf, user.username)
+end
+
+@testset verbose = true "update user admin" begin
+    @ensuremlf
+
+    user = createuser(mlf, "missy", "gala")
+    updateuseradmin(mlf, "missy", true)
+
+    retrieved_user = getuser(mlf, "missy")
+    @test retrieved_user.is_admin == true
+
+    deleteuser(mlf, retrieved_user.username)
+end
+
+@testset verbose = true "delete user" begin
+    @ensuremlf
+
+    user = createuser(mlf, "missy", "gala")
+    deleteuser(mlf, "missy")
+
+    @test_throws ErrorException getuser(mlf, "missy")
+end