From c2334a179f8c72fc554fce79f7f5b38e9ebcfa22 Mon Sep 17 00:00:00 2001
From: Andreas Motl <andreas.motl@crate.io>
Date: Mon, 6 Nov 2023 08:18:20 +0100
Subject: [PATCH] Fix two more uniqueness constraints

- SqlExperimentPermission: "experiment_id", "user_id"
- SqlRegisteredModelPermission: "name", "user_id"
---
 CHANGES.md                           |  5 ++++-
 mlflow_cratedb/patch/mlflow/model.py | 14 +++++++++++++-
 pyproject.toml                       |  4 ++--
 3 files changed, 19 insertions(+), 4 deletions(-)

diff --git a/CHANGES.md b/CHANGES.md
index 430697d..a3b3e92 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -14,7 +14,10 @@
   Note that it is always advised to create backups of your database content.
   This is an excellent opportunity to do that.
 
-- Fix uniqueness constraint with `mlflow.server.auth.db.models.SqlUser.username`.
+- Fix uniqueness constraints
+  - `mlflow.server.auth.db.models.SqlUser.username`.
+  - `m.s.a.d.m.SqlExperimentPermission`: "experiment_id", "user_id"
+  - `m.s.a.d.m.SqlRegisteredModelPermission`: "name", "user_id"
 
 ## 2023-11-01 v2.7.1
 - Fix uniqueness constraint with `SqlRegisteredModel.name`. Thanks, @andnig.
diff --git a/mlflow_cratedb/patch/mlflow/model.py b/mlflow_cratedb/patch/mlflow/model.py
index 052b197..b708e7f 100644
--- a/mlflow_cratedb/patch/mlflow/model.py
+++ b/mlflow_cratedb/patch/mlflow/model.py
@@ -14,12 +14,24 @@ def polyfill_uniqueness_constraints():
           - SqlExperimentPermission: "experiment_id", "user_id"
           - SqlRegisteredModelPermission: "name", "user_id"
     """
-    from mlflow.server.auth.db.models import SqlUser
+    from mlflow.server.auth.db.models import SqlExperimentPermission, SqlRegisteredModelPermission, SqlUser
     from mlflow.store.model_registry.dbmodels.models import SqlRegisteredModel
     from mlflow.store.tracking.dbmodels.models import SqlExperiment
 
     listen(SqlExperiment, "before_insert", check_uniqueness_factory(SqlExperiment, "name"))
+    listen(
+        SqlExperimentPermission,
+        "before_insert",
+        check_uniqueness_factory(SqlExperimentPermission, "experiment_id", "user_id"),
+    )
+
     listen(SqlRegisteredModel, "before_insert", check_uniqueness_factory(SqlRegisteredModel, "name"))
+    listen(
+        SqlRegisteredModelPermission,
+        "before_insert",
+        check_uniqueness_factory(SqlRegisteredModelPermission, "name", "user_id"),
+    )
+
     listen(SqlUser, "before_insert", check_uniqueness_factory(SqlUser, "username"))
 
 
diff --git a/pyproject.toml b/pyproject.toml
index f117738..c42be50 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -78,7 +78,7 @@ dynamic = [
 dependencies = [
   "crash",
   "crate[sqlalchemy]>=0.34",
-  "cratedb-toolkit==0.0.1",
+  "cratedb-toolkit==0.0.2",
   "mlflow==2.8",
   "sqlparse<0.5",
 ]
@@ -93,8 +93,8 @@ develop = [
   "validate-pyproject<0.16",
 ]
 examples = [
+  "pycaret[analysis,models,parallel,test,tuner]==3.1",
   "salesforce-merlion<2.1",
-  "pycaret[analysis,models,tuner,parallel,test]==3.1.0",
   "werkzeug==2.2.3",
 ]
 release = [