mirror of
https://github.com/onyx-dot-app/onyx.git
synced 2026-10-22 11:22:49 -05:00
66 lines
2.5 KiB
Python
66 lines
2.5 KiB
Python
from threading import Thread
|
|
|
|
from sqlalchemy import select
|
|
from sqlalchemy.orm import Session
|
|
|
|
from danswer.access.models import DocumentAccess
|
|
from danswer.db.document import get_acccess_info_for_documents
|
|
from danswer.db.engine import get_sqlalchemy_engine
|
|
from danswer.db.models import Document
|
|
from danswer.document_index.factory import get_default_document_index
|
|
from danswer.document_index.interfaces import UpdateRequest
|
|
from danswer.document_index.vespa.index import VespaIndex
|
|
from danswer.dynamic_configs import get_dynamic_config_store
|
|
from danswer.dynamic_configs.interface import ConfigNotFoundError
|
|
from danswer.utils.logger import setup_logger
|
|
|
|
logger = setup_logger()
|
|
|
|
|
|
_COMPLETED_ACL_UPDATE_KEY = "completed_acl_update"
|
|
|
|
|
|
def set_acl_for_vespa(should_check_if_already_done: bool = False) -> None:
|
|
"""Updates the ACL for all documents based on the state of Postgres."""
|
|
dynamic_config_store = get_dynamic_config_store()
|
|
if should_check_if_already_done:
|
|
try:
|
|
# if entry is found, then we've already done this
|
|
dynamic_config_store.load(_COMPLETED_ACL_UPDATE_KEY)
|
|
return
|
|
except ConfigNotFoundError:
|
|
pass
|
|
|
|
vespa_index = get_default_document_index()
|
|
if not isinstance(vespa_index, VespaIndex):
|
|
raise ValueError("This script is only for Vespa indexes")
|
|
|
|
logger.info("Populating Access Control List fields in Vespa")
|
|
with Session(get_sqlalchemy_engine()) as db_session:
|
|
# for all documents, set the `access_control_list` field apporpriately
|
|
# based on the state of Postgres
|
|
documents = db_session.scalars(select(Document)).all()
|
|
document_access_info = get_acccess_info_for_documents(
|
|
db_session=db_session,
|
|
document_ids=[document.id for document in documents],
|
|
)
|
|
vespa_index.update(
|
|
update_requests=[
|
|
UpdateRequest(
|
|
document_ids=[document_id],
|
|
access=DocumentAccess.build(user_ids, is_public),
|
|
)
|
|
for document_id, user_ids, is_public in document_access_info
|
|
],
|
|
)
|
|
|
|
dynamic_config_store.store(_COMPLETED_ACL_UPDATE_KEY, True)
|
|
|
|
|
|
def set_acl_for_vespa_nonblocking(should_check_if_already_done: bool = False) -> None:
|
|
"""Kick off the ACL update in a separate thread so that other work can continue."""
|
|
Thread(
|
|
target=set_acl_for_vespa,
|
|
args=[should_check_if_already_done],
|
|
).start()
|