[Feature] Integration Tool for Neon Serverless Postgres

[Boburmirzo]

Description

We would like to see a new tool in Arcade that allows AI agents to interact with Neon, the serverless Postgres platform. This tool would let agents create and manage Neon databases, branches, and perform basic operations. It would be useful for AI workflows that rely on structured data storage, such as RAG pipelines, multi-tenant apps, and version-controlled AI experiments. Developers building AI agents could fully automate their database operations in real time.

With Neon, AI agents can create a new database or branch in milliseconds, making it ideal for dynamic and ephemeral workloads driven by AI logic.

Proposed Feature

A new neon_tool that supports:

• Creating and deleting Neon projects and databases
• Creating and managing branches (Neon branching guide)
• Running SQL queries
• Viewing project metadata

The tool should work with the Neon API and accept API key as credentials.

Alternatives Considered

One option is to use custom Python tools inside the agent logic or use the Neon MCP server, but it adds a lot of boilerplate and security concerns. Developers would need to manually handle auth, API calls, and error handling. A built-in Arcade tool would simplify everything and make it reusable for many apps.

Implementation Strategy

Create a new neon_tool using Arcade’s existing tool SDK

Wrap Neon’s REST API for common operations

Use environment variables for API keys

Allow tool parameters like branch_name, sql=..., etc.

Return structured output to agents for decision making

Optional: add logging and response summaries to make agent decisions more transparent.

Use Case

Imagine an AI agent that spins up a temporary database branch when a user starts a new project. After processing or testing, it automatically deletes the branch to save costs. Another example: an agent that runs analytics by querying data directly from a Neon database and replying in natural language. I also wrote about why to use a separate database per agent in this article.

Example 1: An AI agent instantly creates a new Neon database for a user session and deletes it when the task is done, optimizing cost and isolation.
Example 2: An AI assistant managing multi-tenant apps can create isolated branches per customer for analytics and experimentation.

With built-in Neon support, these workflows become much easier and more secure for developers.

Additional Context
Neon is a fast-growing serverless Postgres provider used in many AI and SaaS apps. Adding Arcade integration opens up many automation and AI-driven data workflows. This would also be the first database management tool in Arcade’s toolset, enabling a whole new category of use cases.

Thanks for considering this feature! Happy to contribute to community tool or test prototypes if needed.

[Spartee]

Tool Storage is absolutely something we've got on the roadmap so I'm very happy to see this idea come up.

I don't think the @tool decorator would be easy or clean to abstract though. Decorators in general are not really meant for this and I've regretted doing it in the past.

This said, what about starting with a base toolkit for Neon?

Something that an Arcade Tool Developer could build on for their agent use case or Agent developers if the use case was truly to manage Neon with an agent.

important note: I make this distinction because, in most cases besides someone deploying or managing infra, the agent doens't have access to project_id in the context.

But if we provide this as a pip package it's really easy for people to build on top and remove the need for the agent to know this by using requires_secrets. Usually this also helps with making the parameter annotations and descriptions for the tool closer to the business logic the agent needs to perform the right task.

Here's an example of what I mean

instead of

@tool(
    requires_secrets=["NEON_API_KEY"]
)
async def create_neon_project(
    context: ToolContext,
    project_name: Annotated[str, "Name for the new Neon project"],
    region: Annotated[str, "AWS region for the project (e.g., 'us-east-1')"] = "us-east-1",
    postgres_version: Annotated[int, "PostgreSQL version (14, 15, or 16)"] = 16
) -> Annotated[Dict[str, Any], "Created project information including connection details"]:
    """
    Create a new Neon project with a Postgres database.
    
    <use_case>
    Use this tool when someone is trying to create a database. This provisions
    a new serverless Postgres cluster that can be used for data storage,
    analytics, or as a backend for AI applications.
    </use_case>
    
    <important_notes>
    The tool will automatically create a default database and provide connection
    details. The connection string can be used immediately with any PostgreSQL client.
    </important_notes>
    """

You can have something easier for the agent to grok.

