init

db_models.py

@@ -0,0 +1,81 @@
+import datetime
+from collections.abc import Sequence
+from dataclasses import dataclass
+
+from sqlalchemy import Date, DateTime, Float, Integer, String, UniqueConstraint, select
+from sqlalchemy.orm import DeclarativeBase, Mapped, Session, mapped_column
+from sqlalchemy.sql import func
+
+
+class Base(DeclarativeBase):
+    pass
+
+
+class WeatherForecast(Base):
+    __tablename__ = "weather_forecast"
+    __table_args__ = (
+        UniqueConstraint("city", "forecast_date", name="unique_city_forecast_date"),
+    )
+
+    id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
+    city: Mapped[str] = mapped_column(String, nullable=False)
+    forecast_date: Mapped[datetime.date] = mapped_column(Date, nullable=False)
+    temp_min: Mapped[float] = mapped_column(Float, nullable=False)
+    temp_max: Mapped[float] = mapped_column(Float, nullable=False)
+    humidity: Mapped[int] = mapped_column(Integer, nullable=False)
+    wind_speed: Mapped[float] = mapped_column(Float, nullable=False)
+    description: Mapped[str | None] = mapped_column(String, nullable=True)
+    created_at: Mapped[datetime.datetime] = mapped_column(
+        DateTime, default=func.now(), nullable=False
+    )
+
+    def __repr__(self) -> str:
+        return (
+            f"WeatherForecast(id={self.id!r}, city={self.city!r}, "
+            f"forecast_date={self.forecast_date!r}, temp_min={self.temp_min!r}, "
+            f"temp_max={self.temp_max!r}, humidity={self.humidity!r}, "
+            f"wind_speed={self.wind_speed!r}, description={self.description!r})"
+        )
+
+
+@dataclass(frozen=True)
+class DailyForecast:
+    city: str
+    forecast_date: datetime.date
+    temp_min: float
+    temp_max: float
+    humidity: int
+    wind_speed: float
+    description: str
+
+
+def save_forecasts(session: Session, forecasts: Sequence[DailyForecast]) -> int:
+    if not forecasts:
+        return 0
+
+    known: dict[str, set[datetime.date]] = {}
+    for city in {f.city for f in forecasts}:
+        known[city] = set(
+            session.scalars(
+                select(WeatherForecast.forecast_date).where(
+                    WeatherForecast.city == city
+                )
+            ).all()
+        )
+
+    add = [
+        WeatherForecast(
+            city=f.city,
+            forecast_date=f.forecast_date,
+            temp_min=f.temp_min,
+            temp_max=f.temp_max,
+            humidity=f.humidity,
+            wind_speed=f.wind_speed,
+            description=f.description,
+        )
+        for f in forecasts
+        if f.forecast_date not in known[f.city]
+    ]
+    session.add_all(add)
+    session.commit()
+    return len(add)