Skip to content

Multi-Tenancy

Mirabelle supports multiple customers (tenants) on a single server deployment. Each tenant has a completely isolated PostgreSQL database. A shared master database holds the account registry and the JWT secret.

Overview

flowchart TD
    subgraph master["forecastai_master"]
        ACC["master.accounts\nid · display_name · db_name\nschema_name · connection_params"]
        SA["master.superadmins\nCross-tenant admin users"]
        MP["master.parameters\nJWT secret · token expiry"]
        UA["master.user_accounts\nemail ↔ account mapping"]
    end

    subgraph cache["In-Process Cache (main.py)"]
        AC["_account_cache: dict\naccount_id (UUID) → DB config"]
    end

    subgraph req["Per-Request (asyncio context)"]
        MW["JWTAuthMiddleware\nDecode JWT → extract account_id\nLook up _account_cache\nSet _account_ctx ContextVar"]
        CTX["_account_ctx\nContextVar[dict | None]"]
    end

    subgraph tenants["Tenant Databases"]
        T1["thebicycle\nscenario schema"]
        T2["aerospace\nscenario schema"]
        T3["thegreenery\nscenario schema"]
    end

    ACC --> AC
    AC --> MW
    MW --> CTX
    CTX --> T1
    CTX --> T2
    CTX --> T3

The Master Database

forecastai_master is a separate PostgreSQL database that acts as the control plane. It contains only four tables:

  • master.accounts — one row per tenant with the database name, schema name, and optional connection overrides.
  • master.superadmins — cross-tenant administrator accounts.
  • master.parameters — global configuration including the JWT secret.
  • master.user_accounts — maps user emails to the accounts they can access.

Master DB separate from tenant DBs

The master database must be accessible at all times. If it is unreachable, _build_account_cache() fails silently with a warning and the account cache remains empty, causing all multi-tenant authentication to fail. The API will still operate in single-tenant mode if config/config.yaml is present.

Account Cache (_account_cache)

At startup, main.py calls _build_account_cache():

def _build_account_cache() -> None:
    """Load all active accounts from master DB into _account_cache."""
    global _account_cache
    from db.master_db import get_all_accounts, get_master_pg_config
    master_pg = get_master_pg_config()
    new_entries: dict = {}
    for acc in get_all_accounts():
        cp = acc.get("connection_params") or {}
        new_entries[acc["id"]] = {
            "host":         cp.get("host", master_pg["host"]),
            "port":         cp.get("port", master_pg["port"]),
            "database":     acc["db_name"],
            "user":         cp.get("user", master_pg["user"]),
            "password":     cp.get("password", master_pg["password"]),
            "schema":       acc["schema_name"],
            "sslmode":      cp.get("sslmode", "disable"),
            "display_name": acc.get("display_name", ""),
        }
    _account_cache.clear()
    _account_cache.update(new_entries)

The resulting dict maps each account UUID (as a string) to its full DB connection config. This dict is passed by reference into JWTAuthMiddleware and into admin.py's set_account_cache_ref(), so changes (e.g. after a new tenant is provisioned) are immediately reflected in the middleware without a restart.

Cache Refresh

The cache is refreshed by calling POST /api/admin/cache/refresh (SuperAdmin only). This endpoint calls _build_account_cache() again and the middleware instantly picks up the new entries via the shared mutable dict reference.

Per-Request Context (ContextVar)

Python's contextvars.ContextVar provides async-safe per-request isolation. The implementation in files/db/db.py:

_account_ctx: ContextVar[Optional[dict]] = ContextVar("_account_ctx", default=None)

def set_account_context(cfg: dict):
    return _account_ctx.set(cfg)   # returns a Token for later reset

def reset_account_context(token) -> None:
    _account_ctx.reset(token)      # restores previous value

And in get_conn():

def get_conn() -> psycopg2.extensions.connection:
    pg = _account_ctx.get() or _get_pg_config()  # tenant DB or config.yaml fallback
    conn = psycopg2.connect(
        host=pg["host"], port=pg["port"], dbname=pg["database"],
        user=pg["user"], password=pg["password"],
        sslmode=pg.get("sslmode", "disable"),
        options=f"-c search_path={pg['schema']},public",
    )
    conn.autocommit = False
    return conn

This means every DB query inside a request handler automatically uses the correct tenant's database without any explicit parameter threading.

JWT and Multi-Tenancy

Every JWT payload includes the account_id claim:

{
  "jti": "uuid",
  "sub": "user-uuid",
  "email": "user@example.com",
  "role": "admin",
  "display_name": "Alice",
  "auth_provider": "local",
  "allowed_segments": [1, 2],
  "can_run_process": true,
  "can_create_override": true,
  "is_superadmin": false,
  "account_id": "tenant-account-uuid",
  "account_name": "The Bicycle Company",
  "exp": 1744000000
}

The account_id field is the UUID from master.accounts.id. It tells the middleware which tenant's database to use.

SuperAdmin tokens carry is_superadmin: true and always include an account_id indicating which tenant account they are currently operating as. SuperAdmins can switch accounts by logging in again with a different account_id.

Middleware Flow (Per Request)

