Skip to content

Commit

Permalink
update
Browse files Browse the repository at this point in the history
  • Loading branch information
pablonyx committed Feb 13, 2025
1 parent 895b032 commit 5a709b5
Show file tree
Hide file tree
Showing 2 changed files with 31 additions and 20 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,6 @@
from redis import Redis
from redis.exceptions import LockError
from redis.lock import Lock as RedisLock
from sqlalchemy.exc import IntegrityError
from sqlalchemy.orm import Session

from ee.onyx.db.connector_credential_pair import get_all_auto_sync_cc_pairs
Expand Down Expand Up @@ -463,20 +462,11 @@ def update_external_document_permissions_task(

try:
with get_session_with_tenant(tenant_id) as db_session:
try:
# Add the users to the DB if they don't exist
batch_add_ext_perm_user_if_not_exists(
db_session=db_session,
emails=list(external_access.external_user_emails),
)
except IntegrityError:
db_session.rollback()
task_logger.error(
"IntegrityError (likely 'UniqueViolation') occurred while adding users. "
f"Emails: {list(external_access.external_user_emails)}. Retrying once."
)
return False

batch_add_ext_perm_user_if_not_exists(
db_session=db_session,
emails=list(external_access.external_user_emails),
continue_on_error=True,
)
# Then upsert the document's external permissions
created_new_doc = upsert_document_external_perms(
db_session=db_session,
Expand Down
31 changes: 26 additions & 5 deletions backend/onyx/db/users.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
from fastapi_users.password import PasswordHelper
from sqlalchemy import func
from sqlalchemy import select
from sqlalchemy.exc import IntegrityError
from sqlalchemy.orm import Session
from sqlalchemy.sql import expression
from sqlalchemy.sql.elements import ColumnElement
Expand All @@ -20,8 +21,11 @@
from onyx.db.models import SamlAccount
from onyx.db.models import User
from onyx.db.models import User__UserGroup
from onyx.setup import setup_logger
from onyx.utils.variable_functionality import fetch_ee_implementation_or_noop

logger = setup_logger()


def validate_user_role_update(requested_role: UserRole, current_role: UserRole) -> None:
"""
Expand Down Expand Up @@ -274,7 +278,7 @@ def _generate_ext_permissioned_user(email: str) -> User:


def batch_add_ext_perm_user_if_not_exists(
db_session: Session, emails: list[str]
db_session: Session, emails: list[str], continue_on_error: bool = False
) -> list[User]:
lower_emails = [email.lower() for email in emails]
found_users, missing_lower_emails = _get_users_by_emails(db_session, lower_emails)
Expand All @@ -283,10 +287,27 @@ def batch_add_ext_perm_user_if_not_exists(
for email in missing_lower_emails:
new_users.append(_generate_ext_permissioned_user(email=email))

db_session.add_all(new_users)
db_session.commit()

return found_users + new_users
try:
db_session.add_all(new_users)
db_session.commit()
except IntegrityError:
db_session.rollback()
logger.warning(
"IntegrityError occurred during batch insert. Falling back to individual inserts."
)
if not continue_on_error:
raise
for user in new_users:
try:
db_session.add(user)
db_session.commit()
except IntegrityError:
db_session.rollback()
logger.warning(f"Skipping duplicate user: {user.email}")
continue
# Fetch all users again to ensure we have the most up-to-date list
all_users, _ = _get_users_by_emails(db_session, lower_emails)
return all_users


def delete_user_from_db(
Expand Down

0 comments on commit 5a709b5

Please sign in to comment.