import fcntl import os from pathlib import Path from time import sleep import duckdb import polars as pl import polars.selectors as cs from polars.exceptions import ComputeError from src.utils import get_last_modified, logger def should_rebuild(db_path: Path, parquet_path: str) -> bool: db_path = Path(db_path) if not db_path.exists(): return True dev = os.getenv("DEVELOPMENT", "False").lower() == "true" force = os.getenv("REBUILD_DUCKDB", "False").lower() == "true" if dev and not force: return False last_modified: float = get_last_modified(parquet_path) return last_modified > db_path.stat().st_mtime def _load_source_frame() -> pl.DataFrame: """Read the source parquet and apply the row-level transforms. Kept here (not in utils.py) so src.db has no dependency on utils. Mirrors the behavior previously in utils.get_decp_data(). """ parquet_path: str = os.getenv("DATA_FILE_PARQUET_PATH", "") if not (parquet_path.startswith("http")): assert os.path.exists(parquet_path) try: lff: pl.LazyFrame = pl.scan_parquet(str(parquet_path)) except ComputeError: logger.info("Lecture du parquet échouée, nouvelle tentative dans 10s...") sleep(10) lff = pl.scan_parquet(str(parquet_path)) lff = lff.sort(by=["dateNotification", "uid"], descending=True, nulls_last=True) lff = lff.filter(pl.col("donneesActuelles")).drop("donneesActuelles") # booleans_to_strings: true → "oui", false → "non" lff = lff.with_columns( pl.col(cs.Boolean) .cast(pl.String) .str.replace("true", "oui") .str.replace("false", "non") ) for col in ["acheteur_nom", "titulaire_nom"]: lff = lff.with_columns( pl.when(pl.col(col).is_null()) .then(pl.lit("[Identifiant non reconnu dans la base INSEE]")) .otherwise(pl.col(col)) .name.keep() ) return lff.collect() def build_database(db_path: Path) -> None: """Build the DuckDB database atomically under an exclusive lock. Caller MUST hold the fcntl.flock on the .lock file. """ db_path = Path(db_path) tmp_path = db_path.with_suffix(".duckdb.tmp") staging_parquet = db_path.with_suffix(".staging.parquet") if tmp_path.exists(): tmp_path.unlink() logger.info( f"Construction de la base DuckDB à partir de {os.getenv('DATA_FILE_PARQUET_PATH', '')}..." ) frame = _load_source_frame() # Write transformed frame as parquet so DuckDB can read it natively # (avoids pyarrow dependency for the Polars→DuckDB handoff) frame.write_parquet(str(staging_parquet)) try: with duckdb.connect(str(tmp_path)) as w: w.execute( f"CREATE TABLE decp AS SELECT * FROM read_parquet('{staging_parquet}')" ) w.execute( "CREATE TABLE acheteurs_marches AS " "SELECT DISTINCT uid, objet, acheteur_id FROM decp " "ORDER BY acheteur_id" ) w.execute( "CREATE TABLE titulaires_marches AS " "SELECT DISTINCT uid, objet, titulaire_id FROM decp " "ORDER BY titulaire_id" ) w.execute( "CREATE TABLE acheteurs_departement AS " "SELECT DISTINCT acheteur_id, acheteur_nom, acheteur_departement_code " "FROM decp ORDER BY acheteur_nom" ) w.execute( "CREATE TABLE titulaires_departement AS " "SELECT DISTINCT titulaire_id, titulaire_nom, titulaire_departement_code " "FROM decp ORDER BY titulaire_nom" ) finally: if staging_parquet.exists(): staging_parquet.unlink() os.replace(tmp_path, db_path) logger.info(f"Base DuckDB construite : {db_path}") def _ensure_database() -> Path: db_path = Path(os.getenv("DUCKDB_PATH", "./decp.duckdb")) parquet_path = os.getenv("DATA_FILE_PARQUET_PATH", "") lock_path = db_path.with_suffix(".duckdb.lock") db_exists = db_path.exists() with open(lock_path, "w") as lock_fd: fcntl.flock(lock_fd, fcntl.LOCK_EX) try: if should_rebuild(db_path, parquet_path): build_database(db_path) else: logger.debug("Base de données déjà disponible et à jour.") except Exception as e: if db_exists and db_path.exists(): logger.error( f"Bootstrap données KO ({e}). " f"Réutilisation du DuckDB existant : {db_path}" ) else: logger.critical("Aucune base DuckDB et reconstruction impossible.") raise return db_path DB_PATH = _ensure_database() conn: duckdb.DuckDBPyConnection = duckdb.connect(str(DB_PATH), read_only=True) schema: pl.Schema = conn.execute("SELECT * FROM decp LIMIT 0").pl().schema def get_cursor() -> duckdb.DuckDBPyConnection: """Return a per-request cursor that shares the process-wide connection.""" return conn.cursor() def query_marches( where_sql: str = "TRUE", params: tuple | list = (), columns: list[str] | None = None, order_by: str | None = None, limit: int | None = None, offset: int | None = None, ) -> pl.DataFrame: """Run a parameterized SELECT against the decp table and return Polars. `where_sql` and `order_by` are trusted SQL fragments (callers are internal code, never user input). `params` values are passed through DuckDB's parameter binding. """ cols = ", ".join(columns) if columns else "*" sql = f"SELECT {cols} FROM decp WHERE {where_sql}" if order_by: sql += f" ORDER BY {order_by}" if limit is not None: sql += f" LIMIT {int(limit)}" if offset is not None: sql += f" OFFSET {int(offset)}" logger.debug("query_marches: " + sql.replace("?", "{}").format(*params)) return get_cursor().execute(sql, list(params)).pl() def count_marches(where_sql: str = "TRUE", params: tuple | list = ()) -> int: """Retourne le nombre de lignes correspondant à where_sql.""" sql = f"SELECT COUNT(*) FROM decp WHERE {where_sql}" logger.debug("count_marches: " + sql.replace("?", "{}").format(*params)) result = get_cursor().execute(sql, list(params)).fetchone() return int(result[0]) if result else 0 def count_unique_marches(where_sql: str = "TRUE", params: tuple | list = ()) -> int: """Retourne le nombre de uid distincts correspondant à where_sql.""" sql = f"SELECT COUNT(DISTINCT uid) FROM decp WHERE {where_sql}" logger.debug("count_unique_marches: " + sql.replace("?", "{}").format(*params)) result = get_cursor().execute(sql, list(params)).fetchone() return int(result[0]) if result else 0