diff --git a/surfsense_backend/alembic/versions/23_add_connector_schedules_table.py b/surfsense_backend/alembic/versions/23_add_connector_schedules_table.py index edd8a9cce..57e6702a9 100644 --- a/surfsense_backend/alembic/versions/23_add_connector_schedules_table.py +++ b/surfsense_backend/alembic/versions/23_add_connector_schedules_table.py @@ -42,6 +42,10 @@ def upgrade() -> None: search_space_id INTEGER NOT NULL REFERENCES searchspaces(id) ON DELETE CASCADE, schedule_type scheduletype NOT NULL, cron_expression VARCHAR(100), + daily_time TIME, + weekly_day SMALLINT, + weekly_time TIME, + hourly_minute SMALLINT, is_active BOOLEAN NOT NULL DEFAULT TRUE, last_run_at TIMESTAMPTZ, next_run_at TIMESTAMPTZ, diff --git a/surfsense_backend/app/app.py b/surfsense_backend/app/app.py index 9c882bf85..f814e6f37 100644 --- a/surfsense_backend/app/app.py +++ b/surfsense_backend/app/app.py @@ -1,6 +1,6 @@ import asyncio import logging -from contextlib import asynccontextmanager +from contextlib import asynccontextmanager, suppress from fastapi import Depends, FastAPI from fastapi.middleware.cors import CORSMiddleware @@ -9,6 +9,7 @@ from sqlalchemy.ext.asyncio import AsyncSession from app.config import config from app.db import User, create_db_and_tables, get_async_session from app.routes import router as crud_router +from app.routes.connector_schedules_routes import router as connector_schedules_router from app.routes.scheduler_routes import router as scheduler_router from app.schemas import UserCreate, UserRead, UserUpdate from app.services.connector_scheduler_service import start_scheduler, stop_scheduler @@ -43,10 +44,8 @@ async def lifespan(app: FastAPI): # Cancel the scheduler task if not scheduler_task.done(): scheduler_task.cancel() - try: + with suppress(asyncio.CancelledError): await scheduler_task - except asyncio.CancelledError: - pass logger.info("Application shutdown complete") @@ -98,6 +97,7 @@ if config.AUTH_TYPE == "GOOGLE": ) app.include_router(crud_router, prefix="/api/v1", tags=["crud"]) +app.include_router(connector_schedules_router, prefix="/api/v1", tags=["connector-schedules"]) app.include_router(scheduler_router, prefix="/api/v1", tags=["scheduler"]) diff --git a/surfsense_backend/app/db.py b/surfsense_backend/app/db.py index d2d30afc4..9d295f5fe 100644 --- a/surfsense_backend/app/db.py +++ b/surfsense_backend/app/db.py @@ -16,6 +16,7 @@ from sqlalchemy import ( Integer, String, Text, + Time, UniqueConstraint, text, ) @@ -316,6 +317,11 @@ class ConnectorSchedule(BaseModel, TimestampMixin): ) schedule_type = Column(SQLAlchemyEnum(ScheduleType), nullable=False) cron_expression = Column(String(100), nullable=True) + # Optional schedule config fields (persist user selections) + daily_time = Column(Time(timezone=False), nullable=True) + weekly_day = Column(Integer, nullable=True) # 0=Mon .. 6=Sun + weekly_time = Column(Time(timezone=False), nullable=True) + hourly_minute = Column(Integer, nullable=True) # 0..59 is_active = Column(Boolean, nullable=False, default=True, index=True) last_run_at = Column(TIMESTAMP(timezone=True), nullable=True) next_run_at = Column(TIMESTAMP(timezone=True), nullable=True, index=True) diff --git a/surfsense_backend/app/routes/connector_schedules_routes.py b/surfsense_backend/app/routes/connector_schedules_routes.py index eaa3333e8..0b73ae1bf 100644 --- a/surfsense_backend/app/routes/connector_schedules_routes.py +++ b/surfsense_backend/app/routes/connector_schedules_routes.py @@ -9,7 +9,6 @@ PATCH /connector-schedules/{schedule_id}/toggle - Activate/deactivate a schedule """ import logging -from typing import Any from fastapi import APIRouter, Depends, HTTPException from sqlalchemy.exc import IntegrityError @@ -18,6 +17,7 @@ from sqlalchemy.future import select from app.db import ( ConnectorSchedule, + ScheduleType, SearchSourceConnector, SearchSpace, User, @@ -66,6 +66,13 @@ async def create_connector_schedule( # Verify search space belongs to user await check_ownership(session, SearchSpace, schedule.search_space_id, user) + # Ensure the connector belongs to the same search space + if connector.search_space_id != schedule.search_space_id: + raise HTTPException( + status_code=400, + detail="Connector does not belong to the provided search space", + ) + # Check if schedule already exists for this connector-space pair result = await session.execute( select(ConnectorSchedule).filter( @@ -82,13 +89,29 @@ async def create_connector_schedule( # Calculate next run time next_run_at = calculate_next_run( - schedule.schedule_type, schedule.cron_expression + schedule.schedule_type, + schedule.cron_expression, + schedule.daily_time, + schedule.weekly_day, + schedule.weekly_time, + schedule.hourly_minute ) - # Create schedule - db_schedule = ConnectorSchedule( - **schedule.model_dump(), next_run_at=next_run_at + # Create schedule (only DB columns) + db_data = schedule.model_dump( + include={ + "connector_id", + "search_space_id", + "schedule_type", + "cron_expression", + "daily_time", + "weekly_day", + "weekly_time", + "hourly_minute", + "is_active" + } ) + db_schedule = ConnectorSchedule(**db_data, next_run_at=next_run_at) session.add(db_schedule) await session.commit() await session.refresh(db_schedule) @@ -203,4 +226,139 @@ async def update_connector_schedule( Can update schedule_type, cron_expression, and is_active. If schedule_type changes, next_run_at is recalculated automatically. """ - \ No newline at end of file + try: + # Get the existing schedule + result = await session.execute( + select(ConnectorSchedule).filter(ConnectorSchedule.id == schedule_id) + ) + schedule = result.scalars().first() + + if not schedule: + raise HTTPException(status_code=404, detail="Schedule not found") + + # Verify schedule's connector belongs to user + await check_ownership(session, SearchSourceConnector, schedule.connector_id, user) + + # Update fields that were provided + update_data = schedule_update.model_dump(exclude_unset=True) + + # If schedule_type is being updated, recalculate next_run_at + if "schedule_type" in update_data: + # Use the new schedule_type and existing values for calculation + new_schedule_type = update_data["schedule_type"] + cron_expr = update_data.get("cron_expression", schedule.cron_expression) + daily_time = update_data.get("daily_time", schedule.daily_time) + weekly_day = update_data.get("weekly_day", schedule.weekly_day) + weekly_time = update_data.get("weekly_time", schedule.weekly_time) + hourly_minute = update_data.get("hourly_minute", schedule.hourly_minute) + update_data["next_run_at"] = calculate_next_run( + new_schedule_type, cron_expr, daily_time, weekly_day, weekly_time, hourly_minute + ) + elif "cron_expression" in update_data and schedule.schedule_type == ScheduleType.CUSTOM: + # If only cron_expression is updated for CUSTOM schedule, recalculate next_run_at + update_data["next_run_at"] = calculate_next_run( + schedule.schedule_type, + update_data["cron_expression"], + schedule.daily_time, + schedule.weekly_day, + schedule.weekly_time, + schedule.hourly_minute + ) + + # Apply updates + for field, value in update_data.items(): + setattr(schedule, field, value) + + await session.commit() + await session.refresh(schedule) + + logger.info(f"Updated schedule {schedule_id} with fields: {list(update_data.keys())}") + return schedule + + except HTTPException: + await session.rollback() + raise + except Exception as e: + await session.rollback() + logger.error(f"Failed to update connector schedule: {e!s}") + raise HTTPException( + status_code=500, detail=f"Failed to update connector schedule: {e!s}" + ) from e + + +@router.delete("/connector-schedules/{schedule_id}") +async def delete_connector_schedule( + schedule_id: int, + session: AsyncSession = Depends(get_async_session), + user: User = Depends(current_active_user), +): + """Delete a connector schedule.""" + try: + # Get the existing schedule + result = await session.execute( + select(ConnectorSchedule).filter(ConnectorSchedule.id == schedule_id) + ) + schedule = result.scalars().first() + + if not schedule: + raise HTTPException(status_code=404, detail="Schedule not found") + + # Verify schedule's connector belongs to user + await check_ownership(session, SearchSourceConnector, schedule.connector_id, user) + + # Delete the schedule + await session.delete(schedule) + await session.commit() + + logger.info(f"Deleted schedule {schedule_id}") + return {"message": "Schedule deleted successfully"} + + except HTTPException: + await session.rollback() + raise + except Exception as e: + await session.rollback() + logger.error(f"Failed to delete connector schedule: {e!s}") + raise HTTPException( + status_code=500, detail=f"Failed to delete connector schedule: {e!s}" + ) from e + + +@router.patch("/connector-schedules/{schedule_id}/toggle", response_model=ConnectorScheduleRead) +async def toggle_connector_schedule( + schedule_id: int, + session: AsyncSession = Depends(get_async_session), + user: User = Depends(current_active_user), +): + """Toggle the active status of a connector schedule.""" + try: + # Get the existing schedule + result = await session.execute( + select(ConnectorSchedule).filter(ConnectorSchedule.id == schedule_id) + ) + schedule = result.scalars().first() + + if not schedule: + raise HTTPException(status_code=404, detail="Schedule not found") + + # Verify schedule's connector belongs to user + await check_ownership(session, SearchSourceConnector, schedule.connector_id, user) + + # Toggle the active status + schedule.is_active = not schedule.is_active + await session.commit() + await session.refresh(schedule) + + logger.info(f"Toggled schedule {schedule_id} to {'active' if schedule.is_active else 'inactive'}") + return schedule + + except HTTPException: + await session.rollback() + raise + except Exception as e: + await session.rollback() + logger.error(f"Failed to toggle connector schedule: {e!s}") + raise HTTPException( + status_code=500, detail=f"Failed to toggle connector schedule: {e!s}" + ) from e + diff --git a/surfsense_backend/app/schemas/connector_schedule.py b/surfsense_backend/app/schemas/connector_schedule.py index 960e67c88..559dcb338 100644 --- a/surfsense_backend/app/schemas/connector_schedule.py +++ b/surfsense_backend/app/schemas/connector_schedule.py @@ -1,7 +1,7 @@ from datetime import datetime, time from typing import Optional -from pydantic import BaseModel, ConfigDict, field_validator +from pydantic import BaseModel, ConfigDict, FieldValidationInfo, field_validator from app.db import ScheduleType @@ -23,9 +23,9 @@ class ConnectorScheduleBase(BaseModel): @field_validator("cron_expression") @classmethod - def validate_cron_expression(cls, v: str | None, values: dict) -> str | None: + def validate_cron_expression(cls, v: str | None, info: FieldValidationInfo) -> str | None: """Validate cron expression is provided when schedule_type is CUSTOM.""" - schedule_type = values.data.get("schedule_type") + schedule_type = info.data.get("schedule_type") if schedule_type == ScheduleType.CUSTOM and not v: raise ValueError( "cron_expression is required when schedule_type is CUSTOM" @@ -38,9 +38,9 @@ class ConnectorScheduleBase(BaseModel): @field_validator("daily_time") @classmethod - def validate_daily_time(cls, v: time | None, values: dict) -> time | None: + def validate_daily_time(cls, v: time | None, info: FieldValidationInfo) -> time | None: """Validate daily_time is only provided for DAILY schedule type.""" - schedule_type = values.data.get("schedule_type") + schedule_type = info.data.get("schedule_type") if v is not None and schedule_type != ScheduleType.DAILY: raise ValueError( "daily_time should only be provided for DAILY schedule_type" @@ -49,9 +49,9 @@ class ConnectorScheduleBase(BaseModel): @field_validator("weekly_day") @classmethod - def validate_weekly_day(cls, v: int | None, values: dict) -> int | None: + def validate_weekly_day(cls, v: int | None, info: FieldValidationInfo) -> int | None: """Validate weekly_day is only provided for WEEKLY schedule type.""" - schedule_type = values.data.get("schedule_type") + schedule_type = info.data.get("schedule_type") if v is not None and schedule_type != ScheduleType.WEEKLY: raise ValueError( "weekly_day should only be provided for WEEKLY schedule_type" @@ -62,9 +62,9 @@ class ConnectorScheduleBase(BaseModel): @field_validator("weekly_time") @classmethod - def validate_weekly_time(cls, v: time | None, values: dict) -> time | None: + def validate_weekly_time(cls, v: time | None, info: FieldValidationInfo) -> time | None: """Validate weekly_time is only provided for WEEKLY schedule type.""" - schedule_type = values.data.get("schedule_type") + schedule_type = info.data.get("schedule_type") if v is not None and schedule_type != ScheduleType.WEEKLY: raise ValueError( "weekly_time should only be provided for WEEKLY schedule_type" @@ -73,9 +73,9 @@ class ConnectorScheduleBase(BaseModel): @field_validator("hourly_minute") @classmethod - def validate_hourly_minute(cls, v: int | None, values: dict) -> int | None: + def validate_hourly_minute(cls, v: int | None, info: FieldValidationInfo) -> int | None: """Validate hourly_minute is only provided for HOURLY schedule type.""" - schedule_type = values.data.get("schedule_type") + schedule_type = info.data.get("schedule_type") if v is not None and schedule_type != ScheduleType.HOURLY: raise ValueError( "hourly_minute should only be provided for HOURLY schedule_type" @@ -100,9 +100,9 @@ class ConnectorScheduleUpdate(BaseModel): @field_validator("cron_expression") @classmethod - def validate_cron_expression_update(cls, v: str | None, values: dict) -> str | None: + def validate_cron_expression_update(cls, v: str | None, info: FieldValidationInfo) -> str | None: """Validate cron expression for updates.""" - schedule_type = values.data.get("schedule_type") + schedule_type = info.data.get("schedule_type") if schedule_type == ScheduleType.CUSTOM and v is None: raise ValueError( "cron_expression is required when schedule_type is CUSTOM" diff --git a/surfsense_backend/app/services/connector_scheduler_service.py b/surfsense_backend/app/services/connector_scheduler_service.py index 24ed8e55e..b119c122e 100644 --- a/surfsense_backend/app/services/connector_scheduler_service.py +++ b/surfsense_backend/app/services/connector_scheduler_service.py @@ -7,7 +7,7 @@ It runs as a background service to check for due schedules and execute them. import asyncio import logging -from datetime import datetime, timedelta +from datetime import datetime, timedelta, timezone from typing import Dict, List, Optional from sqlalchemy import select, update @@ -45,7 +45,6 @@ class ConnectorSchedulerService: """Service for managing automated connector scheduling and execution.""" def __init__(self): - self.task_logger = TaskLoggingService() self.running = False self.check_interval = 60 # Check every 60 seconds self.max_concurrent_jobs = 5 # Maximum concurrent indexing jobs @@ -127,7 +126,7 @@ class ConnectorSchedulerService: async def _get_due_schedules(self, session: AsyncSession) -> List[ConnectorSchedule]: """Get all schedules that are due for execution.""" - now = datetime.utcnow() + now = datetime.now(datetime.utc) query = ( select(ConnectorSchedule) @@ -204,9 +203,10 @@ class ConnectorSchedulerService: try: # Create a task log entry - log_entry = await self.task_logger.log_task_start( - session, + task_logger = TaskLoggingService(session, schedule.search_space_id) + log_entry = await task_logger.log_task_start( f"scheduled_sync_{connector.connector_type.value}", + "connector_scheduler", f"Scheduled sync for {connector.name}", { "schedule_id": schedule_id, @@ -224,7 +224,7 @@ class ConnectorSchedulerService: "%Y-%m-%d" ) - end_date = datetime.utcnow().strftime("%Y-%m-%d") + end_date = datetime.now(datetime.utc).strftime("%Y-%m-%d") # Execute the indexer function documents_processed, error_message = await indexer_func( @@ -238,7 +238,7 @@ class ConnectorSchedulerService: ) if error_message: - await self.task_logger.log_task_failure( + await task_logger.log_task_failure( log_entry, f"Scheduled sync failed for {connector.name}", error_message, @@ -248,7 +248,7 @@ class ConnectorSchedulerService: f"Scheduled sync failed for connector {connector.name}: {error_message}" ) else: - await self.task_logger.log_task_success( + await task_logger.log_task_success( log_entry, f"Scheduled sync completed for {connector.name}", {"documents_processed": documents_processed}, @@ -279,7 +279,7 @@ class ConnectorSchedulerService: await session.execute( update(ConnectorSchedule) .filter(ConnectorSchedule.id == schedule_id) - .values(last_run_at=datetime.utcnow()) + .values(last_run_at=datetime.now(timezone.utc)) ) await session.commit() @@ -288,7 +288,12 @@ class ConnectorSchedulerService: ): """Update the next_run_at timestamp for a schedule.""" next_run = calculate_next_run( - schedule.schedule_type, schedule.cron_expression + schedule.schedule_type, + schedule.cron_expression, + schedule.daily_time, + schedule.weekly_day, + schedule.weekly_time, + schedule.hourly_minute ) await session.execute( @@ -302,7 +307,7 @@ class ConnectorSchedulerService: f"Updated next run time for schedule {schedule.id} to {next_run}" ) - async def get_scheduler_status(self) -> Dict: + async def get_scheduler_status(self) -> dict: """Get the current status of the scheduler service.""" return { "running": self.running,