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, group_ids: list[int] | None,
db_session: Session, db_session: Session,
) -> None: ) -> 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( db_session.query(Persona__User).filter(
Persona__User.persona_id == persona_id Persona__User.persona_id == persona_id
).delete(synchronize_session="fetch") ).delete(synchronize_session="fetch")
@@ -23,19 +26,22 @@ def make_persona_private(
).delete(synchronize_session="fetch") ).delete(synchronize_session="fetch")
if user_ids: if user_ids:
for user_uuid in user_ids: user_ids_set = set(user_ids)
db_session.add(Persona__User(persona_id=persona_id, user_id=user_uuid)) for user_id in user_ids_set:
db_session.add(Persona__User(persona_id=persona_id, user_id=user_id))
create_notification( create_notification(
user_id=user_uuid, user_id=user_id,
notif_type=NotificationType.PERSONA_SHARED, notif_type=NotificationType.PERSONA_SHARED,
db_session=db_session, db_session=db_session,
additional_data=PersonaSharedNotificationData( additional_data=PersonaSharedNotificationData(
persona_id=persona_id, persona_id=persona_id,
).model_dump(), ).model_dump(),
) )
if group_ids: 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( db_session.add(
Persona__UserGroup(persona_id=persona_id, user_group_id=group_id) Persona__UserGroup(persona_id=persona_id, user_group_id=group_id)
) )

View File

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

View File

@@ -58,6 +58,7 @@ def test_persona_permissions(reset: None) -> None:
description="A persona created by basic user", description="A persona created by basic user",
is_public=False, is_public=False,
groups=[], groups=[],
users=[admin_user.id],
user_performing_action=basic_user, user_performing_action=basic_user,
) )
PersonaManager.verify(basic_user_persona, 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""" """Test admin permissions"""
# Admin can edit any persona # 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( PersonaManager.edit(
persona=basic_user_persona, 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, user_performing_action=admin_user,
) )
PersonaManager.verify(basic_user_persona, user_performing_action=admin_user) PersonaManager.verify(basic_user_persona, user_performing_action=admin_user)