Database & Data Warehouse Agents

9 min read Module 5 of 10 Topic 14 of 30

What you'll learn

  • Apply the principle of least privilege by creating read-only database roles scoped to agent query needs
  • Enforce parameterized queries across all agent-generated SQL to eliminate SQL injection risk
  • Implement BigQuery dry-run cost estimation and Snowflake query tagging for cost attribution
  • Cache repeated warehouse queries in Redis with TTL-based invalidation to reduce cloud spend
Building this at your company? For enterprise and company teams taking this to production: book a 30-minute session with our AI engineers for architecture guidance, code review, and a rollout plan for your use case.
Book a Team Session

Agents that can query your production databases are among the most powerful, and most dangerous, components in an enterprise system. A misconfigured agent with write access and poorly formed SQL can corrupt data, leak PII, or run a query that costs thousands of dollars in warehouse compute. This lesson builds the safety architecture around database agents: read-only roles, parameterized queries, cost controls for Snowflake and BigQuery, and Redis caching to avoid redundant expensive queries.

Principle of Least Privilege for Agent Database Roles

Never connect an agent to a database with more permissions than the narrowest possible definition of its job. For a query agent, that means a read-only role scoped to specific schemas: not the entire database, and certainly not a superuser or admin role. Create dedicated roles per agent type, not per agent instance, so permission management remains tractable.

-- PostgreSQL: create a read-only role for the analytics agent
CREATE ROLE analytics_agent_ro;
GRANT CONNECT ON DATABASE production TO analytics_agent_ro;
GRANT USAGE ON SCHEMA analytics, reporting TO analytics_agent_ro;
GRANT SELECT ON ALL TABLES IN SCHEMA analytics TO analytics_agent_ro;
GRANT SELECT ON ALL TABLES IN SCHEMA reporting TO analytics_agent_ro;

-- Crucially: deny access to PII schemas, write operations, and DDL
-- PostgreSQL denies everything not explicitly granted, so this
-- role has no access to the 'users', 'payments', or 'internal' schemas.

-- For future tables: ensure the role auto-inherits SELECT on new tables
ALTER DEFAULT PRIVILEGES IN SCHEMA analytics GRANT SELECT ON TABLES TO analytics_agent_ro;

-- Create a login user bound to this role
CREATE USER analytics_agent_user WITH PASSWORD '...';
GRANT analytics_agent_ro TO analytics_agent_user;

Parameterized Queries and SQL Injection Prevention

Agent-generated SQL introduces a unique injection risk: the agent itself constructs the query, potentially incorporating user-supplied values. String formatting SQL, f"SELECT * FROM orders WHERE customer_id = '{customer_id}'", is never safe. Always use parameterized queries, and add a validator that rejects non-SELECT statements before they reach the database.

import asyncpg
import re
from typing import Any

# Compile once at module load: cheap to call repeatedly
_ALLOWED_STATEMENT = re.compile(r"^\s*(SELECT|WITH)\b", re.IGNORECASE)
_DANGEROUS_KEYWORDS = re.compile(
    r"\b(INSERT|UPDATE|DELETE|DROP|ALTER|CREATE|TRUNCATE|EXECUTE|GRANT|REVOKE)\b",
    re.IGNORECASE,
)

def validate_agent_sql(sql: str) -> None:
    """
    Enforce that agent-generated SQL is read-only before it reaches the database.
    This is a defense-in-depth layer, the read-only role is the primary control,
    but this catches logic errors in the agent's SQL generation early.
    """
    if not _ALLOWED_STATEMENT.match(sql):
        raise ValueError(f"Agent SQL must begin with SELECT or WITH, got: {sql[:80]}")
    if _DANGEROUS_KEYWORDS.search(sql):
        raise ValueError("Agent SQL contains prohibited data-manipulation keywords")

async def run_agent_query(
    pool: asyncpg.Pool,
    sql: str,
    params: list[Any],
    max_rows: int = 10_000,
    timeout_seconds: float = 30.0,
) -> list[dict]:
    """
    Execute an agent-generated query with all safety controls active.
    params are always passed separately, asyncpg binds them server-side,
    making SQL injection structurally impossible regardless of param content.
    """
    validate_agent_sql(sql)

    # Append a LIMIT clause to prevent unbounded result sets.
    # We do this at the Python layer so the agent cannot override it.
    bounded_sql = f"SELECT * FROM ({sql}) AS _agent_query LIMIT {max_rows + 1}"

    async with pool.acquire() as conn:
        rows = await conn.fetch(
            bounded_sql,
            *params,
            timeout=timeout_seconds,  # asyncpg enforces this at the connection level
        )

    if len(rows) > max_rows:
        raise ValueError(
            f"Query returned more than {max_rows} rows. Refine the query or add filters."
        )

    return [dict(row) for row in rows]

BigQuery Cost Controls with Dry-Run

BigQuery charges per byte scanned. A single unguarded query against a large table can cost hundreds of dollars. Use the BigQuery dry-run feature to estimate cost before execution, the API performs query parsing and statistics estimation without executing the query, returning the number of bytes that would be processed.