This can be set in the dashboard for all users or you can use the username from the context as a prefix to a project id if you need it to be user isolated. Either way, Agents tend to do much better with something higher level. Building a Neon Toolkit would enable everyone to build that for their use case.

@tool(
    requires_secrets=[
        "NEON_API_KEY", 
        "NEON_PROJECT_NAME_PREFIX", 
        "NEON_REGION", 
        "NEON_POSTGRES_VERSION"
    ]
)
async def create_agent_session(
    context: ToolContext,
    agent_type: Annotated[str, "Type of agent (e.g., 'rag', 'analytics', 'chatbot')"] = "general"
) -> Annotated[Dict[str, Any], "Created agent session id with new database to use for storage of tool results and connection details"]:
    """
    Create an isolated database environment for an AI agent session.
    
    Use this tool when you need to:
    - Process user data that requires isolation from other users or sessions
    - Store temporary results, embeddings, or computed values during complex workflows
    - Create a dedicated workspace for data analysis, RAG pipelines, or multi-step operations
    - Ensure data privacy and security for sensitive user interactions
    - Set up ephemeral storage that can be easily cleaned up after the session ends
    
    This tool automatically provisions a fresh database instance, eliminating setup overhead
    and ensuring each agent session has its own clean data environment.
    """

Implementation of the Base Neon Toolkit

It's important to note that this was vibe coded and posted by Arcade.dev

"""
Neon Toolkit for Arcade
Provides serverless Postgres database management capabilities
"""

import asyncio
import json
import logging
import os
import uuid
from contextlib import asynccontextmanager
from typing import Annotated, Dict, List, Optional, Any, Union
import httpx
import asyncpg
from datetime import datetime

from arcade.sdk import ToolContext, tool
from arcade.sdk.errors import RetryableToolError

logger = logging.getLogger(__name__)

# Global connection pool for reuse across tools
_connection_pools: Dict[str, asyncpg.Pool] = {}
_http_client: Optional[httpx.AsyncClient] = None

class NeonStorage:
    """Enhanced storage interface for Neon toolkit state management"""
    
    def __init__(self, context: ToolContext):
        self.context = context
        self._storage = {}
        self._migrations = {}
        self._tuning_sessions = {}
    
    def set_connection_string(self, project_id: str, connection_string: str):
        """Store connection string for a project"""
        self._storage[f"conn_{project_id}"] = connection_string
    
    def get_connection_string(self, project_id: str) -> Optional[str]:
        """Retrieve connection string for a project"""
        return self._storage.get(f"conn_{project_id}")
    
    def set_branch_info(self, branch_id: str, info: Dict[str, Any]):
        """Store branch metadata"""
        self._storage[f"branch_{branch_id}"] = info
    
    def get_branch_info(self, branch_id: str) -> Optional[Dict[str, Any]]:
        """Retrieve branch metadata"""
        return self._storage.get(f"branch_{branch_id}")
    
    def store_migration(self, migration_id: str, migration_data: Dict[str, Any]):
        """Store migration state for later completion"""
        self._migrations[migration_id] = migration_data
    
    def get_migration(self, migration_id: str) -> Optional[Dict[str, Any]]:
        """Retrieve migration state"""
        return self._migrations.get(migration_id)
    
    def store_tuning_session(self, tuning_id: str, session_data: Dict[str, Any]):
        """Store query tuning session data"""
        self._tuning_sessions[tuning_id] = session_data
    
    def get_tuning_session(self, tuning_id: str) -> Optional[Dict[str, Any]]:
        """Retrieve query tuning session data"""
        return self._tuning_sessions.get(tuning_id)

async def startup_hook():
    """Initialize HTTP client and connection pools"""
    global _http_client
    
    logger.info("Initializing Neon toolkit...")
    
    # Initialize HTTP client for Neon API calls
    _http_client = httpx.AsyncClient(
        timeout=httpx.Timeout(30.0),
        headers={
            "Content-Type": "application/json",
            "User-Agent": "Arcade-Neon-Toolkit/1.0"
        }
    )
    
    logger.info("Neon toolkit initialized successfully")

