Only run set_acl_for_vespa once

This commit is contained in:
Weves 2023-09-30 14:11:22 -07:00 committed by Chris Weaver
parent 22f2398269
commit aa9071e441
2 changed files with 19 additions and 2 deletions

View File

@ -216,7 +216,8 @@ def get_application() -> FastAPI:
# TODO: remove this once everyone is migrated to ACL
logger.info("Populating Access Control List fields in Vespa")
set_acl_for_vespa()
# does nothing if this has been successfully run before
set_acl_for_vespa(should_check_if_already_done=True)
return application

View File

@ -8,13 +8,27 @@ from danswer.datastores.vespa.store import VespaIndex
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.dynamic_configs import get_dynamic_config_store
from danswer.dynamic_configs.interface import ConfigNotFoundError
from danswer.utils.logger import setup_logger
logger = setup_logger()
def set_acl_for_vespa() -> None:
_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")
@ -36,3 +50,5 @@ def set_acl_for_vespa() -> None:
for document_id, user_ids, is_public in document_access_info
],
)
dynamic_config_store.store(_COMPLETED_ACL_UPDATE_KEY, True)