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 annotations1011import re12from datetime import date13from typing import Any1415import duckdb16import pyarrow as pa17from findata_core.db import connection18from findata_schema import BARS_SCHEMA1920from findata_ingest.lake import Lake, live_files212223def _con(read_only: bool = True) -> duckdb.DuckDBPyConnection:24 import resource25 try: # scans over many Parquet files need file descriptors (macOS default soft limit is 256)26 soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE)27 if soft < 65536:28 resource.setrlimit(resource.RLIMIT_NOFILE, (min(65536, hard), hard))29 except (ValueError, OSError):30 pass31 con = duckdb.connect(database=":memory:", config={"threads": 4, "memory_limit": "4GB", "TimeZone": "UTC"})32 return con333435def _arrow(con: duckdb.DuckDBPyConnection, sql: str, params: dict[str, Any] | None) -> pa.Table:36 try:37 rel = con.sql(sql, params=params) if params else con.sql(sql)38 return rel.to_arrow_table()39 except TypeError: # older duckdb without params= on sql()40 return con.execute(sql, params).fetch_arrow_table()414243def empty_bars() -> pa.Table:44 return pa.table({f.name: pa.array([], type=f.type) for f in BARS_SCHEMA}, schema=BARS_SCHEMA)454647def read_bars(dataset_id: str, symbols: list[str], start: date | None, end: date | None, version: int | None = None,48 limit: int | None = None, columns: list[str] | None = None, lake: Lake | None = None) -> tuple[pa.Table, dict[str, Any]]:49 lake = lake or Lake()50 with connection() as conn:51 files = live_files(conn, dataset_id, symbols, start, end, version)52 ver = conn.execute("SELECT current_version FROM datasets WHERE id=%s", (dataset_id,)).fetchone()53 meta = {"dataset_id": dataset_id, "version": version if version is not None else (ver["current_version"] if ver else None),54 "files_scanned": len(files)}55 if not files:56 return empty_bars(), meta57 paths = [str(lake.resolve(f["path"])) for f in files]58 con = _con()59 con.register("paths", pa.table({"p": paths}))60 cols = ", ".join(columns) if columns else "*"61 where = ["symbol_as_reported IN (SELECT UNNEST($syms))"]62 params: dict[str, Any] = {"syms": symbols}63 if start:64 where.append("session_date >= $start")65 params["start"] = start66 if end:67 where.append("session_date <= $end")68 params["end"] = end69 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"70 params["paths"] = paths71 if limit:72 sql += f" LIMIT {int(limit) + 1}"73 tbl = _arrow(con, sql, params)74 meta["truncated"] = bool(limit and tbl.num_rows > limit)75 if limit and tbl.num_rows > limit:76 tbl = tbl.slice(0, limit)77 if not columns:78 tbl = tbl.cast(BARS_SCHEMA)79 return tbl, meta808182def apply_adjustment(table: pa.Table, adjustment: str, lake: Lake | None = None) -> tuple[pa.Table, list[str]]:83 """Return an adjusted copy of raw bars (`adjustment_type` column updated). Never mutates the lake."""84 if adjustment in ("raw", None) or table.num_rows == 0:85 return table, []86 import polars as pl87 df = pl.from_arrow(table)88 syms = df.get_column("symbol_as_reported").unique().to_list()89 iids = [f"firstrate:{df.filter(pl.col('symbol_as_reported') == s).get_column('asset_class')[0]}:{s}" for s in syms]90 with connection() as conn:91 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",92 (iids,)).fetchall()93 notes: list[str] = []94 if not acts:95 notes.append("no corporate actions on file for requested symbols")96 return df.with_columns(pl.lit(adjustment).alias("adjustment_type")).to_arrow().cast(BARS_SCHEMA), notes97 out_frames = []98 for s in syms:99 sub = df.filter(pl.col("symbol_as_reported") == s).sort("timestamp_utc")100 iid = f"firstrate:{sub.get_column('asset_class')[0]}:{s}"101 factor = pl.Series("f", [1.0] * sub.height)102 sd = sub.get_column("session_date")103 for a in acts:104 if a["instrument_id"] != iid:105 continue106 if a["action_type"] == "split":107 f = 1.0 / float(a["value"])108 factor = factor * pl.Series([f if d < a["ex_date"] else 1.0 for d in sd])109 elif a["action_type"] == "dividend" and adjustment == "split_div":110 prev = sub.filter(pl.col("session_date") < a["ex_date"]).get_column("close")111 if prev.len() == 0:112 notes.append(f"{s}: dividend {a['ex_date']} skipped (no prior close in range)")113 continue114 f = 1.0 - float(a["value"]) / float(prev[-1])115 factor = factor * pl.Series([f if d < a["ex_date"] else 1.0 for d in sd])116 sub = sub.with_columns(factor.alias("_f")).with_columns(117 (pl.col("open") * pl.col("_f")).alias("open"), (pl.col("high") * pl.col("_f")).alias("high"),118 (pl.col("low") * pl.col("_f")).alias("low"), (pl.col("close") * pl.col("_f")).alias("close"),119 (pl.col("volume") / pl.col("_f")).alias("volume"), pl.lit(adjustment).alias("adjustment_type")).drop("_f")120 out_frames.append(sub)121 res = pl.concat(out_frames).sort(["symbol_as_reported", "timestamp_utc"])122 return res.to_arrow().cast(BARS_SCHEMA), notes123124125def run_sql(sql: str, max_rows: int, timeout_s: int, views: dict[str, Any] | None = None) -> tuple[pa.Table, dict[str, Any]]:126 """Sandboxed read-only SQL: DuckDB in-memory, external access limited to the registered dataset views."""127 import re128 import threading129130 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)131 if forbidden.search(sql) or ";" in sql.strip().rstrip(";"):132 raise ValueError("only a single read-only SELECT statement is allowed")133 if not re.match(r"^\s*(select|with|describe|show|summarize|explain)\b", sql, re.IGNORECASE):134 raise ValueError("statement must start with SELECT / WITH")135 is_explain = bool(re.match(r"^\s*explain\b", sql, re.I))136 con = _con()137 scoped_used: list[str] = []138 dirs = []139 for r in Lake().roots():140 rr = str(r.resolve()).replace("'", "''")141 dirs += [f"'{rr}/curated'", f"'{rr}/_retired'"]142 con.execute(f"SET allowed_directories = [{', '.join(dirs)}]")143 for name, spec in (views or {}).items():144 if not re.fullmatch(r"[a-z][a-z0-9_]*", name):145 raise ValueError(f"bad view name {name}")146 if name not in _referenced_views(sql, {name: spec}):147 continue # only materialize views the statement uses (directory globs are not free)148 if isinstance(spec, dict) and spec.get("sql"):149 fast = None150 if spec.get("dataset_id"):151 try:152 fast = scoped_view_sql(spec["dataset_id"], spec.get("kind", ""), extract_scope(sql), Lake())153 except Exception:154 fast = None155 con.execute(f"CREATE VIEW {name} AS {fast or spec['sql']}")156 scoped_used.append(name if fast else f"{name} (full scan)")157 elif isinstance(spec, list) and spec: # legacy explicit file list158 lst = ", ".join("'" + g.replace("'", "''") + "'" for g in spec)159 con.execute(f"CREATE VIEW {name} AS SELECT * FROM read_parquet([{lst}], union_by_name=true, hive_partitioning=false)")160 else:161 con.execute(f"CREATE VIEW {name} AS SELECT * FROM (VALUES (NULL)) t(empty) WHERE false")162 con.execute("SET enable_external_access = false")163 timer = threading.Timer(timeout_s, con.interrupt)164 timer.start()165 try:166 body = sql.strip().rstrip(";")167 tbl = _arrow(con, body if is_explain else f"SELECT * FROM ({body}) q LIMIT {int(max_rows) + 1}", None)168 except duckdb.InterruptException as e:169 raise TimeoutError(f"query exceeded {timeout_s}s") from e170 finally:171 timer.cancel()172 truncated = tbl.num_rows > max_rows173 return tbl.slice(0, max_rows), {"truncated": truncated, "views": scoped_used}174175176_views_cache: dict[str, Any] = {"at": 0.0, "value": None}177178179def dataset_views() -> dict[str, dict[str, Any]]:180 """SQL views for the playground/assistant: one per dataset with data, built from Hive-style directory globs181 (`curated/<dataset>/{base,delta}/**`) so that DuckDB prunes partitions (`symbol=`, `year=`, `session_date=`,182 `underlying=`, `period=`) from the path instead of receiving millions of file names. The directory tree is the183 live set (retired files are moved out); files of an in-progress publication may be visible for a few minutes.184 Cached 5 minutes."""185 import time186 if _views_cache["value"] is not None and time.monotonic() - _views_cache["at"] < 300:187 return _views_cache["value"]188 lake = Lake()189 out: dict[str, dict[str, Any]] = {}190 with connection() as conn:191 ds = conn.execute("SELECT id, kind FROM datasets WHERE kind IN ('market_bars','derived','options') AND current_version > 0").fetchall()192 for d in ds:193 root = lake.root_for(d["id"])194 layouts = {}195 for layout in ("base", "delta"):196 ldir = root / "curated" / d["id"] / layout197 if ldir.exists() and next(ldir.rglob("*.parquet"), None) is not None: # skip emptied layouts198 layouts[layout] = str(ldir / "**" / "*.parquet").replace("'", "''")199 if not layouts:200 continue201 parts = []202 if d["kind"] == "options":203 parts.append(f"SELECT * FROM read_parquet('{layouts['base']}', hive_partitioning=true)")204 else:205 if "base" in layouts: # partition columns symbol=, year= come from the path (pruned before any file is opened)206 parts.append(f"SELECT * FROM read_parquet('{layouts['base']}', hive_partitioning=true)")207 if "delta" in layouts: # session files: derive the same columns so UNION ALL BY NAME lines up208 parts.append(f"SELECT *, symbol_as_reported AS symbol, year(session_date)::BIGINT AS year "209 f"FROM read_parquet('{layouts['delta']}', hive_partitioning=false)")210 out[d["id"].replace(".", "_")] = {"sql": " UNION ALL BY NAME ".join(parts), "layouts": sorted(layouts), "root": str(root),211 "kind": d["kind"], "dataset_id": d["id"]}212 _views_cache.update(at=time.monotonic(), value=out)213 return out214215216_SYM_RE = re.compile(r"\b(?:symbol|symbol_as_reported|underlying)\s*=\s*'([A-Za-z0-9.\-_]+)'", re.I)217_SYM_IN_RE = re.compile(r"\b(?:symbol|symbol_as_reported|underlying)\s+IN\s*\(([^)]+)\)", re.I)218_YEAR_RE = re.compile(r"\byear\s*=\s*(\d{4})\b", re.I)219_YEAR_IN_RE = re.compile(r"\byear\s+IN\s*\(([^)]+)\)", re.I)220_DATE_RE = re.compile(r"\b(?:session_date|trade_date)\s*(>=|<=|=|>|<)\s*(?:DATE\s*)?'(\d{4}-\d{2}-\d{2})'", re.I)221_BETWEEN_RE = re.compile(r"\b(?:session_date|trade_date)\s+BETWEEN\s+(?:DATE\s*)?'(\d{4}-\d{2}-\d{2})'\s+AND\s+(?:DATE\s*)?'(\d{4}-\d{2}-\d{2})'", re.I)222223224def extract_scope(sql: str) -> dict[str, Any]:225 """Pull simple equality/IN/range predicates on partition columns out of a statement so that the file list can226 be resolved from the catalog (indexed) instead of globbing hundreds of thousands of files."""227 from datetime import date as _date228 syms: set[str] = {m.group(1).upper() for m in _SYM_RE.finditer(sql)}229 for m in _SYM_IN_RE.finditer(sql):230 syms |= {x.strip().strip("'\"").upper() for x in m.group(1).split(",") if x.strip()}231 years: set[int] = {int(m.group(1)) for m in _YEAR_RE.finditer(sql)}232 for m in _YEAR_IN_RE.finditer(sql):233 years |= {int(x) for x in re.findall(r"\d{4}", m.group(1))}234 start = end = None235 for op, d in _DATE_RE.findall(sql):236 dd = _date.fromisoformat(d)237 if op in (">=", ">", "="):238 start = dd if start is None else max(start, dd)239 if op in ("<=", "<", "="):240 end = dd if end is None else min(end, dd)241 for a, b in _BETWEEN_RE.findall(sql):242 start = max(start, _date.fromisoformat(a)) if start else _date.fromisoformat(a)243 end = min(end, _date.fromisoformat(b)) if end else _date.fromisoformat(b)244 if years:245 ys, ye = _date(min(years), 1, 1), _date(max(years), 12, 31)246 start = ys if start is None else max(start, ys)247 end = ye if end is None else min(end, ye)248 return {"symbols": sorted(syms), "start": start, "end": end}249250251def scoped_view_sql(dataset_id: str, kind: str, scope: dict[str, Any], lake: Lake) -> str | None:252 """Explicit file list from the catalog for symbol-scoped statements (fast path). None → use the glob view."""253 if not scope["symbols"] and not (scope["start"] or scope["end"]):254 return None255 with connection() as conn:256 files = live_files(conn, dataset_id, scope["symbols"] or None, scope["start"], scope["end"])257 if len(files) > 20000:258 return None # too broad for an explicit list: directory scan instead259 if not files:260 return "SELECT * FROM (VALUES (NULL)) t(empty) WHERE false"261 lst = ", ".join("'" + str(lake.resolve(f["path"])).replace("'", "''") + "'" for f in files)262 if kind == "options":263 return (f"SELECT * EXCLUDE (filename), regexp_extract(filename, 'period=([^/]+)', 1) AS period "264 f"FROM read_parquet([{lst}], union_by_name=true, hive_partitioning=false, filename=true)")265 return (f"SELECT *, symbol_as_reported AS symbol, year(session_date)::BIGINT AS year "266 f"FROM read_parquet([{lst}], union_by_name=true, hive_partitioning=false)")267268269def _referenced_views(sql: str, views: dict[str, Any]) -> list[str]:270 import re as _re271 low = sql.lower()272 return [name for name in views if _re.search(r"\b" + _re.escape(name) + r"\b", low)]273