cosmo/backend/app/services/nasa_worker.py

122 lines
5.3 KiB
Python

import logging
import asyncio
from datetime import datetime
from sqlalchemy.ext.asyncio import AsyncSession
from typing import List
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
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,
db=db,
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))