modsetter--surfsense/surfsense_backend/app/tasks/celery_tasks/schedule_checker_task.py

"""Meta-scheduler task that checks for connectors needing periodic indexing."""

import logging
from datetime import UTC, datetime

from sqlalchemy.future import select

from app.celery_app import celery_app
from app.db import SearchSourceConnector, SearchSourceConnectorType
from app.notifications.persistence import Notification
from app.tasks.celery_tasks import get_celery_session_maker, run_async_celery_task
from app.utils.indexing_locks import is_connector_indexing_locked

logger = logging.getLogger(__name__)


@celery_app.task(name="check_periodic_schedules")
def check_periodic_schedules_task():
    """
    Check all connectors for periodic indexing that's due.
    This task runs every minute and triggers indexing for any connector
    whose next_scheduled_at time has passed.
    """
    return run_async_celery_task(_check_and_trigger_schedules)


async def _check_and_trigger_schedules():
    """Check database for connectors that need indexing and trigger their tasks."""
    async with get_celery_session_maker()() as session:
        try:
            # Find all connectors with periodic indexing enabled that are due
            now = datetime.now(UTC)
            result = await session.execute(
                select(SearchSourceConnector).filter(
                    SearchSourceConnector.periodic_indexing_enabled == True,  # noqa: E712
                    SearchSourceConnector.next_scheduled_at <= now,
                )
            )
            due_connectors = result.scalars().all()

            if not due_connectors:
                logger.debug("No connectors due for periodic indexing")
                return

            logger.info(f"Found {len(due_connectors)} connectors due for indexing")

            # Import indexing tasks for KB connectors only.
            # Live connectors (Linear, Slack, Jira, ClickUp, Airtable, Discord,
            # Teams, Gmail, Calendar, Luma) use real-time tools instead.
            from app.tasks.celery_tasks.connector_tasks import (
                index_confluence_pages_task,
                index_elasticsearch_documents_task,
                index_github_repos_task,
                index_google_drive_files_task,
                index_notion_pages_task,
            )

            task_map = {
                SearchSourceConnectorType.NOTION_CONNECTOR: index_notion_pages_task,
                SearchSourceConnectorType.GITHUB_CONNECTOR: index_github_repos_task,
                SearchSourceConnectorType.CONFLUENCE_CONNECTOR: index_confluence_pages_task,
                SearchSourceConnectorType.ELASTICSEARCH_CONNECTOR: index_elasticsearch_documents_task,
                SearchSourceConnectorType.GOOGLE_DRIVE_CONNECTOR: index_google_drive_files_task,
                SearchSourceConnectorType.COMPOSIO_GOOGLE_DRIVE_CONNECTOR: index_google_drive_files_task,
            }

            from app.services.mcp_oauth.registry import (
                DEPRECATED_INDEXING_CONNECTOR_TYPES,
                LIVE_CONNECTOR_TYPES,
            )

            # Disable obsolete periodic indexing in one batch: live connectors
            # (now real-time agent tools) and deprecated-indexing connectors
            # (KB is files/notes/uploads only) no longer index on a schedule.
            live_disabled = []
            for connector in due_connectors:
                if (
                    connector.connector_type in LIVE_CONNECTOR_TYPES
                    or connector.connector_type in DEPRECATED_INDEXING_CONNECTOR_TYPES
                ):
                    connector.periodic_indexing_enabled = False
                    connector.next_scheduled_at = None
                    live_disabled.append(connector)
            if live_disabled:
                await session.commit()
                for c in live_disabled:
                    logger.info(
                        "Disabled obsolete periodic indexing for connector %s (%s)",
                        c.id,
                        c.connector_type.value,
                    )

            # Trigger indexing for each due connector
            for connector in due_connectors:
                if connector in live_disabled:
                    continue

                # Primary guard: Redis lock indicates a task is currently running.
                if is_connector_indexing_locked(connector.id):
                    logger.info(
                        f"Skipping periodic indexing for connector {connector.id} "
                        "(Redis lock indicates indexing is already in progress)"
                    )
                    continue

                # Skip scheduling if a sync for this connector is already in progress.
                # This prevents duplicate tasks from piling up under slow/rate-limited providers.
                in_progress_result = await session.execute(
                    select(Notification.id).where(
                        Notification.type == "connector_indexing",
                        Notification.notification_metadata["connector_id"].astext
                        == str(connector.id),
                        Notification.notification_metadata["status"].astext
                        == "in_progress",
                    )
                )
                if in_progress_result.first():
                    logger.info(
                        f"Skipping periodic indexing for connector {connector.id} "
                        "(already has in-progress indexing notification)"
                    )
                    continue

                task = task_map.get(connector.connector_type)
                if task:
                    logger.info(
                        f"Triggering periodic indexing for connector {connector.id} "
                        f"({connector.connector_type.value})"
                    )

                    # Special handling for Google Drive (native and Composio) - uses config for folder/file selection
                    if connector.connector_type in [
                        SearchSourceConnectorType.GOOGLE_DRIVE_CONNECTOR,
                        SearchSourceConnectorType.COMPOSIO_GOOGLE_DRIVE_CONNECTOR,
                    ]:
                        connector_config = connector.config or {}
                        selected_folders = connector_config.get("selected_folders", [])
                        selected_files = connector_config.get("selected_files", [])
                        indexing_options = connector_config.get(
                            "indexing_options",
                            {
                                "max_files_per_folder": 100,
                                "incremental_sync": True,
                                "include_subfolders": True,
                            },
                        )

                        if selected_folders or selected_files:
                            task.delay(
                                connector.id,
                                connector.workspace_id,
                                str(connector.user_id),
                                {
                                    "folders": selected_folders,
                                    "files": selected_files,
                                    "indexing_options": indexing_options,
                                },
                            )
                        else:
                            # No folders/files selected - skip indexing but still update next_scheduled_at
                            # to prevent checking every minute
                            logger.info(
                                f"Google Drive connector {connector.id} has no folders or files selected, "
                                "skipping periodic indexing (will check again at next scheduled time)"
                            )
                            from datetime import timedelta

                            connector.next_scheduled_at = now + timedelta(
                                minutes=connector.indexing_frequency_minutes
                            )
                            await session.commit()
                            continue

                    else:
                        task.delay(
                            connector.id,
                            connector.workspace_id,
                            str(connector.user_id),
                            None,  # start_date - uses last_indexed_at
                            None,  # end_date - uses now
                        )

                    # Update next_scheduled_at for next run
                    from datetime import timedelta

                    connector.next_scheduled_at = now + timedelta(
                        minutes=connector.indexing_frequency_minutes
                    )
                    await session.commit()
                else:
                    logger.warning(
                        f"No task found for connector type {connector.connector_type}"
                    )

        except Exception as e:
            logger.error(f"Error checking periodic schedules: {e!s}", exc_info=True)
            await session.rollback()