server_b_ingest_sql.py

import asyncio
import json
import os
import time
from datetime import datetime, timezone, timedelta
from typing import Any, Dict

import psycopg2
from psycopg2.extras import execute_values
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field

# =========================
# CONFIG
# =========================
FLUSH_INTERVAL_SEC = 20.0
MAX_BATCH_SIZE = 200
MAX_ROWS = 100000
RESET_N = 10
MAX_QUEUE_SIZE = 50000
CLEANUP_INTERVAL_SEC = 6 * 3600  # every 6 hours instead of every 30 min

IST = timezone(timedelta(hours=5, minutes=30))

# =========================
# DATA MODEL
# =========================
class IngestMsg(BaseModel):
    ts: str
    plant_id: str = Field(..., min_length=1, max_length=50)
    payload: Dict[str, Any]


# =========================
# DB CONNECTION
# =========================
def make_conn():
    return psycopg2.connect(
        os.environ["DATABASE_URL"],
        connect_timeout=5,
    )

# =========================
# TIMESTAMP PARSE
# =========================
def parse_ts_ist(ts_str: str) -> datetime:
    try:
        dt = datetime.fromisoformat(ts_str.replace("Z", "+00:00"))
    except Exception:
        raise ValueError(f"Invalid ts format: {ts_str}")

    if dt.tzinfo is None:
        dt = dt.replace(tzinfo=IST)

    return dt


# =========================
# RESET LOGIC
# =========================
def reset_last_n_old_records(cursor, conn, n=10):
    cursor.execute("""
        SELECT id FROM dbo.raw_plant_data
        ORDER BY ts DESC
        LIMIT %s
    """, (n,))
    
    rows = cursor.fetchall()
    if not rows:
        return

    now = datetime.utcnow()

    for i, (row_id,) in enumerate(rows, start=1):
        new_ts = now - timedelta(seconds=i)

        cursor.execute("""
            UPDATE dbo.raw_plant_data
            SET ts = %s
            WHERE id = %s
        """, (new_ts, row_id))

    conn.commit()
    print(f"Reset {len(rows)} old records timestamps")


# =========================
# CLEANUP LOGIC
# =========================
def cleanup_old_rows(cursor, conn):
    cursor.execute("""
        DELETE FROM dbo.raw_plant_data
        WHERE ts < NOW() - INTERVAL '3 days'
    """)
    
    deleted = cursor.rowcount
    conn.commit()

    if deleted > 0:
        print(f"Deleted {deleted} old rows")

# =========================
# APP
# =========================
app = FastAPI(title="SWaT Ingest API")

queue: asyncio.Queue[IngestMsg] = asyncio.Queue(
    maxsize=MAX_QUEUE_SIZE
)

stats = {
    "received": 0,
    "inserted": 0,
    "bad_ts": 0,
}


# =========================
# STARTUP
# =========================
@app.on_event("startup")
async def startup():
    try:
        conn = make_conn()
        cursor = conn.cursor()

        try:
            cleanup_old_rows(cursor, conn)
            print("Startup DB cleanup completed")

        finally:
            cursor.close()
            conn.close()

    except Exception as e:
        print(f"Startup DB unavailable: {e}")

    # ALWAYS start flush loop
    asyncio.create_task(flush_loop())

    print("Application startup complete")

# =========================
# ROUTES
# =========================
@app.get("/")
async def root():
    return {"status": "running"}

@app.api_route("/health", methods=["GET", "POST", "HEAD"])
async def health():
    return {"status": "ok"}
    
@app.get("/ingest")
async def ingest_health():
    return {"ok": True, "queue_size": queue.qsize()}


@app.post("/ingest")
async def ingest(msg: IngestMsg):
    if not msg.payload:
        raise HTTPException(
            status_code=400,
            detail="payload required"
        )

    try:
        parse_ts_ist(msg.ts)
    except ValueError as e:
        stats["bad_ts"] += 1
        raise HTTPException(
            status_code=400,
            detail=str(e)
        )

    try:
        queue.put_nowait(msg)
    except asyncio.QueueFull:
        raise HTTPException(
            status_code=503,
            detail="queue full"
        )

    stats["received"] += 1

    return {"ok": True}

# =========================
# FLUSH LOOP
# =========================
async def flush_loop():
    conn = None
    cursor = None

    def ensure_sql():
        nonlocal conn, cursor
    
        if conn is not None:
            return
    
        try:
            conn = make_conn()
            cursor = conn.cursor()
    
            print("Database connection established")
    
        except Exception as e:
            print(f"Database unavailable: {e}")
    
            conn = None
            cursor = None
    
            raise

    def close_sql():
        # Close the DB connection between flushes so it isn't held open
        # idly. On scale-to-zero Neon plans, an idle-but-open connection
        # can prevent the compute from suspending and racking up cost.
        nonlocal conn, cursor

        try:
            if cursor:
                cursor.close()
        except:
            pass

        try:
            if conn:
                conn.close()
        except:
            pass

        conn = None
        cursor = None

    batch = []
    last_flush = time.time()
    # Cleanup every CLEANUP_INTERVAL_SEC seconds
    last_cleanup = time.time()

    while True:
        try:
            try:
                # Wait longer for a message when idle instead of polling
                # every 0.1s — this cuts wake-ups ~50x while keeping
                # responsiveness for the 20s flush interval.
                msg = await asyncio.wait_for(queue.get(), timeout=5.0)
                batch.append(msg)
            except asyncio.TimeoutError:
                pass

            now = time.time()
            time_due = (now - last_flush) >= FLUSH_INTERVAL_SEC
            size_due = len(batch) >= MAX_BATCH_SIZE

            if batch and (time_due or size_due):
                ensure_sql()
                if time.time() - last_cleanup > CLEANUP_INTERVAL_SEC:
                    cleanup_old_rows(cursor, conn)
                    last_cleanup = time.time()
                rows = []
                for m in batch:
                    dt = parse_ts_ist(m.ts)
                    payload_json = json.dumps(m.payload)
                    rows.append((dt, m.plant_id, payload_json))

                # execute_values sends one multi-row INSERT statement
                # instead of N separate INSERTs (executemany does NOT
                # batch under the hood for psycopg2) — much less DB CPU.
                execute_values(
                    cursor,
                    "INSERT INTO dbo.raw_plant_data (ts, plant_id, payload_json) VALUES %s",
                    rows,
                )
                conn.commit()

                stats["inserted"] += len(rows)
                print(f"Inserted {len(rows)} rows")

                batch.clear()
                last_flush = now

                # Close the connection now that the write is done, rather
                # than keeping it open until the next flush in ~20s.
                close_sql()

        except Exception as e:
            print(f"Flush error: {e}")

            try:
                if conn:
                    conn.rollback()
            except:
                pass

            close_sql()

            await asyncio.sleep(1)


# =========================
# MAIN
# =========================
if __name__ == "__main__":
    import uvicorn
    uvicorn.run("ml_api:app", host="0.0.0.0", port=7860)