From 42a0f45a96d3d917478937093c6c9ff6a29a1651 Mon Sep 17 00:00:00 2001 From: pablodanswer Date: Sat, 1 Feb 2025 16:02:49 -0800 Subject: [PATCH] update --- backend/ee/onyx/server/oauth.py | 2 ++ backend/onyx/auth/users.py | 2 ++ backend/onyx/background/celery/tasks/monitoring/tasks.py | 2 +- backend/onyx/background/celery/tasks/shared/tasks.py | 1 + backend/onyx/background/indexing/run_indexing.py | 4 +++- backend/onyx/connectors/google_utils/resources.py | 1 + backend/onyx/db/document_set.py | 1 + backend/onyx/onyxbot/slack/listener.py | 6 +++++- .../google/test_google_drive_permission_sync.py | 1 + 9 files changed, 17 insertions(+), 3 deletions(-) diff --git a/backend/ee/onyx/server/oauth.py b/backend/ee/onyx/server/oauth.py index 4b79750f0..de239f65a 100644 --- a/backend/ee/onyx/server/oauth.py +++ b/backend/ee/onyx/server/oauth.py @@ -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) diff --git a/backend/onyx/auth/users.py b/backend/onyx/auth/users.py index da9865550..62c3e9e80 100644 --- a/backend/onyx/auth/users.py +++ b/backend/onyx/auth/users.py @@ -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") diff --git a/backend/onyx/background/celery/tasks/monitoring/tasks.py b/backend/onyx/background/celery/tasks/monitoring/tasks.py index 244b4abd2..fe25c1d80 100644 --- a/backend/onyx/background/celery/tasks/monitoring/tasks.py +++ b/backend/onyx/background/celery/tasks/monitoring/tasks.py @@ -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 diff --git a/backend/onyx/background/celery/tasks/shared/tasks.py b/backend/onyx/background/celery/tasks/shared/tasks.py index a1bc12252..84f7e12c6 100644 --- a/backend/onyx/background/celery/tasks/shared/tasks.py +++ b/backend/onyx/background/celery/tasks/shared/tasks.py @@ -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}" diff --git a/backend/onyx/background/indexing/run_indexing.py b/backend/onyx/background/indexing/run_indexing.py index fd0ab0b35..72fd8e4fe 100644 --- a/backend/onyx/background/indexing/run_indexing.py +++ b/backend/onyx/background/indexing/run_indexing.py @@ -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(): diff --git a/backend/onyx/connectors/google_utils/resources.py b/backend/onyx/connectors/google_utils/resources.py index c93c58097..48bd981c2 100644 --- a/backend/onyx/connectors/google_utils/resources.py +++ b/backend/onyx/connectors/google_utils/resources.py @@ -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) diff --git a/backend/onyx/db/document_set.py b/backend/onyx/db/document_set.py index bf9d6d38d..0a5397eb8 100644 --- a/backend/onyx/db/document_set.py +++ b/backend/onyx/db/document_set.py @@ -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, diff --git a/backend/onyx/onyxbot/slack/listener.py b/backend/onyx/onyxbot/slack/listener.py index e8b26f006..ae6274adb 100644 --- a/backend/onyx/onyxbot/slack/listener.py +++ b/backend/onyx/onyxbot/slack/listener.py @@ -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) diff --git a/backend/tests/integration/connector_job_tests/google/test_google_drive_permission_sync.py b/backend/tests/integration/connector_job_tests/google/test_google_drive_permission_sync.py index 58f58a6c5..30283097c 100644 --- a/backend/tests/integration/connector_job_tests/google/test_google_drive_permission_sync.py +++ b/backend/tests/integration/connector_job_tests/google/test_google_drive_permission_sync.py @@ -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 = {