This commit is contained in:
pablodanswer 2025-02-01 16:02:49 -08:00
parent 125e5eaab1
commit 42a0f45a96
9 changed files with 17 additions and 3 deletions

View File

@ -286,6 +286,7 @@ def prepare_authorization_request(
oauth_state = (
base64.urlsafe_b64encode(oauth_uuid.bytes).rstrip(b"=").decode("utf-8")
)
session: str
if connector == DocumentSource.SLACK:
oauth_url = SlackOAuth.generate_oauth_url(oauth_state)
@ -554,6 +555,7 @@ def handle_google_drive_oauth_callback(
)
session_json = session_json_bytes.decode("utf-8")
session: GoogleDriveOAuth.OAuthSession
try:
session = GoogleDriveOAuth.parse_session(session_json)

View File

@ -1047,6 +1047,8 @@ async def api_key_dep(
if AUTH_TYPE == AuthType.DISABLED:
return None
user: User | None = None
hashed_api_key = get_hashed_api_key_from_request(request)
if not hashed_api_key:
raise HTTPException(status_code=401, detail="Missing API key")

View File

@ -747,7 +747,7 @@ def cloud_check_alembic() -> bool | None:
revision_counts: dict[str, int] = {}
out_of_date_tenants: dict[str, str | None] = {}
top_revision: str = ""
tenant_ids: list[str] = []
tenant_ids: list[str] | list[None] = []
try:
# map each tenant_id to its revision

View File

@ -168,6 +168,7 @@ def document_by_cc_pair_cleanup_task(
task_logger.info(f"SoftTimeLimitExceeded exception. doc={document_id}")
return False
except Exception as ex:
e: Exception | None = None
if isinstance(ex, RetryError):
task_logger.warning(
f"Tenacity retry failed: num_attempts={ex.last_attempt.attempt_number}"

View File

@ -256,6 +256,8 @@ def _run_indexing(
document_count = 0
chunk_count = 0
run_end_dt = None
tracer_counter: int
for ind, (window_start, window_end) in enumerate(
get_time_windows_for_index_attempt(
last_successful_run=datetime.fromtimestamp(
@ -266,6 +268,7 @@ def _run_indexing(
):
cc_pair_loop: ConnectorCredentialPair | None = None
index_attempt_loop: IndexAttempt | None = None
tracer_counter = 0
try:
window_start = max(
@ -290,7 +293,6 @@ def _run_indexing(
tenant_id=tenant_id,
)
tracer_counter = 0
if INDEXING_TRACER_INTERVAL > 0:
tracer.snap()
for doc_batch in connector_runner.run():

View File

@ -26,6 +26,7 @@ def _get_google_service(
creds: ServiceAccountCredentials | OAuthCredentials,
user_email: str | None = None,
) -> GoogleDriveService | GoogleDocsService | AdminService | GmailService:
service: Resource
if isinstance(creds, ServiceAccountCredentials):
creds = creds.with_subject(user_email)
service = build(service_name, service_version, credentials=creds)

View File

@ -222,6 +222,7 @@ def insert_document_set(
)
new_document_set_row: DocumentSetDBModel
ds_cc_pairs: list[DocumentSet__ConnectorCredentialPair]
try:
new_document_set_row = DocumentSetDBModel(
name=document_set_creation_request.name,

View File

@ -5,6 +5,7 @@ import sys
import threading
import time
from collections.abc import Callable
from contextvars import Token
from threading import Event
from types import FrameType
from typing import Any
@ -250,6 +251,8 @@ class SlackbotHandler:
"""
all_tenants = get_all_tenant_ids()
token: Token[str]
# 1) Try to acquire locks for new tenants
for tenant_id in all_tenants:
if (
@ -771,6 +774,7 @@ def process_message(
client=client.web_client, channel_id=channel
)
token: Token[str] | None = None
# Set the current tenant ID at the beginning for all DB calls within this thread
if client.tenant_id:
logger.info(f"Setting tenant ID to {client.tenant_id}")
@ -825,7 +829,7 @@ def process_message(
if notify_no_answer:
apologize_for_fail(details, client)
finally:
if client.tenant_id:
if token:
CURRENT_TENANT_ID_CONTEXTVAR.reset(token)

View File

@ -54,6 +54,7 @@ def google_drive_test_env_setup() -> (
service_account_key = os.environ["FULL_CONTROL_DRIVE_SERVICE_ACCOUNT"]
drive_id: str | None = None
drive_service: GoogleDriveService | None = None
try:
credentials = {