chore(auth): strengthen check constraints related to auth method (#4704)

Arize-ai · Sep 21, 2024 · afe991c · afe991c
1 parent 87227b7
commit afe991c
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 2 deletions.
diff --git a/src/phoenix/db/migrations/versions/cd164e83824f_users_and_tokens.py b/src/phoenix/db/migrations/versions/cd164e83824f_users_and_tokens.py
@@ -66,7 +66,18 @@ def upgrade() -> None:
             sa.TIMESTAMP(timezone=True),
             nullable=True,
         ),
-        sa.CheckConstraint("password_hash is null or password_salt is not null", name="salt"),
+        sa.CheckConstraint(
+            "(password_hash IS NULL) = (password_salt IS NULL)",
+            name="password_hash_and_salt",
+        ),
+        sa.CheckConstraint(
+            "(oauth2_client_id IS NULL) = (oauth2_user_id IS NULL)",
+            name="oauth2_client_id_and_user_id",
+        ),
+        sa.CheckConstraint(
+            "password_hash IS NULL or oauth2_client_id IS NULL",
+            name="at_most_one_auth_method",
+        ),
         sa.UniqueConstraint(
             "oauth2_client_id",
             "oauth2_user_id",

diff --git a/src/phoenix/db/models.py b/src/phoenix/db/models.py
@@ -690,7 +690,18 @@ def _auth_method_expression(cls) -> ColumnElement[Optional[str]]:
         )
 
     __table_args__ = (
-        CheckConstraint("password_hash is null or password_salt is not null", name="salt"),
+        CheckConstraint(
+            "(password_hash IS NULL) = (password_salt IS NULL)",
+            name="password_hash_and_salt",
+        ),
+        CheckConstraint(
+            "(oauth2_client_id IS NULL) = (oauth2_user_id IS NULL)",
+            name="oauth2_client_id_and_user_id",
+        ),
+        CheckConstraint(
+            "password_hash IS NULL or oauth2_client_id IS NULL",
+            name="at_most_one_auth_method",
+        ),
         UniqueConstraint(
             "oauth2_client_id",
             "oauth2_user_id",