feat: rebuild UI (#3)

mlflow_oidc_auth/app.py

@@ -14,71 +14,45 @@
 app.secret_key = app.config["SECRET_KEY"].encode("utf8")
 app.template_folder = template_dir
 
-# Add new routes
+# OIDC routes
 app.add_url_rule(rule=routes.LOGIN, methods=["GET"], view_func=views.login)
-app.add_url_rule(rule=routes.LOGOUT, methods=["GET", "POST"], view_func=views.logout)
-app.add_url_rule(
-    rule=routes.CALLBACK, methods=["GET", "POST"], view_func=views.callback
-)
+app.add_url_rule(rule=routes.LOGOUT, methods=["GET"], view_func=views.logout)
+app.add_url_rule(rule=routes.CALLBACK, methods=["GET"], view_func=views.callback)
+
+# UI routes
 app.add_url_rule(rule=routes.STATIC, methods=["GET"], view_func=views.oidc_static)
 app.add_url_rule(rule=routes.UI, methods=["GET"], view_func=views.oidc_ui)
 app.add_url_rule(rule=routes.UI_ROOT, methods=["GET"], view_func=views.oidc_ui)
-app.add_url_rule(rule=routes.OIDC_HOME, methods=["GET"], view_func=views.oidc_home)
-app.add_url_rule(
-    rule=routes.SEARCH_MODEL, methods=["GET"], view_func=views.search_model
-)
-app.add_url_rule(
-    rule=routes.SEARCH_EXPERIMENT, methods=["POST"], view_func=views.create_user
-)
-# app.add_url_rule(routes.LOGIN_MLFLOW, methods=['POST'], view_func=views.login_mlflow)
-app.add_url_rule(rule=routes.PERMISSIONS, methods=["GET"], view_func=views.permissions)
-app.add_url_rule(
-    rule=routes.PERMISSIONS_USERS, methods=["GET"], view_func=views.permissions_users
-)
-app.add_url_rule(
-    rule=routes.PERMISSIONS_EXPERIMENTS,
-    methods=["GET"],
-    view_func=views.permissions_experiments,
-)
-app.add_url_rule(
-    rule=routes.PERMISSIONS_MODELS, methods=["GET"], view_func=views.permissions_models
-)
-app.add_url_rule(
-    rule=routes.PERMISSIONS_USER_DETAILS,
-    methods=["GET"],
-    view_func=views.permissions_user_details,
-)
-app.add_url_rule(
-    rule=routes.PERMISSIONS_EXPERIMENT_DETAILS,
-    methods=["GET"],
-    view_func=views.permissions_experiment_details,
-)
-app.add_url_rule(
-    rule=routes.PERMISSIONS_MODEL_DETAILS,
-    methods=["GET"],
-    view_func=views.permissions_model_details,
-)
-app.add_url_rule(
-    rule=routes.GET_EXPERIMENT_PERMISSION,
-    methods=["GET"],
-    view_func=views.get_experiment_permission,
-)
-app.add_url_rule(
-    rule=routes.UPDATE_USER_PASSWORD,
-    methods=["GET"],
-    view_func=views.update_username_password,
-)
+
+# User token
+app.add_url_rule(rule=routes.CREATE_ACCESS_TOKEN, methods=["GET"], view_func=views.create_access_token)
+app.add_url_rule(rule=routes.GET_CURRENT_USER, methods=["GET"], view_func=views.get_current_user)
+
+# UI routes support
+app.add_url_rule(rule=routes.GET_EXPERIMENTS, methods=["GET"], view_func=views.get_experiments)
+app.add_url_rule(rule=routes.GET_MODELS, methods=["GET"], view_func=views.get_models)
+app.add_url_rule(rule=routes.GET_USERS, methods=["GET"], view_func=views.get_users)
+app.add_url_rule(rule=routes.GET_USER_EXPERIMENTS, methods=["GET"], view_func=views.get_user_experiments)
+app.add_url_rule(rule=routes.GET_USER_MODELS, methods=["GET"], view_func=views.get_user_models)
+app.add_url_rule(rule=routes.GET_EXPERIMENT_USERS, methods=["GET"], view_func=views.get_experiment_users)
+app.add_url_rule(rule=routes.GET_MODEL_USERS, methods=["GET"], view_func=views.get_model_users)
+
+# User management
+app.add_url_rule(rule=routes.CREATE_USER, methods=["POST"], view_func=views.create_user)
 app.add_url_rule(rule=routes.GET_USER, methods=["GET"], view_func=views.get_user)
-app.add_url_rule(
-    rule=routes.UPDATE_EXPERIMENT_PERMISSION,
-    methods=["POST"],
-    view_func=views.update_experiment_permission,
-)
-app.add_url_rule(
-    rule=routes.UPDATE_REGISTERED_MODEL_PERMISSION,
-    methods=["POST"],
-    view_func=views.update_model_permission,
-)
+app.add_url_rule(rule=routes.UPDATE_USER_PASSWORD, methods=["GET"], view_func=views.update_username_password)
+app.add_url_rule(rule=routes.UPDATE_USER_ADMIN, methods=["GET"], view_func=views.update_user_admin)
+app.add_url_rule(rule=routes.DELETE_USER, methods=["GET"], view_func=views.delete_user)
+
+# permission management
+app.add_url_rule(rule=routes.CREATE_EXPERIMENT_PERMISSION, methods=["POST"], view_func=views.create_experiment_permission)
+app.add_url_rule(rule=routes.GET_EXPERIMENT_PERMISSION, methods=["GET"], view_func=views.get_experiment_permission)
+app.add_url_rule(rule=routes.UPDATE_EXPERIMENT_PERMISSION, methods=["POST"], view_func=views.update_experiment_permission)
+app.add_url_rule(rule=routes.DELETE_EXPERIMENT_PERMISSION, methods=["POST"], view_func=views.delete_experiment_permission)
+app.add_url_rule(rule=routes.CREATE_REGISTERED_MODEL_PERMISSION, methods=["POST"], view_func=views.create_model_permission)
+app.add_url_rule(rule=routes.GET_REGISTERED_MODEL_PERMISSION, methods=["GET"], view_func=views.get_model_permission)
+app.add_url_rule(rule=routes.UPDATE_REGISTERED_MODEL_PERMISSION, methods=["POST"], view_func=views.update_model_permission)
+app.add_url_rule(rule=routes.DELETE_REGISTERED_MODEL_PERMISSION, methods=["POST"], view_func=views.delete_model_permission)
 
 # Add new hooks
 app.before_request(views.before_request_hook)

mlflow_oidc_auth/client.py

@@ -1,10 +1,10 @@
 import mlflow
-from mlflow.server.auth.entities import (
+from mlflow_oidc_auth.entities import (
     ExperimentPermission,
     RegisteredModelPermission,
     User,
 )
-from mlflow.server.auth.routes import (
+from mlflow_oidc_auth.routes import (
     CREATE_EXPERIMENT_PERMISSION,
     CREATE_REGISTERED_MODEL_PERMISSION,
     CREATE_USER,
@@ -151,9 +151,7 @@ def delete_user(self, username: str):
         except mlflow.exceptions.RestException as e:
             raise e
 
-    def create_experiment_permission(
-        self, experiment_id: str, username: str, permission: str
-    ):
+    def create_experiment_permission(self, experiment_id: str, username: str, permission: str):
         """
         Create a permission on an experiment for a user.
 
@@ -168,9 +166,7 @@ def create_experiment_permission(
         if not username:
             raise ValueError("Username must not be empty.")
         if permission not in ["READ", "EDIT", "MANAGE", "NO_PERMISSIONS"]:
-            raise ValueError(
-                "Permission must be one of 'READ', 'EDIT', 'MANAGE', or 'NO_PERMISSIONS'."
-            )
+            raise ValueError("Permission must be one of 'READ', 'EDIT', 'MANAGE', or 'NO_PERMISSIONS'.")
 
         try:
             resp = self._request(
@@ -212,9 +208,7 @@ def get_experiment_permission(self, experiment_id: str, username: str):
 
         return ExperimentPermission.from_json(resp["experiment_permission"])
 
-    def update_experiment_permission(
-        self, experiment_id: str, username: str, permission: str
-    ):
+    def update_experiment_permission(self, experiment_id: str, username: str, permission: str):
         """
         Update an existing experiment permission for a user.
 
@@ -229,9 +223,7 @@ def update_experiment_permission(
         if not username:
             raise ValueError("Username must not be empty.")
         if permission not in ["READ", "EDIT", "MANAGE", "NO_PERMISSIONS"]:
-            raise ValueError(
-                "Permission must be one of 'READ', 'EDIT', 'MANAGE', or 'NO_PERMISSIONS'."
-            )
+            raise ValueError("Permission must be one of 'READ', 'EDIT', 'MANAGE', or 'NO_PERMISSIONS'.")
 
         try:
             self._request(
@@ -268,9 +260,7 @@ def delete_experiment_permission(self, experiment_id: str, username: str):
         except mlflow.exceptions.RestException as e:
             raise e
 
-    def create_registered_model_permission(
-        self, name: str, username: str, permission: str
-    ):
+    def create_registered_model_permission(self, name: str, username: str, permission: str):
         """
         Create a permission on an registered model for a user.
 
@@ -285,9 +275,7 @@ def create_registered_model_permission(
         if not username:
             raise ValueError("Username must not be empty.")
         if permission not in ["READ", "EDIT", "MANAGE", "NO_PERMISSIONS"]:
-            raise ValueError(
-                "Permission must be one of 'READ', 'EDIT', 'MANAGE', or 'NO_PERMISSIONS'."
-            )
+            raise ValueError("Permission must be one of 'READ', 'EDIT', 'MANAGE', or 'NO_PERMISSIONS'.")
 
         try:
             resp = self._request(
@@ -324,9 +312,7 @@ def get_registered_model_permission(self, name: str, username: str):
 
         return RegisteredModelPermission.from_json(resp["registered_model_permission"])
 
-    def update_registered_model_permission(
-        self, name: str, username: str, permission: str
-    ):
+    def update_registered_model_permission(self, name: str, username: str, permission: str):
         """
         Update an existing registered model permission for a user.
 
@@ -341,9 +327,7 @@ def update_registered_model_permission(
         if not username:
             raise ValueError("Username must not be empty.")
         if permission not in ["READ", "EDIT", "MANAGE", "NO_PERMISSIONS"]:
-            raise ValueError(
-                "Permission must be one of 'READ', 'EDIT', 'MANAGE', or 'NO_PERMISSIONS'."
-            )
+            raise ValueError("Permission must be one of 'READ', 'EDIT', 'MANAGE', or 'NO_PERMISSIONS'.")
 
         try:
             self._request(

mlflow_oidc_auth/db/cli.py

@@ -1,7 +1,7 @@
 import click
 import sqlalchemy
 
-from mlflow.server.auth.db import utils
+from mlflow_oidc_auth.db import utils
 
 
 @click.group(name="db")

mlflow_oidc_auth/db/migrations/env.py

@@ -3,7 +3,7 @@
 from alembic import context
 from sqlalchemy import engine_from_config, pool
 
-from mlflow.server.auth.db.models import Base
+from mlflow_oidc_auth.db.models import Base
 
 # this is the Alembic Config object, which provides
 # access to the values within the .ini file in use.

mlflow_oidc_auth/db/migrations/versions/8606fa83a998_initial_migration.py

@@ -5,6 +5,7 @@
 Create Date: 2023-07-07 23:30:50.921970
 
 """
+
 import sqlalchemy as sa
 from alembic import op
 
@@ -20,6 +21,7 @@ def upgrade() -> None:
         "users",
         sa.Column("id", sa.Integer(), nullable=False, primary_key=True),
         sa.Column("username", sa.String(length=255), nullable=True),
+        sa.Column("display_name", sa.String(length=255), nullable=True),
         sa.Column("password_hash", sa.String(length=255), nullable=True),
         sa.Column("is_admin", sa.Boolean(), nullable=True),
         sa.UniqueConstraint("username"),

mlflow_oidc_auth/db/models.py

@@ -8,7 +8,7 @@
 )
 from sqlalchemy.orm import declarative_base, relationship
 
-from mlflow.server.auth.entities import (
+from mlflow_oidc_auth.entities import (
     ExperimentPermission,
     RegisteredModelPermission,
     User,
@@ -21,25 +21,21 @@ class SqlUser(Base):
     __tablename__ = "users"
     id = Column(Integer(), primary_key=True)
     username = Column(String(255), unique=True)
+    display_name = Column(String(255))
     password_hash = Column(String(255))
     is_admin = Column(Boolean, default=False)
     experiment_permissions = relationship("SqlExperimentPermission", backref="users")
-    registered_model_permissions = relationship(
-        "SqlRegisteredModelPermission", backref="users"
-    )
+    registered_model_permissions = relationship("SqlRegisteredModelPermission", backref="users")
 
     def to_mlflow_entity(self):
         return User(
             id_=self.id,
             username=self.username,
+            display_name=self.display_name,
             password_hash=self.password_hash,
             is_admin=self.is_admin,
-            experiment_permissions=[
-                p.to_mlflow_entity() for p in self.experiment_permissions
-            ],
-            registered_model_permissions=[
-                p.to_mlflow_entity() for p in self.registered_model_permissions
-            ],
+            experiment_permissions=[p.to_mlflow_entity() for p in self.experiment_permissions],
+            registered_model_permissions=[p.to_mlflow_entity() for p in self.registered_model_permissions],
         )
 
 
@@ -49,9 +45,7 @@ class SqlExperimentPermission(Base):
     experiment_id = Column(String(255), nullable=False)
     user_id = Column(Integer, ForeignKey("users.id"), nullable=False)
     permission = Column(String(255))
-    __table_args__ = (
-        UniqueConstraint("experiment_id", "user_id", name="unique_experiment_user"),
-    )
+    __table_args__ = (UniqueConstraint("experiment_id", "user_id", name="unique_experiment_user"),)
 
     def to_mlflow_entity(self):
         return ExperimentPermission(

mlflow_oidc_auth/entities.py

@@ -5,6 +5,7 @@ def __init__(
         username,
         password_hash,
         is_admin,
+        display_name,
         experiment_permissions=None,
         registered_model_permissions=None,
     ):
@@ -14,6 +15,7 @@ def __init__(
         self._is_admin = is_admin
         self._experiment_permissions = experiment_permissions
         self._registered_model_permissions = registered_model_permissions
+        self._display_name = display_name
 
     @property
     def id(self):
@@ -51,33 +53,35 @@ def registered_model_permissions(self):
     def registered_model_permissions(self, registered_model_permissions):
         self._registered_model_permissions = registered_model_permissions
 
+    @property
+    def display_name(self):
+        return self._display_name
+
+    @display_name.setter
+    def display_name(self, display_name):
+        self._display_name = display_name
+
     def to_json(self):
         return {
             "id": self.id,
             "username": self.username,
             "is_admin": self.is_admin,
-            "experiment_permissions": [
-                p.to_json() for p in self.experiment_permissions
-            ],
-            "registered_model_permissions": [
-                p.to_json() for p in self.registered_model_permissions
-            ],
+            "display_name": self.display_name,
+            "experiment_permissions": [p.to_json() for p in self.experiment_permissions],
+            "registered_model_permissions": [p.to_json() for p in self.registered_model_permissions],
         }
 
     @classmethod
     def from_json(cls, dictionary):
         return cls(
             id_=dictionary["id"],
             username=dictionary["username"],
+            display_name=dictionary["display_name"],
             password_hash="REDACTED",
             is_admin=dictionary["is_admin"],
-            experiment_permissions=[
-                ExperimentPermission.from_json(p)
-                for p in dictionary["experiment_permissions"]
-            ],
+            experiment_permissions=[ExperimentPermission.from_json(p) for p in dictionary["experiment_permissions"]],
             registered_model_permissions=[
-                RegisteredModelPermission.from_json(p)
-                for p in dictionary["registered_model_permissions"]
+                RegisteredModelPermission.from_json(p) for p in dictionary["registered_model_permissions"]
             ],
         )