184 lines
6.3 KiB
Python
184 lines
6.3 KiB
Python
"""Persistence operations for celestial position data."""
|
|
|
|
from datetime import datetime
|
|
from typing import Any, Dict, List, Optional
|
|
|
|
from sqlalchemy import and_, delete, func, select
|
|
from sqlalchemy.dialects.postgresql import insert
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.database import AsyncSessionLocal
|
|
from app.models.db import Position
|
|
|
|
|
|
class PositionService:
|
|
"""Store and query time-series position data."""
|
|
|
|
@staticmethod
|
|
async def save_positions(
|
|
body_id: str,
|
|
positions: List[Dict[str, Any]],
|
|
source: str = "nasa_horizons",
|
|
session: Optional[AsyncSession] = None,
|
|
) -> int:
|
|
async def save(active_session: AsyncSession) -> int:
|
|
for position_data in positions:
|
|
statement = insert(Position).values(
|
|
body_id=body_id,
|
|
time=position_data["time"],
|
|
x=position_data["x"],
|
|
y=position_data["y"],
|
|
z=position_data["z"],
|
|
vx=position_data.get("vx"),
|
|
vy=position_data.get("vy"),
|
|
vz=position_data.get("vz"),
|
|
source=source,
|
|
)
|
|
statement = statement.on_conflict_do_update(
|
|
index_elements=["body_id", "time"],
|
|
set_={
|
|
"x": position_data["x"],
|
|
"y": position_data["y"],
|
|
"z": position_data["z"],
|
|
"vx": position_data.get("vx"),
|
|
"vy": position_data.get("vy"),
|
|
"vz": position_data.get("vz"),
|
|
"source": source,
|
|
},
|
|
)
|
|
await active_session.execute(statement)
|
|
|
|
await active_session.commit()
|
|
return len(positions)
|
|
|
|
if session is not None:
|
|
return await save(session)
|
|
async with AsyncSessionLocal() as active_session:
|
|
return await save(active_session)
|
|
|
|
@staticmethod
|
|
async def get_positions(
|
|
body_id: str,
|
|
start_time: Optional[datetime] = None,
|
|
end_time: Optional[datetime] = None,
|
|
session: Optional[AsyncSession] = None,
|
|
) -> List[Position]:
|
|
async def query(active_session: AsyncSession) -> List[Position]:
|
|
statement = select(Position).where(Position.body_id == body_id)
|
|
if start_time and end_time:
|
|
statement = statement.where(
|
|
and_(Position.time >= start_time, Position.time <= end_time)
|
|
)
|
|
elif start_time:
|
|
statement = statement.where(Position.time >= start_time)
|
|
elif end_time:
|
|
statement = statement.where(Position.time <= end_time)
|
|
|
|
result = await active_session.execute(statement.order_by(Position.time))
|
|
return list(result.scalars().all())
|
|
|
|
if session is not None:
|
|
return await query(session)
|
|
async with AsyncSessionLocal() as active_session:
|
|
return await query(active_session)
|
|
|
|
@staticmethod
|
|
async def get_positions_in_range(
|
|
body_id: str,
|
|
start_time: datetime,
|
|
end_time: datetime,
|
|
session: Optional[AsyncSession] = None,
|
|
) -> List[Position]:
|
|
return await PositionService.get_positions(
|
|
body_id, start_time, end_time, session
|
|
)
|
|
|
|
@staticmethod
|
|
async def save_position(
|
|
body_id: str,
|
|
time: datetime,
|
|
x: float,
|
|
y: float,
|
|
z: float,
|
|
source: str = "nasa_horizons",
|
|
vx: Optional[float] = None,
|
|
vy: Optional[float] = None,
|
|
vz: Optional[float] = None,
|
|
session: Optional[AsyncSession] = None,
|
|
) -> Position:
|
|
async def save(active_session: AsyncSession) -> Position:
|
|
result = await active_session.execute(
|
|
select(Position).where(
|
|
and_(Position.body_id == body_id, Position.time == time)
|
|
)
|
|
)
|
|
position = result.scalar_one_or_none()
|
|
if position is None:
|
|
position = Position(body_id=body_id, time=time, x=x, y=y, z=z)
|
|
active_session.add(position)
|
|
|
|
position.x = x
|
|
position.y = y
|
|
position.z = z
|
|
position.vx = vx
|
|
position.vy = vy
|
|
position.vz = vz
|
|
position.source = source
|
|
await active_session.commit()
|
|
await active_session.refresh(position)
|
|
return position
|
|
|
|
if session is not None:
|
|
return await save(session)
|
|
async with AsyncSessionLocal() as active_session:
|
|
return await save(active_session)
|
|
|
|
@staticmethod
|
|
async def delete_old_positions(
|
|
before_time: datetime,
|
|
session: Optional[AsyncSession] = None,
|
|
) -> int:
|
|
async def remove(active_session: AsyncSession) -> int:
|
|
result = await active_session.execute(
|
|
delete(Position).where(Position.time < before_time)
|
|
)
|
|
await active_session.commit()
|
|
return result.rowcount
|
|
|
|
if session is not None:
|
|
return await remove(session)
|
|
async with AsyncSessionLocal() as active_session:
|
|
return await remove(active_session)
|
|
|
|
@staticmethod
|
|
async def get_available_dates(
|
|
body_id: str,
|
|
start_time: datetime,
|
|
end_time: datetime,
|
|
session: Optional[AsyncSession] = None,
|
|
) -> List[datetime]:
|
|
async def query(active_session: AsyncSession) -> List[datetime]:
|
|
date_expression = func.date(Position.time)
|
|
statement = (
|
|
select(date_expression)
|
|
.where(
|
|
and_(
|
|
Position.body_id == body_id,
|
|
Position.time >= start_time,
|
|
Position.time <= end_time,
|
|
)
|
|
)
|
|
.distinct()
|
|
.order_by(date_expression)
|
|
)
|
|
result = await active_session.execute(statement)
|
|
return [row[0] for row in result]
|
|
|
|
if session is not None:
|
|
return await query(session)
|
|
async with AsyncSessionLocal() as active_session:
|
|
return await query(active_session)
|
|
|
|
|
|
position_service = PositionService()
|