ariadne/internal/platform/database/migrate.go
2026-07-24 05:45:04 +09:00

584 lines
15 KiB
Go

package database
import (
"context"
"crypto/sha256"
"embed"
"errors"
"fmt"
"io/fs"
"sort"
"strings"
"time"
"github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgxpool"
"go.uber.org/zap"
)
const migrationLockID int64 = 648093714083
//go:embed migrations/*.up.sql
var migrations embed.FS
type MigrationOptions struct {
DatabaseURL string
RuntimeRole string
IdentityRole string
LockTimeout time.Duration
}
type migrationFile struct {
Name string
Checksum string
Content []byte
}
func RunMigrations(
ctx context.Context,
options MigrationOptions,
logger *zap.Logger,
) error {
if err := options.validate(); err != nil {
return err
}
poolConfig, err := pgxpool.ParseConfig(options.DatabaseURL)
if err != nil {
return fmt.Errorf("parse migration database URL: %w", err)
}
poolConfig.ConnConfig.RuntimeParams["application_name"] = "ariadne-migrate"
pool, err := pgxpool.NewWithConfig(ctx, poolConfig)
if err != nil {
return fmt.Errorf("open migration database: %w", err)
}
defer pool.Close()
connection, err := pool.Acquire(ctx)
if err != nil {
return fmt.Errorf("acquire migration connection: %w", err)
}
defer connection.Release()
if err := connection.Conn().Ping(ctx); err != nil {
return fmt.Errorf("ping migration database: %w", err)
}
if err := validateMigrationRole(ctx, connection, options); err != nil {
return err
}
if err := acquireMigrationLock(ctx, connection, options.LockTimeout); err != nil {
return err
}
defer releaseMigrationLock(connection)
if err := applyMigrations(ctx, connection); err != nil {
return err
}
if err := synchronizeAccess(ctx, connection, options); err != nil {
return err
}
logger.Info(
"database migrations and access policies applied",
zap.String("runtime_role", options.RuntimeRole),
zap.String("identity_role", options.IdentityRole),
)
return nil
}
func (options MigrationOptions) validate() error {
if strings.TrimSpace(options.DatabaseURL) == "" {
return errors.New("ARIADNE_DATABASE_MIGRATION_URL is required")
}
if strings.TrimSpace(options.RuntimeRole) == "" {
return errors.New("ARIADNE_DATABASE_RUNTIME_ROLE is required")
}
if strings.TrimSpace(options.IdentityRole) == "" {
return errors.New("ARIADNE_DATABASE_IDENTITY_ROLE is required")
}
if options.RuntimeRole == options.IdentityRole {
return errors.New("runtime and identity database roles must be different")
}
if options.LockTimeout <= 0 {
return errors.New("ARIADNE_MIGRATION_LOCK_TIMEOUT must be greater than zero")
}
return nil
}
func validateMigrationRole(
ctx context.Context,
connection *pgxpool.Conn,
options MigrationOptions,
) error {
var (
roleName string
isSuperuser bool
canBypassRLS bool
canCreateDB bool
canCreateRole bool
canReplicate bool
canInherit bool
)
if err := connection.QueryRow(ctx, `
SELECT
rolname,
rolsuper,
rolbypassrls,
rolcreatedb,
rolcreaterole,
rolreplication,
rolinherit
FROM pg_roles
WHERE rolname = current_user
`).Scan(
&roleName,
&isSuperuser,
&canBypassRLS,
&canCreateDB,
&canCreateRole,
&canReplicate,
&canInherit,
); err != nil {
return fmt.Errorf("inspect migration role: %w", err)
}
if isSuperuser ||
!canBypassRLS ||
canCreateDB ||
canCreateRole ||
canReplicate ||
canInherit {
return fmt.Errorf(
"migration role %q must be NOSUPERUSER, BYPASSRLS, "+
"NOCREATEDB, NOCREATEROLE, NOREPLICATION and NOINHERIT",
roleName,
)
}
if roleName == options.RuntimeRole || roleName == options.IdentityRole {
return fmt.Errorf("migration role %q must be dedicated", roleName)
}
return nil
}
func acquireMigrationLock(
ctx context.Context,
connection *pgxpool.Conn,
timeout time.Duration,
) error {
lockContext, cancel := context.WithTimeout(ctx, timeout)
defer cancel()
ticker := time.NewTicker(250 * time.Millisecond)
defer ticker.Stop()
for {
var acquired bool
if err := connection.QueryRow(
lockContext,
"SELECT pg_try_advisory_lock($1)",
migrationLockID,
).Scan(&acquired); err != nil {
return fmt.Errorf("try migration lock: %w", err)
}
if acquired {
return nil
}
select {
case <-lockContext.Done():
return fmt.Errorf(
"acquire migration lock within %s: %w",
timeout,
lockContext.Err(),
)
case <-ticker.C:
}
}
}
func releaseMigrationLock(connection *pgxpool.Conn) {
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
_, _ = connection.Exec(ctx, "SELECT pg_advisory_unlock($1)", migrationLockID)
}
func applyMigrations(ctx context.Context, connection *pgxpool.Conn) error {
if _, err := connection.Exec(ctx, `
CREATE TABLE IF NOT EXISTS public.ariadne_schema_migrations (
version text PRIMARY KEY,
checksum text NOT NULL,
applied_at timestamptz NOT NULL DEFAULT now()
)
`); err != nil {
return fmt.Errorf("create migration table: %w", err)
}
manifest, err := embeddedMigrationManifest()
if err != nil {
return err
}
for _, migration := range manifest {
var appliedChecksum string
err := connection.QueryRow(
ctx,
`
SELECT checksum
FROM public.ariadne_schema_migrations
WHERE version = $1
`,
migration.Name,
).Scan(&appliedChecksum)
switch {
case err == nil:
if appliedChecksum != migration.Checksum {
return fmt.Errorf(
"migration %s checksum changed: applied %s, current %s",
migration.Name,
appliedChecksum,
migration.Checksum,
)
}
continue
case !errors.Is(err, pgx.ErrNoRows):
return fmt.Errorf("check migration %s: %w", migration.Name, err)
}
transaction, err := connection.Begin(ctx)
if err != nil {
return fmt.Errorf("begin migration %s: %w", migration.Name, err)
}
if _, err := transaction.Exec(ctx, string(migration.Content)); err != nil {
_ = transaction.Rollback(ctx)
return fmt.Errorf("apply migration %s: %w", migration.Name, err)
}
if _, err := transaction.Exec(
ctx,
`
INSERT INTO public.ariadne_schema_migrations (version, checksum)
VALUES ($1, $2)
`,
migration.Name,
migration.Checksum,
); err != nil {
_ = transaction.Rollback(ctx)
return fmt.Errorf("record migration %s: %w", migration.Name, err)
}
if err := transaction.Commit(ctx); err != nil {
return fmt.Errorf("commit migration %s: %w", migration.Name, err)
}
}
return nil
}
func embeddedMigrationManifest() ([]migrationFile, error) {
entries, err := fs.ReadDir(migrations, "migrations")
if err != nil {
return nil, fmt.Errorf("read migrations: %w", err)
}
names := make([]string, 0, len(entries))
for _, entry := range entries {
if !entry.IsDir() && strings.HasSuffix(entry.Name(), ".up.sql") {
names = append(names, entry.Name())
}
}
sort.Strings(names)
manifest := make([]migrationFile, 0, len(names))
for _, name := range names {
content, err := migrations.ReadFile("migrations/" + name)
if err != nil {
return nil, fmt.Errorf("read migration %s: %w", name, err)
}
manifest = append(manifest, migrationFile{
Name: name,
Checksum: fmt.Sprintf("%x", sha256.Sum256(content)),
Content: content,
})
}
return manifest, nil
}
func validateSchema(ctx context.Context, database databaseQuerier) error {
manifest, err := embeddedMigrationManifest()
if err != nil {
return err
}
rows, err := database.Query(ctx, `
SELECT version, checksum
FROM public.ariadne_schema_migrations
ORDER BY version
`)
if err != nil {
return fmt.Errorf("read applied migrations: %w", err)
}
defer rows.Close()
applied := make(map[string]string, len(manifest))
for rows.Next() {
var version, checksum string
if err := rows.Scan(&version, &checksum); err != nil {
return fmt.Errorf("scan applied migration: %w", err)
}
applied[version] = checksum
}
if err := rows.Err(); err != nil {
return fmt.Errorf("iterate applied migrations: %w", err)
}
if len(applied) != len(manifest) {
return fmt.Errorf(
"database migration count = %d, application expects %d",
len(applied),
len(manifest),
)
}
for _, migration := range manifest {
if checksum, exists := applied[migration.Name]; !exists {
return fmt.Errorf("database migration %s is missing", migration.Name)
} else if checksum != migration.Checksum {
return fmt.Errorf("database migration %s checksum does not match", migration.Name)
}
}
return nil
}
func synchronizeAccess(
ctx context.Context,
connection *pgxpool.Conn,
options MigrationOptions,
) error {
transaction, err := connection.Begin(ctx)
if err != nil {
return fmt.Errorf("begin access policy transaction: %w", err)
}
defer func() {
rollbackContext, cancel := context.WithTimeout(
context.Background(),
5*time.Second,
)
defer cancel()
_ = transaction.Rollback(rollbackContext)
}()
if err := validateTenantTableCatalog(ctx, transaction); err != nil {
return err
}
if _, err := transaction.Exec(
ctx,
"REVOKE CREATE ON SCHEMA public FROM PUBLIC",
); err != nil {
return fmt.Errorf("restrict public schema creation: %w", err)
}
if err := synchronizeRuntimeAccess(ctx, transaction, options.RuntimeRole); err != nil {
return err
}
if err := synchronizeIdentityAccess(ctx, transaction, options.IdentityRole); err != nil {
return err
}
if err := validateRuntimePrivileges(ctx, transaction, options.RuntimeRole); err != nil {
return err
}
if err := validateIdentityPrivileges(ctx, transaction, options.IdentityRole); err != nil {
return err
}
if err := transaction.Commit(ctx); err != nil {
return fmt.Errorf("commit access policies: %w", err)
}
return nil
}
func synchronizeRuntimeAccess(ctx context.Context, transaction pgx.Tx, role string) error {
if err := validateRestrictedRole(ctx, transaction, role, "runtime"); err != nil {
return err
}
roleIdentifier := pgx.Identifier{role}.Sanitize()
if _, err := transaction.Exec(
ctx,
"REVOKE CREATE ON SCHEMA public FROM "+roleIdentifier,
); err != nil {
return fmt.Errorf("restrict runtime schema creation: %w", err)
}
if _, err := transaction.Exec(
ctx,
"GRANT USAGE ON SCHEMA public TO "+roleIdentifier,
); err != nil {
return fmt.Errorf("grant runtime schema access: %w", err)
}
for _, tableName := range orderedRuntimeTables() {
permissions, exists := runtimeTablePermissions(tableName)
if !exists {
return fmt.Errorf("runtime permissions for %q are not registered", tableName)
}
tableIdentifier := pgx.Identifier{"public", tableName}.Sanitize()
if _, err := transaction.Exec(
ctx,
"REVOKE ALL PRIVILEGES ON TABLE "+tableIdentifier+" FROM "+roleIdentifier,
); err != nil {
return fmt.Errorf("reset runtime access to %s: %w", tableName, err)
}
if _, err := transaction.Exec(
ctx,
"GRANT "+permissions+" ON TABLE "+tableIdentifier+" TO "+roleIdentifier,
); err != nil {
return fmt.Errorf("grant runtime access to %s: %w", tableName, err)
}
}
for _, tableName := range []string{"users", "ariadne_schema_migrations"} {
tableIdentifier := pgx.Identifier{"public", tableName}.Sanitize()
if _, err := transaction.Exec(
ctx,
"REVOKE ALL PRIVILEGES ON TABLE "+tableIdentifier+" FROM "+roleIdentifier,
); err != nil {
return fmt.Errorf("reset runtime access to %s: %w", tableName, err)
}
}
if _, err := transaction.Exec(
ctx,
"GRANT SELECT ON TABLE public.ariadne_schema_migrations TO "+roleIdentifier,
); err != nil {
return fmt.Errorf("grant runtime migration metadata access: %w", err)
}
if _, err := transaction.Exec(
ctx,
"GRANT EXECUTE ON FUNCTION public.ariadne_current_organization_id() TO "+
roleIdentifier,
); err != nil {
return fmt.Errorf("grant runtime tenant function access: %w", err)
}
return nil
}
func synchronizeIdentityAccess(ctx context.Context, transaction pgx.Tx, role string) error {
if err := validateRestrictedRole(ctx, transaction, role, "identity"); err != nil {
return err
}
roleIdentifier := pgx.Identifier{role}.Sanitize()
if _, err := transaction.Exec(
ctx,
"REVOKE CREATE ON SCHEMA public FROM "+roleIdentifier,
); err != nil {
return fmt.Errorf("restrict identity schema creation: %w", err)
}
if _, err := transaction.Exec(
ctx,
"GRANT USAGE ON SCHEMA public TO "+roleIdentifier,
); err != nil {
return fmt.Errorf("grant identity schema access: %w", err)
}
for _, tableName := range orderedRuntimeTables() {
tableIdentifier := pgx.Identifier{"public", tableName}.Sanitize()
if _, err := transaction.Exec(
ctx,
"REVOKE ALL PRIVILEGES ON TABLE "+tableIdentifier+" FROM "+roleIdentifier,
); err != nil {
return fmt.Errorf("reset identity access to %s: %w", tableName, err)
}
}
for _, tableName := range []string{"users", "ariadne_schema_migrations"} {
tableIdentifier := pgx.Identifier{"public", tableName}.Sanitize()
if _, err := transaction.Exec(
ctx,
"REVOKE ALL PRIVILEGES ON TABLE "+tableIdentifier+" FROM "+roleIdentifier,
); err != nil {
return fmt.Errorf("reset identity access to %s: %w", tableName, err)
}
}
if _, err := transaction.Exec(
ctx,
"GRANT SELECT ON TABLE public.users, public.organizations, "+
"public.organization_memberships TO "+roleIdentifier,
); err != nil {
return fmt.Errorf("grant identity lookup access: %w", err)
}
policies := []struct {
name string
table string
}{
{name: "organizations_identity_lookup", table: "organizations"},
{name: "organization_memberships_identity_lookup", table: "organization_memberships"},
}
for _, policy := range policies {
policyIdentifier := pgx.Identifier{policy.name}.Sanitize()
tableIdentifier := pgx.Identifier{"public", policy.table}.Sanitize()
if _, err := transaction.Exec(
ctx,
"DROP POLICY IF EXISTS "+policyIdentifier+" ON "+tableIdentifier,
); err != nil {
return fmt.Errorf("drop identity policy %s: %w", policy.name, err)
}
if _, err := transaction.Exec(
ctx,
"CREATE POLICY "+policyIdentifier+" ON "+tableIdentifier+
" FOR SELECT TO "+roleIdentifier+" USING (true)",
); err != nil {
return fmt.Errorf("create identity policy %s: %w", policy.name, err)
}
}
return nil
}
func validateRestrictedRole(
ctx context.Context,
database databaseQuerier,
role string,
purpose string,
) error {
var (
exists bool
isSuperuser bool
canBypassRLS bool
canCreateDB bool
canCreateRole bool
canReplicate bool
canInherit bool
)
if err := database.QueryRow(ctx, `
SELECT
COUNT(*) = 1,
COALESCE(bool_or(rolsuper), false),
COALESCE(bool_or(rolbypassrls), false),
COALESCE(bool_or(rolcreatedb), false),
COALESCE(bool_or(rolcreaterole), false),
COALESCE(bool_or(rolreplication), false),
COALESCE(bool_or(rolinherit), false)
FROM pg_roles
WHERE rolname = $1
`, role).Scan(
&exists,
&isSuperuser,
&canBypassRLS,
&canCreateDB,
&canCreateRole,
&canReplicate,
&canInherit,
); err != nil {
return fmt.Errorf("check %s role: %w", purpose, err)
}
if !exists {
return fmt.Errorf("%s role %q does not exist", purpose, role)
}
if isSuperuser ||
canBypassRLS ||
canCreateDB ||
canCreateRole ||
canReplicate ||
canInherit {
return fmt.Errorf(
"%s role %q must be NOSUPERUSER, NOBYPASSRLS, "+
"NOCREATEDB, NOCREATEROLE, NOREPLICATION and NOINHERIT",
purpose,
role,
)
}
return nil
}