cleveragents-core/robot/helper_database_resources.py

"""Helper utilities for database resource Robot smoke tests.

Each command prints a single ``<command>-ok`` token on success so the
calling Robot test can assert on ``stdout``.  Only uses sqlite3 (stdlib)
for concrete database testing.
"""

from __future__ import annotations

import os
import sqlite3
import sys
import tempfile

from cleveragents.infrastructure.sandbox.factory import SandboxFactory
from cleveragents.infrastructure.sandbox.transaction_sandbox import (
    TransactionSandbox,
)
from cleveragents.resource.handlers.database import (
    DATABASE_TYPE_DEFS,
    DATABASE_TYPE_NAMES,
    DatabaseResourceHandler,
    resolve_connection_args,
    validate_connection,
)
from cleveragents.resource.handlers.protocol import ResourceHandler
from cleveragents.shared.redaction import mask_database_url


def _protocol_check() -> None:
    """Verify DatabaseResourceHandler satisfies ResourceHandler."""
    handler = DatabaseResourceHandler()
    assert isinstance(handler, ResourceHandler), (
        "DatabaseResourceHandler not a ResourceHandler"
    )
    print("protocol-check-ok")


def _type_defs() -> None:
    """Verify all 4 database type definitions are present."""
    names = {td["name"] for td in DATABASE_TYPE_DEFS}
    for expected in ("postgres", "mysql", "sqlite", "duckdb"):
        assert expected in names, f"Missing type def: {expected}"
    assert names == DATABASE_TYPE_NAMES
    print("type-defs-ok")


def _sqlite_validate() -> None:
    """Validate SQLite connection with a temp file."""
    fd, path = tempfile.mkstemp(suffix=".db")
    os.close(fd)
    try:
        success, msg = validate_connection("sqlite", {"path": path})
        assert success, f"SQLite validation failed: {msg}"
        print("sqlite-validate-ok")
    finally:
        os.unlink(path)


def _tx_commit() -> None:
    """Test transaction sandbox commit persists data."""
    fd, path = tempfile.mkstemp(suffix=".db")
    os.close(fd)
    try:
        # Setup table
        conn = sqlite3.connect(path)
        conn.execute("CREATE TABLE items (id INTEGER PRIMARY KEY, name TEXT)")
        conn.commit()
        conn.close()

        # Use sandbox
        sandbox = TransactionSandbox(resource_id="robot-tx-1", original_path=path)
        sandbox.create(plan_id="PLAN-ROBOT-COMMIT")
        sandbox.execute("INSERT INTO items (id, name) VALUES (1, 'committed')")
        result = sandbox.commit("robot commit")
        assert result.success
        sandbox.cleanup()

        # Verify data persisted
        conn = sqlite3.connect(path)
        rows = conn.execute("SELECT name FROM items WHERE id = 1").fetchall()
        conn.close()
        assert len(rows) == 1
        assert rows[0][0] == "committed"
        print("tx-commit-ok")
    finally:
        os.unlink(path)


def _tx_rollback() -> None:
    """Test transaction sandbox rollback discards data."""
    fd, path = tempfile.mkstemp(suffix=".db")
    os.close(fd)
    try:
        # Setup table
        conn = sqlite3.connect(path)
        conn.execute("CREATE TABLE items (id INTEGER PRIMARY KEY, name TEXT)")
        conn.commit()
        conn.close()

        # Use sandbox
        sandbox = TransactionSandbox(resource_id="robot-tx-2", original_path=path)
        sandbox.create(plan_id="PLAN-ROBOT-ROLLBACK")
        sandbox.execute("INSERT INTO items (id, name) VALUES (2, 'rolled-back')")
        # Transition to ACTIVE by executing, then rollback
        sandbox.rollback()
        sandbox.cleanup()

        # Verify data NOT persisted
        conn = sqlite3.connect(path)
        rows = conn.execute("SELECT name FROM items WHERE id = 2").fetchall()
        conn.close()
        assert len(rows) == 0
        print("tx-rollback-ok")
    finally:
        os.unlink(path)


def _cred_mask() -> None:
    """Verify credential masking in validation messages."""
    props = {"connection-string": "postgresql://admin:supersecret@db:5432/app"}
    resolved = resolve_connection_args(props, "postgres")
    _success, msg = validate_connection("postgres", resolved)
    assert "supersecret" not in msg, f"Secret leaked: {msg}"
    # Also test mask_database_url directly
    masked = mask_database_url("postgresql://admin:supersecret@db:5432/app")
    assert "supersecret" not in masked
    print("cred-mask-ok")


def _factory_routing() -> None:
    """Verify SandboxFactory routes transaction_rollback correctly."""
    factory = SandboxFactory()
    sandbox = factory.create_sandbox(
        resource_id="factory-test",
        original_path=":memory:",
        sandbox_strategy="transaction_rollback",
    )
    assert isinstance(sandbox, TransactionSandbox)
    assert SandboxFactory.is_supported("transaction_rollback")
    print("factory-routing-ok")


def _read_only() -> None:
    """Verify read-only sandbox prevents writes."""
    fd, path = tempfile.mkstemp(suffix=".db")
    os.close(fd)
    try:
        conn = sqlite3.connect(path)
        conn.execute("CREATE TABLE items (id INTEGER PRIMARY KEY, name TEXT)")
        conn.commit()
        conn.close()

        sandbox = TransactionSandbox(
            resource_id="ro-test", original_path=path, read_only=True
        )
        sandbox.create(plan_id="PLAN-RO-ROBOT")
        rejected = False
        try:
            sandbox.execute("INSERT INTO items (id, name) VALUES (1, 'forbidden')")
        except Exception:
            rejected = True
        finally:
            sandbox.cleanup()
        assert rejected, "Write was not rejected in read-only mode"
        print("read-only-ok")
    finally:
        os.unlink(path)


_COMMANDS = {
    "protocol-check": _protocol_check,
    "type-defs": _type_defs,
    "sqlite-validate": _sqlite_validate,
    "tx-commit": _tx_commit,
    "tx-rollback": _tx_rollback,
    "cred-mask": _cred_mask,
    "factory-routing": _factory_routing,
    "read-only": _read_only,
}


if __name__ == "__main__":
    if len(sys.argv) < 2 or sys.argv[1] not in _COMMANDS:
        print(
            f"Usage: {sys.argv[0]} <{'|'.join(_COMMANDS)}>",
            file=sys.stderr,
        )
        sys.exit(1)
    _COMMANDS[sys.argv[1]]()