| 1 | import asyncio |
| 2 | from typing import Optional |
| 3 | |
| 4 | from fastapi import APIRouter |
| 5 | from fastapi import Depends |
| 6 | from fastapi import HTTPException |
| 7 | from fastapi import Security |
| 8 | from loguru import logger |
| 9 | from sqlalchemy.ext.asyncio import AsyncSession |
| 10 | from sqlalchemy.future import select |
| 11 | |
| 12 | from app.auth.utils import AuthHandler |
| 13 | from app.db.db_session import get_db |
| 14 | from app.schedulers.models.scheduler import JobMetadata |
| 15 | from app.schedulers.scheduler import get_function_by_name |
| 16 | from app.schedulers.scheduler import get_scheduler_instance |
| 17 | from app.schedulers.scheduler import init_scheduler |
| 18 | from app.schedulers.schema.scheduler import JobsNextRunResponse |
| 19 | from app.schedulers.schema.scheduler import JobsResponse |
| 20 | |
| 21 | scheduler_router = APIRouter( |
| 22 | dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))], |
| 23 | ) |
| 24 | |
| 25 | |
| 26 | async def get_scheduler(): |
| 27 | # Singleton pattern or reference to existing instance |
| 28 | return await init_scheduler() |
| 29 | |
| 30 | |
| 31 | async def find_job_by_id(scheduler, job_id): |
| 32 | """ |
| 33 | Find a job in the scheduler by its ID. |
| 34 | |
| 35 | Args: |
| 36 | scheduler (Scheduler): The scheduler object. |
| 37 | job_id (str): The ID of the job to find. |
| 38 | |
| 39 | Returns: |
| 40 | Job: The job object if found, None otherwise. |
| 41 | """ |
| 42 | for job in scheduler.get_jobs(): |
| 43 | if job.id == job_id: |
| 44 | return job |
| 45 | return None |
| 46 | |
| 47 | |
| 48 | async def manage_job_metadata(session, job_id, action, **kwargs): |
| 49 | """ |
| 50 | Manage job metadata based on the specified action. |
| 51 | |
| 52 | Args: |
| 53 | session (Session): The database session. |
| 54 | job_id (int): The ID of the job. |
| 55 | action (str): The action to perform on the job metadata. Possible values are "update" and "delete". |
| 56 | **kwargs: Additional keyword arguments representing the fields to update and their new values. |
| 57 | |
| 58 | Returns: |
| 59 | JobMetadata: The updated or deleted job metadata. |
| 60 | """ |
| 61 | if action == "add": |
| 62 | job_metadata = JobMetadata( |
| 63 | job_id=job_id, |
| 64 | last_success=None, |
| 65 | time_interval=kwargs["time_interval"], |
| 66 | enabled=True, |
| 67 | extra_data=kwargs["extra_data"], |
| 68 | ) |
| 69 | session.add(job_metadata) |
| 70 | await session.commit() |
| 71 | return job_metadata |
| 72 | |
| 73 | job_metadata = await session.execute(select(JobMetadata).filter_by(job_id=job_id)) |
| 74 | job_metadata = job_metadata.scalars().first() |
| 75 | |
| 76 | if action == "update": |
| 77 | for key, value in kwargs.items(): |
| 78 | setattr(job_metadata, key, value) |
| 79 | await session.commit() |
| 80 | elif action == "delete": |
| 81 | await session.delete(job_metadata) |
| 82 | await session.commit() |
| 83 | |
| 84 | return job_metadata |
| 85 | |
| 86 | |
| 87 | @scheduler_router.get("", response_model=JobsResponse, description="Get all jobs") |
| 88 | async def get_all_jobs(session: AsyncSession = Depends(get_db)) -> JobsResponse: |
| 89 | """ |
| 90 | Retrieve all jobs from the scheduler. |
| 91 | |
| 92 | Args: |
| 93 | session (AsyncSession): The database session. |
| 94 | |
| 95 | Returns: |
| 96 | JobsResponse: The response containing the list of jobs. |
| 97 | |
| 98 | """ |
| 99 | scheduler = await get_scheduler_instance() |
| 100 | jobs = scheduler.get_jobs() |
| 101 | apscheduler_jobs = [] |
| 102 | for job in jobs: |
| 103 | job_metadata = await session.execute( |
| 104 | select(JobMetadata).filter_by(job_id=job.id), |
| 105 | ) |
| 106 | job_metadata = job_metadata.scalars().first() |
| 107 | logger.info(f"job_metadata: {job_metadata}") |
| 108 | apscheduler_jobs.append( |
| 109 | { |
| 110 | "id": job.id, |
| 111 | "name": job.name, |
| 112 | "time_interval": job_metadata.time_interval, |
| 113 | "enabled": job_metadata.enabled, |
| 114 | "description": job_metadata.job_description, |
| 115 | "last_success": job_metadata.last_success, |
| 116 | }, |
| 117 | ) |
| 118 | logger.info(f"apscheduler_jobs: {apscheduler_jobs}") |
| 119 | return JobsResponse( |
| 120 | jobs=apscheduler_jobs, |
| 121 | success=True, |
| 122 | message="Jobs successfully retrieved.", |
| 123 | ) |
| 124 | |
| 125 | |
| 126 | @scheduler_router.get("/next_run/{job_id}", response_model=JobsNextRunResponse, description="Get the next run time of a job") |
| 127 | async def get_next_run(job_id: str) -> JobsNextRunResponse: |
| 128 | """ |
| 129 | Get the next run time of a job. |
| 130 | |
| 131 | Args: |
| 132 | job_id (str): The ID of the job. |
| 133 | |
| 134 | Returns: |
| 135 | dict: A dictionary containing the next run time of the job. |
| 136 | """ |
| 137 | scheduler = await get_scheduler_instance() |
| 138 | job = await find_job_by_id(scheduler, job_id) |
| 139 | if job is None: |
| 140 | raise HTTPException(status_code=404, detail="Job not found") |
| 141 | next_run_time = job.next_run_time |
| 142 | logger.info(f"Next run time for job {job_id}: {next_run_time}") |
| 143 | return JobsNextRunResponse( |
| 144 | next_run_time=next_run_time, |
| 145 | success=True, |
| 146 | message="Next run time successfully retrieved.", |
| 147 | ) |
| 148 | |
| 149 | |
| 150 | @scheduler_router.post("/add", description="Add a job") |
| 151 | async def add_job( |
| 152 | job_id: str, |
| 153 | function_name: str, |
| 154 | time_interval: int, |
| 155 | extra_data: Optional[str] = None, |
| 156 | session: AsyncSession = Depends(get_db), |
| 157 | ): |
| 158 | """ |
| 159 | Add a job to the scheduler. |
| 160 | |
| 161 | Args: |
| 162 | job_id (str): The ID of the job. |
| 163 | function_name (str): The name of the function to be scheduled. |
| 164 | time_interval (int): The time interval for the job in minutes. |
| 165 | extra_data (str, optional): Additional data to be stored with the job metadata. Defaults to None. |
| 166 | session (AsyncSession, optional): The database session. Defaults to Depends(get_db). |
| 167 | |
| 168 | Returns: |
| 169 | dict: A dictionary containing the success status and a message. |
| 170 | """ |
| 171 | scheduler = await get_scheduler_instance() |
| 172 | job_function = get_function_by_name(function_name) |
| 173 | scheduler.add_job( |
| 174 | job_function, |
| 175 | "interval", |
| 176 | minutes=time_interval, |
| 177 | id=job_id, |
| 178 | replace_existing=True, |
| 179 | ) |
| 180 | await manage_job_metadata( |
| 181 | session, |
| 182 | job_id, |
| 183 | "add", |
| 184 | time_interval=time_interval, |
| 185 | extra_data=extra_data, |
| 186 | ) |
| 187 | if not scheduler.running: |
| 188 | scheduler.start() |
| 189 | logger.info(f"Job {job_id} added successfully") |
| 190 | return {"success": True, "message": "Job added successfully"} |
| 191 | |
| 192 | |
| 193 | @scheduler_router.post("/jobs/run/{job_id}", description="Run a job") |
| 194 | async def run_job_manually(job_id: str, session: AsyncSession = Depends(get_db)): |
| 195 | """ |
| 196 | Manually triggers a scheduled job for immediate execution. |
| 197 | |
| 198 | Args: |
| 199 | job_id (str): The identifier of the job to run. |
| 200 | session (AsyncSession): The database session dependency. |
| 201 | |
| 202 | Returns: |
| 203 | A JSON response with the result of the operation. |
| 204 | """ |
| 205 | scheduler = await get_scheduler_instance() # Make sure your scheduler is properly initialized |
| 206 | job = scheduler.get_job(job_id) |
| 207 | |
| 208 | if job is None: |
| 209 | raise HTTPException(status_code=404, detail="Job not found") |
| 210 | |
| 211 | try: |
| 212 | # Retrieve the function associated with the job and run it |
| 213 | job_function = get_function_by_name(job.name) # Ensure this function maps job names to function objects |
| 214 | logger.info(f"Running job {job_id} manually") |
| 215 | if asyncio.iscoroutinefunction(job_function): |
| 216 | logger.info(f"Running async job {job_id}") |
| 217 | result = await job_function() # Execute the function if it's async |
| 218 | else: |
| 219 | logger.info(f"Running sync job {job_id}") |
| 220 | result = job_function() # Execute synchronously if not an async function |
| 221 | |
| 222 | return {"success": True, "message": "Job executed successfully", "result": result} |
| 223 | except Exception as e: |
| 224 | raise HTTPException(status_code=500, detail=str(e)) |
| 225 | |
| 226 | |
| 227 | @scheduler_router.post("/start/{job_id}", description="Start a job") |
| 228 | async def start_job(job_id: str): |
| 229 | """ |
| 230 | Start a job by resuming its execution. |
| 231 | |
| 232 | Args: |
| 233 | job_id (str): The ID of the job to start. |
| 234 | |
| 235 | Returns: |
| 236 | dict: A dictionary containing the success status and a message. |
| 237 | - If the job is found and successfully started, the success status is True and the message is "Job started successfully". |
| 238 | - If the job is not found, the success status is False and the message is "Job not found". |
| 239 | """ |
| 240 | scheduler = await get_scheduler_instance() |
| 241 | job = await find_job_by_id(scheduler, job_id) |
| 242 | if job: |
| 243 | job.resume() |
| 244 | logger.info(f"Job {job_id} started successfully") |
| 245 | return {"success": True, "message": "Job started successfully"} |
| 246 | logger.error(f"Job {job_id} not found for starting") |
| 247 | return {"success": False, "message": "Job not found"} |
| 248 | |
| 249 | |
| 250 | @scheduler_router.post("/pause/{job_id}", description="Pause a job") |
| 251 | async def pause_job(job_id: str): |
| 252 | """ |
| 253 | Pause a job. |
| 254 | |
| 255 | Args: |
| 256 | job_id (str): The ID of the job to be paused. |
| 257 | |
| 258 | Returns: |
| 259 | dict: A dictionary containing the success status and a message. |
| 260 | - If the job is paused successfully, the success status is True and the message is "Job paused successfully". |
| 261 | - If the job is not found, the success status is False and the message is "Job not found". |
| 262 | """ |
| 263 | scheduler = await get_scheduler_instance() |
| 264 | job = await find_job_by_id(scheduler, job_id) |
| 265 | if job: |
| 266 | job.pause() |
| 267 | logger.info(f"Job {job_id} paused successfully") |
| 268 | return {"success": True, "message": "Job paused successfully"} |
| 269 | logger.error(f"Job {job_id} not found for pausing") |
| 270 | return {"success": False, "message": "Job not found"} |
| 271 | |
| 272 | |
| 273 | @scheduler_router.put("/update/{job_id}", description="Update a job") |
| 274 | async def update_job( |
| 275 | job_id: str, |
| 276 | time_interval: int, |
| 277 | extra_data: Optional[str] = None, |
| 278 | session: AsyncSession = Depends(get_db), |
| 279 | ): |
| 280 | """ |
| 281 | Update a job with the specified job_id and time_interval. |
| 282 | |
| 283 | Parameters: |
| 284 | - job_id (str): The ID of the job to be updated. |
| 285 | - time_interval (int): The new time interval for the job in minutes. |
| 286 | - session (AsyncSession, optional): The database session. Defaults to Depends(get_db). |
| 287 | |
| 288 | Returns: |
| 289 | - dict: A dictionary containing the success status and a message. |
| 290 | |
| 291 | Example: |
| 292 | { |
| 293 | "success": True, |
| 294 | "message": "Job updated successfully" |
| 295 | } |
| 296 | """ |
| 297 | scheduler = await get_scheduler_instance() |
| 298 | job = await find_job_by_id(scheduler, job_id) |
| 299 | if job: |
| 300 | job.reschedule(trigger="interval", minutes=time_interval) |
| 301 | await manage_job_metadata( |
| 302 | session, |
| 303 | job_id, |
| 304 | "update", |
| 305 | time_interval=time_interval, |
| 306 | extra_data=extra_data, |
| 307 | ) |
| 308 | logger.info(f"Job {job_id} updated successfully") |
| 309 | # Update the job metadata |
| 310 | await manage_job_metadata( |
| 311 | session, |
| 312 | job_id, |
| 313 | "update", |
| 314 | time_interval=time_interval, |
| 315 | extra_data=extra_data, |
| 316 | ) |
| 317 | |
| 318 | return {"success": True, "message": "Job updated successfully"} |
| 319 | logger.error(f"Job {job_id} not found for updating") |
| 320 | return {"success": False, "message": "Job not found"} |
| 321 | |
| 322 | |
| 323 | @scheduler_router.delete("/{job_id}", description="Delete a job") |
| 324 | async def delete_job(job_id: str, session: AsyncSession = Depends(get_db)): |
| 325 | """ |
| 326 | Delete a job. |
| 327 | |
| 328 | Args: |
| 329 | job_id (str): The ID of the job to be deleted. |
| 330 | session (AsyncSession, optional): The database session. Defaults to Depends(get_db). |
| 331 | |
| 332 | Returns: |
| 333 | dict: A dictionary containing the success status and a message. |
| 334 | """ |
| 335 | scheduler = await get_scheduler_instance() |
| 336 | job = await find_job_by_id(scheduler, job_id) |
| 337 | if job: |
| 338 | scheduler.remove_job(job_id) |
| 339 | await manage_job_metadata(session, job_id, "delete") |
| 340 | logger.info(f"Job {job_id} deleted successfully") |
| 341 | return {"success": True, "message": "Job deleted successfully"} |
| 342 | logger.error(f"Job {job_id} not found for deletion") |
| 343 | return {"success": False, "message": "Job not found"} |