replace pandas to sql with sql alchemy

database/sqlalchemy_store.py

@@ -0,0 +1,111 @@
+"""Transactional SQLAlchemy Core helpers for SQLite feature-store tables."""
+
+from __future__ import annotations
+
+from contextlib import contextmanager
+from datetime import date, datetime
+import re
+
+import pandas as pd
+from sqlalchemy import Boolean, DateTime, Float, Integer, MetaData, Table, Text, Column, create_engine, event, text
+
+
+def _engine(database_url: str):
+    engine = create_engine(database_url, connect_args={"timeout": 30})
+
+    @event.listens_for(engine, "connect")
+    def _configure_sqlite(dbapi_connection, _):
+        cursor = dbapi_connection.cursor()
+        cursor.execute("PRAGMA foreign_keys = ON")
+        cursor.execute("PRAGMA journal_mode = WAL")
+        cursor.execute("PRAGMA busy_timeout = 30000")
+        cursor.close()
+
+    return engine
+
+
+@contextmanager
+def read_connection(database_url: str):
+    engine = _engine(database_url)
+    try:
+        with engine.connect() as connection:
+            yield connection
+    finally:
+        engine.dispose()
+
+
+def _column_type(values: pd.Series):
+    if pd.api.types.is_bool_dtype(values):
+        return Boolean()
+    if pd.api.types.is_integer_dtype(values):
+        return Integer()
+    if pd.api.types.is_float_dtype(values):
+        return Float()
+    if pd.api.types.is_datetime64_any_dtype(values):
+        return DateTime()
+    return Text()
+
+
+def _records(frame: pd.DataFrame) -> list[dict]:
+    clean = frame.where(pd.notna(frame), None)
+    records = []
+    for row in clean.to_dict(orient="records"):
+        records.append({
+            key: (
+                value.to_pydatetime() if isinstance(value, pd.Timestamp)
+                else value.item() if hasattr(value, "item") and not isinstance(value, (str, bytes))
+                else value
+            )
+            for key, value in row.items()
+        })
+    return records
+
+
+def replace_tables(database_url: str, frames: dict[str, pd.DataFrame]) -> None:
+    """Atomically replace dynamic feature tables without exposing an empty table.
+
+    SQLite serializes this `BEGIN IMMEDIATE` transaction across API and
+    scheduler processes. DDL and inserts are rolled back together on failure.
+    """
+    if not frames:
+        return
+    for table_name in frames:
+        if not re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]*", table_name):
+            raise ValueError(f"Unsafe table name: {table_name}")
+
+    engine = _engine(database_url)
+    try:
+        with engine.connect() as connection:
+            connection.execute(text("BEGIN IMMEDIATE"))
+            try:
+                staging: dict[str, str] = {}
+                for name, frame in frames.items():
+                    stage = f"__staging_{name}"
+                    staging[name] = stage
+                    connection.execute(text(f'DROP TABLE IF EXISTS "{stage}"'))
+                    metadata = MetaData()
+                    table = Table(stage, metadata, *[
+                        Column(str(column), _column_type(frame[column])) for column in frame.columns
+                    ])
+                    table.create(connection)
+                    rows = _records(frame)
+                    for offset in range(0, len(rows), 1000):
+                        connection.execute(table.insert(), rows[offset:offset + 1000])
+
+                for name, stage in staging.items():
+                    backup = f"__backup_{name}"
+                    connection.execute(text(f'DROP TABLE IF EXISTS "{backup}"'))
+                    exists = connection.execute(text(
+                        "SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = :name"
+                    ), {"name": name}).scalar() is not None
+                    if exists:
+                        connection.execute(text(f'ALTER TABLE "{name}" RENAME TO "{backup}"'))
+                    connection.execute(text(f'ALTER TABLE "{stage}" RENAME TO "{name}"'))
+                    if exists:
+                        connection.execute(text(f'DROP TABLE "{backup}"'))
+                connection.commit()
+            except Exception:
+                connection.rollback()
+                raise
+    finally:
+        engine.dispose()

jobs/rebuild_feature_stores.py

@@ -3,51 +3,43 @@
 from __future__ import annotations
 
 import pandas as pd
+from sqlalchemy import text
 from scrapper_service.logger import logger
 from app.config import settings
-from database.connection import connect
+from database.sqlalchemy_store import read_connection, replace_tables
 from feature_engine.training_transforms import macro, nifty, ohlcv, vix
 from feature_engine.engine import FeatureEngine
 
 
 def _load_candles(conn, market: str | None = None, symbols: tuple[str, ...] | None = None) -> pd.DataFrame:
     clauses = ["t.active"]
-    params: list[object] = []
+    params: dict[str, object] = {}
     if market is not None:
-        clauses.append("t.market = ?")
-        params.append(market)
+        clauses.append("t.market = :market")
+        params["market"] = market
     if symbols is not None:
