1"""Read side: DuckDB over the live Parquet file set selected from the Postgres manifest (`lake_files`).23Adjustments are applied at query time from `corporate_actions` (explicit, reproducible):4 split : multiply prices by Π(1/ratio) for splits with ex_date > bar session_date; divide volume by the same factor5 split_div : additionally multiply by Π(1 - div/prev_close) for dividends with ex_date > session_date6 (prev_close = last raw close strictly before ex_date in the requested data — if unavailable the7 dividend factor is skipped and reported in `adjustment_notes`).8"""9from __future__ import annotations1011from datetime import date12from typing import Any1314import duckdb15import pyarrow as pa16from findata_core.db import connection17from findata_schema import BARS_SCHEMA1819from findata_ingest.lake import Lake, live_files202122def _con(read_only: bool = True) -> duckdb.DuckDBPyConnection:23 con = duckdb.connect(database=":memory:", config={"threads": 4, "memory_limit": "4GB", "TimeZone": "UTC"})24 return con252627def _arrow(con: duckdb.DuckDBPyConnection, sql: str, params: dict[str, Any] | None) -> pa.Table:28 try:29 rel = con.sql(sql, params=params) if params else con.sql(sql)30 return rel.to_arrow_table()31 except TypeError: # older duckdb without params= on sql()32 return con.execute(sql, params).fetch_arrow_table()333435def empty_bars() -> pa.Table:36 return pa.table({f.name: pa.array([], type=f.type) for f in BARS_SCHEMA}, schema=BARS_SCHEMA)373839def read_bars(dataset_id: str, symbols: list[str], start: date | None, end: date | None, version: int | None = None,40 limit: int | None = None, columns: list[str] | None = None, lake: Lake | None = None) -> tuple[pa.Table, dict[str, Any]]:41 lake = lake or Lake()42 with connection() as conn:43 files = live_files(conn, dataset_id, symbols, start, end, version)44 ver = conn.execute("SELECT current_version FROM datasets WHERE id=%s", (dataset_id,)).fetchone()45 meta = {"dataset_id": dataset_id, "version": version if version is not None else (ver["current_version"] if ver else None),46 "files_scanned": len(files)}47 if not files:48 return empty_bars(), meta49 paths = [str(lake.resolve(f["path"])) for f in files]50 con = _con()51 con.register("paths", pa.table({"p": paths}))52 cols = ", ".join(columns) if columns else "*"53 where = ["symbol_as_reported IN (SELECT UNNEST($syms))"]54 params: dict[str, Any] = {"syms": symbols}55 if start:56 where.append("session_date >= $start")57 params["start"] = start58 if end:59 where.append("session_date <= $end")60 params["end"] = end61 sql = f"SELECT {cols} FROM read_parquet($paths, union_by_name=true, hive_partitioning=false) WHERE {' AND '.join(where)} ORDER BY symbol_as_reported, timestamp_utc"62 params["paths"] = paths63 if limit:64 sql += f" LIMIT {int(limit) + 1}"65 tbl = _arrow(con, sql, params)66 meta["truncated"] = bool(limit and tbl.num_rows > limit)67 if limit and tbl.num_rows > limit:68 tbl = tbl.slice(0, limit)69 if not columns:70 tbl = tbl.cast(BARS_SCHEMA)71 return tbl, meta727374def apply_adjustment(table: pa.Table, adjustment: str, lake: Lake | None = None) -> tuple[pa.Table, list[str]]:75 """Return an adjusted copy of raw bars (`adjustment_type` column updated). Never mutates the lake."""76 if adjustment in ("raw", None) or table.num_rows == 0:77 return table, []78 import polars as pl79 df = pl.from_arrow(table)80 syms = df.get_column("symbol_as_reported").unique().to_list()81 iids = [f"firstrate:{df.filter(pl.col('symbol_as_reported') == s).get_column('asset_class')[0]}:{s}" for s in syms]82 with connection() as conn:83 acts = conn.execute("SELECT instrument_id, action_type, ex_date, value::float8 AS value FROM corporate_actions WHERE instrument_id = ANY(%s) ORDER BY ex_date",84 (iids,)).fetchall()85 notes: list[str] = []86 if not acts:87 notes.append("no corporate actions on file for requested symbols")88 return df.with_columns(pl.lit(adjustment).alias("adjustment_type")).to_arrow().cast(BARS_SCHEMA), notes89 out_frames = []90 for s in syms:91 sub = df.filter(pl.col("symbol_as_reported") == s).sort("timestamp_utc")92 iid = f"firstrate:{sub.get_column('asset_class')[0]}:{s}"93 factor = pl.Series("f", [1.0] * sub.height)94 sd = sub.get_column("session_date")95 for a in acts:96 if a["instrument_id"] != iid:97 continue98 if a["action_type"] == "split":99 f = 1.0 / float(a["value"])100 factor = factor * pl.Series([f if d < a["ex_date"] else 1.0 for d in sd])101 elif a["action_type"] == "dividend" and adjustment == "split_div":102 prev = sub.filter(pl.col("session_date") < a["ex_date"]).get_column("close")103 if prev.len() == 0:104 notes.append(f"{s}: dividend {a['ex_date']} skipped (no prior close in range)")105 continue106 f = 1.0 - float(a["value"]) / float(prev[-1])107 factor = factor * pl.Series([f if d < a["ex_date"] else 1.0 for d in sd])108 sub = sub.with_columns(factor.alias("_f")).with_columns(109 (pl.col("open") * pl.col("_f")).alias("open"), (pl.col("high") * pl.col("_f")).alias("high"),110 (pl.col("low") * pl.col("_f")).alias("low"), (pl.col("close") * pl.col("_f")).alias("close"),111 (pl.col("volume") / pl.col("_f")).alias("volume"), pl.lit(adjustment).alias("adjustment_type")).drop("_f")112 out_frames.append(sub)113 res = pl.concat(out_frames).sort(["symbol_as_reported", "timestamp_utc"])114 return res.to_arrow().cast(BARS_SCHEMA), notes115116117def run_sql(sql: str, max_rows: int, timeout_s: int, views: dict[str, list[str]] | None = None) -> tuple[pa.Table, dict[str, Any]]:118 """Sandboxed read-only SQL: DuckDB in-memory, external access limited to the registered dataset views."""119 import re120 import threading121122 forbidden = re.compile(r"\b(attach|copy|export|import|install|load|create|insert|update|delete|drop|alter|pragma|set|call|read_csv|read_json|read_text|glob|httpfs)\b", re.IGNORECASE)123 if forbidden.search(sql) or ";" in sql.strip().rstrip(";"):124 raise ValueError("only a single read-only SELECT statement is allowed")125 if not re.match(r"^\s*(select|with|describe|show|summarize|explain)\b", sql, re.IGNORECASE):126 raise ValueError("statement must start with SELECT / WITH")127 is_explain = bool(re.match(r"^\s*explain\b", sql, re.I))128 con = _con()129 dirs = []130 for r in Lake().roots():131 rr = str(r.resolve()).replace("'", "''")132 dirs += [f"'{rr}/curated'", f"'{rr}/_retired'"]133 con.execute(f"SET allowed_directories = [{', '.join(dirs)}]")134 for name, paths in (views or {}).items():135 if not re.fullmatch(r"[a-z][a-z0-9_]*", name):136 raise ValueError(f"bad view name {name}")137 if paths:138 lst = ", ".join("'" + p.replace("'", "''") + "'" for p in paths)139 con.execute(f"CREATE VIEW {name} AS SELECT * FROM read_parquet([{lst}], union_by_name=true, hive_partitioning=false)")140 else:141 con.execute(f"CREATE VIEW {name} AS SELECT * FROM (VALUES (NULL)) t(empty) WHERE false")142 con.execute("SET enable_external_access = false")143 timer = threading.Timer(timeout_s, con.interrupt)144 timer.start()145 try:146 body = sql.strip().rstrip(";")147 tbl = _arrow(con, body if is_explain else f"SELECT * FROM ({body}) q LIMIT {int(max_rows) + 1}", None)148 except duckdb.InterruptException as e:149 raise TimeoutError(f"query exceeded {timeout_s}s") from e150 finally:151 timer.cancel()152 truncated = tbl.num_rows > max_rows153 return tbl.slice(0, max_rows), {"truncated": truncated}154155156def dataset_views() -> dict[str, list[str]]:157 """Map SQL view names (dataset ids with dots → underscores) to live file lists."""158 lake = Lake()159 out: dict[str, list[str]] = {}160 with connection() as conn:161 ds = conn.execute("SELECT id FROM datasets WHERE kind IN ('market_bars','derived','options') AND current_version > 0").fetchall()162 for d in ds:163 files = live_files(conn, d["id"])164 out[d["id"].replace(".", "_")] = [str(lake.resolve(f["path"])) for f in files]165 return out166