sentiment-engine/MALKHUT/malkhut/storage/ch_store.py

"""
ClickHouse persistence layer for MALKHUT.

Stores: replay traces, fulfilment decisions, self-play episodes,
policy snapshots, policy pool, live discrepancies.

ClickHouse credentials: dolphin:dolphin_ch_2026 @ localhost:8123
Database: dolphin_malkhut
"""
from __future__ import annotations

import json
import time
from typing import Any, Mapping, Optional, Sequence

import urllib.request
import urllib.parse


CH_URL = "http://localhost:8123/"
CH_USER = "dolphin"
CH_PASS = "dolphin_ch_2026"
CH_DB = "dolphin_malkhut"


def _ch_query(sql: str, data: str = "") -> str:
    """Execute a ClickHouse query via HTTP POST."""
    qparams = {
        "query": sql,
        "user": CH_USER,
        "password": CH_PASS,
        "database": CH_DB,
    }
    url = f"{CH_URL}?{urllib.parse.urlencode(qparams)}"
    body = data.encode("utf-8") if data else b""
    req = urllib.request.Request(url, data=body, method="POST")
    try:
        with urllib.request.urlopen(req, timeout=10) as resp:
            return resp.read().decode("utf-8")
    except urllib.error.HTTPError as e:
        err_body = e.read().decode("utf-8", errors="replace") if e.fp else ""
        raise RuntimeError(f"CH query failed: {e.code} {err_body}") from e


def _ch_insert(table: str, rows_json: str) -> None:
    """Insert JSONEachRow into ClickHouse."""
    _ch_query(f"INSERT INTO {table} FORMAT JSONEachRow", rows_json)


class MalkhutCHStore:
    """
    ClickHouse storage for MALKHUT.

    Tables are created on first use (idempotent).
    """

    def __init__(self) -> None:
        self._ensure_db()

    def _ensure_db(self) -> None:
        _ch_query(f"CREATE DATABASE IF NOT EXISTS {CH_DB}")

    def ensure_tables(self) -> None:
        _ch_query("""
            CREATE TABLE IF NOT EXISTS replay_steps (
                ts_ns Int64,
                symbol String,
                step_index UInt32,
                before_state String,
                joint_action String,
                after_state String,
                inserted_at DateTime DEFAULT now()
            ) ENGINE = MergeTree()
            ORDER BY (symbol, ts_ns, step_index)
        """)

        _ch_query("""
            CREATE TABLE IF NOT EXISTS fulfilment_decisions (
                ts_ns Int64,
                exchange String,
                symbol String,
                intent_id String,
                state_hash String,
                selected_action String,
                root_distribution String,
                risk_decision String,
                policy_version String,
                latency_ms Float64,
                inserted_at DateTime DEFAULT now()
            ) ENGINE = MergeTree()
            ORDER BY (exchange, symbol, ts_ns, intent_id)
        """)

        _ch_query("""
            CREATE TABLE IF NOT EXISTS self_play_episodes (
                ts_ns Int64,
                policy_version String,
                scenario_id String,
                seed Int64,
                pnl_bps Float64,
                max_drawdown_bps Float64,
                fill_ratio Float64,
                adverse_fill_ratio Float64,
                avg_slippage_bps Float64,
                liq_near_misses UInt32,
                cancel_count UInt32,
                diagnostics String,
                inserted_at DateTime DEFAULT now()
            ) ENGINE = MergeTree()
            ORDER BY (policy_version, scenario_id, seed)
        """)

        _ch_query("""
            CREATE TABLE IF NOT EXISTS policy_snapshots (
                version String,
                score Float64,
                created_ts_ns Int64,
                params String,
                cma_vector String,
                evaluation_summary String,
                git_hash String,
                inserted_at DateTime DEFAULT now()
            ) ENGINE = MergeTree()
            ORDER BY (version)
        """)

        _ch_query("""
            CREATE TABLE IF NOT EXISTS live_discrepancies (
                ts_ns Int64,
                exchange String,
                symbol String,
                predicted String,
                actual String,
                severity String,
                inserted_at DateTime DEFAULT now()
            ) ENGINE = MergeTree()
            ORDER BY (exchange, symbol, ts_ns)
        """)

    def store_replay_step(
        self, symbol: str, ts_ns: int, step_index: int,
        before: str, action: str, after: str,
    ) -> None:
        row = json.dumps({
            "ts_ns": ts_ns, "symbol": symbol, "step_index": step_index,
            "before_state": before, "joint_action": action, "after_state": after,
        })
        _ch_insert("replay_steps", row)

    def store_fulfilment_decision(
        self, ts_ns: int, exchange: str, symbol: str, intent_id: str,
        state_hash: str, selected_action: str, root_distribution: str,
        risk_decision: str, policy_version: str, latency_ms: float,
    ) -> None:
        row = json.dumps({
            "ts_ns": ts_ns, "exchange": exchange, "symbol": symbol,
            "intent_id": intent_id, "state_hash": state_hash,
            "selected_action": selected_action, "root_distribution": root_distribution,
            "risk_decision": risk_decision, "policy_version": policy_version,
            "latency_ms": latency_ms,
        })
        _ch_insert("fulfilment_decisions", row)

    def store_episode(
        self, policy_version: str, scenario_id: str, seed: int,
        pnl_bps: float, max_drawdown_bps: float, fill_ratio: float,
        adverse_fill_ratio: float, avg_slippage_bps: float,
        liq_near_misses: int, cancel_count: int, diagnostics: str,
    ) -> None:
        row = json.dumps({
            "ts_ns": time.time_ns(), "policy_version": policy_version,
            "scenario_id": scenario_id, "seed": seed,
            "pnl_bps": pnl_bps, "max_drawdown_bps": max_drawdown_bps,
            "fill_ratio": fill_ratio, "adverse_fill_ratio": adverse_fill_ratio,
            "avg_slippage_bps": avg_slippage_bps, "liq_near_misses": liq_near_misses,
            "cancel_count": cancel_count, "diagnostics": diagnostics,
        })
        _ch_insert("self_play_episodes", row)

    def store_policy_snapshot(
        self, version: str, score: float, params_str: str,
        cma_vector: str = "", evaluation_summary: str = "", git_hash: str = "",
    ) -> None:
        row = json.dumps({
            "version": version, "score": score,
            "created_ts_ns": time.time_ns(), "params": params_str,
            "cma_vector": cma_vector, "evaluation_summary": evaluation_summary,
            "git_hash": git_hash,
        })
        _ch_insert("policy_snapshots", row)

    def store_discrepancy(
        self, ts_ns: int, exchange: str, symbol: str,
        predicted: str, actual: str, severity: str,
    ) -> None:
        row = json.dumps({
            "ts_ns": ts_ns, "exchange": exchange, "symbol": symbol,
            "predicted": predicted, "actual": actual, "severity": severity,
        })
        _ch_insert("live_discrepancies", row)

    def query(self, sql: str) -> str:
        return _ch_query(sql)