async def shutdown_hook():
    """Clean up connections and HTTP client"""
    global _http_client, _connection_pools
    
    logger.info("Shutting down Neon toolkit...")
    
    # Close all database connection pools
    for pool_name, pool in _connection_pools.items():
        logger.info(f"Closing connection pool: {pool_name}")
        await pool.close()
    
    _connection_pools.clear()
    
    # Close HTTP client
    if _http_client:
        await _http_client.aclose()
        _http_client = None
    
    logger.info("Neon toolkit shutdown complete")

@asynccontextmanager
async def get_database_connection(connection_string: str):
    """Get a database connection with proper cleanup"""
    pool_key = connection_string
    
    if pool_key not in _connection_pools:
        # Create new connection pool
        _connection_pools[pool_key] = await asyncpg.create_pool(
            connection_string,
            min_size=1,
            max_size=10,
            command_timeout=60
        )
    
    pool = _connection_pools[pool_key]
    async with pool.acquire() as connection:
        yield connection

async def make_neon_api_request(
    api_key: str, 
    method: str, 
    endpoint: str, 
    data: Optional[Dict] = None,
    retry_count: int = 1
) -> Dict[str, Any]:
    """Make authenticated request to Neon API with retry logic"""
    if not _http_client:
        raise RetryableToolError("HTTP client not initialized")
    
    url = f"https://console.neon.tech/api/v2/{endpoint}"
    headers = {"Authorization": f"Bearer {api_key}"}
    
    for attempt in range(retry_count + 1):
        try:
            if method == "GET":
                response = await _http_client.get(url, headers=headers)
            elif method == "POST":
                response = await _http_client.post(url, headers=headers, json=data)
            elif method == "DELETE":
                response = await _http_client.delete(url, headers=headers)
            else:
                raise ValueError(f"Unsupported HTTP method: {method}")
            
            response.raise_for_status()
            return response.json()
        
        except httpx.HTTPStatusError as e:
            if attempt < retry_count and e.response.status_code >= 500:
                logger.warning(f"Attempt {attempt + 1} failed, retrying...")
                await asyncio.sleep(2 ** attempt)  # Exponential backoff
                continue
            elif e.response.status_code >= 500:
                raise RetryableToolError(f"Neon API server error: {e}")
            else:
                raise Exception(f"Neon API client error: {e}")
        except httpx.RequestError as e:
            if attempt < retry_count:
                logger.warning(f"Request attempt {attempt + 1} failed, retrying...")
                await asyncio.sleep(2 ** attempt)
                continue
            else:
                raise RetryableToolError(f"Neon API request failed: {e}")

def split_sql_statements(sql: str) -> List[str]:
    """Split SQL string into individual statements"""
    # Simple implementation - in production, use a proper SQL parser
    statements = [stmt.strip() for stmt in sql.split(';') if stmt.strip()]
    return statements

@tool(
    requires_secrets=["NEON_API_KEY"]
)
async def list_neon_projects(
    context: ToolContext,
    limit: Annotated[int, "Maximum number of projects to return"] = 10
) -> Annotated[List[Dict[str, Any]], "List of Neon projects"]:
    """
    Lists the first N Neon projects in your account.
    
    If you can't find the project you're looking for, increase the limit
    by passing a higher value to the limit parameter.
    """
    api_key = context.get_secret("NEON_API_KEY")
    
    try:
        result = await make_neon_api_request(
            api_key, "GET", f"projects?limit={limit}"
        )
        
        projects = result.get("projects", [])
        logger.info(f"Retrieved {len(projects)} Neon projects")
        
        return [
            {
                "id": project["id"],
                "name": project["name"],
                "region": project["region_id"],
                "created_at": project["created_at"],
                "pg_version": project["pg_version"]
            }
            for project in projects
        ]
        
    except Exception as e:
        raise RetryableToolError(f"Failed to list Neon projects: {str(e)}")

