mirror of
https://github.com/danswer-ai/danswer.git
synced 2025-05-24 02:30:15 +02:00
b9c29f2a36
* Fix pagination for index attempts table * fixed index attempts pagination * fixed query history table * query clearnup * fixed test * fixed weird tests???
107 lines
4.0 KiB
Python
107 lines
4.0 KiB
Python
from datetime import datetime
|
|
from datetime import timedelta
|
|
from urllib.parse import urlencode
|
|
|
|
import requests
|
|
|
|
from onyx.db.engine import get_session_context_manager
|
|
from onyx.db.enums import IndexModelStatus
|
|
from onyx.db.models import IndexAttempt
|
|
from onyx.db.models import IndexingStatus
|
|
from onyx.db.search_settings import get_current_search_settings
|
|
from onyx.server.documents.models import IndexAttemptSnapshot
|
|
from onyx.server.documents.models import PaginatedReturn
|
|
from tests.integration.common_utils.constants import API_SERVER_URL
|
|
from tests.integration.common_utils.constants import GENERAL_HEADERS
|
|
from tests.integration.common_utils.test_models import DATestIndexAttempt
|
|
from tests.integration.common_utils.test_models import DATestUser
|
|
|
|
|
|
class IndexAttemptManager:
|
|
@staticmethod
|
|
def create_test_index_attempts(
|
|
num_attempts: int,
|
|
cc_pair_id: int,
|
|
from_beginning: bool = False,
|
|
status: IndexingStatus = IndexingStatus.SUCCESS,
|
|
new_docs_indexed: int = 10,
|
|
total_docs_indexed: int = 10,
|
|
docs_removed_from_index: int = 0,
|
|
error_msg: str | None = None,
|
|
base_time: datetime | None = None,
|
|
) -> list[DATestIndexAttempt]:
|
|
if base_time is None:
|
|
base_time = datetime.now()
|
|
|
|
attempts = []
|
|
with get_session_context_manager() as db_session:
|
|
# Get the current search settings
|
|
search_settings = get_current_search_settings(db_session)
|
|
if (
|
|
not search_settings
|
|
or search_settings.status != IndexModelStatus.PRESENT
|
|
):
|
|
raise ValueError("No current search settings found with PRESENT status")
|
|
|
|
for i in range(num_attempts):
|
|
time_created = base_time - timedelta(hours=i)
|
|
|
|
index_attempt = IndexAttempt(
|
|
connector_credential_pair_id=cc_pair_id,
|
|
from_beginning=from_beginning,
|
|
status=status,
|
|
new_docs_indexed=new_docs_indexed,
|
|
total_docs_indexed=total_docs_indexed,
|
|
docs_removed_from_index=docs_removed_from_index,
|
|
error_msg=error_msg,
|
|
time_created=time_created,
|
|
time_started=time_created,
|
|
time_updated=time_created,
|
|
search_settings_id=search_settings.id,
|
|
)
|
|
|
|
db_session.add(index_attempt)
|
|
db_session.flush() # To get the ID
|
|
|
|
attempts.append(
|
|
DATestIndexAttempt(
|
|
id=index_attempt.id,
|
|
status=index_attempt.status,
|
|
new_docs_indexed=index_attempt.new_docs_indexed,
|
|
total_docs_indexed=index_attempt.total_docs_indexed,
|
|
docs_removed_from_index=index_attempt.docs_removed_from_index,
|
|
error_msg=index_attempt.error_msg,
|
|
time_started=index_attempt.time_started,
|
|
time_updated=index_attempt.time_updated,
|
|
)
|
|
)
|
|
|
|
db_session.commit()
|
|
|
|
return attempts
|
|
|
|
@staticmethod
|
|
def get_index_attempt_page(
|
|
cc_pair_id: int,
|
|
page: int = 0,
|
|
page_size: int = 10,
|
|
user_performing_action: DATestUser | None = None,
|
|
) -> PaginatedReturn[IndexAttemptSnapshot]:
|
|
query_params: dict[str, str | int] = {
|
|
"page_num": page,
|
|
"page_size": page_size,
|
|
}
|
|
|
|
response = requests.get(
|
|
url=f"{API_SERVER_URL}/manage/admin/cc-pair/{cc_pair_id}/index-attempts?{urlencode(query_params, doseq=True)}",
|
|
headers=user_performing_action.headers
|
|
if user_performing_action
|
|
else GENERAL_HEADERS,
|
|
)
|
|
response.raise_for_status()
|
|
data = response.json()
|
|
return PaginatedReturn(
|
|
items=[IndexAttemptSnapshot(**item) for item in data["items"]],
|
|
total_items=data["total_items"],
|
|
)
|