Merge pull request #840 from airweave-ai/fix/unified-temporal-scheduling

fix: unified temporal scheduling

Author: orhanrauf

.vscode/settings.json

@@ -9,6 +9,15 @@
         "${workspaceFolder}/backend"
     ],
     "python.defaultInterpreterPath": "${command:python.interpreterPath}",
+    "terminal.integrated.env.osx": {
+        "PYTHONPATH": "${workspaceFolder}/backend:${env:PYTHONPATH}"
+    },
+    "terminal.integrated.env.linux": {
+        "PYTHONPATH": "${workspaceFolder}/backend:${env:PYTHONPATH}"
+    },
+    "terminal.integrated.env.windows": {
+        "PYTHONPATH": "${workspaceFolder}/backend;${env:PYTHONPATH}"
+    },
     "python.testing.pytestPath": "${command:python.interpreterPath}",
     "python.testing.autoTestDiscoverOnSaveEnabled": true,
     "python.debugOptions": [

backend/airweave/core/source_connection_service.py

@@ -56,6 +56,68 @@ def _get_default_cron_schedule(self, ctx: ApiContext) -> str:
         )
         return cron_schedule
 
+    def _validate_cron_schedule_for_source(
+        self, cron_schedule: str, source: schemas.Source, ctx: ApiContext
+    ) -> None:
+        """Validate CRON schedule based on source capabilities.
+
+        Args:
+            cron_schedule: The CRON expression to validate
+            source: The source model
+            ctx: API context
+
+        Raises:
+            HTTPException: If the schedule is invalid for the source
+        """
+        import re
+
+        if not cron_schedule:
+            return
+
+        # Parse the CRON expression to check if it's minute-level
+        # We need to distinguish between:
+        # - "*/N * * * *" where N < 60 - runs every N minutes (minute-level)
+        # - "* * * * *" - runs every minute (minute-level)
+        # - "0 * * * *" - runs at minute 0 of every hour (hourly, not minute-level)
+        # - "30 2 * * *" - runs at 2:30 AM daily (daily, not minute-level)
+
+        # Check for patterns that run more frequently than hourly
+        # Pattern 1: */N where N < 60 (e.g., */5, */15, */30)
+        interval_pattern = r"^\*/([1-5]?[0-9]) \* \* \* \*$"
+        match = re.match(interval_pattern, cron_schedule)
+
+        if match:
+            interval = int(match.group(1))
+            if interval < 60:
+                # This is sub-hourly (minute-level)
+                if not source.supports_continuous:
+                    raise HTTPException(
+                        status_code=400,
+                        detail=f"Source '{source.short_name}' does not support continuous syncs. "
+                        f"Minimum schedule interval is 1 hour (e.g., '0 * * * *' for hourly).",
+                    )
+                # For continuous sources, sub-hourly is allowed
+                ctx.logger.info(
+                    f"Source '{source.short_name}' supports continuous syncs, "
+                    f"allowing minute-level schedule: {cron_schedule}"
+                )
+                return
+
+        # Pattern 2: * * * * * (every minute)
+        if cron_schedule == "* * * * *":
+            if not source.supports_continuous:
+                raise HTTPException(
+                    status_code=400,
+                    detail=f"Source '{source.short_name}' does not support continuous syncs. "
+                    f"Minimum schedule interval is 1 hour (e.g., '0 * * * *' for hourly).",
+                )
+            ctx.logger.info(
+                f"Source '{source.short_name}' supports continuous syncs, "
+                f"allowing every-minute schedule: {cron_schedule}"
+            )
+
+        # All other patterns (including "0 * * * *" for hourly) are allowed
+
     """Clean service with automatic auth method inference.
 
     Key improvements:
@@ -225,11 +287,15 @@ async def update(
         ctx: ApiContext,
     ) -> SourceConnection:
         """Update a source connection."""
+        # First check if the source connection exists
         source_conn = await crud.source_connection.get(db, id=id, ctx=ctx)
         if not source_conn:
             raise HTTPException(status_code=404, detail="Source connection not found")
 
         async with UnitOfWork(db) as uow:
+            # Re-fetch the source_conn within the UoW session to avoid session mismatch
+            source_conn = await crud.source_connection.get(uow.session, id=id, ctx=ctx)
+
             # Update fields
             update_data = obj_in.model_dump(exclude_unset=True)
 
@@ -242,20 +308,32 @@ async def update(
                 del update_data["config"]
 
             # Handle schedule update
-            if "schedule" in update_data and update_data["schedule"]:
+            if (
+                "schedule" in update_data and update_data["schedule"]
+            ):  # TODO: only if actually different
                 if source_conn.sync_id:
+                    new_cron = update_data["schedule"].get("cron")
+                    if new_cron:
+                        # Get the source to validate schedule
+                        source = await self._get_and_validate_source(
+                            uow.session, source_conn.short_name
+                        )
+                        self._validate_cron_schedule_for_source(new_cron, source, ctx)
                     await self._update_sync_schedule(
                         uow.session,
                         source_conn.sync_id,
-                        update_data["schedule"].get("cron"),
+                        new_cron,
                         ctx,
                         uow,
                     )
                     del update_data["schedule"]
 
             # Handle credential update (direct auth only)
             if "credentials" in update_data:
-                auth_method = self._determine_auth_method(source_conn)
+                # Use the schema function that works with database models
+                from airweave.schemas.source_connection import determine_auth_method
+
+                auth_method = determine_auth_method(source_conn)
                 if auth_method != AuthenticationMethod.DIRECT:
                     raise HTTPException(
                         status_code=400,
@@ -365,6 +443,9 @@ async def _create_with_direct_auth(
             if cron_schedule is None:
                 cron_schedule = self._get_default_cron_schedule(ctx)
 
+            # Validate the schedule for this source
+            self._validate_cron_schedule_for_source(cron_schedule, source, ctx)
+
             sync, sync_job = await self._create_sync_without_schedule(
                 uow.session,
                 obj_in.name,
@@ -400,6 +481,7 @@ async def _create_with_direct_auth(
                     cron_schedule=cron_schedule,
                     db=uow.session,
                     ctx=ctx,
+                    uow=uow,
                 )
 
             # Convert to schemas while still in session
@@ -582,6 +664,9 @@ async def _create_with_oauth_token(
             if cron_schedule is None:
                 cron_schedule = self._get_default_cron_schedule(ctx)
 
+            # Validate the schedule for this source
+            self._validate_cron_schedule_for_source(cron_schedule, source, ctx)
+
             sync, sync_job = await self._create_sync_without_schedule(
                 uow.session,
                 obj_in.name,
@@ -617,6 +702,7 @@ async def _create_with_oauth_token(
                     cron_schedule=cron_schedule,
                     db=uow.session,
                     ctx=ctx,
+                    uow=uow,
                 )
 
             # Convert to schemas while still in session
@@ -728,6 +814,9 @@ async def _create_with_auth_provider(
             if cron_schedule is None:
                 cron_schedule = self._get_default_cron_schedule(ctx)
 
+            # Validate the schedule for this source
+            self._validate_cron_schedule_for_source(cron_schedule, source, ctx)
+
             sync, sync_job = await self._create_sync_without_schedule(
                 uow.session,
                 obj_in.name,
@@ -765,6 +854,7 @@ async def _create_with_auth_provider(
                     cron_schedule=cron_schedule,
                     db=uow.session,
                     ctx=ctx,
+                    uow=uow,
                 )
 
             # Convert to schemas while still in session

backend/airweave/core/source_connection_service_helpers.py

@@ -800,9 +800,9 @@ async def update_sync_schedule(
         self,
         db: AsyncSession,
         sync_id: UUID,
-        cron_schedule: str,
+        cron_schedule: Optional[str],
         ctx: ApiContext,
-        uow: Any,
+        uow: UnitOfWork,
     ) -> None:
         """Update sync schedule in database and Temporal."""
         sync = await crud.sync.get(db, id=sync_id, ctx=ctx)
@@ -811,15 +811,23 @@ async def update_sync_schedule(
             sync_update = schemas.SyncUpdate(cron_schedule=cron_schedule)
             await crud.sync.update(db, db_obj=sync, obj_in=sync_update, ctx=ctx, uow=uow)
 
-            # Update in Temporal if a schedule exists
+            # Update in Temporal
             from airweave.platform.temporal.schedule_service import temporal_schedule_service
 
-            await temporal_schedule_service.create_or_update_schedule(
-                sync_id=sync_id,
-                cron_schedule=cron_schedule,
-                db=db,
-                ctx=ctx,
-            )
+            if cron_schedule is None:
+                # If cron_schedule is None, delete the Temporal schedule
+                await temporal_schedule_service.delete_all_schedules_for_sync(
+                    sync_id=sync_id, db=db, ctx=ctx
+                )
+            else:
+                # Otherwise create or update the schedule
+                await temporal_schedule_service.create_or_update_schedule(
+                    sync_id=sync_id,
+                    cron_schedule=cron_schedule,
+                    db=db,
+                    ctx=ctx,
+                    uow=uow,
+                )
 
     async def update_auth_fields(
         self,
@@ -1184,6 +1192,7 @@ async def complete_oauth_connection(
                     cron_schedule=cron_schedule,
                     db=uow.session,
                     ctx=ctx,
+                    uow=uow,
                 )
 
             # Mark init session complete

backend/airweave/core/sync_service.py

@@ -84,6 +84,7 @@ async def _create_and_run_with_uow(
                 cron_schedule=sync_in.cron_schedule,
                 db=db,
                 ctx=ctx,
+                uow=uow,
             )
 
         # Run immediately if requested

backend/airweave/main.py

@@ -4,6 +4,7 @@
 and unhandled exceptions.
 """
 
