Bugfix/dedupe ids (#3952)

* dedupe make_private_persona and update test

* add comment

* comments, and just have duplicate user id's for the test instead of modifying edit

* found the magic word

---------

Co-authored-by: Richard Kuo (Danswer) <rkuo@onyx.app>
This commit is contained in:
rkuo-danswer
2025-02-10 18:27:55 -08:00
committed by GitHub
parent da5c83a96d
commit b56877cc2e
3 changed files with 20 additions and 7 deletions

View File

@@ -15,6 +15,9 @@ def make_persona_private(
group_ids: list[int] | None,
db_session: Session,
) -> None:
"""NOTE(rkuo): This function batches all updates into a single commit. If we don't
dedupe the inputs, the commit will exception."""
db_session.query(Persona__User).filter(
Persona__User.persona_id == persona_id
).delete(synchronize_session="fetch")
@@ -23,19 +26,22 @@ def make_persona_private(
).delete(synchronize_session="fetch")
if user_ids:
for user_uuid in user_ids:
db_session.add(Persona__User(persona_id=persona_id, user_id=user_uuid))
user_ids_set = set(user_ids)
for user_id in user_ids_set:
db_session.add(Persona__User(persona_id=persona_id, user_id=user_id))
create_notification(
user_id=user_uuid,
user_id=user_id,
notif_type=NotificationType.PERSONA_SHARED,
db_session=db_session,
additional_data=PersonaSharedNotificationData(
persona_id=persona_id,
).model_dump(),
)
if group_ids:
for group_id in group_ids:
group_ids_set = set(group_ids)
for group_id in group_ids_set:
db_session.add(
Persona__UserGroup(persona_id=persona_id, user_group_id=group_id)
)

View File

@@ -66,7 +66,7 @@ class PersonaManager:
response = requests.post(
f"{API_SERVER_URL}/persona",
json=persona_creation_request.model_dump(),
json=persona_creation_request.model_dump(mode="json"),
headers=user_performing_action.headers
if user_performing_action
else GENERAL_HEADERS,
@@ -119,6 +119,7 @@ class PersonaManager:
) -> DATestPersona:
system_prompt = system_prompt or f"System prompt for {persona.name}"
task_prompt = task_prompt or f"Task prompt for {persona.name}"
persona_update_request = PersonaUpsertRequest(
name=name or persona.name,
description=description or persona.description,
@@ -146,7 +147,7 @@ class PersonaManager:
response = requests.patch(
f"{API_SERVER_URL}/persona/{persona.id}",
json=persona_update_request.model_dump(),
json=persona_update_request.model_dump(mode="json"),
headers=user_performing_action.headers
if user_performing_action
else GENERAL_HEADERS,

View File

@@ -58,6 +58,7 @@ def test_persona_permissions(reset: None) -> None:
description="A persona created by basic user",
is_public=False,
groups=[],
users=[admin_user.id],
user_performing_action=basic_user,
)
PersonaManager.verify(basic_user_persona, user_performing_action=basic_user)
@@ -139,9 +140,14 @@ def test_persona_permissions(reset: None) -> None:
"""Test admin permissions"""
# Admin can edit any persona
# the persona was shared with the admin user on creation
# this edit call will simulate having the same user in the list twice.
# The server side should dedupe and handle this correctly (prior bug)
PersonaManager.edit(
persona=basic_user_persona,
description="Updated by admin",
description="Updated by admin 2",
users=[admin_user.id, admin_user.id],
user_performing_action=admin_user,
)
PersonaManager.verify(basic_user_persona, user_performing_action=admin_user)