244 lines
8.6 KiB
Go
244 lines
8.6 KiB
Go
package pgstreambroker
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"strings"
|
|
|
|
"github.com/centrifugal/centrifugo/v6/internal/pgschema"
|
|
)
|
|
|
|
// reconcileShardLock ensures the shard_lock table (for both variants) has
|
|
// exactly one row for each shard in [0, NumShards). Called on every
|
|
// EnsureSchema including the fast path so a NumShards change between runs
|
|
// is picked up. A stale shard_lock is load-bearing, not cosmetic: a missing
|
|
// row makes `PERFORM 1 FROM shard_lock WHERE shard_id = X FOR UPDATE`
|
|
// lock nothing, breaking per-shard publish serialization, which in turn
|
|
// lets stream IDs commit out of order and the outbox cursor skip rows.
|
|
// bothVariantsPresent reports whether the primary table exists for both the
|
|
// jsonb and binary variants, so EnsureSchema's fast path doesn't trust a partial
|
|
// install where only one variant was created.
|
|
func (e *PostgresStreamBroker) bothVariantsPresent(ctx context.Context) bool {
|
|
for _, prefix := range []string{e.names.jsonbPrefix, e.names.binaryPrefix} {
|
|
table := strings.TrimRight(prefix, "_")
|
|
if _, err := e.pool.Exec(ctx, fmt.Sprintf(`SELECT 1 FROM %s LIMIT 0`, table)); err != nil {
|
|
return false
|
|
}
|
|
}
|
|
return true
|
|
}
|
|
|
|
func (e *PostgresStreamBroker) reconcileShardLock(ctx context.Context) error {
|
|
for _, prefix := range []string{e.names.jsonbPrefix, e.names.binaryPrefix} {
|
|
shardLock := prefix + "shard_lock"
|
|
if _, err := e.pool.Exec(ctx, fmt.Sprintf(
|
|
`INSERT INTO %s (shard_id) SELECT generate_series(0, $1 - 1) ON CONFLICT DO NOTHING`,
|
|
shardLock), e.conf.NumShards); err != nil {
|
|
return &SchemaError{
|
|
Object: SchemaObject{Type: "table", Name: shardLock},
|
|
Op: "create",
|
|
Err: fmt.Errorf("populate shard_lock: %w", err),
|
|
}
|
|
}
|
|
if _, err := e.pool.Exec(ctx, fmt.Sprintf(
|
|
`DELETE FROM %s WHERE shard_id >= $1`,
|
|
shardLock), e.conf.NumShards); err != nil {
|
|
return &SchemaError{
|
|
Object: SchemaObject{Type: "table", Name: shardLock},
|
|
Op: "create",
|
|
Err: fmt.Errorf("trim shard_lock: %w", err),
|
|
}
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// versionTables returns the per-variant schema_version table names — used as
|
|
// the targets for every UPDATE schema_version performed during the schema
|
|
// lifecycle (specifically the final SetSchemaVersion after DDL).
|
|
func (e *PostgresStreamBroker) versionTables() []string {
|
|
return []string{
|
|
e.names.jsonbPrefix + "schema_version",
|
|
e.names.binaryPrefix + "schema_version",
|
|
}
|
|
}
|
|
|
|
// migrationVariants renders `template` once per variant (jsonb + binary) and
|
|
// pairs each rendered SQL with its variant's schema_version table.
|
|
func (e *PostgresStreamBroker) migrationVariants(template string) []pgschema.MigrationVariant {
|
|
return []pgschema.MigrationVariant{
|
|
{
|
|
SQL: renderSchemaTemplate(template, e.names.jsonbPrefix, false),
|
|
VersionTable: e.names.jsonbPrefix + "schema_version",
|
|
},
|
|
{
|
|
SQL: renderSchemaTemplate(template, e.names.binaryPrefix, true),
|
|
VersionTable: e.names.binaryPrefix + "schema_version",
|
|
},
|
|
}
|
|
}
|
|
|
|
// runMigrationsUnderLock — see pgmapbroker.runMigrationsUnderLock; identical
|
|
// shape and rationale, scoped to this broker's prefix tables.
|
|
func (e *PostgresStreamBroker) runMigrationsUnderLock(ctx context.Context, label string) error {
|
|
release, err := pgschema.AcquireMigrationLock(ctx, e.pool, label)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer release()
|
|
|
|
dbVersion, isFresh, err := pgschema.ReadSchemaVersion(ctx, e.pool, e.names.schemaVersion)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if isFresh {
|
|
return nil
|
|
}
|
|
if err := pgschema.CheckDowngrade(label, dbVersion, schemaVersion); err != nil {
|
|
return err
|
|
}
|
|
for v := dbVersion + 1; v <= schemaVersion; v++ {
|
|
sql, ok := schemaMigrations[v]
|
|
if !ok {
|
|
return fmt.Errorf("%s: missing schemaMigrations[%d] at runtime (schemaVersion=%d)", label, v, schemaVersion)
|
|
}
|
|
if err := pgschema.ApplyMigrationInTx(ctx, e.pool, label, v, e.migrationVariants(sql)); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// EnsureSchema creates all required database objects idempotently.
|
|
//
|
|
// Flow:
|
|
// 1. Read schema_version with error discrimination (transient errors must
|
|
// propagate — see pgschema.ReadSchemaVersion).
|
|
// 2. Fast path: version matches AND probe ok AND partitioned shape verified
|
|
// → reconcile shard_lock + lookahead partitions and return.
|
|
// 3. Reject downgrade.
|
|
// 4. Run pending migrations BEFORE DDL, each in its own transaction with
|
|
// atomic schema_version bump.
|
|
// 5. Render and run DDL for both variants (CREATE OR REPLACE picks up the
|
|
// latest function bodies; CREATE TABLE IF NOT EXISTS is a no-op when
|
|
// migrations have already brought the shape current).
|
|
// 6. Reconcile shard_lock; verify partitioned shape; pre-create partitions.
|
|
// 7. Final schema_version write — fatal on failure.
|
|
func (e *PostgresStreamBroker) EnsureSchema(ctx context.Context) error {
|
|
const label = "pgstreambroker"
|
|
|
|
dbVersion, isFresh, err := pgschema.ReadSchemaVersion(ctx, e.pool, e.names.schemaVersion)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
// Fast path: version current, both variants' history tables exist, partition
|
|
// shape ok. Probe BOTH variants (not just the active one): the two variants
|
|
// are created by separate Exec calls, so a partial fresh install can commit
|
|
// one and fail the other transiently. Trusting only the active variant would
|
|
// take the fast path and then fail in reconcileShardLock (which touches both)
|
|
// on the missing tables, before reaching the idempotent DDL that recreates
|
|
// them — wedging the broker on every restart. If a variant is missing, fall
|
|
// through and re-apply DDL to self-heal.
|
|
if !isFresh && dbVersion == schemaVersion {
|
|
if e.bothVariantsPresent(ctx) {
|
|
if err := e.verifyPartitionedShape(ctx); err != nil {
|
|
return err
|
|
}
|
|
if err := e.reconcileShardLock(ctx); err != nil {
|
|
return err
|
|
}
|
|
return e.ensureInitialPartitions(ctx)
|
|
}
|
|
}
|
|
|
|
if !isFresh {
|
|
if err := pgschema.CheckDowngrade(label, dbVersion, schemaVersion); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
|
|
// Migrations run BEFORE DDL, under a Postgres advisory lock so concurrent
|
|
// rolling-deploy upgrades can't run the same migration chain twice.
|
|
// See pgmapbroker.EnsureSchema for the full rationale.
|
|
if !isFresh {
|
|
if err := e.runMigrationsUnderLock(ctx, label); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
|
|
// Render and run DDL for both variants.
|
|
jsonbSQL := renderSchema(e.names.jsonbPrefix, false)
|
|
binarySQL := renderSchema(e.names.binaryPrefix, true)
|
|
jsonbDDL, jsonbFuncs := splitSchemaSQL(jsonbSQL)
|
|
binaryDDL, binaryFuncs := splitSchemaSQL(binarySQL)
|
|
for _, sql := range []string{jsonbDDL, binaryDDL, jsonbFuncs, binaryFuncs} {
|
|
if sql == "" {
|
|
continue
|
|
}
|
|
if err := e.execSchemaWithRetry(ctx, sql); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
|
|
if err := e.reconcileShardLock(ctx); err != nil {
|
|
return err
|
|
}
|
|
if err := e.verifyPartitionedShape(ctx); err != nil {
|
|
return err
|
|
}
|
|
if err := e.ensureInitialPartitions(ctx); err != nil {
|
|
return err
|
|
}
|
|
|
|
// Final schema_version write — see pgmapbroker for the rationale on
|
|
// "fatal on failure".
|
|
return pgschema.SetSchemaVersion(ctx, e.pool, label, schemaVersion, e.versionTables())
|
|
}
|
|
|
|
// verifyPartitionedShape probes pg_partitioned_table for the active variant's
|
|
// history table. If a non-partitioned table exists (from an earlier build of
|
|
// pgstreambroker), fail loudly with a clear "drop the legacy tables" message.
|
|
func (e *PostgresStreamBroker) verifyPartitionedShape(ctx context.Context) error {
|
|
var isPartitioned bool
|
|
err := e.pool.QueryRow(ctx, `
|
|
SELECT EXISTS(
|
|
SELECT 1 FROM pg_partitioned_table
|
|
WHERE partrelid = $1::regclass
|
|
)
|
|
`, e.names.stream).Scan(&isPartitioned)
|
|
if err != nil {
|
|
// Table doesn't exist yet (or other lookup error). Fall through —
|
|
// the schema template's CREATE TABLE will run later.
|
|
return nil
|
|
}
|
|
if !isPartitioned {
|
|
return &SchemaError{
|
|
Object: SchemaObject{Type: "table", Name: e.names.stream},
|
|
Op: "verify",
|
|
Err: fmt.Errorf(
|
|
"%s exists but is not partitioned. pgstreambroker schema requires "+
|
|
"a partitioned table and this build does not include migration logic. "+
|
|
"Drop the existing tables manually with: DROP TABLE IF EXISTS %s, %s, %s, %s, %s CASCADE",
|
|
e.names.stream,
|
|
e.names.stream, e.names.meta, e.names.idempotency, e.names.shardLock, e.names.schemaVersion,
|
|
),
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// ensureInitialPartitions pre-creates today's and lookahead partitions
|
|
// via pgoutbox.Partitioner.EnsureLookaheadPartitions.
|
|
func (e *PostgresStreamBroker) ensureInitialPartitions(ctx context.Context) error {
|
|
p := e.newPartitioner()
|
|
if err := p.EnsureLookaheadPartitions(ctx); err != nil {
|
|
return &SchemaError{
|
|
Object: SchemaObject{Type: "table", Name: e.names.stream},
|
|
Op: "create",
|
|
Err: err,
|
|
}
|
|
}
|
|
return nil
|
|
}
|