diff --git a/README.md b/README.md index a72cd9e..4fd8dd8 100644 --- a/README.md +++ b/README.md @@ -235,7 +235,7 @@ true for Founders/Pro/Business/Backer, false for Standard by default). | `GET /api/drive/connections/` | Caller's personal connections + their company's company connections | | `DELETE /api/drive/connections//` | Disconnect (owner for personal, company manager for company) — deactivates + clears tokens, keeps history | | `POST /api/drive/connections//resources/` | `{ "resource_ids": [...], "resource_labels": [...] }` — folder/shared-drive/site ids to sync; empty = provider root | -| `POST /api/drive/connections//sync/` | Sync now (`chat_backend/services/drive_sync.py::sync_connection`) | +| `POST /api/drive/connections//sync/` | Enqueue sync now (`chat_backend/drive_tasks.py`) — returns **202** with `connection.last_sync_status=pending`; poll `GET /api/drive/connections/` for `ok` / `error` + `last_sync_error` | **Provider scope differences:** - Google: same `drive.readonly` scope for personal and company; company sync @@ -243,10 +243,12 @@ true for Founders/Pro/Business/Backer, false for Standard by default). - Microsoft: personal uses `Files.Read`; company uses `Files.Read.All Sites.Read.All` and syncs SharePoint sites (`selected_resource_ids` = site ids). -**Workers / webhooks (#52):** -- `python manage.py sync_drive_connections [--connection-id N]` — cron/worker entry point. +**Workers / webhooks (#52, #57):** +- `python manage.py sync_drive_connections [--connection-id N]` — enqueue sync tasks (default). +- `python manage.py sync_drive_connections --sync-now` — run sync inline in this process (cron/debug). +- Django 6 `TASKS` (see `settings.py`): default `ImmediateBackend` runs tasks in-process; Sync now still returns 202 by dispatching on a background thread. Swap `TASKS` to a durable queue + worker for production scale. - `POST /api/drive/webhooks/google/` / `POST /api/drive/webhooks/microsoft/` — - provider push-notification stubs (`AllowAny`); acknowledge `200` and call + provider push-notification stubs (`AllowAny`); acknowledge `200` and enqueue `sync_connection` when the notification's `connection_id` is resolvable, else just `200` (no-op). Microsoft's subscription-creation `validationToken` handshake is echoed back as `text/plain`. diff --git a/llm_be/chat_backend/drive_tasks.py b/llm_be/chat_backend/drive_tasks.py new file mode 100644 index 0000000..47a7d44 --- /dev/null +++ b/llm_be/chat_backend/drive_tasks.py @@ -0,0 +1,95 @@ +"""Drive sync background tasks (#57). + +Uses Django 6 Tasks API (same pattern as ``email_tasks``). With the default +``ImmediateBackend``, enqueue still runs in-process — we dispatch that call on +a daemon thread after commit so ``POST .../sync/`` can return 202 Pending +without waiting for Drive list/download/ingest. Swap ``TASKS`` to a durable +queue + worker later; enqueue path stays the same. +""" + +from __future__ import annotations + +import logging +import threading +from functools import partial + +from django.conf import settings +from django.db import transaction +from django.tasks import task + +from chat_backend.models import DriveConnection +from chat_backend.services.drive_sync import sync_connection + +logger = logging.getLogger(__name__) + + +def _uses_immediate_backend() -> bool: + backend = ( + (getattr(settings, "TASKS", {}) or {}) + .get("default", {}) + .get("BACKEND", "") + ) + return "ImmediateBackend" in str(backend) + + +@task +def run_drive_connection_sync(connection_id: int) -> dict: + """Load a DriveConnection and run ``sync_connection`` (#57).""" + try: + connection = DriveConnection.objects.get(pk=connection_id, is_active=True) + except DriveConnection.DoesNotExist: + logger.error( + "DriveConnection %s missing or inactive; cannot sync", connection_id + ) + return {"error": "connection_not_found"} + return sync_connection(connection) + + +def _dispatch_sync(connection_id: int) -> None: + """Enqueue (or run) the sync task; ImmediateBackend runs off the request thread.""" + try: + if _uses_immediate_backend(): + # ImmediateBackend executes during enqueue — keep HTTP snappy. + threading.Thread( + target=run_drive_connection_sync.enqueue, + kwargs={"connection_id": connection_id}, + daemon=True, + name=f"drive-sync-{connection_id}", + ).start() + else: + run_drive_connection_sync.enqueue(connection_id=connection_id) + except Exception: + logger.exception( + "Failed to dispatch Drive sync for connection=%s", connection_id + ) + DriveConnection.objects.filter(pk=connection_id).update( + last_sync_status=DriveConnection.SyncStatus.ERROR, + last_sync_error="Failed to enqueue Drive sync task.", + ) + + +def enqueue_drive_sync( + connection: DriveConnection, *, force: bool = False +) -> tuple[DriveConnection, bool]: + """Mark connection pending and enqueue sync after DB commit. + + Returns ``(connection, enqueued)``. If already ``pending`` and ``force`` is + false, does not enqueue a duplicate job. + """ + connection.refresh_from_db() + if ( + not force + and connection.last_sync_status == DriveConnection.SyncStatus.PENDING + ): + return connection, False + + connection.last_sync_status = DriveConnection.SyncStatus.PENDING + connection.last_sync_error = "" + connection.save( + update_fields=["last_sync_status", "last_sync_error", "last_modified"] + ) + + transaction.on_commit( + partial(_dispatch_sync, connection_id=connection.id) + ) + return connection, True diff --git a/llm_be/chat_backend/management/commands/sync_drive_connections.py b/llm_be/chat_backend/management/commands/sync_drive_connections.py index d5a65d5..12ece1b 100644 --- a/llm_be/chat_backend/management/commands/sync_drive_connections.py +++ b/llm_be/chat_backend/management/commands/sync_drive_connections.py @@ -1,20 +1,21 @@ -"""Worker entry point for scheduled Drive sync (#52). +"""Worker entry point for scheduled Drive sync (#52 / #57). Usage: python manage.py sync_drive_connections python manage.py sync_drive_connections --connection-id 42 + python manage.py sync_drive_connections --sync-now """ from __future__ import annotations from django.core.management.base import BaseCommand, CommandError +from chat_backend.drive_tasks import enqueue_drive_sync, run_drive_connection_sync from chat_backend.models import DriveConnection -from chat_backend.services.drive_sync import sync_connection class Command(BaseCommand): - help = "Sync active Drive connections (Google Drive / OneDrive / SharePoint) into Documents." + help = "Enqueue (or run) sync for active Drive connections into Documents." def add_arguments(self, parser): parser.add_argument( @@ -23,9 +24,18 @@ class Command(BaseCommand): default=None, help="Sync only the DriveConnection with this id.", ) + parser.add_argument( + "--sync-now", + action="store_true", + help=( + "Run sync inline in this process instead of enqueueing a " + "background task (useful for cron/debugging)." + ), + ) def handle(self, *args, **options): connection_id = options.get("connection_id") + sync_now = options.get("sync_now") queryset = DriveConnection.objects.filter(is_active=True) if connection_id is not None: queryset = queryset.filter(id=connection_id) @@ -41,15 +51,25 @@ class Command(BaseCommand): for connection in connections: self.stdout.write( - f"Syncing connection={connection.id} " + f"{'Syncing' if sync_now else 'Enqueueing'} connection={connection.id} " f"provider={connection.provider} kind={connection.kind}..." ) - result = sync_connection(connection) - if result.get("error"): - self.stderr.write(f" connection={connection.id} failed: {result['error']}") + if sync_now: + result = run_drive_connection_sync.call(connection_id=connection.id) + if result.get("error"): + self.stderr.write( + f" connection={connection.id} failed: {result['error']}" + ) + else: + self.stdout.write( + f" connection={connection.id} added={result.get('added', 0)} " + f"updated={result.get('updated', 0)} " + f"removed={result.get('removed', 0)} " + f"failed={len(result.get('failed') or [])}" + ) else: + _, enqueued = enqueue_drive_sync(connection, force=True) self.stdout.write( - f" connection={connection.id} added={result['added']} " - f"updated={result['updated']} removed={result['removed']} " - f"failed={len(result['failed'])}" + f" connection={connection.id} " + f"{'queued' if enqueued else 'already pending'}" ) diff --git a/llm_be/chat_backend/tests/test_drive_tasks.py b/llm_be/chat_backend/tests/test_drive_tasks.py new file mode 100644 index 0000000..ce8af0b --- /dev/null +++ b/llm_be/chat_backend/tests/test_drive_tasks.py @@ -0,0 +1,74 @@ +"""Tests for Drive sync background tasks (#57).""" + +from __future__ import annotations + +from unittest import mock + +from django.test import TestCase, TransactionTestCase + +from chat_backend.drive_tasks import enqueue_drive_sync, run_drive_connection_sync +from chat_backend.models import DriveConnection +from chat_backend.tests.factories import make_company, make_drive_connection, make_user + + +class RunDriveConnectionSyncTestCase(TestCase): + def setUp(self): + self.company = make_company() + self.user = make_user(company=self.company) + self.connection = make_drive_connection( + self.company, kind=DriveConnection.Kind.PERSONAL, user=self.user + ) + + @mock.patch("chat_backend.drive_tasks.sync_connection") + def test_runs_sync_for_active_connection(self, mock_sync): + mock_sync.return_value = {"added": 1, "updated": 0, "removed": 0, "failed": []} + + result = run_drive_connection_sync.call(connection_id=self.connection.id) + + mock_sync.assert_called_once() + self.assertEqual(mock_sync.call_args.args[0].id, self.connection.id) + self.assertEqual(result["added"], 1) + + def test_missing_connection_returns_error(self): + result = run_drive_connection_sync.call(connection_id=999999) + self.assertEqual(result, {"error": "connection_not_found"}) + + +class EnqueueDriveSyncTestCase(TransactionTestCase): + def setUp(self): + self.company = make_company() + self.user = make_user(company=self.company) + self.connection = make_drive_connection( + self.company, kind=DriveConnection.Kind.PERSONAL, user=self.user + ) + + @mock.patch("chat_backend.drive_tasks._dispatch_sync") + def test_marks_pending_and_dispatches(self, mock_dispatch): + connection, enqueued = enqueue_drive_sync(self.connection) + + self.assertTrue(enqueued) + connection.refresh_from_db() + self.assertEqual(connection.last_sync_status, DriveConnection.SyncStatus.PENDING) + self.assertEqual(connection.last_sync_error, "") + mock_dispatch.assert_called_once_with(connection_id=self.connection.id) + + @mock.patch("chat_backend.drive_tasks._dispatch_sync") + def test_skips_duplicate_while_pending(self, mock_dispatch): + self.connection.last_sync_status = DriveConnection.SyncStatus.PENDING + self.connection.save(update_fields=["last_sync_status"]) + + connection, enqueued = enqueue_drive_sync(self.connection) + + self.assertFalse(enqueued) + mock_dispatch.assert_not_called() + self.assertEqual(connection.last_sync_status, DriveConnection.SyncStatus.PENDING) + + @mock.patch("chat_backend.drive_tasks._dispatch_sync") + def test_force_reenqueues_pending(self, mock_dispatch): + self.connection.last_sync_status = DriveConnection.SyncStatus.PENDING + self.connection.save(update_fields=["last_sync_status"]) + + _, enqueued = enqueue_drive_sync(self.connection, force=True) + + self.assertTrue(enqueued) + mock_dispatch.assert_called_once() diff --git a/llm_be/chat_backend/tests/test_management_drive_sync.py b/llm_be/chat_backend/tests/test_management_drive_sync.py index 5118975..fc53093 100644 --- a/llm_be/chat_backend/tests/test_management_drive_sync.py +++ b/llm_be/chat_backend/tests/test_management_drive_sync.py @@ -1,4 +1,4 @@ -"""Tests for the sync_drive_connections management command (#52).""" +"""Tests for the sync_drive_connections management command (#52 / #57).""" from __future__ import annotations @@ -17,9 +17,11 @@ class SyncDriveConnectionsCommandTestCase(TestCase): def setUp(self): self.company = make_company() - @mock.patch("chat_backend.management.commands.sync_drive_connections.sync_connection") - def test_syncs_all_active_connections(self, mock_sync): - mock_sync.return_value = {"added": 1, "updated": 0, "removed": 0, "failed": []} + @mock.patch( + "chat_backend.management.commands.sync_drive_connections.enqueue_drive_sync" + ) + def test_enqueues_all_active_connections(self, mock_enqueue): + mock_enqueue.return_value = (mock.Mock(), True) active = make_drive_connection(self.company) make_drive_connection( self.company, @@ -30,24 +32,55 @@ class SyncDriveConnectionsCommandTestCase(TestCase): out = StringIO() call_command("sync_drive_connections", stdout=out) - mock_sync.assert_called_once_with(active) + mock_enqueue.assert_called_once() + self.assertEqual(mock_enqueue.call_args.args[0].id, active.id) + self.assertIn("queued", out.getvalue()) + + @mock.patch("chat_backend.drive_tasks.sync_connection") + def test_sync_now_runs_inline(self, mock_sync): + mock_sync.return_value = { + "added": 1, + "updated": 0, + "removed": 0, + "failed": [], + } + active = make_drive_connection(self.company) + + out = StringIO() + call_command("sync_drive_connections", "--sync-now", stdout=out) + + mock_sync.assert_called_once() + self.assertEqual(mock_sync.call_args.args[0].id, active.id) self.assertIn("added=1", out.getvalue()) - @mock.patch("chat_backend.management.commands.sync_drive_connections.sync_connection") - def test_syncs_single_connection_by_id(self, mock_sync): - mock_sync.return_value = {"added": 0, "updated": 1, "removed": 0, "failed": []} + @mock.patch( + "chat_backend.management.commands.sync_drive_connections.enqueue_drive_sync" + ) + def test_enqueues_single_connection_by_id(self, mock_enqueue): + mock_enqueue.return_value = (mock.Mock(), True) target = make_drive_connection(self.company) - other = make_drive_connection( + make_drive_connection( self.company, provider=DriveConnection.Provider.MICROSOFT ) - call_command("sync_drive_connections", "--connection-id", str(target.id), stdout=StringIO()) + call_command( + "sync_drive_connections", + "--connection-id", + str(target.id), + stdout=StringIO(), + ) - mock_sync.assert_called_once_with(target) + mock_enqueue.assert_called_once() + self.assertEqual(mock_enqueue.call_args.args[0].id, target.id) def test_unknown_connection_id_raises(self): with self.assertRaises(CommandError): - call_command("sync_drive_connections", "--connection-id", "999999", stdout=StringIO()) + call_command( + "sync_drive_connections", + "--connection-id", + "999999", + stdout=StringIO(), + ) def test_no_connections_reports_and_exits_cleanly(self): out = StringIO() diff --git a/llm_be/chat_backend/tests/test_views_drive.py b/llm_be/chat_backend/tests/test_views_drive.py index d44a983..f18fb0c 100644 --- a/llm_be/chat_backend/tests/test_views_drive.py +++ b/llm_be/chat_backend/tests/test_views_drive.py @@ -189,35 +189,38 @@ class DriveConnectionSyncViewTestCase(APITestCase): ) self.url = reverse("drive_connection_sync", kwargs={"connection_id": self.connection.id}) - @mock.patch("chat_backend.views_drive.sync_connection") - def test_triggers_sync_for_owner(self, mock_sync): - mock_sync.return_value = {"added": 2, "updated": 0, "removed": 0, "failed": []} + @mock.patch("chat_backend.views_drive.enqueue_drive_sync") + def test_triggers_sync_for_owner(self, mock_enqueue): + self.connection.last_sync_status = DriveConnection.SyncStatus.PENDING + self.connection.save(update_fields=["last_sync_status"]) + mock_enqueue.return_value = (self.connection, True) response = self.client.post(self.url) - self.assertEqual(response.status_code, status.HTTP_200_OK) - mock_sync.assert_called_once_with(self.connection) - self.assertEqual(response.data["result"]["added"], 2) + self.assertEqual(response.status_code, status.HTTP_202_ACCEPTED) + mock_enqueue.assert_called_once() + self.assertTrue(response.data["queued"]) + self.assertEqual(response.data["connection"]["last_sync_status"], "pending") - @mock.patch("chat_backend.views_drive.sync_connection") - def test_non_owner_forbidden(self, mock_sync): + @mock.patch("chat_backend.views_drive.enqueue_drive_sync") + def test_non_owner_forbidden(self, mock_enqueue): other = make_user(email="other@example.com", company=self.company) self.client.force_authenticate(user=other) response = self.client.post(self.url) self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN) - mock_sync.assert_not_called() + mock_enqueue.assert_not_called() - @mock.patch("chat_backend.views_drive.sync_connection") - def test_disconnected_connection_rejected(self, mock_sync): + @mock.patch("chat_backend.views_drive.enqueue_drive_sync") + def test_disconnected_connection_rejected(self, mock_enqueue): self.connection.is_active = False self.connection.save(update_fields=["is_active"]) response = self.client.post(self.url) self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST) - mock_sync.assert_not_called() + mock_enqueue.assert_not_called() class DriveWebhookViewTestCase(APITestCase): @@ -227,33 +230,34 @@ class DriveWebhookViewTestCase(APITestCase): self.company, kind=DriveConnection.Kind.PERSONAL, provider=DriveConnection.Provider.GOOGLE ) - @mock.patch("chat_backend.views_drive.sync_connection") - def test_google_webhook_syncs_matching_connection(self, mock_sync): + @mock.patch("chat_backend.views_drive.enqueue_drive_sync") + def test_google_webhook_syncs_matching_connection(self, mock_enqueue): url = reverse("drive_webhook_google") response = self.client.post(f"{url}?connection_id={self.connection.id}") self.assertEqual(response.status_code, status.HTTP_200_OK) - mock_sync.assert_called_once_with(self.connection) + mock_enqueue.assert_called_once() + self.assertEqual(mock_enqueue.call_args.args[0].id, self.connection.id) - @mock.patch("chat_backend.views_drive.sync_connection") - def test_google_webhook_without_connection_id_is_a_noop(self, mock_sync): + @mock.patch("chat_backend.views_drive.enqueue_drive_sync") + def test_google_webhook_without_connection_id_is_a_noop(self, mock_enqueue): url = reverse("drive_webhook_google") response = self.client.post(url) self.assertEqual(response.status_code, status.HTTP_200_OK) - mock_sync.assert_not_called() + mock_enqueue.assert_not_called() - @mock.patch("chat_backend.views_drive.sync_connection") - def test_microsoft_webhook_validation_handshake(self, mock_sync): + @mock.patch("chat_backend.views_drive.enqueue_drive_sync") + def test_microsoft_webhook_validation_handshake(self, mock_enqueue): url = reverse("drive_webhook_microsoft") response = self.client.post(f"{url}?validationToken=abc123") self.assertEqual(response.status_code, status.HTTP_200_OK) self.assertEqual(response.content.decode(), "abc123") - mock_sync.assert_not_called() + mock_enqueue.assert_not_called() - @mock.patch("chat_backend.views_drive.sync_connection") - def test_microsoft_webhook_syncs_via_client_state(self, mock_sync): + @mock.patch("chat_backend.views_drive.enqueue_drive_sync") + def test_microsoft_webhook_syncs_via_client_state(self, mock_enqueue): ms_connection = make_drive_connection( self.company, kind=DriveConnection.Kind.PERSONAL, @@ -269,4 +273,5 @@ class DriveWebhookViewTestCase(APITestCase): ) self.assertEqual(response.status_code, status.HTTP_200_OK) - mock_sync.assert_called_once_with(ms_connection) + mock_enqueue.assert_called_once() + self.assertEqual(mock_enqueue.call_args.args[0].id, ms_connection.id) diff --git a/llm_be/chat_backend/views_drive.py b/llm_be/chat_backend/views_drive.py index eb1721a..7b73147 100644 --- a/llm_be/chat_backend/views_drive.py +++ b/llm_be/chat_backend/views_drive.py @@ -12,9 +12,9 @@ from rest_framework.views import APIView from finance.services.quotas import FeatureNotAllowed, assert_feature_allowed +from .drive_tasks import enqueue_drive_sync from .models import DriveConnection from .serializers import DriveConnectionResourcesSerializer, DriveConnectionSerializer -from .services.drive_sync import sync_connection logger = logging.getLogger(__name__) @@ -138,10 +138,14 @@ class DriveConnectionSyncView(APIView): status=status.HTTP_400_BAD_REQUEST, ) - result = sync_connection(connection) + connection, enqueued = enqueue_drive_sync(connection) connection.refresh_from_db() return Response( - {"result": result, "connection": DriveConnectionSerializer(connection).data} + { + "queued": enqueued, + "connection": DriveConnectionSerializer(connection).data, + }, + status=status.HTTP_202_ACCEPTED, ) @@ -161,10 +165,11 @@ def _maybe_sync_from_webhook(connection_id, provider: str) -> None: if connection is None: return try: - sync_connection(connection) + enqueue_drive_sync(connection) except Exception: logger.exception( - "Webhook-triggered Drive sync failed for connection=%s", connection_id + "Webhook-triggered Drive sync enqueue failed for connection=%s", + connection_id, )