+import os
 import subprocess
 from contextlib import asynccontextmanager
 
@@ -45,12 +46,22 @@
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
-    """Lifespan context manager for startup and shutdown events."""
-    # Startup
+    """Lifespan context manager for startup and shutdown events.
+
+    Runs alembic migrations and syncs platform components.
+    """
     async with AsyncSessionLocal() as db:
         if settings.RUN_ALEMBIC_MIGRATIONS:
             logger.info("Running alembic migrations...")
-            subprocess.run(["alembic", "upgrade", "head"], check=True)
+            env = os.environ.copy()
+            backend_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+            env["PYTHONPATH"] = backend_dir
+            subprocess.run(
+                ["alembic", "upgrade", "head"],
+                check=True,
+                cwd=backend_dir,
+                env=env,
+            )
         if settings.RUN_DB_SYNC:
             # Ensure all FileEntity subclasses have their parent and chunk models created
             ensure_file_entity_models()

backend/airweave/models/source.py

@@ -33,6 +33,9 @@ class Source(Base):
     output_entity_definition_ids = Column(JSON, nullable=False)
     config_schema = Column(JSON, nullable=True)  # JSON Schema for configuration
     labels: Mapped[list[str]] = mapped_column(JSON, nullable=True, default=list)
+    supports_continuous: Mapped[bool] = mapped_column(
+        Boolean, default=False, nullable=False, server_default="false"
+    )
 
     __table_args__ = (UniqueConstraint("name", "organization_id", name="uq_source_name_org"),)
 