@tool(
    requires_secrets=["NEON_API_KEY"]
)
async def create_neon_project(
    context: ToolContext,
    project_name: Annotated[str, "Name for the new Neon project"],
    region: Annotated[str, "AWS region for the project (e.g., 'us-east-1')"] = "us-east-1",
    postgres_version: Annotated[int, "PostgreSQL version (14, 15, or 16)"] = 16
) -> Annotated[Dict[str, Any], "Created project information including connection details"]:
    """
    Create a new Neon project with a Postgres database.
    
    <use_case>
    Use this tool when someone is trying to create a database. This provisions
    a new serverless Postgres cluster that can be used for data storage,
    analytics, or as a backend for AI applications.
    </use_case>
    
    <important_notes>
    The tool will automatically create a default database and provide connection
    details. The connection string can be used immediately with any PostgreSQL client.
    </important_notes>
    """
    api_key = context.get_secret("NEON_API_KEY")
    storage = NeonStorage(context)
    
    project_data = {
        "project": {
            "name": project_name,
            "region_id": region,
            "pg_version": postgres_version
        }
    }
    
    try:
        result = await make_neon_api_request(
            api_key, "POST", "projects", project_data, retry_count=1
        )
        
        project = result["project"]
        branch = result["branch"]
        databases = result["databases"]
        
        # Get connection string for the new project
        connection_uri_result = await make_neon_api_request(
            api_key, "GET", 
            f"projects/{project['id']}/connection_uri?" +
            f"database_name={databases[0]['name']}&role_name={databases[0]['owner_name']}"
        )
        
        connection_uri = connection_uri_result["uri"]
        
        # Store connection information
        storage.set_connection_string(project["id"], connection_uri)
        
        logger.info(f"Created Neon project: {project['id']}")
        
        return {
            "project_id": project["id"],
            "name": project["name"],
            "region": project["region_id"],
            "postgres_version": project["pg_version"],
            "branch_id": branch["id"],
            "branch_name": branch["name"],
            "database_name": databases[0]["name"],
            "connection_uri": connection_uri,
            "created_at": project["created_at"],
            "dashboard_url": f"https://console.neon.tech/app/projects/{project['id']}"
        }
        
    except Exception as e:
        if "project limit" in str(e).lower():
            raise RetryableToolError(
                "You have reached your Neon project limit. "
                "Please upgrade your account at https://console.neon.tech/app/billing"
            )
        raise RetryableToolError(f"Failed to create Neon project: {str(e)}")

@tool(
    requires_secrets=["NEON_API_KEY"]
)
async def create_database_branch(
    context: ToolContext,
    project_id: Annotated[str, "Neon project ID"],
    branch_name: Annotated[Optional[str], "Name for the new branch"] = None,
    parent_branch_id: Annotated[Optional[str], "Parent branch ID (defaults to main)"] = None
) -> Annotated[Dict[str, Any], "Created branch information"]:
    """
    Create a new database branch from an existing branch.
    
    <use_case>
    Branches provide instant point-in-time copies of your database,
    perfect for testing, experiments, or isolated user sessions.
    Leverages Neon's branching feature for development, testing, or migrations.
    </use_case>
    
    <workflow>
    1. Creates a new branch with optional name
    2. Sets up compute endpoint for the branch
    3. Returns branch details for immediate use
    </workflow>
    """
    api_key = context.get_secret("NEON_API_KEY")
    storage = NeonStorage(context)
    
    # Generate branch name if not provided
    if not branch_name:
        branch_name = f"branch-{datetime.now().strftime('%Y%m%d-%H%M%S')}"
    
    branch_data = {
        "branch": {
            "name": branch_name
        },
        "endpoints": [
            {
                "type": "read_write",
                "autoscaling_limit_min_cu": 0.25,
                "autoscaling_limit_max_cu": 0.25,
                "provisioner": "k8s-neonvm"
            }
        ]
    }
    
    if parent_branch_id:
        branch_data["branch"]["parent_id"] = parent_branch_id
    
    try:
        result = await make_neon_api_request(
            api_key, "POST", f"projects/{project_id}/branches", branch_data
        )
        
        branch = result["branch"]
        endpoints = result.get("endpoints", [])
        
        # Store branch information
        storage.set_branch_info(branch["id"], {
            "project_id": project_id,
            "name": branch_name,
            "created_at": branch["created_at"],
            "endpoints": endpoints
        })
        
        logger.info(f"Created branch: {branch['id']} in project {project_id}")
        
        return {
            "branch_id": branch["id"],
            "name": branch["name"],
            "project_id": project_id,
            "parent_branch_id": branch.get("parent_id"),
            "created_at": branch["created_at"],
            "current_state": branch.get("current_state", "ready"),
            "endpoints": endpoints
        }
        
    except Exception as e:
        raise RetryableToolError(f"Failed to create database branch: {str(e)}")

