Janady13/missionmatch / internal/platform/database/database.go

Source identity

Field	Recorded value
Source type	GitHub source
Original record	Janady13/missionmatch / internal/platform/database/database.go
Date	Not specified in source
Repository	Janady13/missionmatch
Revision	720d47fae476b0d38477ae4401697d372b1ea2bc
Branch	feature/opportunity-to-outcome
Path or archive member	internal/platform/database/database.go
Original SHA-256	0ea84cec52036595458a12e68af1f7d359c0d8c3f2611a72f8a331bd9553e96c
Git blob SHA	00c6dc5deceefc94bff6edfac6fd5d3ff4a8dfb4
Reading-copy SHA-256	0ea84cec52036595458a12e68af1f7d359c0d8c3f2611a72f8a331bd9553e96c

Original source ↗

Every recorded occurrence

These 2 records contain identical original source bytes. Their paths and revisions remain part of the archive.

Original record	Branch or archive	Revision / identifier	Original location
Janady13/missionmatch / internal/platform/database/database.go	feature/opportunity-to-outcome	720d47fae476b0d38477ae4401697d372b1ea2bc	Open original
Janady13/missionmatch / internal/platform/database/database.go	feature/aioe-upgrade	2178e7298d0e5d9af83afa12402fd587f0ba6ed8	Open original

Recovered source transcript

Download complete source text ↓128 lines · 5.8 KB

// Package database constructs the application's Postgres connection pool.
//
// Two things are wired here that the rest of the app relies on (see docs/STACK_DECISIONS.md §2/§3/§5):
//
//  1. pgvector codecs are registered on EVERY physical connection via AfterConnect, so vector
//     columns bind and scan as pgvector.Vector. Registering once on a single conn is a bug under
//     pool churn, it must be AfterConnect.
//  2. The pool is expected to authenticate as the NOBYPASSRLS application role (never the migrator),
//     so Postgres Row-Level Security is always in force. Which role is used is a property of the
//     DSN, not of this code.
package database

import (
	"context"
	"fmt"
	"time"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
	pgxvec "github.com/pgvector/pgvector-go/pgx"
)

// NewPool parses dsn, registers the pgvector codecs per connection, and returns a ready pool.
// The caller owns the returned pool's Close(). dsn MUST point at the NOBYPASSRLS application role.
func NewPool(ctx context.Context, dsn string) (*pgxpool.Pool, error) {
	return newPool(ctx, dsn, nil)
}

// NewBoundedPool is NewPool for a per tenant pool: same pgvector registration and same role
// expectation, but sized small because the database per tenant model keeps MANY pools alive at once and
// the product of pool count and per pool MaxConns must stay under the server's max_connections.
//
// Three deliberate differences from the single shared pool:
//
//  1. MaxConns is capped by the caller rather than left at the pgx default of runtime.NumCPU, which
//     would let a handful of tenants exhaust the server's connection slots on a large host.
//  2. MinConns is zero, so an idle tenant holds no connection at all. A shared pool benefits from warm
//     connections; a tenant pool that may sit unused for hours should cost nothing while it does.
//  3. A nonzero MaxConnLifetimeJitter, so the connections of many tenant pools opened at the same
//     moment (a restart, a traffic spike) do not all expire together and reconnect in a thundering
//     herd. Jitter defaults to a tenth of the configured lifetime, or five minutes when the DSN sets
//     no lifetime at all.
func NewBoundedPool(ctx context.Context, dsn string, maxConns int32) (*pgxpool.Pool, error) {
	return newPool(ctx, dsn, func(cfg *pgxpool.Config) {
		if maxConns > 0 {
			cfg.MaxConns = maxConns
		}
		cfg.MinConns = 0
		if cfg.MaxConnLifetime > 0 && cfg.MaxConnLifetimeJitter == 0 {
			cfg.MaxConnLifetimeJitter = cfg.MaxConnLifetime / 10
		} else if cfg.MaxConnLifetimeJitter == 0 {
			cfg.MaxConnLifetimeJitter = 5 * time.Minute
		}
	})
}

// newPool is the one place a pool is actually constructed, so the pgvector registration that both the
// shared pool and every tenant pool depend on cannot drift between them. tune may be nil.
func newPool(ctx context.Context, dsn string, tune func(*pgxpool.Config)) (*pgxpool.Pool, error) {
	cfg, err := pgxpool.ParseConfig(dsn)
	if err != nil {
		return nil, fmt.Errorf("parse pool config: %w", err)
	}

	// AfterConnect runs once per physical connection the pool opens, the only correct place to
	// register vector/halfvec/sparsevec so pgvector.Vector works as a bind param and Scan target.
	cfg.AfterConnect = func(ctx context.Context, conn *pgx.Conn) error {
		if err := pgxvec.RegisterTypes(ctx, conn); err != nil {
			return fmt.Errorf("register pgvector types: %w", err)
		}
		return nil
	}

	if tune != nil {
		tune(cfg)
	}

	pool, err := pgxpool.NewWithConfig(ctx, cfg)
	if err != nil {
		return nil, fmt.Errorf("create pool: %w", err)
	}
	return pool, nil
}

// AssertRole fails closed unless the pool authenticates as wantUser with rolbypassrls=false. It is
// the runtime guardrail behind the whole RLS and EEO isolation story: if an operator pastes a
// BYPASSRLS or superuser DSN (for example the migrator DSN) into the app or compliance variable, RLS
// would be silently bypassed and every tenant's rows, including protected demographic attributes,
// would leak cross org. This converts that misconfiguration into a hard boot failure.
func AssertRole(ctx context.Context, pool *pgxpool.Pool, wantUser string) error {
	var user string
	var bypass bool
	err := pool.QueryRow(ctx,
		`SELECT current_user, rolbypassrls FROM pg_roles WHERE rolname = current_user`).Scan(&user, &bypass)
	if err != nil {
		return fmt.Errorf("role assertion query: %w", err)
	}
	if user != wantUser {
		return fmt.Errorf("pool authenticates as %q, expected %q", user, wantUser)
	}
	if bypass {
		return fmt.Errorf("pool role %q has BYPASSRLS, which would bypass row level security", user)
	}
	return nil
}

// AssertMigratorRole is the mirror of AssertRole for the privileged batch jobs (the sync and import
// commands). Those jobs need the migrator role for one of two reasons: import-candidates writes a
// tenant scoped row with an explicit organization_id and no app.current_org_id, so a NOBYPASSRLS role
// would be rejected by the RLS WITH CHECK; the catalog syncs (SAM, awards, wages) write GLOBAL
// reference tables that the application role is granted read only, so a NOBYPASSRLS role would be
// denied INSERT. Both failures are caught here, at boot, by requiring a BYPASSRLS role. If an operator
// pastes the application DSN into a job's MM_MIGRATOR_DATABASE_URL this becomes a hard, legible boot
// failure instead of an RLS rejection or a permission denied error partway through the run.
func AssertMigratorRole(ctx context.Context, pool *pgxpool.Pool) error {
	var user string
	var bypass bool
	err := pool.QueryRow(ctx,
		`SELECT current_user, rolbypassrls FROM pg_roles WHERE rolname = current_user`).Scan(&user, &bypass)
	if err != nil {
		return fmt.Errorf("role assertion query: %w", err)
	}
	if !bypass {
		return fmt.Errorf("pool role %q lacks BYPASSRLS; this job needs the migrator DSN (MM_MIGRATOR_DATABASE_URL)", user)
	}
	return nil
}