fix tests

Arize-ai · Dec 24, 2024 · e522957 · e522957
1 parent 74eba58
commit e522957
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/src/phoenix/server/api/mutations/prompt_mutations.py b/src/phoenix/server/api/mutations/prompt_mutations.py
@@ -2,7 +2,8 @@
 
 import strawberry
 from pydantic import ValidationError
-from sqlean.dbapi2 import IntegrityError  # type: ignore[import-untyped]
+from sqlalchemy.exc import IntegrityError as PostgresIntegrityError
+from sqlean.dbapi2 import IntegrityError as SQLiteIntegrityError  # type: ignore[import-untyped]
 from strawberry.types import Info
 
 from phoenix.db import models
@@ -47,6 +48,6 @@ async def create_prompt(self, info: Info[Context, None], input: CreatePromptInpu
             session.add(prompt)
             try:
                 await session.commit()
-            except IntegrityError:
+            except (PostgresIntegrityError, SQLiteIntegrityError):
                 raise Conflict(f"A prompt named '{input.name}' already exists")
         return to_gql_prompt(prompt)