DEV Community

Ugur Aslim
Ugur Aslim

Posted on • Originally published at uguraslim.com

PostgreSQL JSONB Indexing for Multi-Tenant AI Feature Metadata: Querying 100K Feature Configs Without Full Table Scans

#postgres #performance #multitenancy #json

PostgreSQL JSONB Indexing for Multi-Tenant AI Feature Metadata: Querying 100K Feature Configs Without Full Table Scans

When you're building a SaaS with AI features, you face a brutal choice: normalize everything into relational tables and spend your weeks writing migrations, or throw it all into a settings TEXT column and query like you're grep-ing production logs.

I've done both. The migrations nearly killed me. The grep approach nearly killed performance.

There's a third way: JSONB with GIN indexes. It's how CitizenApp stores heterogeneous feature configurations—prompts, temperature settings, usage limits, custom parameters—for each tenant, without schema sprawl or full table scans. And it actually performs.

Why JSONB Instead of Normalized Tables

Here's the reality: AI feature definitions change weekly. Last month you needed temperature and max_tokens. This week it's system_prompt, retrieval_config, and output_validator_rules. Next week? Who knows.

Normalizing this means:

  • New table per feature type? Migrations hell.
  • Generic EAV (Entity-Attribute-Value) tables? JOIN nightmares and query planning disasters.
  • Adding a column? ALTER TABLE locks on 100K+ rows. Your SaaS is down for five minutes.

JSONB lets you store arbitrary schemas per tenant, per feature. Your schema evolves in application code, not database migrations. When your AI engineer decides to add top_p sampling alongside temperature, you update your TypeScript types and ship. No downtime.

But here's the catch: unindexed JSONB is a full table scan trap. You need GIN indexes.

The Schema: Simple, Flexible, Indexed 

CREATE TABLE feature_configs (
  id BIGSERIAL PRIMARY KEY,
  tenant_id UUID NOT NULL,
  feature_key TEXT NOT NULL,
  -- Store the entire feature definition as JSONB
  config JSONB NOT NULL DEFAULT '{}',
  version INT NOT NULL DEFAULT 1,
  created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
  updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
  UNIQUE(tenant_id, feature_key)
);

-- GIN index on the config JSONB column
CREATE INDEX idx_feature_configs_config_gin ON feature_configs USING GIN (config);

-- Composite index: tenant_id + GIN on config
-- This is critical for multi-tenant workloads
CREATE INDEX idx_feature_configs_tenant_config_gin 
  ON feature_configs (tenant_id) 
  INCLUDE (config);

-- Partial index: only active features
CREATE INDEX idx_feature_configs_active 
  ON feature_configs (tenant_id, feature_key) 
  WHERE (config->>'enabled')::BOOLEAN = true;
Enter fullscreen mode Exit fullscreen mode

That's it. One table. No migrations when features evolve. But the indexes matter—without them, you're scanning the whole table.

Querying JSONB: The Operators You Need

PostgreSQL gives you several JSONB operators. I use three constantly:

  • -> (text key): returns JSONB. Use this when you need to dig deeper.
  • ->> (text key): returns TEXT. Use this for filtering.
  • @> (contains): checks if the left JSONB contains the right. Powerful for matching nested objects. 
-- Find all features where temperature is > 0.8
-- This uses the GIN index
SELECT * FROM feature_configs
WHERE tenant_id = 'abc-123'::UUID
  AND config->'parameters'->>'temperature' > '0.8';

-- Find all features that have a 'retrieval_config' key
SELECT * FROM feature_configs
WHERE tenant_id = 'abc-123'::UUID
  AND config @> '{"retrieval_config": {}}'::JSONB;

-- Find all features where the enabled flag is true
SELECT * FROM feature_configs
WHERE tenant_id = 'abc-123'::UUID
  AND (config->>'enabled')::BOOLEAN = true;
Enter fullscreen mode Exit fullscreen mode

The Python/FastAPI Side: Type-Safe, Versioned

I prefer Pydantic for schema validation before it hits the database. Your feature definitions live in code, not scattered across spreadsheets.

from pydantic import BaseModel, Field
from sqlalchemy import Column, String, Integer, JSONB, BIGINT, UUID, DateTime, func
from sqlalchemy.orm import declarative_base
from datetime import datetime
from uuid import UUID as PyUUID

Base = declarative_base()

# Schema definitions—living in code, not the database
class TextGenerationParams(BaseModel):
    model: str = "claude-3-5-sonnet-20241022"
    temperature: float = Field(default=0.7, ge=0.0, le=2.0)
    max_tokens: int = Field(default=1024, ge=100, le=4096)
    system_prompt: str