@tool(
    requires_secrets=["NEON_API_KEY"]
)
async def execute_sql_query(
    context: ToolContext,
    project_id: Annotated[str, "Neon project ID"],
    query: Annotated[str, "SQL query to execute"],
    branch_id: Annotated[Optional[str], "Branch ID (defaults to main)"] = None,
    database_name: Annotated[Optional[str], "Database name"] = None
) -> Annotated[Dict[str, Any], "Query results and metadata"]:
    """
    Execute a SQL query against a Neon database.
    
    <use_case>
    Use this tool to execute a single SQL statement against a Neon database.
    Supports SELECT, INSERT, UPDATE, DELETE and other SQL operations
    with automatic connection management and error handling.
    </use_case>
    
    <important_notes>
    If you have a temporary branch from a prior step, you MUST:
    1. Pass the branch ID to this tool unless explicitly told otherwise
    2. Tell the user that you are using the temporary branch with ID [branch_id]
    </important_notes>
    """
    api_key = context.get_secret("NEON_API_KEY")
    storage = NeonStorage(context)
    
    try:
        # Get connection string
        connection_string = await get_connection_string_internal(
            api_key, project_id, branch_id, database_name
        )
        
        # Execute query with connection pool
        async with get_database_connection(connection_string) as conn:
            start_time = datetime.now()
            
            # Determine query type
            query_type = query.strip().upper().split()[0]
            
            if query_type == "SELECT":
                rows = await conn.fetch(query)
                result_data = [dict(row) for row in rows]
                row_count = len(result_data)
            else:
                result = await conn.execute(query)
                result_data = None
                # Extract row count from result string (e.g., "UPDATE 5")
                row_count = int(result.split()[-1]) if result.split()[-1].isdigit() else 0
            
            execution_time = (datetime.now() - start_time).total_seconds()
            
            logger.info(f"Executed {query_type} query on project {project_id}: {row_count} rows affected")
            
            return {
                "query_type": query_type,
                "rows_affected": row_count,
                "execution_time_seconds": execution_time,
                "data": result_data,
                "project_id": project_id,
                "branch_id": branch_id,
                "database_name": database_name
            }
            
    except Exception as e:
        raise RetryableToolError(f"Failed to execute SQL query: {str(e)}")

@tool(
    requires_secrets=["NEON_API_KEY"]
)
async def execute_sql_transaction(
    context: ToolContext,
    project_id: Annotated[str, "Neon project ID"],
    sql_statements: Annotated[List[str], "List of SQL statements to execute in transaction"],
    branch_id: Annotated[Optional[str], "Branch ID (defaults to main)"] = None,
    database_name: Annotated[Optional[str], "Database name"] = None
) -> Annotated[List[Any], "Transaction results"]:
    """
    Execute multiple SQL statements in a single transaction.
    
    <use_case>
    Use this tool to execute a SQL transaction against a Neon database,
    should be used for multiple SQL statements that need to be executed atomically.
    </use_case>
    
    <important_notes>
    If you have a temporary branch from a prior step, you MUST:
    1. Pass the branch ID to this tool unless explicitly told otherwise
    2. Tell the user that you are using the temporary branch with ID [branch_id]
    </important_notes>
    """
    api_key = context.get_secret("NEON_API_KEY")
    
    try:
        # Get connection string
        connection_string = await get_connection_string_internal(
            api_key, project_id, branch_id, database_name
        )
        
        # Execute transaction
        async with get_database_connection(connection_string) as conn:
            async with conn.transaction():
                results = []
                for sql in sql_statements:
                    if sql.strip().upper().startswith("SELECT"):
                        rows = await conn.fetch(sql)
                        results.append([dict(row) for row in rows])
                    else:
                        result = await conn.execute(sql)
                        results.append(result)
                
                logger.info(f"Executed transaction with {len(sql_statements)} statements")
                return results
        
    except Exception as e:
        raise RetryableToolError(f"Failed to execute SQL transaction: {str(e)}")