flowchart TD
    A["Agent SQL Query"] --> B["Redis Cache Lookup\nkey = hash(sql + params)"]
    B -->|Cache Hit| C["Return Cached Result\n(zero compute cost)"]
    B -->|Cache Miss| D["BigQuery Dry-Run\nestimate bytes scanned"]
    D --> E{"Cost > $threshold?"}
    E -->|Yes| F["Reject Query\nReturn cost estimate to agent"]
    E -->|No| G["Execute Query\nBigQuery Job API"]
    G --> H["Store Result in Redis\nTTL = freshness window"]
    H --> I["Return Result"]

    style A fill:#EEF0F7,stroke:#6366F1,color:#0F172A
    style B fill:#fff7ed,stroke:#f59e0b,color:#0F172A
    style C fill:#f0fdf9,stroke:#0D9488,color:#0F172A
    style D fill:#EEF0F7,stroke:#6366F1,color:#0F172A
    style E fill:#fff7ed,stroke:#f59e0b,color:#0F172A
    style F fill:#EEF0F7,stroke:#6366F1,color:#0F172A
    style G fill:#EEF0F7,stroke:#6366F1,color:#0F172A
    style H fill:#f0fdf9,stroke:#0D9488,color:#0F172A
    style I fill:#f0fdf9,stroke:#0D9488,color:#0F172A

Snowflake Query Tagging and Timeout Enforcement

Snowflake’s query history is queryable: but only useful for cost attribution if every agent query is tagged with the agent type, session ID, and task ID. Set query tags via ALTER SESSION SET QUERY_TAG before executing each query. This lets you run INFORMATION_SCHEMA.QUERY_HISTORY aggregations to see which agent types drive the most cost, which queries are slowest, and whether a specific task triggered a warehouse surge.

import snowflake.connector
import json

async def run_snowflake_query(
    conn: snowflake.connector.SnowflakeConnection,
    sql: str,
    params: dict,
    agent_id: str,
    task_id: str,
    timeout_seconds: int = 30,
) -> list[dict]:
    """
    Execute a Snowflake query with tagging and timeout.
    QUERY_TAG is stored in QUERY_HISTORY and is the primary mechanism
    for per-agent cost attribution in Snowflake's billing views.
    """
    tag = json.dumps({"agent_id": agent_id, "task_id": task_id, "service": "agent-network"})

    with conn.cursor(snowflake.connector.DictCursor) as cur:
        # Tag this session: every subsequent query carries this tag
        cur.execute(f"ALTER SESSION SET QUERY_TAG = '{tag}'")

        # Enforce a hard statement timeout: prevents runaway full-table scans
        cur.execute(f"ALTER SESSION SET STATEMENT_TIMEOUT_IN_SECONDS = {timeout_seconds}")

        # Parameterized execution via %(name)s syntax: Snowflake binds server-side
        cur.execute(sql, params)
        return cur.fetchmany(10_000)  # Never fetch unlimited rows

Caching Repeated Warehouse Queries

sequenceDiagram
    participant AG1 as Agent Instance 1
    participant AG2 as Agent Instance 2
    participant RC as Redis Cache
    participant BQ as BigQuery / Snowflake

    AG1->>RC: GET cache:sql_hash_abc123
    RC-->>AG1: MISS
    AG1->>BQ: Execute Query ($12 cost)
    BQ-->>AG1: 8,000 rows
    AG1->>RC: SET cache:sql_hash_abc123 TTL=3600s
    AG2->>RC: GET cache:sql_hash_abc123
    RC-->>AG2: HIT, return cached rows ($0 cost)

    style AG1 fill:#EEF0F7,stroke:#6366F1,color:#0F172A
    style AG2 fill:#EEF0F7,stroke:#6366F1,color:#0F172A
    style RC fill:#fff7ed,stroke:#f59e0b,color:#0F172A
    style BQ fill:#f0fdf9,stroke:#0D9488,color:#0F172A

The cache key should be a stable hash of the normalized SQL plus bound parameter values. Normalize by stripping comments and collapsing whitespace before hashing, this ensures that semantically identical queries from different agents share a cache entry. Set TTL based on how frequently the underlying data changes: hourly for aggregated reporting data, 5 minutes for near-real-time operational queries, 24 hours for historical analytics that never change.

In the next lesson, we shift from consuming enterprise APIs to exposing agent capabilities: building custom MCP servers that package your enterprise tools as standardized, versioned capabilities any agent in the network can discover and call.

Knowledge Check

3 questions to test your understanding

1 An agent generates SQL from user natural language input and executes it against a production PostgreSQL database. The agent's database user has full read/write access. What is the primary risk?

2 Your BigQuery agent runs a complex analytical query that costs $45 in slot time. The same query is requested 200 times per day by different agents. What is the correct cost control?

3 An agent is querying Snowflake. A runaway query with a missing WHERE clause returns 50 million rows. What mechanisms prevent this from causing an incident?

Discussion

Questions and notes from learners on this topic

Loading discussion…

Go further with expert guidance

Ready to build production AI?
Talk to our R&D team.

These courses give you the foundation. Our embedded AI teams take you from prototype to production in 30–90 days, with your team, your codebase, your goals. Book a free strategy call to see how we can accelerate your AI initiative.

30 minutes · No obligation · Expert AI engineers, not sales reps

AI Architecture Review

Audit your current stack and identify high-impact improvements

Project Review

Get expert feedback on your AI implementation and codebase

Team Mentoring

Upskill your engineers with hands-on AI coaching sessions

AI Strategy

Define your AI roadmap, prioritization, and implementation plan