diff --git a/scripts/rain_model_common.py b/scripts/rain_model_common.py index 6126e14..87fe00b 100644 --- a/scripts/rain_model_common.py +++ b/scripts/rain_model_common.py @@ -49,7 +49,16 @@ def fetch_ws90(conn, site: str, start: str, end: str) -> pd.DataFrame: AND (%s = '' OR ts <= %s::timestamptz) ORDER BY ts ASC """ - return pd.read_sql_query(sql, conn, params=(site, start, start, end, end), parse_dates=["ts", "received_at"]) + with conn.cursor() as cur: + cur.execute(sql, (site, start, start, end, end)) + rows = cur.fetchall() + cols = [d.name for d in cur.description] + + df = pd.DataFrame.from_records(rows, columns=cols) + if not df.empty: + df["ts"] = pd.to_datetime(df["ts"], utc=True) + df["received_at"] = pd.to_datetime(df["received_at"], utc=True) + return df def fetch_baro(conn, site: str, start: str, end: str) -> pd.DataFrame: @@ -61,7 +70,16 @@ def fetch_baro(conn, site: str, start: str, end: str) -> pd.DataFrame: AND (%s = '' OR ts <= %s::timestamptz) ORDER BY ts ASC """ - return pd.read_sql_query(sql, conn, params=(site, start, start, end, end), parse_dates=["ts", "received_at"]) + with conn.cursor() as cur: + cur.execute(sql, (site, start, start, end, end)) + rows = cur.fetchall() + cols = [d.name for d in cur.description] + + df = pd.DataFrame.from_records(rows, columns=cols) + if not df.empty: + df["ts"] = pd.to_datetime(df["ts"], utc=True) + df["received_at"] = pd.to_datetime(df["received_at"], utc=True) + return df def build_dataset(