chore: bootstrap housing research platform

src/shanghai_housing/db.py

@@ -0,0 +1,109 @@
+from __future__ import annotations
+
+import csv
+import sqlite3
+from pathlib import Path
+from typing import Iterable, Mapping, Union
+
+
+PROJECT_ROOT = Path(__file__).resolve().parents[2]
+DEFAULT_DB_PATH = PROJECT_ROOT / "data" / "shanghai_housing.sqlite"
+SCHEMA_PATH = Path(__file__).with_name("schema.sql")
+SAMPLE_DIR = PROJECT_ROOT / "data" / "sample"
+
+
+def connect(db_path: Union[Path, str] = DEFAULT_DB_PATH) -> sqlite3.Connection:
+    conn = sqlite3.connect(str(db_path))
+    conn.row_factory = sqlite3.Row
+    conn.execute("PRAGMA foreign_keys = ON")
+    return conn
+
+
+def init_db(db_path: Union[Path, str] = DEFAULT_DB_PATH) -> Path:
+    path = Path(db_path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with connect(path) as conn:
+        conn.executescript(SCHEMA_PATH.read_text(encoding="utf-8"))
+    return path
+
+
+def load_sample_data(db_path: Union[Path, str] = DEFAULT_DB_PATH) -> None:
+    init_db(db_path)
+    with connect(db_path) as conn:
+        clear_tables(
+            conn,
+            [
+                "neighborhood_monthly_metrics",
+                "neighborhoods",
+                "area_monthly_metrics",
+                "policy_events",
+                "areas",
+            ],
+        )
+        load_csv(conn, "areas", SAMPLE_DIR / "areas.csv")
+        load_csv(conn, "area_monthly_metrics", SAMPLE_DIR / "area_monthly_metrics.csv")
+        load_csv(conn, "neighborhoods", SAMPLE_DIR / "neighborhoods.csv")
+        load_csv(
+            conn,
+            "neighborhood_monthly_metrics",
+            SAMPLE_DIR / "neighborhood_monthly_metrics.csv",
+        )
+
+
+def clear_tables(conn: sqlite3.Connection, tables: Iterable[str]) -> None:
+    for table in tables:
+        conn.execute(f"DELETE FROM {table}")
+
+
+def load_csv(conn: sqlite3.Connection, table: str, path: Path) -> int:
+    with path.open("r", encoding="utf-8-sig", newline="") as handle:
+        rows = list(csv.DictReader(handle))
+    if not rows:
+        return 0
+
+    columns = list(rows[0].keys())
+    placeholders = ", ".join(["?"] * len(columns))
+    column_sql = ", ".join(columns)
+    update_sql = ", ".join([f"{column}=excluded.{column}" for column in columns])
+    sql = (
+        f"INSERT INTO {table} ({column_sql}) VALUES ({placeholders}) "
+        f"ON CONFLICT DO UPDATE SET {update_sql}"
+    )
+    conn.executemany(sql, [tuple(row[column] for column in columns) for row in rows])
+    return len(rows)
+
+
+def fetch_area_metrics(conn: sqlite3.Connection, month: str) -> list[sqlite3.Row]:
+    return conn.execute(
+        """
+        SELECT
+            a.area_id,
+            a.name,
+            a.district,
+            a.segment,
+            m.*
+        FROM area_monthly_metrics m
+        JOIN areas a ON a.area_id = m.area_id
+        WHERE m.month = ?
+        ORDER BY a.district, a.name
+        """,
+        (month,),
+    ).fetchall()
+
+
+def fetch_area_history(
+    conn: sqlite3.Connection, area_id: str, through_month: str
+) -> list[sqlite3.Row]:
+    return conn.execute(
+        """
+        SELECT *
+        FROM area_monthly_metrics
+        WHERE area_id = ? AND month <= ?
+        ORDER BY month
+        """,
+        (area_id, through_month),
+    ).fetchall()
+
+
+def row_to_dict(row: Union[sqlite3.Row, Mapping[str, object]]) -> dict[str, object]:
+    return {key: row[key] for key in row.keys()} if isinstance(row, sqlite3.Row) else dict(row)