From b56877cc2ebe7f4cd86dfea1470ed60faf7d341f Mon Sep 17 00:00:00 2001 From: rkuo-danswer Date: Mon, 10 Feb 2025 18:27:55 -0800 Subject: [PATCH] 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) --- backend/ee/onyx/db/persona.py | 14 ++++++++++---- .../integration/common_utils/managers/persona.py | 5 +++-- .../tests/permissions/test_persona_permissions.py | 8 +++++++- 3 files changed, 20 insertions(+), 7 deletions(-) diff --git a/backend/ee/onyx/db/persona.py b/backend/ee/onyx/db/persona.py index b5773588b7a..9fab07a58ea 100644 --- a/backend/ee/onyx/db/persona.py +++ b/backend/ee/onyx/db/persona.py @@ -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) ) diff --git a/backend/tests/integration/common_utils/managers/persona.py b/backend/tests/integration/common_utils/managers/persona.py index e662427d736..3f85456b930 100644 --- a/backend/tests/integration/common_utils/managers/persona.py +++ b/backend/tests/integration/common_utils/managers/persona.py @@ -66,7 +66,7 @@ def create( 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 @@ def edit( ) -> 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 @@ def edit( 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, diff --git a/backend/tests/integration/tests/permissions/test_persona_permissions.py b/backend/tests/integration/tests/permissions/test_persona_permissions.py index 5f63e5fc21b..baea7b456bd 100644 --- a/backend/tests/integration/tests/permissions/test_persona_permissions.py +++ b/backend/tests/integration/tests/permissions/test_persona_permissions.py @@ -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)