584 lines
15 KiB
Go
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
|
|
}
|