Function execute_readonly_query

web/pgadmin/llm/tools/database.py:312–395

Execute a read-only SQL query against a PostgreSQL database. This function: 1. Opens a new connection with LLM-specific application_name 2. Starts a READ ONLY transaction 3. Executes the query 4. Returns results (limited to max_rows) 5. Rolls back and closes the connect

(
    sid: int,
    did: int,
    query: str,
    max_rows: int = 1000
)

def execute_readonly_query(
    sid: int,
    did: int,
    query: str,
    max_rows: int = 1000
) -> dict:
    """
    Execute a read-only SQL query against a PostgreSQL database.

    This function:
    1. Opens a new connection with LLM-specific application_name
    2. Starts a READ ONLY transaction
    3. Executes the query
    4. Returns results (limited to max_rows)
    5. Rolls back and closes the connection

    Args:
        sid: Server ID from the Object Explorer
        did: Database ID (OID) from the Object Explorer
        query: SQL query to execute (should be SELECT or read-only)
        max_rows: Maximum number of rows to return (default 1000)

    Returns:
        Dictionary containing:
        - columns: List of column names
        - rows: List of row data (as lists)
        - row_count: Number of rows returned
        - truncated: True if results were limited

    Raises:
        DatabaseToolError: If the query is rejected by validation, the
            query fails at runtime, or a connection cannot be
            established.
    """
    # Validate the LLM-supplied query before allocating a connection.
    # The BEGIN TRANSACTION READ ONLY wrapper below is only effective
    # if the query is a single read-only statement; see
    # _validate_readonly_query for the threat model and rules.
    _validate_readonly_query(query)

    # Generate unique connection ID for this LLM query
    conn_id = f"llm_{secrets.choice(range(1, 9999999))}"

    manager = None
    conn = None

    try:
        # Get connection manager and connection object
        manager, conn = _get_connection(sid, did, conn_id)

        # Connect with read-only settings
        status, error = _connect_readonly(manager, conn, conn_id)
        if not status:
            raise DatabaseToolError(
                f"Connection failed: {error}",
                code="CONNECTION_ERROR"
            )