"""Helper utilities for database resource Robot smoke tests. Each command prints a single ``-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]]()