main.py

from fastapi import FastAPI, Depends, HTTPException
from sqlalchemy import create_engine, Column, Integer, String
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker, Session
from pydantic import BaseModel
import os
import json
import asyncpg
from pgqueuer.queries import Queries
from pgqueuer.db import AsyncpgDriver

DATABASE_URL = os.getenv("DATABASE_URL", "postgresql+psycopg://fastapi_user:fastapi_password@localhost:5432/fastapi_db")
ASYNC_DATABASE_URL = DATABASE_URL.replace("postgresql+psycopg://", "postgresql://")

engine = create_engine(DATABASE_URL)
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
Base = declarative_base()

class User(Base):
    __tablename__ = "users"
    
    id = Column(Integer, primary_key=True, index=True)
    name = Column(String, index=True)

# Create SQLAlchemy tables
Base.metadata.create_all(bind=engine)

app = FastAPI()

# Initialize pgqueuer tables 
async def init_pgqueuer():
    """Initialize pgqueuer tables on startup (idempotent)"""
    conn = await asyncpg.connect(ASYNC_DATABASE_URL)
    try:
        # Check if pgqueuer tables already exist
        table_exists = await conn.fetchval(
            "SELECT EXISTS (SELECT FROM information_schema.tables WHERE table_name = 'pgqueuer')"
        )
        
        if not table_exists:
            driver = AsyncpgDriver(conn)
            from pgqueuer.queries import Queries
            queries = Queries(driver)
            await queries.install()
            await queries.upgrade()
            print("pgqueuer tables initialized successfully!")
        else:
            print("pgqueuer tables already exist, skipping initialization")
            
    except Exception as e:
        print(f"pgqueuer initialization: {e}")
    finally:
        await conn.close()

@app.on_event("startup")
async def startup_event():
    await init_pgqueuer()

class UserCreate(BaseModel):
    name: str

class UserResponse(BaseModel):
    id: int
    name: str
    
    class Config:
        orm_mode = True

def get_db():
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()

@app.post("/users/", response_model=UserResponse)
async def create_user(user: UserCreate):
    """
    Create user and audit job atomically in the same transaction.
    Both operations succeed or both fail together.
    """
    conn = await asyncpg.connect(ASYNC_DATABASE_URL)
    try:
        async with conn.transaction():
            # 1. Create user using raw SQL in the same transaction
            user_id = await conn.fetchval(
                "INSERT INTO users (name) VALUES ($1) RETURNING id",
                user.name
            )
            
            # 2. Create audit job in the same transaction
            audit_payload = {
                "event_type": "user_created",
                "user_id": user_id,
                "user_name": user.name,
                "timestamp": str(user_id)
            }
            
            # 3. Enqueue job using pgqueuer in same transaction
            driver = AsyncpgDriver(conn)
            queries = Queries(driver)
            
            await queries.enqueue(
                ["audit_events"],  # Entrypoint names
                [json.dumps(audit_payload).encode()],  # Job payload as bytes
                [0]  # Priority
            )
            
            # Both operations committed together automatically when transaction exits
            
        # Return the created user
        return UserResponse(id=user_id, name=user.name)
        
    except Exception as e:
        # Transaction automatically rolled back on exception
        raise HTTPException(status_code=500, detail=f"Failed to create user: {str(e)}")
    finally:
        await conn.close()

@app.get("/")
def read_root():
    return {"message": "FastAPI with SQLAlchemy and PostgreSQL"}