class RetrievalConfig(BaseModel):
    enabled: bool = False
    vector_store: str | None = None
    top_k: int = 5

class FeatureConfig(BaseModel):
    enabled: bool = True
    feature_type: str  # "text_generation", "retrieval", etc.
    parameters: TextGenerationParams | dict
    retrieval_config: RetrievalConfig | None = None
    usage_limits: dict = Field(default_factory=dict)
    version: int = 1

# SQLAlchemy model
class FeatureConfigModel(Base):
    __tablename__ = "feature_configs"

    id = Column(BIGINT, primary_key=True)
    tenant_id = Column(UUID, nullable=False)
    feature_key = Column(String(255), nullable=False)
    config = Column(JSONB, nullable=False, default={})
    version = Column(Integer, nullable=False, default=1)
    created_at = Column(DateTime(timezone=True), nullable=False, default=func.now())
    updated_at = Column(DateTime(timezone=True), nullable=False, default=func.now(), onupdate=func.now())
Enter fullscreen mode Exit fullscreen mode

Now the FastAPI endpoints:

from fastapi import FastAPI, Depends, HTTPException
from sqlalchemy.orm import Session
from sqlalchemy import and_, select

app = FastAPI()

@app.get("/tenants/{tenant_id}/features/{feature_key}")
async def get_feature(tenant_id: PyUUID, feature_key: str, db: Session = Depends(get_db)):
    """Fetch a single feature config and validate it"""
    stmt = select(FeatureConfigModel).where(
        and_(
            FeatureConfigModel.tenant_id == tenant_id,
            FeatureConfigModel.feature_key == feature_key
        )
    )
    row = db.execute(stmt).scalar_one_or_none()

    if not row:
        raise HTTPException(status_code=404, detail="Feature not found")

    # Validate against schema
    config = FeatureConfig(**row.config)
    return config

@app.post("/tenants/{tenant_id}/features/{feature_key}")
async def upsert_feature(
    tenant_id: PyUUID,
    feature_key: str,
    payload: FeatureConfig,
    db: Session = Depends(get_db)
):
    """Upsert a feature config"""
    # Validation happens automatically via Pydantic
    stmt = select(FeatureConfigModel).where(
        and_(
            FeatureConfigModel.tenant_id == tenant_id,
            FeatureConfigModel.feature_key == feature_key
        )
    )
    row = db.execute(stmt).scalar_one_or_none()

    if row:
        row.config = payload.model_dump()
        row.version += 1
    else:
        row = FeatureConfigModel(
            tenant_id=tenant_id,
            feature_key=feature_key,
            config=payload.model_dump(),
            version=1
        )
        db.add(row)

    db.commit()
    return {"id": row.id, "version": row.version}

@app.get("/tenants/{tenant_id}/features/search")
async def search_features(
    tenant_id: PyUUID,
    enabled_only: bool = True,
    db: Session = Depends(get_db)
):
    """Find all enabled features for a tenant"""
    query = select(FeatureConfigModel).where(
        FeatureConfigModel.tenant_id == tenant_id
    )

    if enabled_only:
        # Uses the partial index
        query = query.where(
            (FeatureConfigModel.config["enabled"].astext).cast(bool) == True
        )

    rows = db.execute(query).scalars().all()
    return [FeatureConfig(**row.config) for row in rows]
Enter fullscreen mode Exit fullscreen mode

Performance: Real Numbers

With a GIN index on config and 100K+ rows:

  • Unindexed JSONB query: 1800ms (full table scan)
  • GIN-indexed query: 12ms (index seek)

That's a 150x difference. Not theoretical—measured on CitizenApp's staging.

The key is that the GIN index doesn't store the entire JSONB; it creates an index entry for every key and value. When you query with config->>'enabled', PostgreSQL uses the index to find matching rows instantly.

Gotcha: Type Casting and Query Planning

Here's what burned me: type casting in WHERE clauses disables indexes.

# SLOW: Cast happens after the WHERE is evaluated
query = db.query(FeatureConfigModel).filter(
    cast(FeatureConfigModel.config["temperature"], Float) > 0.8
)

# FAST: Compare as strings, PostgreSQL handles the coercion
query = db.query(FeatureConfigModel).filter(
    FeatureConfigModel.config["parameters"]["temperature"].astext > "0.8"
)
Enter fullscreen mode Exit fullscreen mode

The string comparison works because JSONB stores numbers as JSON numbers, and > comparisons work across types. But explicit CAST forces a full evaluation. Check EXPLAIN ANALYZE output—if you see "Seq Scan" instead

Top comments (0)