237 lines
9.8 KiB
Python
237 lines
9.8 KiB
Python
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()
|
|
)
|