api-versioning-with-neondb-branching/
│
├── app/
│   ├── main.py                   
│   ├── neon_db_setup.py          
|   ├── db_connection.py         
│── data/
│   ├── schema_v1.sql              # SQL schema for v1
│   ├── schema_v2.sql              # SQL schema for v2
│── .env
|── requirements.txt 
└── README.md

NEON_API_KEY=your_neon_api_key
NEON_PROJECT_ID=your_neon_project_id

-- schema_v1.sql
CREATE TABLE users (
    id SERIAL PRIMARY KEY,
    name TEXT,
    email TEXT
);

INSERT INTO users (name, email) VALUES 
('Alice', 'alice@example.com'), 
('Bob', 'bob@example.com');

Define Schema for API v2:

-- schema_v2.sql with additional fields
CREATE TABLE users (
    id SERIAL PRIMARY KEY,
    name TEXT,
    email TEXT,
    age INT,
    city TEXT
);

INSERT INTO users (name, email, age, city) VALUES 
('Alice', 'alice@example.com', 25, 'New York'), 
('Bob', 'bob@example.com', 30, 'San Francisco');

import os
import requests
import psycopg2
from dotenv import load_dotenv

load_dotenv()

NEON_API_KEY = os.getenv("NEON_API_KEY")
NEON_PROJECT_ID = os.getenv("NEON_PROJECT_ID")
NEON_API_URL = "<https://console.neon.tech/api/v2>"
DEFAULT_DATABASE_NAME = "neondb"
DEFAULT_DATABASE_ROLE = "neondb_owner"

SCHEMA_FILES = {
    "v1": "./data/schema_v1.sql",
    "v2": "./data/schema_v2.sql",
}

versions = ["v1", "v2"]

def create_or_get_branch(branch_name):
    """
    Checks if a branch exists. If not, creates it.

    :param branch_name: The branch to check or create.
    :return: The branch ID.
    """
    existing_branches = list_existing_branches()

    if branch_name in existing_branches:
        print(f"✅ Branch '{branch_name}' already exists.")
        return existing_branches[branch_name]

    url = f"{NEON_API_URL}/projects/{NEON_PROJECT_ID}/branches"
    headers = {
        "Authorization": f"Bearer {NEON_API_KEY}",
        "Accept": "application/json",
        "Content-type": "application/json",
    }
    data = {"branch": {"name": branch_name}, "endpoints": [{"type": "read_write"}]}

    response = requests.post(url, headers=headers, json=data)

    if response.status_code == 201:
        branch_data = response.json()
        branch_id = branch_data.get("branch", {}).get("id")
        print(f"✅ Branch '{branch_name}' created successfully with ID: {branch_id}")
        return branch_id
    else:
        print(f"❌ Failed to create branch: {response.text}")
        return None

def list_existing_branches():
    """
    Fetches the list of existing branches in the Neon project.

    :return: List of existing branch names.
    """
    url = f"{NEON_API_URL}/projects/{NEON_PROJECT_ID}/branches"
    headers = {"Authorization": f"Bearer {NEON_API_KEY}"}

    response = requests.get(url, headers=headers)

    if response.status_code == 200:
        branches = response.json().get("branches", [])
        return {branch["name"]: branch["id"] for branch in branches}
    else:
        print(f"❌ Failed to fetch existing branches: {response.text}")
        return {}

def get_connection_uri(branch_id):
    """
    Fetches the database connection URI for a specific Neon branch.

    :param branch_id: The ID of the branch.
    :return: The connection string URI.
    """
    url = f"{NEON_API_URL}/projects/{NEON_PROJECT_ID}/connection_uri"
    params = {
        "database_name": DEFAULT_DATABASE_NAME,
        "role_name": DEFAULT_DATABASE_ROLE,
        "branch_id": branch_id,  # Pass the branch ID as a query parameter
    }
    headers = {"Authorization": f"Bearer {NEON_API_KEY}"}

    response = requests.get(url, headers=headers, params=params)

    if response.status_code == 200:
        return response.json().get("uri")
    else:
        print(f"❌ Failed to fetch connection URI: {response.text}")
        return None

def initialize_connection_strings():
    """
    Initializes and caches database connection strings for different API versions.
    This function is called once at application startup.
    """
    connection_strings = {}

    for version in versions:
        branch_name = f"api_{version}_branch"
        branch_id = create_or_get_branch(branch_name)

        if branch_id:
            connection_uri = get_connection_uri(branch_id)
            if connection_uri:
                connection_strings[version] = connection_uri
                print(f"✅ Cached connection for {version}: {connection_uri}")

                # Initialize database schema for this version
                schema_file = SCHEMA_FILES.get(version)
                if schema_file:
                    initialize_database_schema(connection_uri, schema_file)
            else:
                print(f"❌ Failed to get connection URI for {version}")
        else:
            print(f"❌ Failed to retrieve branch ID for {version}")

    return connection_strings

def initialize_database_schema(connection_uri, schema_file):
    """Reads SQL schema from file and applies it to the database."""
    try:
        # Read schema file
        with open(schema_file, "r") as file:
            schema_sql = file.read()

        # Split SQL statements by semicolon
        sql_statements = [
            stmt.strip() for stmt in schema_sql.split(";") if stmt.strip()
        ]

        # Connect to the database
        conn = psycopg2.connect(connection_uri)
        cursor = conn.cursor()

        # Execute SQL statements one by one
        for statement in sql_statements:
            print(f"📌 Executing SQL: {statement}")
            cursor.execute(statement)

        conn.commit()
        cursor.close()
        conn.close()
        print("✅ Database schema initialized successfully.")

    except Exception as e:
        print(f"❌ Error initializing database schema: {e}")
        raise

import psycopg2
from dotenv import load_dotenv
from app.neon_db_setup import initialize_connection_strings

load_dotenv()

# Store connection strings
CONNECTION_STRINGS = initialize_connection_strings()

def get_db_connection(version: str):
    """Returns a database connection based on the API version."""
    if version not in CONNECTION_STRINGS:
        raise Exception(f"No cached connection string found for version {version}")

    return psycopg2.connect(CONNECTION_STRINGS[version])

from fastapi import FastAPI
from app.db_connection import get_db_connection

app = FastAPI()

@app.get("/v1/users")
def get_users_v1():
    """Fetch users from v1 schema (Basic users table)."""
    conn = get_db_connection("v1")
    cur = conn.cursor()
    cur.execute("SELECT id, name, email FROM users;")
    users = cur.fetchall()
    conn.close()
    return [{"id": u[0], "name": u[1], "email": u[2]} for u in users]

@app.get("/v2/users")
def get_users_v2():
    """Fetch users from v2 schema (With additional columns)."""
    conn = get_db_connection("v2")
    cur = conn.cursor()
    cur.execute("SELECT id, name, email, age, city FROM users;")
    users = cur.fetchall()
    conn.close()
    return [
        {"id": u[0], "name": u[1], "email": u[2], "age": u[3], "city": u[4]}
        for u in users
    ]