156 lines
4.8 KiB
Python
156 lines
4.8 KiB
Python
from __future__ import annotations
|
|
|
|
import csv
|
|
import sqlite3
|
|
from pathlib import Path
|
|
from typing import Iterable, Mapping, Union
|
|
|
|
from .import_templates import ImportTemplateError, load_tabular_rows, validate_import_rows
|
|
|
|
|
|
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"
|
|
|
|
SQLITE_IMPORT_TABLES = {
|
|
"area_monthly_metrics": "area_monthly_metrics",
|
|
"neighborhood_monthly_metrics": "neighborhood_monthly_metrics",
|
|
"policy_events": "policy_events",
|
|
}
|
|
|
|
|
|
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 import_template_file(
|
|
db_path: Union[Path, str],
|
|
template_name: str,
|
|
file_path: Path,
|
|
) -> int:
|
|
table = SQLITE_IMPORT_TABLES.get(template_name)
|
|
if table is None:
|
|
supported = ", ".join(sorted(SQLITE_IMPORT_TABLES))
|
|
raise ImportTemplateError(
|
|
f"template '{template_name}' cannot be imported into SQLite; "
|
|
f"supported templates: {supported}"
|
|
)
|
|
|
|
rows = load_tabular_rows(file_path)
|
|
result = validate_import_rows(template_name, rows)
|
|
if not result.valid:
|
|
first_error = result.errors[0]
|
|
raise ImportTemplateError(
|
|
f"validation failed at row {first_error.row}, "
|
|
f"column {first_error.column}: {first_error.message}"
|
|
)
|
|
if not rows:
|
|
return 0
|
|
|
|
init_db(db_path)
|
|
with connect(db_path) as conn:
|
|
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)
|