-        clauses.append(f"t.symbol IN ({', '.join('?' * len(symbols))})")
-        params.extend(symbols)
+        placeholders = []
+        for index, symbol in enumerate(symbols):
+            key = f"symbol_{index}"
+            placeholders.append(f":{key}")
+            params[key] = symbol
+        clauses.append(f"t.symbol IN ({', '.join(placeholders)})")
     query = f"""
         SELECT o.trading_date AS timestamp, t.symbol, o.open, o.high, o.low, o.close, o.volume
         FROM all_market_ohlcv o JOIN market_tickers t ON t.id = o.ticker_id
         WHERE {" AND ".join(clauses)}
         ORDER BY t.symbol, o.trading_date
     """
-    with conn.cursor() as cur:
-        cur.execute(query, params)
-        return pd.DataFrame(cur.fetchall())
-
-
-def _write_table(df: pd.DataFrame, conn, table_name: str) -> None:
-    try:
-        df.to_sql(table_name, conn, if_exists="replace", index=False)
-        logger.info(
-            f"Successfully wrote feature table {table_name}: {len(df)} rows"
-        )
-    except Exception:
-        logger.exception(f"Failed to write feature table {table_name}")
-        raise
+    return pd.read_sql(text(query), conn, params=params)
 
 
 def _build_labels(conn) -> pd.DataFrame:
     """Create auditable D+1-open / next-five-closes training labels in SQLite."""
-    raw = pd.read_sql_query("""
+    raw = pd.read_sql(text("""
         SELECT t.symbol, o.trading_date AS timestamp, o.open, o.close
         FROM ohlcv o JOIN market_tickers t ON t.id = o.ticker_id
         WHERE t.active ORDER BY t.symbol, o.trading_date
-    """, conn, parse_dates=["timestamp"])
+    """), conn, parse_dates=["timestamp"])
     # SQLite data imported from older schema versions may retain numeric
     # values as TEXT. Labels must never compare those strings to a number.
     raw["open"] = pd.to_numeric(raw["open"], errors="coerce")
@@ -69,31 +61,7 @@ def rebuild_feature_stores() -> dict[str, int]:
         "Rebuilding feature stores from SQLite yfinance OHLCV"
     )
 
-    with connect(settings.DATABASE_URL) as conn:
-
-        # -------------------------------------------------
-        # Database identity
-        # -------------------------------------------------
-
-        with conn.cursor() as cur:
-
-            cur.execute("PRAGMA database_list")
-
-            logger.info(
-                f"SQLite database: {cur.fetchall()}"
-            )
-
-            cur.execute("""
-                SELECT name
-                FROM sqlite_master
-                WHERE type = 'table'
-                ORDER BY name
-            """)
-
-            logger.info(
-                f"Tables BEFORE rebuild: "
-                f"{[row['name'] for row in cur.fetchall()]}"
-            )
+    with read_connection(settings.DATABASE_URL) as conn:
 
         # -------------------------------------------------
         # Load raw data
@@ -204,43 +172,23 @@ def rebuild_feature_stores() -> dict[str, int]:
 
         try:
 
-            _write_table(
-                nse_features,
-                conn,
-                settings.OHLCV_FEATURES_TABLE,
-            )
-
-            _write_table(
-                us_features,
-                conn,
-                settings.US_FEATURES_TABLE,
-            )
-
-            _write_table(
-                nifty_features,
-                conn,
-                settings.NIFTY_FEATURES_TABLE,
-            )
-
-            _write_table(
-                vix_features,
-                conn,
-                settings.VIX_FEATURES_TABLE,
-            )
-
-            _write_table(
-                macro_features,
-                conn,
-                settings.MACRO_FEATURES_TABLE,
-            )
-
-            # The merged matrix is also persisted in SQLite so serving never
-            # needs an in-memory DuckDB query.
-            merged_features = FeatureEngine().build_all_features(conn)
-            _write_table(merged_features, conn, "merged_features")
-            _write_table(_build_labels(conn), conn, "training_labels")
+            replace_tables(settings.DATABASE_URL, {
+                settings.OHLCV_FEATURES_TABLE: nse_features,
+                settings.US_FEATURES_TABLE: us_features,
+                settings.NIFTY_FEATURES_TABLE: nifty_features,
+                settings.VIX_FEATURES_TABLE: vix_features,
+                settings.MACRO_FEATURES_TABLE: macro_features,
+            })
 
-            conn.commit()
+            # These read the committed source stores, then replace the two
+            # derived tables in their own all-or-nothing transaction.
+            merged_features = FeatureEngine().build_all_features()
+            with read_connection(settings.DATABASE_URL) as label_conn:
+                labels = _build_labels(label_conn)
+            replace_tables(settings.DATABASE_URL, {
+                "merged_features": merged_features,
+                "training_labels": labels,
+            })
 
         except Exception:
 
@@ -256,19 +204,7 @@ def rebuild_feature_stores() -> dict[str, int]:
         # Verify AFTER commit
         # -------------------------------------------------
 
-        with conn.cursor() as cur:
-
-            cur.execute("""
-                SELECT name
-                FROM sqlite_master
-                WHERE type = 'table'
-                ORDER BY name
-            """)
-
-            tables = [
-                row["name"]
-                for row in cur.fetchall()
-            ]
+        tables = pd.read_sql(text("SELECT name FROM sqlite_master WHERE type = 'table' ORDER BY name"), conn)["name"].tolist()
 
             logger.info(
                 f"Tables AFTER rebuild: {tables}"