Spaces:
Sleeping
Sleeping
| """Database abstraction layer for TransitPulse. | |
| Supports local DuckDB (direct Parquet querying) and BigQuery in cloud mode. | |
| """ | |
| from __future__ import annotations | |
| import os | |
| from pathlib import Path | |
| from typing import Any | |
| import duckdb | |
| import pandas as pd | |
| from config import CFG | |
| class TransitPulseDB: | |
| """Manages queries to either local DuckDB or BigQuery.""" | |
| def __init__(self, mode: str = "local"): | |
| self.mode = mode | |
| self.duck_conn = None | |
| if self.mode == "local": | |
| # Initialize DuckDB local in-memory or file database | |
| self.duck_conn = duckdb.connect(database=":memory:") | |
| self.init_local_tables() | |
| else: | |
| # Cloud mode: initialize BigQuery client | |
| from google.cloud import bigquery | |
| self.bq_client = bigquery.Client(project=CFG.gcp_project) | |
| print(f"BigQuery serving layer initialized in project: {CFG.gcp_project}") | |
| def init_local_tables(self) -> None: | |
| """Create views in DuckDB over the output parquet files.""" | |
| scores_path = CFG.output_dir / "route_scores.parquet" | |
| segment_path = CFG.output_dir / "daily_segment_metrics.parquet" | |
| anomaly_path = CFG.output_dir / "anomaly_events.parquet" | |
| if not scores_path.exists(): | |
| # If parquet doesn't exist, create empty tables or register mock schemas | |
| print("WARNING: Parquet outputs not found. DuckDB views might fail until pipeline runs.") | |
| # We can create dummy files so startup doesn't crash | |
| return | |
| # Register views pointing to parquet files directly | |
| # This keeps the DuckDB database dynamically in sync with parquet outputs | |
| self.duck_conn.execute(f"CREATE OR REPLACE VIEW route_scores AS SELECT * FROM read_parquet('{scores_path}')") | |
| self.duck_conn.execute(f"CREATE OR REPLACE VIEW daily_segment_metrics AS SELECT * FROM read_parquet('{segment_path}')") | |
| self.duck_conn.execute(f"CREATE OR REPLACE VIEW anomaly_events AS SELECT * FROM read_parquet('{anomaly_path}')") | |
| print("DuckDB views created successfully.") | |
| def run_query(self, query: str, params: dict[str, Any] | None = None) -> pd.DataFrame: | |
| """Helper to run SQL query on DuckDB or BigQuery.""" | |
| if self.mode == "local": | |
| if self.duck_conn is None: | |
| # Late-bind if connection was missed | |
| self.duck_conn = duckdb.connect(database=":memory:") | |
| self.init_local_tables() | |
| # DuckDB supports named parameters with $name or standard python interpolation | |
| # Let's execute and return dataframe | |
| if params: | |
| df = self.duck_conn.execute(query, params).df() | |
| else: | |
| df = self.duck_conn.execute(query).df() | |
| else: | |
| # Cloud mode (BigQuery) | |
| # Standard BigQuery query | |
| job_config = None | |
| if params: | |
| from google.cloud import bigquery | |
| query_params = [ | |
| bigquery.ScalarQueryParameter(name, "STRING" if isinstance(val, str) else "FLOAT", val) | |
| for name, val in params.items() | |
| ] | |
| job_config = bigquery.QueryJobConfig(query_parameters=query_params) | |
| # Map simple table names to dataset qualified names | |
| bq_query = query.replace("route_scores", f"`{CFG.gcp_project}.{CFG.bq_dataset}.route_scores`") | |
| bq_query = bq_query.replace("daily_segment_metrics", f"`{CFG.gcp_project}.{CFG.bq_dataset}.daily_segment_metrics`") | |
| bq_query = bq_query.replace("anomaly_events", f"`{CFG.gcp_project}.{CFG.bq_dataset}.anomaly_events`") | |
| query_job = self.bq_client.query(bq_query, job_config=job_config) | |
| df = query_job.to_dataframe() | |
| import numpy as np | |
| return df.replace({np.nan: None}) | |
| # ── Database Endpoint Operations ─────────────────────────────────── | |
| def get_route_scores(self, sort_by: str = "reliability_score", ascending: bool = True) -> list[dict[str, Any]]: | |
| """Returns sorted route reliability scores.""" | |
| direction = "ASC" if ascending else "DESC" | |
| # Validate column name to prevent SQL injection | |
| valid_cols = ["route_id", "reliability_score", "mean_headway", "bunching_count", "gap_count", "mean_dwell_sec", "wow_trend", "date", "route_base_boardings"] | |
| if sort_by not in valid_cols: | |
| sort_by = "reliability_score" | |
| # Get latest date available in the route_scores table | |
| try: | |
| latest_date_df = self.run_query("SELECT MAX(date) as max_date FROM route_scores") | |
| latest_date = latest_date_df.iloc[0]["max_date"] | |
| except Exception: | |
| latest_date = None | |
| if latest_date: | |
| query = f"SELECT * FROM route_scores WHERE date = $latest_date ORDER BY {sort_by} {direction}" | |
| params = {"latest_date": str(latest_date)} | |
| else: | |
| query = f"SELECT * FROM route_scores ORDER BY {sort_by} {direction}" | |
| params = {} | |
| df = self.run_query(query, params) | |
| return df.to_dict(orient="records") | |
| def get_route_timeline(self, route_id: str) -> list[dict[str, Any]]: | |
| """Returns historical scores timeline for a specific route.""" | |
| query = "SELECT date, reliability_score, mean_headway, mean_dwell_sec, bunching_count, gap_count FROM route_scores WHERE route_id = $route_id ORDER BY date ASC" | |
| df = self.run_query(query, {"route_id": route_id}) | |
| return df.to_dict(orient="records") | |
| def get_worst_segments(self, limit: int = 10) -> list[dict[str, Any]]: | |
| """Returns the bottom N route segments sorted by reliability score.""" | |
| try: | |
| latest_date_df = self.run_query("SELECT MAX(date) as max_date FROM daily_segment_metrics") | |
| latest_date = latest_date_df.iloc[0]["max_date"] | |
| except Exception: | |
| latest_date = None | |
| if latest_date: | |
| query = f"SELECT route_id, stop_id, reliability_score, total_trips, bunching_rate, gap_rate, mean_dwell_sec, route_base_boardings FROM daily_segment_metrics WHERE date = $latest_date ORDER BY reliability_score ASC LIMIT {limit}" | |
| params = {"latest_date": str(latest_date)} | |
| else: | |
| query = f"SELECT route_id, stop_id, reliability_score, total_trips, bunching_rate, gap_rate, mean_dwell_sec, route_base_boardings FROM daily_segment_metrics ORDER BY reliability_score ASC LIMIT {limit}" | |
| params = {} | |
| df = self.run_query(query, params) | |
| return df.to_dict(orient="records") | |
| def get_anomalies(self, route_id: str | None = None, anomaly_type: str | None = None, limit: int = 100) -> list[dict[str, Any]]: | |
| """Returns recent anomaly events, optionally filtered.""" | |
| where_clauses = [] | |
| params = {} | |
| if route_id: | |
| where_clauses.append("route_id = $route_id") | |
| params["route_id"] = route_id | |
| if anomaly_type: | |
| where_clauses.append("anomaly_type = $anomaly_type") | |
| params["anomaly_type"] = anomaly_type | |
| where_str = "" | |
| if where_clauses: | |
| where_str = "WHERE " + " AND ".join(where_clauses) | |
| query = f"SELECT timestamp_str, vehicle_id, route_id, stop_id, stop_sequence, headway_sec, scheduled_headway_sec, anomaly_type FROM anomaly_events {where_str} ORDER BY timestamp_str DESC LIMIT {limit}" | |
| df = self.run_query(query, params) | |
| return df.to_dict(orient="records") | |
| def compare_periods(self, route_id: str, date1: str, date2: str) -> dict[str, Any]: | |
| """Compares route performance metrics between two dates.""" | |
| query = "SELECT reliability_score, mean_headway, bunching_count, gap_count FROM route_scores WHERE route_id = $route_id AND date IN ($date1, $date2)" | |
| df = self.run_query(query, {"route_id": route_id, "date1": date1, "date2": date2}) | |
| return df.to_dict(orient="records") | |
| # Database Singleton | |
| DB = TransitPulseDB(mode=CFG.mode) | |