@tool(
    requires_secrets=["NEON_API_KEY"]
)
async def prepare_database_migration(
    context: ToolContext,
    project_id: Annotated[str, "Neon project ID"],
    migration_sql: Annotated[str, "SQL statements for the migration"],
    database_name: Annotated[Optional[str], "Database name (uses default if not provided)"] = None
) -> Annotated[Dict[str, Any], "Migration preparation result"]:
    """
    Prepare a database schema migration by creating a temporary branch and applying changes.
    
    <use_case>
    This tool performs database schema migrations by automatically generating and executing DDL statements.
    
    Supported operations:
    - CREATE operations: Add new columns, tables, constraints
    - ALTER operations: Modify column types, rename columns, add/modify indexes
    - DROP operations: Remove columns, tables, constraints
    
    The tool will:
    1. Parse your migration SQL
    2. Create a temporary branch for safety
    3. Execute in the temporary branch
    4. Return migration details for verification
    </use_case>
    
    <workflow>
    1. Creates a temporary branch
    2. Applies the migration SQL in that branch
    3. Returns migration details for verification
    </workflow>
    
    <important_notes>
    After executing this tool, you MUST:
    1. Test the migration in the temporary branch using 'execute_sql_query'
    2. Ask for confirmation before proceeding
    3. Use 'complete_database_migration' tool to apply changes to main branch
    </important_notes>
    
    <error_handling>
    On error, the tool will:
    1. Automatically attempt ONE retry of the exact same operation
    2. If the retry fails, terminate execution and return error details
    3. Clean up temporary branch automatically
    </error_handling>
    """
    api_key = context.get_secret("NEON_API_KEY")
    storage = NeonStorage(context)
    
    migration_id = str(uuid.uuid4())
    temp_branch = None
    
    try:
        # Create temporary branch for migration
        branch_result = await create_database_branch(
            context, project_id, f"migration-{migration_id[:8]}"
        )
        temp_branch = branch_result
        
        # Execute migration in temporary branch
        statements = split_sql_statements(migration_sql)
        migration_result = await execute_sql_transaction(
            context, project_id, statements, temp_branch["branch_id"], database_name
        )
        
        # Store migration information
        migration_data = {
            "migration_sql": migration_sql,
            "database_name": database_name,
            "project_id": project_id,
            "temp_branch": temp_branch,
            "statements": statements,
            "created_at": datetime.now().isoformat()
        }
        storage.store_migration(migration_id, migration_data)
        
        logger.info(f"Prepared migration {migration_id} in temporary branch {temp_branch['branch_id']}")
        
        return {
            "migration_id": migration_id,
            "temporary_branch_name": temp_branch["name"],
            "temporary_branch_id": temp_branch["branch_id"],
            "migration_result": "success",
            "statements_executed": len(statements),
            "execution_results": migration_result
        }
        
    except Exception as e:
        # Clean up temporary branch on error
        if temp_branch:
            try:
                await delete_database_branch(
                    context, project_id, temp_branch["branch_id"]
                )
            except:
                pass  # Best effort cleanup
        
        raise RetryableToolError(f"Failed to prepare database migration: {str(e)}")

