SQL Read-Only Policy

from brane import CapabilityDeniedError, Decision, Effect, Runtime

runtime = Runtime(
    agent_id="support-agent",
    environment="prod",
    tenant_id="tenant_acme",
)

@runtime.capability(
    name="execute_sql",
    type="database",
    risk="high",
    effect=Effect(type="database_query", reversible=True),
    data_namespace="customer.records",
    owner="data-team",
)
def execute_sql(query: str, params: dict | None = None):
    return {"rows": [], "count": 0}

@runtime.before_capability(
    "execute_sql",
    name="sql_read_only",
    version="1.0",
    description="Only SELECT statements are allowed",
    priority=100,
)
def sql_read_only(ctx):
    query = ctx.arg("query", "").strip().lower()
    allowed_prefixes = ("select", "with", "explain")
    if not any(query.startswith(p) for p in allowed_prefixes):
        return Decision(
            type="deny",
            reason=f"Only SELECT queries are allowed. Got: {query[:40]}",
        )
    return Decision(type="allow")

try:
    result = execute_sql("SELECT * FROM customers WHERE id = 42")
    print(result)

    result = execute_sql("WITH cte AS (SELECT ...) SELECT * FROM cte")
    print(result)

    execute_sql("DELETE FROM customers WHERE id = 42")
except CapabilityDeniedError as e:
    print(f"Blocked: {e.reason}")
    print(f"Policy: {e.policy_name}")

@runtime.before_capability("execute_sql")
def sql_read_only_strict(ctx):
    query = ctx.arg("query", "").strip().lower()
    if query.startswith("explain analyze"):
        return Decision(type="deny", reason="EXPLAIN ANALYZE is not allowed")
    if not any(query.startswith(p) for p in ("select", "with", "explain")):
        return Decision(type="deny", reason="Only SELECT queries are allowed")
    return Decision(type="allow")

@runtime.before_capability("execute_sql")
def sql_policy(ctx):
    query = ctx.arg("query", "").strip().lower()
    if ctx.is_prod:
        if not query.startswith("select"):
            return Decision(type="deny", reason="Only SELECT in prod")
    else:
        blocked = ("drop", "truncate")
        if any(query.startswith(b) for b in blocked):
            return Decision(type="deny", reason="DROP and TRUNCATE are never allowed")
    return Decision(type="allow")