admin / Syanpse-Vanguard

"""Embedded DuckDB engine — the zero-config default backend.

DuckDB is single-writer and not natively async, so all calls are dispatched to
a thread and serialized behind a lock. Batches are converted to a columnar Arrow
table and appended in one INSERT for throughput."""
from __future__ import annotations

import asyncio
import os
from datetime import datetime
from importlib import resources
from typing import Any, AsyncIterator, Sequence

import duckdb
import pyarrow as pa

from sv.models.event import STORAGE_COLUMNS, EventModel
from sv.storage.base import Query, StorageEngine


class DuckDBEngine(StorageEngine):
    def __init__(self, path: str, cold_dir: str):
        if path != ":memory:":
            os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True)
        os.makedirs(cold_dir, exist_ok=True)
        self._con = duckdb.connect(path)
        self._cold_dir = cold_dir
        self._lock = asyncio.Lock()

    async def _run(self, fn):
        async with self._lock:
            return await asyncio.to_thread(fn)

    # Idempotent, additive migrations for databases created by an earlier
    # version (e.g. a reused v1 volume where CREATE TABLE IF NOT EXISTS is a
    # no-op). Each statement is safe to run repeatedly.
    _MIGRATIONS = [
        "ALTER TABLE agents ADD COLUMN IF NOT EXISTS last_ip VARCHAR",
        "ALTER TABLE alert_rules ADD COLUMN IF NOT EXISTS webhook_id VARCHAR",
    ]

    async def init_schema(self) -> None:
        ddl = resources.files("sv.storage").joinpath("schema.sql").read_text()

        def _do():
            self._con.execute(ddl)
            for stmt in self._MIGRATIONS:
                try:
                    self._con.execute(stmt)
                except Exception:  # noqa: BLE001 — column already present / older engine
                    pass

        await self._run(_do)

    async def write_batch(self, events: Sequence[EventModel]) -> int:
        if not events:
            return 0
        table = _events_to_arrow(events)

        def _do():
            self._con.register("_batch", table)
            self._con.execute(
                f"INSERT INTO events ({', '.join(STORAGE_COLUMNS)}) "
                f"SELECT {', '.join(STORAGE_COLUMNS)} FROM _batch"
            )
            self._con.unregister("_batch")
            return table.num_rows

        return await self._run(_do)

    async def execute(self, query: Query) -> pa.Table:
        return await self._run(
            lambda: self._con.execute(query.text, query.params).fetch_arrow_table()
        )

    async def stream(self, query: Query, batch_rows: int = 10_000
                     ) -> AsyncIterator[pa.RecordBatch]:
        reader = await self._run(
            lambda: self._con.execute(query.text, query.params)
            .fetch_record_batch(batch_rows)
        )
        while True:
            batch = await self._run(lambda: next(reader, None))
            if batch is None:
                break
            yield batch

    async def retention_sweep(self, older_than: datetime) -> int:
        def _do():
            self._con.execute(
                "COPY (SELECT * FROM events WHERE ts < ?) TO ? "
                "(FORMAT parquet, PARTITION_BY (tenant_id), APPEND)",
                [older_than, self._cold_dir],
            )
            self._con.execute("DELETE FROM events WHERE ts < ?", [older_than])
            return self._con.execute(
                "SELECT changes()"  # rows affected by the last statement
            ).fetchone()[0]

        return await self._run(_do)

    async def fetchone(self, query: Query) -> dict | None:
        def _do():
            cur = self._con.execute(query.text, query.params)
            row = cur.fetchone()
            if row is None:
                return None
            cols = [d[0] for d in cur.description]
            return dict(zip(cols, row))

        return await self._run(_do)

    async def fetchall(self, query: Query) -> list[dict]:
        def _do():
            cur = self._con.execute(query.text, query.params)
            rows = cur.fetchall()
            cols = [d[0] for d in cur.description]
            return [dict(zip(cols, r)) for r in rows]

        return await self._run(_do)

    async def run(self, query: Query) -> None:
        await self._run(lambda: self._con.execute(query.text, query.params))

    async def close(self) -> None:
        await self._run(self._con.close)


def _events_to_arrow(events: Sequence[EventModel]) -> pa.Table:
    cols: dict[str, list[Any]] = {c: [] for c in STORAGE_COLUMNS}
    for e in events:
        row = e.to_storage_row()
        for c in STORAGE_COLUMNS:
            cols[c].append(row[c])
    return pa.table(cols)