121 lines
5.2 KiB
Python
121 lines
5.2 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 = 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))
|