PR #17542 - onyx-dot-app/onyx

Description

The Google Drive connector now fetches about 2.2 times as many documents a minute against a real Google Workspace. This PR also adds a mock Google Drive server and a benchmark, so that throughput can be measured and kept without a Workspace.

Results

Live test on the onyx.app Workspace. Both runs used the same Onyx instance, service-account credential and settings: the entire Drive (45 users, 16 shared drives), permission sync on, no docprocessing worker (so the runs measure docfetching only), and a 90-minute window from the start of the index attempt.

Run	Code	Documents in 90 min	Docs/min	Rate limits	Connector errors
Before	`fcabb73a85`	35,195	391.1	0	0
After	this PR	77,508	861.2	0	0

The after run passed the before run's 90-minute total at 36 minutes. In the last 140 seconds of the after window, the host ran out of memory and docfetching stored only one batch (load average 23, 55 s of clock drift between Celery workers); up to that stall the after rate was 884 docs/min. The table keeps the full window. In the before run, the time inside the connector was about 5,831 s and storing batches about 24 s, so the fetch from Google sets the rate.

Permissions. Permission sync ran in both runs. Each of the 38,795 documents that the before run stored is also in the after run, with the same user emails, group ids and public flag. No document in either run has missing or empty permissions. The after run skipped one Google Doc: files.export returned 403 insufficientFilePermissions for its retriever, but Drive says that user can download it. A plain API client outside the connector gets the same 403, so this change did not cause it. The before run's documents do not include that file.

Hill climb against the mock Drive server. 2,700 documents (8 users, 3 shared drives), realistic latency (100 ms metadata, 250 ms listings and Docs API, 300 ms downloads and exports), Google's units quota (1,000,000 units a minute per project, 325,000 per user).

Step	Change	Workers	Docs/min	Share of quota	Rate-limited calls
00	Baseline	8	568	11.5%	0
02	Token reuse, full pages, rolling window, folder prefetch	8	777	15.7%	0
03	Same	16	1,084	21.9%	64
04	Same	32	1,100	22.2%	159
05	+ interleave by user, spread drive principals, light PDF child	16	1,566	31.6%	0
06	Same	32	2,086	42.1%	0

Every mock run indexed all 2,700 documents, with no duplicates and no failures.

The changes

A rolling conversion window with 32 workers. Before, files converted in sub-batches of 50 with 8 threads, and each sub-batch waited for its slowest file; nothing converted while the listing and the folder walk ran. Now BoundedParallelRunner (onyx/utils/threadpool_concurrency.py) keeps up to DRIVE_CONVERSION_WORKERS (32) conversions running while the main thread goes on listing. It copies the caller's context variables to each thread, as the module's other helpers do. Throughput is about the requests in flight divided by the time per document, and the old code averaged about 4 in flight.
One token and one session per user. Before, each new thread's clients fetched their own OAuth token, and each Docs API call opened a new AuthorizedSession and fetched one more (1,412 token requests for 2,700 mock documents; in the live before run, 4 timed out). Now the gateway impersonates each user once and shares that credential (GoogleDriveSourceOperations._impersonated), and keeps the Drive client, directory client and Docs API session per thread. The mock counted 8 token requests.
Full list pages. Every files.list asks for 1,000 files (DRIVE_LIST_PAGE_SIZE) instead of Google's default 100. A call costs the same quota at any page size.
A parallel folder prefetch. Before, each unknown parent folder was fetched one at a time on the main thread while conversions waited, and the retriever and admin were tried in a random order (a set). Now _prefetch_ancestor_folders fetches a sub-batch's unknown folders in parallel, a tree level at a time, into a per-call cache; the walk reads the cache, retriever first. A failed fetch is not cached, so the walk still sees the error.
Load spread over users. Partitions converted one after another, so all concurrency acted as one user and met Google's per-user quotas (300 Docs API reads a minute). Now _interleave_by_user hands files to the pool round-robin by user, reading up to 2,000 listed files ahead, and select_drive_organizer(busy=...) prefers an organizer that lists no other drive. A call converts everything it lists before it returns, so the order does not change what the checkpoint records.
A light PDF subprocess. Each PDF's isolated child imported extract_file_text for about 0.9 CPU seconds. PDFium extraction moved to onyx/file_processing/pdfium_text.py, which imports psutil and pypdfium2; a child now costs about 0.04 CPU seconds (mock: 367 → 23 CPU seconds of PDF children per run). extract_file_text re-exports the names.
Rate-limit retries past the minute window. Downloads and exports gave up after the client library's 3 retries (about 7 s), so a rate limit failed the document (154 of 770 lost under a tight mock quota). Now they retry with backoff up to about 95 s, and the Docs API retries 7 times with waits up to 32 s; the same mock test lost none.

Test tools

Mock Drive server (backend/tests/integration/mock_services/mock_google_drive_server/): serves the token, Drive, Admin SDK, Docs and batch calls with Google's visibility rules, quotas and error shapes, from a corpus generated from a seed. See its README.
GOOGLE_API_ENDPOINT_OVERRIDE (testing only): sends every Google call to another URL. Clients are built from the bundled discovery document with its root URL replaced, so batch requests follow too. It is in the env inventory baseline, not in env.template.
backend/scripts/benchmark_google_drive_connector.py: runs the connector's checkpoint loop as docfetching does and reports docs/min, quota use, requests in flight, CPU and every file that did not become a document.

What limits the connector now

Parsing CPU in one process. Past 32 workers the docfetching process uses about one core under the GIL; DOCX conversion (markitdown → mammoth) is the heaviest per file.
The Docs API quota per user. One user who owns a whole partition caps near 850 docs/min. Lifting it needs Docs API reads spread over a drive's members, or an HTML export instead of the Docs API call (only if Google's HTML export keeps heading ids, which needs a real-Workspace check).

Issues found, not fixed here

doc_conversion.py (generic extraction, about line 455) catches every exception around the download, so a failed download drops the file with no ConnectorFailure.
18 live Google Docs had headings without headingId; section_extraction.py raised KeyError and fell back to plain text.
After a pause, docfetching ran on for about 8 minutes until its watchdog saw the canceled attempt.

How Has This Been Tested?

Unit tests: backend/tests/unit/onyx/connectors, file_processing and utils (5,912 passed). New tests cover BoundedParallelRunner (in-flight bound, errors, context variables), _interleave_by_user, and the conversion bounds; existing batching and Docs API tests were updated to the new seams.
backend/tests/external_dependency_unit/connectors/google_drive/test_mock_drive_indexing.py indexes a small mock corpus end to end and checks that every file becomes a document.
The mock hill climb and the live A/B test above. The live scripts drove a local Onyx instance and are not part of this PR.

Docs

Customer-facing change: I updated documentation/ (text and screenshots).
No docs change needed (internal, tests, CI, or multi-tenant cloud only).

The connector's settings do not change. GOOGLE_API_ENDPOINT_OVERRIDE is for testing only.

Additional Options

[Optional] Please cherry-pick this PR to the latest release version.

🤖 Generated with Claude Code