from __future__ import annotations from contextlib import contextmanager from dataclasses import dataclass from datetime import datetime, timedelta, timezone from os import environ from typing import Iterator from gibil.classes.models import ( WeatherForecastPoint, WeatherForecastRun, WeatherResolvedTruth, WeatherStationObservation, ) from astrape_oracle.classes.datasets import ( WeatherDisplayDataset, WeatherStationDisplayDataset, ) class WeatherStoreConfigurationError(RuntimeError): pass @dataclass(frozen=True) class WeatherStoreConfig: database_url: str @classmethod def from_env(cls) -> "WeatherStoreConfig": database_url = environ.get("ASTRAPE_DATABASE_URL") if not database_url: raise WeatherStoreConfigurationError( "ASTRAPE_DATABASE_URL is required for weather storage" ) return cls(database_url=database_url) class WeatherStore: """Persists external weather forecasts and resolved truth in TimescaleDB.""" def __init__(self, config: WeatherStoreConfig) -> None: self.config = config @classmethod def from_env(cls) -> "WeatherStore": return cls(WeatherStoreConfig.from_env()) def initialize(self) -> None: with self._connection() as connection: with connection.cursor() as cursor: cursor.execute("CREATE EXTENSION IF NOT EXISTS timescaledb") cursor.execute( """ CREATE TABLE IF NOT EXISTS weather_forecast_points ( issued_at TIMESTAMPTZ NOT NULL, target_at TIMESTAMPTZ NOT NULL, horizon_hours INTEGER NOT NULL, source TEXT NOT NULL, latitude DOUBLE PRECISION NOT NULL, longitude DOUBLE PRECISION NOT NULL, temperature_c DOUBLE PRECISION, shortwave_radiation_w_m2 DOUBLE PRECISION, cloud_cover_pct DOUBLE PRECISION, inserted_at TIMESTAMPTZ NOT NULL DEFAULT now(), PRIMARY KEY (issued_at, target_at, source) ) """ ) cursor.execute( """ SELECT create_hypertable( 'weather_forecast_points', 'target_at', if_not_exists => TRUE ) """ ) cursor.execute( """ CREATE TABLE IF NOT EXISTS weather_resolved_truth ( resolved_at TIMESTAMPTZ NOT NULL, source TEXT NOT NULL, temperature_c DOUBLE PRECISION, shortwave_radiation_w_m2 DOUBLE PRECISION, cloud_cover_pct DOUBLE PRECISION, inserted_at TIMESTAMPTZ NOT NULL DEFAULT now(), PRIMARY KEY (resolved_at, source) ) """ ) cursor.execute( """ ALTER TABLE weather_resolved_truth ADD COLUMN IF NOT EXISTS cloud_cover_pct DOUBLE PRECISION """ ) cursor.execute( """ SELECT create_hypertable( 'weather_resolved_truth', 'resolved_at', if_not_exists => TRUE ) """ ) cursor.execute( """ CREATE TABLE IF NOT EXISTS weather_station_observations ( observed_at TIMESTAMPTZ NOT NULL, received_at TIMESTAMPTZ NOT NULL, source TEXT NOT NULL, temperature_c DOUBLE PRECISION, relative_humidity_pct DOUBLE PRECISION, dew_point_c DOUBLE PRECISION, pressure_hpa DOUBLE PRECISION, wind_speed_m_s DOUBLE PRECISION, wind_gust_m_s DOUBLE PRECISION, wind_max_m_s DOUBLE PRECISION, wind_direction_deg DOUBLE PRECISION, solar_radiation_w_m2 DOUBLE PRECISION, uv_index DOUBLE PRECISION, rain_rate_mm_h DOUBLE PRECISION, rain_event_mm DOUBLE PRECISION, rain_hourly_mm DOUBLE PRECISION, rain_daily_mm DOUBLE PRECISION, rain_weekly_mm DOUBLE PRECISION, rain_monthly_mm DOUBLE PRECISION, rain_yearly_mm DOUBLE PRECISION, indoor_temperature_c DOUBLE PRECISION, indoor_humidity_pct DOUBLE PRECISION, absolute_pressure_hpa DOUBLE PRECISION, relative_pressure_hpa DOUBLE PRECISION, metadata JSONB NOT NULL DEFAULT '{}'::jsonb, inserted_at TIMESTAMPTZ NOT NULL DEFAULT now(), PRIMARY KEY (observed_at, source) ) """ ) cursor.execute( """ SELECT create_hypertable( 'weather_station_observations', 'observed_at', if_not_exists => TRUE ) """ ) connection.commit() def save_forecast_run(self, forecast_run: WeatherForecastRun) -> int: rows = [ ( point.issued_at, point.target_at, point.horizon_hours, forecast_run.source, forecast_run.latitude, forecast_run.longitude, point.temperature_c, point.shortwave_radiation_w_m2, point.cloud_cover_pct, ) for point in forecast_run.points ] if not rows: return 0 with self._connection() as connection: with connection.cursor() as cursor: cursor.executemany( """ INSERT INTO weather_forecast_points ( issued_at, target_at, horizon_hours, source, latitude, longitude, temperature_c, shortwave_radiation_w_m2, cloud_cover_pct ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s) ON CONFLICT (issued_at, target_at, source) DO UPDATE SET horizon_hours = EXCLUDED.horizon_hours, latitude = EXCLUDED.latitude, longitude = EXCLUDED.longitude, temperature_c = EXCLUDED.temperature_c, shortwave_radiation_w_m2 = EXCLUDED.shortwave_radiation_w_m2, cloud_cover_pct = EXCLUDED.cloud_cover_pct, inserted_at = now() """, rows, ) connection.commit() return len(rows) def save_resolved_truth(self, truth_points: list[WeatherResolvedTruth]) -> int: rows = [ ( point.resolved_at, point.source, point.temperature_c, point.shortwave_radiation_w_m2, point.cloud_cover_pct, ) for point in truth_points ] if not rows: return 0 with self._connection() as connection: with connection.cursor() as cursor: cursor.executemany( """ INSERT INTO weather_resolved_truth ( resolved_at, source, temperature_c, shortwave_radiation_w_m2, cloud_cover_pct ) VALUES (%s, %s, %s, %s, %s) ON CONFLICT (resolved_at, source) DO UPDATE SET temperature_c = EXCLUDED.temperature_c, shortwave_radiation_w_m2 = EXCLUDED.shortwave_radiation_w_m2, cloud_cover_pct = EXCLUDED.cloud_cover_pct, inserted_at = now() """, rows, ) connection.commit() return len(rows) def save_zero_hour_forecast_as_truth( self, forecast_run: WeatherForecastRun ) -> int: truth_points = [ WeatherResolvedTruth( resolved_at=point.issued_at, source="open_meteo_zero_hour", temperature_c=point.temperature_c, shortwave_radiation_w_m2=point.shortwave_radiation_w_m2, cloud_cover_pct=point.cloud_cover_pct, ) for point in forecast_run.points if point.horizon_hours == 0 ] return self.save_resolved_truth(truth_points) def save_station_observation( self, observation: WeatherStationObservation, ) -> int: with self._connection() as connection: with connection.cursor() as cursor: try: from psycopg.types.json import Jsonb except ImportError as error: raise WeatherStoreConfigurationError( "Install dependencies with `python3 -m pip install -r requirements.txt`" ) from error cursor.execute( """ INSERT INTO weather_station_observations ( observed_at, received_at, source, temperature_c, relative_humidity_pct, dew_point_c, pressure_hpa, wind_speed_m_s, wind_gust_m_s, wind_max_m_s, wind_direction_deg, solar_radiation_w_m2, uv_index, rain_rate_mm_h, rain_event_mm, rain_hourly_mm, rain_daily_mm, rain_weekly_mm, rain_monthly_mm, rain_yearly_mm, indoor_temperature_c, indoor_humidity_pct, absolute_pressure_hpa, relative_pressure_hpa, metadata ) VALUES ( %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s ) ON CONFLICT (observed_at, source) DO UPDATE SET received_at = EXCLUDED.received_at, temperature_c = EXCLUDED.temperature_c, relative_humidity_pct = EXCLUDED.relative_humidity_pct, dew_point_c = EXCLUDED.dew_point_c, pressure_hpa = EXCLUDED.pressure_hpa, wind_speed_m_s = EXCLUDED.wind_speed_m_s, wind_gust_m_s = EXCLUDED.wind_gust_m_s, wind_max_m_s = EXCLUDED.wind_max_m_s, wind_direction_deg = EXCLUDED.wind_direction_deg, solar_radiation_w_m2 = EXCLUDED.solar_radiation_w_m2, uv_index = EXCLUDED.uv_index, rain_rate_mm_h = EXCLUDED.rain_rate_mm_h, rain_event_mm = EXCLUDED.rain_event_mm, rain_hourly_mm = EXCLUDED.rain_hourly_mm, rain_daily_mm = EXCLUDED.rain_daily_mm, rain_weekly_mm = EXCLUDED.rain_weekly_mm, rain_monthly_mm = EXCLUDED.rain_monthly_mm, rain_yearly_mm = EXCLUDED.rain_yearly_mm, indoor_temperature_c = EXCLUDED.indoor_temperature_c, indoor_humidity_pct = EXCLUDED.indoor_humidity_pct, absolute_pressure_hpa = EXCLUDED.absolute_pressure_hpa, relative_pressure_hpa = EXCLUDED.relative_pressure_hpa, metadata = EXCLUDED.metadata, inserted_at = now() """, ( observation.observed_at, observation.received_at, observation.source, observation.temperature_c, observation.relative_humidity_pct, observation.dew_point_c, observation.pressure_hpa, observation.wind_speed_m_s, observation.wind_gust_m_s, observation.wind_max_m_s, observation.wind_direction_deg, observation.solar_radiation_w_m2, observation.uv_index, observation.rain_rate_mm_h, observation.rain_event_mm, observation.rain_hourly_mm, observation.rain_daily_mm, observation.rain_weekly_mm, observation.rain_monthly_mm, observation.rain_yearly_mm, observation.indoor_temperature_c, observation.indoor_humidity_pct, observation.absolute_pressure_hpa, observation.relative_pressure_hpa, Jsonb(observation.metadata), ), ) connection.commit() return 1 def load_latest_forecast_points( self, start_at: datetime, end_at: datetime, ) -> list[WeatherForecastPoint]: with self._connection() as connection: with connection.cursor() as cursor: cursor.execute( """ SELECT issued_at, target_at, horizon_hours, source, temperature_c, shortwave_radiation_w_m2, cloud_cover_pct FROM ( SELECT issued_at, target_at, horizon_hours, source, temperature_c, shortwave_radiation_w_m2, cloud_cover_pct, ROW_NUMBER() OVER ( PARTITION BY target_at ORDER BY issued_at DESC ) as rn FROM weather_forecast_points WHERE target_at >= %s AND target_at <= %s ) as ranked WHERE rn = 1 ORDER BY target_at LIMIT 5000 """, (start_at, end_at), ) rows = cursor.fetchall() return [ WeatherForecastPoint( issued_at=row[0], target_at=row[1], horizon_hours=row[2], source=row[3], temperature_c=row[4], shortwave_radiation_w_m2=row[5], cloud_cover_pct=row[6], ) for row in rows ] def load_forecast_history( self, start_at: datetime, end_at: datetime, ) -> dict[datetime, list[dict[str, object]]]: """Every forecast issued for each target in the window, oldest first. Unlike `load_latest_forecast_points`, this keeps all vintages of a forecast, so a training example can be shown the forecast that was actually available at the time it was issued rather than the best one known in hindsight. """ with self._connection() as connection: with connection.cursor() as cursor: cursor.execute( """ SELECT issued_at, target_at, temperature_c, shortwave_radiation_w_m2, cloud_cover_pct FROM weather_forecast_points WHERE target_at >= %s AND target_at <= %s ORDER BY target_at, issued_at """, (start_at, end_at), ) rows = cursor.fetchall() by_target: dict[datetime, list[dict[str, object]]] = {} for row in rows: by_target.setdefault(row[1], []).append( { "issued_at": row[0], "target_at": row[1], "temperature_c": row[2], "shortwave_radiation_w_m2": row[3], "cloud_cover_pct": row[4], } ) return by_target def load_display_dataset( self, start_at: datetime | None = None, end_at: datetime | None = None, ) -> WeatherDisplayDataset: now = datetime.now(timezone.utc) if start_at is None: start_at = now - timedelta(hours=24) if end_at is None: end_at = now + timedelta(hours=48) with self._connection() as connection: with connection.cursor() as cursor: cursor.execute( """ SELECT issued_at, target_at, horizon_hours, source, temperature_c, shortwave_radiation_w_m2, cloud_cover_pct FROM ( SELECT issued_at, target_at, horizon_hours, source, temperature_c, shortwave_radiation_w_m2, cloud_cover_pct, ROW_NUMBER() OVER (PARTITION BY target_at, horizon_hours ORDER BY issued_at DESC) as rn FROM weather_forecast_points WHERE target_at >= %s AND target_at <= %s ) as ranked WHERE rn = 1 ORDER BY target_at, horizon_hours LIMIT 5000 """, (start_at, end_at), ) forecast_rows = cursor.fetchall() cursor.execute( """ SELECT resolved_at, source, temperature_c, shortwave_radiation_w_m2, cloud_cover_pct FROM weather_resolved_truth WHERE resolved_at >= %s AND resolved_at <= %s ORDER BY resolved_at LIMIT 5000 """, (start_at, now), ) truth_rows = cursor.fetchall() return WeatherDisplayDataset( forecast_points=[ WeatherForecastPoint( issued_at=row[0], target_at=row[1], horizon_hours=row[2], source=row[3], temperature_c=row[4], shortwave_radiation_w_m2=row[5], cloud_cover_pct=row[6], ) for row in forecast_rows ], resolved_truth=[ WeatherResolvedTruth( resolved_at=row[0], source=row[1], temperature_c=row[2], shortwave_radiation_w_m2=row[3], cloud_cover_pct=row[4], ) for row in truth_rows ], ) def load_station_display_dataset( self, start_at: datetime | None = None, end_at: datetime | None = None, ) -> WeatherStationDisplayDataset: now = datetime.now(timezone.utc) if start_at is None: start_at = now - timedelta(days=7) if end_at is None: end_at = now with self._connection() as connection: with connection.cursor() as cursor: cursor.execute( """ SELECT observed_at, received_at, source, temperature_c, relative_humidity_pct, dew_point_c, pressure_hpa, wind_speed_m_s, wind_gust_m_s, wind_max_m_s, wind_direction_deg, solar_radiation_w_m2, uv_index, rain_rate_mm_h, rain_event_mm, rain_hourly_mm, rain_daily_mm, rain_weekly_mm, rain_monthly_mm, rain_yearly_mm, indoor_temperature_c, indoor_humidity_pct, absolute_pressure_hpa, relative_pressure_hpa, metadata FROM weather_station_observations WHERE observed_at >= %s AND observed_at <= %s ORDER BY observed_at LIMIT 100000 """, (start_at, end_at), ) rows = cursor.fetchall() return WeatherStationDisplayDataset( observations=[ WeatherStationObservation( observed_at=row[0], received_at=row[1], source=row[2], temperature_c=row[3], relative_humidity_pct=row[4], dew_point_c=row[5], pressure_hpa=row[6], wind_speed_m_s=row[7], wind_gust_m_s=row[8], wind_max_m_s=row[9], wind_direction_deg=row[10], solar_radiation_w_m2=row[11], uv_index=row[12], rain_rate_mm_h=row[13], rain_event_mm=row[14], rain_hourly_mm=row[15], rain_daily_mm=row[16], rain_weekly_mm=row[17], rain_monthly_mm=row[18], rain_yearly_mm=row[19], indoor_temperature_c=row[20], indoor_humidity_pct=row[21], absolute_pressure_hpa=row[22], relative_pressure_hpa=row[23], metadata=row[24] or {}, ) for row in rows ] ) @contextmanager def _connection(self) -> Iterator[object]: try: import psycopg except ImportError as error: raise WeatherStoreConfigurationError( "Install dependencies with `python3 -m pip install -r requirements.txt`" ) from error with psycopg.connect(self.config.database_url) as connection: yield connection