backend/airweave/models/sync.py

@@ -33,7 +33,6 @@ class Sync(OrganizationBase, UserMixin):
     )
     temporal_schedule_id: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)
     sync_type: Mapped[str] = mapped_column(String(50), default="full")
-    minute_level_cron_schedule: Mapped[Optional[str]] = mapped_column(String(100), nullable=True)
     sync_metadata: Mapped[Optional[dict]] = mapped_column(JSON, nullable=True)
 
     jobs: Mapped[list["SyncJob"]] = relationship(
@@ -95,8 +94,9 @@ def delete_temporal_schedules_after_sync_delete(mapper, connection, target):
         import asyncio
 
         schedule_ids = [
-            f"minute-sync-{target.id}",
-            f"daily-cleanup-{target.id}",
+            f"sync-{target.id}",  # Regular schedule
+            f"minute-sync-{target.id}",  # Minute-level schedule
+            f"daily-cleanup-{target.id}",  # Daily cleanup schedule
         ]
 
         async def _cleanup():

backend/airweave/platform/db_sync.py

@@ -374,6 +374,7 @@ async def _sync_sources(
             class_name=source_class.__name__,
             output_entity_definition_ids=output_entity_ids,
             labels=getattr(source_class, "_labels", []),
+            supports_continuous=getattr(source_class, "_supports_continuous", False),
         )
         source_definitions.append(source_def)
 

backend/airweave/platform/decorators.py

@@ -18,6 +18,7 @@ def source(
     auth_config_class: Optional[Type[BaseModel]] = None,
     config_class: Optional[Type[BaseModel]] = None,
     labels: Optional[List[str]] = None,
+    supports_continuous: bool = False,
 ) -> Callable[[type], type]:
     """Enhanced source decorator with OAuth type tracking.
 
@@ -30,6 +31,7 @@ def source(
         auth_config_class: Pydantic model for auth configuration (for DIRECT auth only)
         config_class: Pydantic model for source configuration
         labels: Tags for categorization (e.g., "CRM", "Database")
+        supports_continuous: Whether source supports cursor-based continuous syncing (default False)
 
     Example:
         # OAuth source (no auth config)
@@ -66,6 +68,7 @@ def decorator(cls: type) -> type:
         cls._auth_config_class = auth_config_class
         cls._config_class = config_class
         cls._labels = labels or []
+        cls._supports_continuous = supports_continuous
 
         # Add validation method if not present
         if not hasattr(cls, "validate"):

backend/airweave/platform/sources/asana.py

@@ -32,6 +32,7 @@
     auth_config_class=None,
     config_class="AsanaConfig",
     labels=["Project Management"],
+    supports_continuous=False,
 )
 class AsanaSource(BaseSource):
     """Asana source connector integrates with the Asana API to extract and synchronize data.