mirror of
https://github.com/danswer-ai/danswer.git
synced 2025-05-03 08:20:40 +02:00
9ea2ae267e
* Initial scaffolding for metrics * iterate * more * More metrics + SyncRecord concept * Add indices, standardize timing * Small cleanup * Address comments
105 lines
3.7 KiB
Python
105 lines
3.7 KiB
Python
from typing import cast
|
|
|
|
from redis import Redis
|
|
from sqlalchemy.orm import Session
|
|
|
|
from ee.onyx.db.user_group import delete_user_group
|
|
from ee.onyx.db.user_group import fetch_user_group
|
|
from ee.onyx.db.user_group import mark_user_group_as_synced
|
|
from ee.onyx.db.user_group import prepare_user_group_for_deletion
|
|
from onyx.background.celery.apps.app_base import task_logger
|
|
from onyx.db.enums import SyncStatus
|
|
from onyx.db.enums import SyncType
|
|
from onyx.db.sync_record import update_sync_record_status
|
|
from onyx.redis.redis_usergroup import RedisUserGroup
|
|
from onyx.utils.logger import setup_logger
|
|
|
|
logger = setup_logger()
|
|
|
|
|
|
def monitor_usergroup_taskset(
|
|
tenant_id: str | None, key_bytes: bytes, r: Redis, db_session: Session
|
|
) -> None:
|
|
"""This function is likely to move in the worker refactor happening next."""
|
|
fence_key = key_bytes.decode("utf-8")
|
|
usergroup_id_str = RedisUserGroup.get_id_from_fence_key(fence_key)
|
|
if not usergroup_id_str:
|
|
task_logger.warning(f"Could not parse usergroup id from {fence_key}")
|
|
return
|
|
|
|
try:
|
|
usergroup_id = int(usergroup_id_str)
|
|
except ValueError:
|
|
task_logger.exception(f"usergroup_id ({usergroup_id_str}) is not an integer!")
|
|
raise
|
|
|
|
rug = RedisUserGroup(tenant_id, usergroup_id)
|
|
if not rug.fenced:
|
|
return
|
|
|
|
initial_count = rug.payload
|
|
if initial_count is None:
|
|
return
|
|
|
|
count = cast(int, r.scard(rug.taskset_key))
|
|
task_logger.info(
|
|
f"User group sync progress: usergroup_id={usergroup_id} remaining={count} initial={initial_count}"
|
|
)
|
|
if count > 0:
|
|
update_sync_record_status(
|
|
db_session=db_session,
|
|
entity_id=usergroup_id,
|
|
sync_type=SyncType.USER_GROUP,
|
|
sync_status=SyncStatus.IN_PROGRESS,
|
|
num_docs_synced=count,
|
|
)
|
|
return
|
|
|
|
user_group = fetch_user_group(db_session=db_session, user_group_id=usergroup_id)
|
|
if user_group:
|
|
usergroup_name = user_group.name
|
|
try:
|
|
if user_group.is_up_for_deletion:
|
|
# this prepare should have been run when the deletion was scheduled,
|
|
# but run it again to be sure we're ready to go
|
|
mark_user_group_as_synced(db_session, user_group)
|
|
prepare_user_group_for_deletion(db_session, usergroup_id)
|
|
delete_user_group(db_session=db_session, user_group=user_group)
|
|
|
|
update_sync_record_status(
|
|
db_session=db_session,
|
|
entity_id=usergroup_id,
|
|
sync_type=SyncType.USER_GROUP,
|
|
sync_status=SyncStatus.SUCCESS,
|
|
num_docs_synced=initial_count,
|
|
)
|
|
|
|
task_logger.info(
|
|
f"Deleted usergroup: name={usergroup_name} id={usergroup_id}"
|
|
)
|
|
else:
|
|
mark_user_group_as_synced(db_session=db_session, user_group=user_group)
|
|
|
|
update_sync_record_status(
|
|
db_session=db_session,
|
|
entity_id=usergroup_id,
|
|
sync_type=SyncType.USER_GROUP,
|
|
sync_status=SyncStatus.SUCCESS,
|
|
num_docs_synced=initial_count,
|
|
)
|
|
|
|
task_logger.info(
|
|
f"Synced usergroup. name={usergroup_name} id={usergroup_id}"
|
|
)
|
|
except Exception as e:
|
|
update_sync_record_status(
|
|
db_session=db_session,
|
|
entity_id=usergroup_id,
|
|
sync_type=SyncType.USER_GROUP,
|
|
sync_status=SyncStatus.FAILED,
|
|
num_docs_synced=initial_count,
|
|
)
|
|
raise e
|
|
|
|
rug.reset()
|