Designing Workflow Data Models with PostgreSQL and SQLAlchemy

PostgreSQL + SQLAlchemy Model Design — Workflow Data Structure

This post covers the model design process for storing XGen platform workflow, node, and edge data in PostgreSQL, making heavy use of SQLAlchemy 2.0's async support.

Database Connection Setup

from sqlalchemy.ext.asyncio import (
    create_async_engine,
    AsyncSession,
    async_sessionmaker,
)
from sqlalchemy.orm import DeclarativeBase

DATABASE_URL = "postgresql+asyncpg://user:pass@localhost:5432/xgen"

engine = create_async_engine(DATABASE_URL, echo=False, pool_size=20)
async_session_maker = async_sessionmaker(engine, class_=AsyncSession)

class Base(DeclarativeBase):
    pass

Core Model Design

Workflow Model

from sqlalchemy import Column, String, DateTime, ForeignKey, JSON, Boolean
from sqlalchemy.orm import relationship
from datetime import datetime
import uuid

class Workflow(Base):
    __tablename__ = "workflows"

    id = Column(String, primary_key=True, default=lambda: str(uuid.uuid4()))
    name = Column(String(255), nullable=False)
    description = Column(String(1000), default="")
    user_id = Column(String, ForeignKey("users.id"), nullable=False)
    is_deployed = Column(Boolean, default=False)
    created_at = Column(DateTime, default=datetime.utcnow)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)

    nodes = relationship("Node", back_populates="workflow", cascade="all, delete-orphan")
    edges = relationship("Edge", back_populates="workflow", cascade="all, delete-orphan")
    user = relationship("User", back_populates="workflows")

Node and Edge Models

class Node(Base):
    __tablename__ = "nodes"

    id = Column(String, primary_key=True, default=lambda: str(uuid.uuid4()))
    workflow_id = Column(String, ForeignKey("workflows.id"), nullable=False)
    node_type = Column(String(50), nullable=False)
    position_x = Column(Float, default=0)
    position_y = Column(Float, default=0)
    params = Column(JSON, default=dict)
    label = Column(String(255), default="")

    workflow = relationship("Workflow", back_populates="nodes")

class Edge(Base):
    __tablename__ = "edges"

    id = Column(String, primary_key=True, default=lambda: str(uuid.uuid4()))
    workflow_id = Column(String, ForeignKey("workflows.id"), nullable=False)
    source_node_id = Column(String, ForeignKey("nodes.id"), nullable=False)
    target_node_id = Column(String, ForeignKey("nodes.id"), nullable=False)
    source_handle = Column(String(50), default="output")
    target_handle = Column(String(50), default="input")

    workflow = relationship("Workflow", back_populates="edges")

Migration Management

Schema changes are managed with Alembic.

# alembic/env.py
from models import Base
target_metadata = Base.metadata

# Generate a migration
alembic revision --autogenerate -m "add_workflow_tables"
# Apply the migration
alembic upgrade head

CRUD Service

class WorkflowService:
    def __init__(self, db: AsyncSession):
        self.db = db

    async def create_workflow(self, user_id: str, data: dict) -> Workflow:
        workflow = Workflow(user_id=user_id, **data)
        self.db.add(workflow)
        await self.db.flush()
        return workflow

    async def get_workflow_with_graph(self, workflow_id: str) -> dict:
        stmt = (
            select(Workflow)
            .options(selectinload(Workflow.nodes), selectinload(Workflow.edges))
            .where(Workflow.id == workflow_id)
        )
        result = await self.db.execute(stmt)
        workflow = result.scalar_one_or_none()
        if not workflow:
            raise AppException(404, "Workflow not found", "NOT_FOUND")
        return workflow

Design Decisions

The key decision was modeling each node's params as a JSON column. Parameters vary completely across node types, so a normalized table would require a new migration every time a node type was added. The JSON column gives us the flexibility we need while still letting us leverage PostgreSQL's JSON query capabilities.