lambdagent/agentpaas/db/models.py

"""
agentpaas.db.models — SQLAlchemy ORM models.

Core tables: tenants, agents, agent_versions, runs, jobs, api_keys, usage_records.
Principle P3: Works with both SQLite and PostgreSQL.
"""
from __future__ import annotations
import uuid
import hashlib
import json
import time
from datetime import datetime, timezone
from typing import Optional, Dict, Any, List

import sqlite3


def now_utc():
    return datetime.now(timezone.utc).isoformat()


def gen_id(prefix: str = "") -> str:
    return f"{prefix}{uuid.uuid4().hex[:12]}"


class Database:
    """
    Lightweight database abstraction. Uses SQLite directly for MVP.
    Designed for easy migration to SQLAlchemy/PostgreSQL later.
    """

    def __init__(self, db_url: str = "sqlite:///./agentpaas.db"):
        if "sqlite" in db_url:
            db_path = db_url.replace("sqlite:///", "")
            self.conn = sqlite3.connect(db_path, check_same_thread=False)
            self.conn.row_factory = sqlite3.Row
        else:
            raise NotImplementedError("Only SQLite supported in MVP. Set AGENTPAAS_DATABASE_URL=sqlite:///./agentpaas.db")
        self._init_tables()

    def _init_tables(self):
        c = self.conn.cursor()
        c.executescript("""
            CREATE TABLE IF NOT EXISTS tenants (
                id TEXT PRIMARY KEY,
                name TEXT NOT NULL,
                plan TEXT DEFAULT 'free',
                quota_tokens_monthly INTEGER DEFAULT 1000000,
                quota_concurrency INTEGER DEFAULT 5,
                quota_agents INTEGER DEFAULT 20,
                status TEXT DEFAULT 'active',
                created_at TEXT
            );

            CREATE TABLE IF NOT EXISTS users (
                id TEXT PRIMARY KEY,
                tenant_id TEXT REFERENCES tenants(id),
                email TEXT,
                role TEXT DEFAULT 'developer',
                created_at TEXT
            );

            CREATE TABLE IF NOT EXISTS api_keys (
                id TEXT PRIMARY KEY,
                tenant_id TEXT REFERENCES tenants(id),
                user_id TEXT REFERENCES users(id),
                key_hash TEXT NOT NULL UNIQUE,
                key_prefix TEXT NOT NULL,
                name TEXT,
                scopes TEXT DEFAULT '[]',
                rate_limit INTEGER DEFAULT 60,
                expires_at TEXT,
                last_used_at TEXT,
                status TEXT DEFAULT 'active',
                created_at TEXT
            );

            CREATE TABLE IF NOT EXISTS agents (
                id TEXT PRIMARY KEY,
                tenant_id TEXT REFERENCES tenants(id),
                name TEXT NOT NULL,
                description TEXT DEFAULT '',
                current_version INTEGER DEFAULT 1,
                tags TEXT DEFAULT '[]',
                environment TEXT DEFAULT 'production',
                traffic_rules TEXT,
                status TEXT DEFAULT 'active',
                created_at TEXT,
                updated_at TEXT
            );

            CREATE TABLE IF NOT EXISTS agent_versions (
                agent_id TEXT REFERENCES agents(id),
                version INTEGER NOT NULL,
                config TEXT NOT NULL,
                config_hash TEXT NOT NULL,
                changelog TEXT DEFAULT '',
                created_by TEXT,
                created_at TEXT,
                PRIMARY KEY (agent_id, version)
            );

            CREATE TABLE IF NOT EXISTS runs (
                id TEXT PRIMARY KEY,
                agent_id TEXT REFERENCES agents(id),
                agent_version INTEGER,
                tenant_id TEXT REFERENCES tenants(id),
                job_id TEXT,
                input TEXT,
                output TEXT,
                status TEXT DEFAULT 'running',
                input_tokens INTEGER DEFAULT 0,
                output_tokens INTEGER DEFAULT 0,
                steps INTEGER DEFAULT 0,
                duration_ms INTEGER DEFAULT 0,
                error TEXT,
                trace_id TEXT,
                created_at TEXT,
                completed_at TEXT
            );

            CREATE TABLE IF NOT EXISTS jobs (
                id TEXT PRIMARY KEY,
                tenant_id TEXT REFERENCES tenants(id),
                agent_id TEXT REFERENCES agents(id),
                priority INTEGER DEFAULT 5,
                status TEXT DEFAULT 'queued',
                input TEXT,
                worker_id TEXT,
                attempts INTEGER DEFAULT 0,
                max_attempts INTEGER DEFAULT 3,
                started_at TEXT,
                completed_at TEXT,
                created_at TEXT
            );

            CREATE TABLE IF NOT EXISTS usage_records (
                id TEXT PRIMARY KEY,
                tenant_id TEXT REFERENCES tenants(id),
                agent_id TEXT REFERENCES agents(id),
                run_id TEXT REFERENCES runs(id),
                model TEXT,
                input_tokens INTEGER DEFAULT 0,
                output_tokens INTEGER DEFAULT 0,
                tool_calls INTEGER DEFAULT 0,
                duration_ms INTEGER DEFAULT 0,
                cost_usd REAL DEFAULT 0.0,
                created_at TEXT
            );
        """)
        self.conn.commit()

    def execute(self, sql: str, params: tuple = ()) -> sqlite3.Cursor:
        return self.conn.execute(sql, params)

    def commit(self):
        self.conn.commit()

    def fetchone(self, sql: str, params: tuple = ()) -> Optional[dict]:
        row = self.conn.execute(sql, params).fetchone()
        return dict(row) if row else None

    def fetchall(self, sql: str, params: tuple = ()) -> List[dict]:
        return [dict(r) for r in self.conn.execute(sql, params).fetchall()]