import logging import asyncio from datetime import datetime from sqlalchemy.ext.asyncio import AsyncSession from typing import List, Optional from app.database import AsyncSessionLocal from app.services.task_service import task_service from app.services.db_service import celestial_body_service, position_service from app.services.horizons import horizons_service from app.services.orbit_service import orbit_service logger = logging.getLogger(__name__) async def download_positions_task(task_id: int, body_ids: List[str], dates: List[str]): """ Background task worker for downloading NASA positions """ logger.info(f"Task {task_id}: Starting download for {len(body_ids)} bodies and {len(dates)} dates") async with AsyncSessionLocal() as db: try: # Mark as running await task_service.update_progress(db, task_id, 0, "running") total_operations = len(body_ids) * len(dates) current_op = 0 success_count = 0 failed_count = 0 results = [] for body_id in body_ids: # Check body body = await celestial_body_service.get_body_by_id(body_id, db) if not body: results.append({"body_id": body_id, "error": "Body not found"}) failed_count += len(dates) current_op += len(dates) continue body_result = { "body_id": body_id, "body_name": body.name, "dates": [] } for date_str in dates: try: target_date = datetime.strptime(date_str, "%Y-%m-%d") # Check existing existing = await position_service.get_positions( body_id=body_id, start_time=target_date, end_time=target_date.replace(hour=23, minute=59, second=59), session=db ) if existing and len(existing) > 0: body_result["dates"].append({"date": date_str, "status": "skipped"}) success_count += 1 else: # Download positions = await horizons_service.get_body_positions( body_id=body_id, start_time=target_date, end_time=target_date, step="1d" ) if positions and len(positions) > 0: pos_data = [{ "time": target_date, "x": positions[0].x, "y": positions[0].y, "z": positions[0].z, "vx": getattr(positions[0], 'vx', None), "vy": getattr(positions[0], 'vy', None), "vz": getattr(positions[0], 'vz', None), }] await position_service.save_positions( body_id=body_id, positions=pos_data, source="nasa_horizons", session=db ) body_result["dates"].append({"date": date_str, "status": "success"}) success_count += 1 else: body_result["dates"].append({"date": date_str, "status": "failed", "error": "No data"}) failed_count += 1 # Sleep slightly to prevent rate limiting and allow context switching # await asyncio.sleep(0.1) except Exception as e: logger.error(f"Error processing {body_id} on {date_str}: {e}") body_result["dates"].append({"date": date_str, "status": "error", "error": str(e)}) failed_count += 1 # Update progress current_op += 1 progress = int((current_op / total_operations) * 100) # Only update DB every 5% or so to reduce load, but update Redis frequently # For now, update every item for simplicity await task_service.update_progress(db, task_id, progress) results.append(body_result) # Complete final_result = { "total_success": success_count, "total_failed": failed_count, "details": results } await task_service.complete_task(db, task_id, final_result) logger.info(f"Task {task_id} completed successfully") except Exception as e: logger.error(f"Task {task_id} failed critically: {e}") await task_service.fail_task(db, task_id, str(e)) async def generate_orbits_task(task_id: int, body_ids: Optional[List[str]] = None): """ Background task to generate orbits Args: task_id: ID of the task record to update body_ids: List of body IDs to generate. If None, generates for all bodies with orbital params. """ logger.info(f"🚀 Starting background orbit generation task {task_id}") async with AsyncSessionLocal() as db: try: # Update task to running await task_service.update_task( db, task_id, status="running", started_at=datetime.utcnow(), progress=0 ) # 1. Determine bodies to process bodies_to_process = [] if body_ids: # Fetch specific bodies requested for bid in body_ids: body = await celestial_body_service.get_body_by_id(bid, db) if body: bodies_to_process.append(body) else: # Fetch all bodies bodies_to_process = await celestial_body_service.get_all_bodies(db) # 2. Filter for valid orbital parameters valid_bodies = [] for body in bodies_to_process: extra_data = body.extra_data or {} # Must have orbit_period_days to generate an orbit if extra_data.get("orbit_period_days"): valid_bodies.append(body) elif body_ids and body.id in body_ids: # If explicitly requested but missing period, log warning logger.warning(f"Body {body.name} ({body.id}) missing 'orbit_period_days', skipping.") total_bodies = len(valid_bodies) if total_bodies == 0: await task_service.update_task( db, task_id, status="completed", progress=100, result={"message": "No bodies with 'orbit_period_days' found to process"} ) return # 3. Process success_count = 0 failure_count = 0 results = [] for i, body in enumerate(valid_bodies): try: # Update progress progress = int((i / total_bodies) * 100) await task_service.update_task(db, task_id, progress=progress) extra_data = body.extra_data or {} period = float(extra_data.get("orbit_period_days")) color = extra_data.get("orbit_color", "#CCCCCC") # Generate orbit orbit = await orbit_service.generate_orbit( body_id=body.id, body_name=body.name_zh or body.name, period_days=period, color=color, session=db, horizons_service=horizons_service ) results.append({ "body_id": body.id, "body_name": body.name_zh or body.name, "status": "success", "num_points": orbit.num_points }) success_count += 1 except Exception as e: logger.error(f"Failed to generate orbit for {body.name}: {e}") results.append({ "body_id": body.id, "body_name": body.name_zh or body.name, "status": "failed", "error": str(e) }) failure_count += 1 # Finish task await task_service.update_task( db, task_id, status="completed", progress=100, completed_at=datetime.utcnow(), result={ "total": total_bodies, "success": success_count, "failed": failure_count, "details": results } ) logger.info(f"🏁 Orbit generation task {task_id} completed") except Exception as e: logger.error(f"Task {task_id} failed: {e}") await task_service.update_task( db, task_id, status="failed", error_message=str(e), completed_at=datetime.utcnow() )