Model/repositories/solar/solar_postgres_repository.py

from __future__ import annotations

from typing import Any, Optional

from sqlmodel import Session, select

from infrastructure.postgres.solar_table import SolarRow
from repositories.solar.solar_repository import SolarRepository


class SolarPostgresRepository(SolarRepository):
    def __init__(self, session: Session) -> None:
        self._session = session

    def save(
        self, uprn: int, *, longitude: float, latitude: float, insights: dict[str, Any]
    ) -> None:
        existing = self._session.exec(
            select(SolarRow).where(SolarRow.uprn == uprn)
        ).first()
        if existing is None:
            self._session.add(
                SolarRow(
                    uprn=uprn,
                    longitude=longitude,
                    latitude=latitude,
                    google_api_response=insights,
                )
            )
        else:
            existing.longitude = longitude
            existing.latitude = latitude
            existing.google_api_response = insights
            self._session.add(existing)

    def get(self, uprn: int) -> Optional[dict[str, Any]]:
        row = self._session.exec(
            select(SolarRow).where(SolarRow.uprn == uprn)
        ).first()
        return row.google_api_response if row is not None else None