@tool(
    requires_secrets=["NEON_API_KEY"]
)
async def complete_database_migration(
    context: ToolContext,
    migration_id: Annotated[str, "Migration ID from prepare_database_migration"]
) -> Annotated[Dict[str, Any], "Migration completion result"]:
    """
    Complete a database migration by applying changes to the main branch and cleaning up.
    
    <use_case>
    Complete a database migration when the user confirms the migration is ready
    to be applied to the main branch. This tool also cleans up the temporary
    branch created by the prepare_database_migration tool.
    </use_case>
    
    <workflow>
    1. Retrieves the prepared migration
    2. Applies the migration SQL to the main branch
    3. Deletes the temporary branch
    4. Returns completion status
    </workflow>
    """
    storage = NeonStorage(context)
    
    try:
        # Get migration data
        migration_data = storage.get_migration(migration_id)
        if not migration_data:
            raise Exception(f"Migration {migration_id} not found")
        
        # Apply migration to main branch
        main_result = await execute_sql_transaction(
            context,
            migration_data["project_id"],
            migration_data["statements"],
            database_name=migration_data["database_name"]
        )
        
        # Clean up temporary branch
        cleanup_result = await delete_database_branch(
            context,
            migration_data["project_id"],
            migration_data["temp_branch"]["branch_id"]
        )
        
        logger.info(f"Completed migration {migration_id}")
        
        return {
            "migration_id": migration_id,
            "status": "completed",
            "main_branch_result": main_result,
            "temporary_branch_deleted": cleanup_result["status"],
            "completed_at": datetime.now().isoformat()
        }
        
    except Exception as e:
        raise RetryableToolError(f"Failed to complete database migration: {str(e)}")

@tool(
    requires_secrets=["NEON_API_KEY"]
)
async def get_database_tables(
    context: ToolContext,
    project_id: Annotated[str, "Neon project ID"],
    branch_id: Annotated[Optional[str], "Branch ID (defaults to main)"] = None,
    database_name: Annotated[Optional[str], "Database name"] = None
) -> Annotated[List[Dict[str, Any]], "List of tables in the database"]:
    """
    Get all tables in a Neon database.
    
    <use_case>
    Lists all tables within a specified Neon database, including schema
    information and table types. Useful for database exploration and
    understanding the current structure.
    </use_case>
    """
    try:
        tables_query = """
        SELECT 
            table_schema,
            table_name,
            table_type
        FROM information_schema.tables 
        WHERE table_schema NOT IN ('pg_catalog', 'information_schema')
        ORDER BY table_schema, table_name;
        """
        
        result = await execute_sql_query(
            context, project_id, tables_query, branch_id, database_name
        )
        
        return result["data"]
        
    except Exception as e:
        raise RetryableToolError(f"Failed to get database tables: {str(e)}")

@tool(
    requires_secrets=["NEON_API_KEY"]
)
async def describe_table_schema(
    context: ToolContext,
    project_id: Annotated[str, "Neon project ID"],
    table_name: Annotated[str, "Name of the table to describe"],
    branch_id: Annotated[Optional[str], "Branch ID (defaults to main)"] = None,
    database_name: Annotated[Optional[str], "Database name"] = None
) -> Annotated[Dict[str, Any], "Table schema information"]:
    """
    Describe the schema of a table in a Neon database.
    
    <use_case>
    Retrieves detailed schema information for a specific table,
    including columns, data types, constraints, and indexes.
    </use_case>
    """
    try:
        # Extract simple table name if schema-qualified
        simple_table_name = table_name.split('.')[-1]
        
        schema_query = f"""
        SELECT 
            column_name,
            data_type,
            is_nullable,
            column_default,
            character_maximum_length,
            numeric_precision,
            numeric_scale
        FROM information_schema.columns
        WHERE table_name = '{simple_table_name}'
        ORDER BY ordinal_position;
        """
        
        result = await execute_sql_query(
            context, project_id, schema_query, branch_id, database_name
        )
        
        return {
            "table_name": table_name,
            "columns": result["data"],
            "column_count": len(result["data"])
        }
        
    except Exception as e:
        raise RetryableToolError(f"Failed to describe table schema: {str(e)}")