sequenceDiagram
    participant Client
    participant FastAPI
    participant JWTAuthMiddleware
    participant account_cache
    participant _account_ctx
    participant Handler
    participant get_conn

    Client->>FastAPI: GET /api/forecasts (Authorization: Bearer <jwt>)
    FastAPI->>JWTAuthMiddleware: dispatch(request)
    JWTAuthMiddleware->>JWTAuthMiddleware: Decode JWT (HS256)
    JWTAuthMiddleware->>account_cache: Look up account_id
    account_cache->>JWTAuthMiddleware: DB config dict
    JWTAuthMiddleware->>_account_ctx: set_account_context(cfg) → token
    JWTAuthMiddleware->>JWTAuthMiddleware: Check revoked_tokens
    JWTAuthMiddleware->>FastAPI: request.state.user = {...}
    FastAPI->>Handler: Execute endpoint
    Handler->>get_conn: get_conn()
    get_conn->>_account_ctx: .get() → DB config
    get_conn->>Handler: psycopg2 connection to tenant DB
    Handler->>FastAPI: Return response
    FastAPI->>JWTAuthMiddleware: (finally) reset_account_context(token)
    FastAPI->>Client: HTTP response

The try/finally in the middleware guarantees reset_account_context(token) is always called, even if the handler raises an exception. This prevents context leakage between requests in the same asyncio event loop.

Token Revocation

When a user logs out, their JWT's jti (JWT ID) is inserted into {schema}.revoked_tokens. On every subsequent request, the middleware checks this table:

cur.execute(
    f"SELECT 1 FROM {effective_schema}.revoked_tokens WHERE jti = %s",
    (jti,),
)
if cur.fetchone():
    return JSONResponse(status_code=401, content={"detail": "Token has been revoked"})

Revocation is per-tenant

The revocation check runs against the tenant's own revoked_tokens table (derived from effective_schema). SuperAdmin tokens are checked against whichever tenant schema their account_id resolves to.

Role Hierarchy

superadmin  (master.superadmins — global)
    │  Full access to all tenants.
    │  Accesses via /api/admin/* endpoints.
    │  JWT includes is_superadmin=true.
    │
    ▼
  admin     (scenario.users where role='admin')
    │  Full access within their tenant.
    │  Can manage users, parameters, segments.
    │
    ▼
  user      (scenario.users where role='user')
         Limited by:
           allowed_segments      — which segments they can view
           allowed_segments_edit — which segments they can edit
           can_run_process       — whether they can trigger pipeline runs
           can_create_override   — whether they can apply forecast overrides

Role checks in the API use the get_current_user dependency injector which reads request.state.user (set by the middleware). Guard functions:

def _user_has_segment_access(user: dict, segment_id: int, edit: bool = False) -> bool:
    if user.get("role") == "admin":
        return True
    allowed = user.get("allowed_segments_edit" if edit else "allowed_segments", []) or []
    return segment_id in allowed

def _user_can_run_process(user: dict) -> bool:
    return user.get("role") in ("admin", "superadmin") or user.get("can_run_process", False)

Provisioning a New Tenant

Step 1: Create the PostgreSQL Database

-- Connect as postgres superuser
CREATE DATABASE newclient;
\c newclient
\i files/DDL/schema.sql
\i files/DDL/supply_schema.sql
\i files/DDL/migrations/exception_workflow.sql

Step 2: Register in Master DB

\c forecastai_master
INSERT INTO master.accounts (display_name, db_name, schema_name)
VALUES ('New Client Ltd', 'newclient', 'scenario')
RETURNING id;

Note the returned UUID — this becomes the account_id in all JWTs for this tenant.

Step 3: Refresh the Account Cache

Either restart the API server, or call the cache refresh endpoint (SuperAdmin token required):

curl -X POST http://localhost:8002/api/admin/cache/refresh \
  -H "Authorization: Bearer <superadmin-token>"

Step 4: Create the First Admin User

Use the Admin Panel in the UI (as SuperAdmin) or call the API:

curl -X POST http://localhost:8002/api/auth/users \
  -H "Authorization: Bearer <superadmin-token>" \
  -H "Content-Type: application/json" \
  -d '{
    "email": "admin@newclient.com",
    "display_name": "Admin User",
    "password": "secure-password",
    "role": "admin"
  }'

Step 5: Seed Parameters

Navigate to Settings as the new admin and configure the data_source parameter set to point to the new client's ERP/demand database.

Connection Pooling Considerations

Mirabelle does not use a connection pool like PgBouncer or SQLAlchemy's pool. Every get_conn() call opens a new psycopg2 connection and closes it in a finally block. This is safe for the expected request volume (dozens to hundreds of concurrent users) but may become a bottleneck at high concurrency.

For production deployments with many concurrent users, consider: 1. Running PgBouncer in transaction-mode pooling in front of PostgreSQL. 2. Implementing connection pooling in files/db/db.py using psycopg2.pool.ThreadedConnectionPool.

ContextVar and Dask workers

Dask workers run in separate OS processes and do not inherit the middleware's ContextVar. Pipeline scripts that launch Dask tasks must pass DB credentials explicitly via environment variables (DB_HOST, DB_PORT, DB_NAME, DB_USER, DB_PASSWORD, DB_SCHEMA) or a temporary config file.