@tool(
    requires_secrets=["NEON_API_KEY"]
)
async def delete_database_branch(
    context: ToolContext,
    project_id: Annotated[str, "Neon project ID"],
    branch_id: Annotated[str, "Branch ID to delete"]
) -> Annotated[Dict[str, Any], "Deletion confirmation"]:
    """
    Delete a database branch to free up resources and optimize costs.
    
    <important_notes>
    This operation is irreversible. Ensure you have backups
    of any important data before deletion.
    </important_notes>
    """
    api_key = context.get_secret("NEON_API_KEY")
    storage = NeonStorage(context)
    
    try:
        # Get branch info before deletion
        branch_info = storage.get_branch_info(branch_id)
        
        await make_neon_api_request(
            api_key, "DELETE", f"projects/{project_id}/branches/{branch_id}"
        )
        
        logger.info(f"Deleted branch: {branch_id} from project {project_id}")
        
        return {
            "branch_id": branch_id,
            "project_id": project_id,
            "deleted_at": datetime.now().isoformat(),
            "branch_name": branch_info.get("name") if branch_info else None,
            "status": "deleted"
        }
        
    except Exception as e:
        raise RetryableToolError(f"Failed to delete database branch: {str(e)}")

@tool(
    requires_secrets=["NEON_API_KEY"]
)
async def get_project_metrics(
    context: ToolContext,
    project_id: Annotated[str, "Neon project ID"]
) -> Annotated[Dict[str, Any], "Project usage metrics and information"]:
    """
    Get usage metrics, performance data, and project information.
    
    <use_case>
    Provides insights into database usage, costs, and performance
    to help optimize resource allocation and budgeting.
    </use_case>
    """
    api_key = context.get_secret("NEON_API_KEY")
    
    try:
        # Get project details
        project_result = await make_neon_api_request(
            api_key, "GET", f"projects/{project_id}"
        )
        
        project = project_result["project"]
        
        logger.info(f"Retrieved metrics for project {project_id}")
        
        return {
            "project_id": project["id"],
            "name": project["name"],
            "region": project["region_id"],
            "postgres_version": project["pg_version"],
            "created_at": project["created_at"],
            "updated_at": project["updated_at"],
            "current_state": project.get("current_state", "ready"),
            "branch_count": len(project.get("branches", [])),
            "dashboard_url": f"https://console.neon.tech/app/projects/{project['id']}"
        }
        
    except Exception as e:
        raise RetryableToolError(f"Failed to get project metrics: {str(e)}")

async def get_connection_string_internal(
    api_key: str,
    project_id: str,
    branch_id: Optional[str] = None,
    database_name: Optional[str] = None,
    role_name: Optional[str] = None
) -> str:
    """Internal helper to get connection string"""
    
    # Get project info if branch_id not provided
    if not branch_id:
        project_result = await make_neon_api_request(
            api_key, "GET", f"projects/{project_id}/branches"
        )
        branches = project_result["branches"]
        default_branch = next((b for b in branches if b.get("primary")), branches[0])
        branch_id = default_branch["id"]
    
    # Build connection URI request
    params = {"branch_id": branch_id}
    if database_name:
        params["database_name"] = database_name
    if role_name:
        params["role_name"] = role_name
    
    query_params = "&".join([f"{k}={v}" for k, v in params.items()])
    
    connection_result = await make_neon_api_request(
        api_key, "GET", f"projects/{project_id}/connection_uri?{query_params}"
    )
    
    return connection_result["uri"]

# Export hooks for worker lifecycle management
__all__ = [
    "list_neon_projects",
    "create_neon_project",
    "create_database_branch", 
    "execute_sql_query",
    "execute_sql_transaction",
    "prepare_database_migration",
    "complete_database_migration",
    "get_database_tables",
    "describe_table_schema",
    "delete_database_branch",
    "get_project_metrics",
    "startup_hook",
    "shutdown_hook",
    "NeonStorage"
]

Use Cases for Toolkits Build on the Base Neon Toolkit

Some ideas for abstractions built on this.

Dynamic Multi-Tenant Architecture

Agents can create isolated database branches for each customer or session, ensuring data privacy and enabling per-tenant customization without complex schema management.

Ephemeral Data Processing

For data-intensive tasks, agents can create temporary branches, process data, extract insights, and automatically clean up resources, optimizing both performance and costs.

AI Experiment Management

Agents can create branches for different model training runs, A/B tests, or feature experiments, maintaining data versioning and enabling easy rollbacks.

RAG Pipeline Optimization

Agents can dynamically create optimized database branches with specific indexes and schemas tailored to different types of queries or knowledge domains.