91 lines
3.6 KiB
Go
91 lines
3.6 KiB
Go
package migrations
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"regexp"
|
|
"strings"
|
|
|
|
"github.com/jackc/pgx/v5"
|
|
)
|
|
|
|
var rolePattern = regexp.MustCompile(`^[a-z_][a-z0-9_]{0,62}$`)
|
|
|
|
var appTables = []struct {
|
|
name string
|
|
privileges string
|
|
}{
|
|
{"assets", "SELECT, INSERT, DELETE"},
|
|
{"generation_jobs", "SELECT, INSERT, UPDATE, DELETE"},
|
|
{"seedream_layer_compositions", "SELECT, INSERT, UPDATE, DELETE"},
|
|
{"usage_events", "SELECT, INSERT, UPDATE"},
|
|
{"projects", "SELECT, UPDATE"},
|
|
{"image_templates", "SELECT, INSERT, UPDATE, DELETE"},
|
|
{"platform_organizations", "SELECT, INSERT, UPDATE, DELETE"},
|
|
{"platform_users", "SELECT, INSERT, UPDATE, DELETE"},
|
|
{"platform_account_migrations", "SELECT, INSERT, UPDATE"},
|
|
{"platform_runtime_settings", "SELECT, INSERT, UPDATE"},
|
|
{"billing_price_rules", "SELECT, INSERT, UPDATE"},
|
|
{"billing_wallets", "SELECT, INSERT, UPDATE"},
|
|
{"billing_ledger", "SELECT, INSERT"},
|
|
{"payment_orders", "SELECT, INSERT, UPDATE"},
|
|
{"payment_refunds", "SELECT, INSERT, UPDATE"},
|
|
}
|
|
|
|
func provisionRole(ctx context.Context, conn *pgx.Conn, role string) error {
|
|
// rolePattern has already validated the identifier. Quote it consistently.
|
|
quotedRole := `"` + role + `"`
|
|
managed := make([]string, len(appTables))
|
|
for i, table := range appTables {
|
|
managed[i] = `public."` + table.name + `"`
|
|
}
|
|
tx, err := conn.Begin(ctx)
|
|
if err != nil {
|
|
return safeDBError("begin application role provisioning", "", err)
|
|
}
|
|
commands := []string{
|
|
"REVOKE CREATE ON SCHEMA public FROM PUBLIC",
|
|
"GRANT USAGE ON SCHEMA public TO " + quotedRole,
|
|
"REVOKE ALL ON TABLE " + strings.Join(managed, ", ") + " FROM " + quotedRole,
|
|
}
|
|
for _, table := range appTables {
|
|
commands = append(commands, fmt.Sprintf("GRANT %s ON TABLE public.\"%s\" TO %s", table.privileges, table.name, quotedRole))
|
|
}
|
|
commands = append(commands,
|
|
"REVOKE ALL ON FUNCTION public.claim_generation_jobs(text, integer, integer) FROM PUBLIC",
|
|
"REVOKE ALL ON FUNCTION public.billing_post_wallet_entry(text, text, text, text, text, bigint, text, text, text, jsonb) FROM PUBLIC",
|
|
"GRANT EXECUTE ON FUNCTION public.claim_generation_jobs(text, integer, integer) TO "+quotedRole,
|
|
"GRANT EXECUTE ON FUNCTION public.billing_post_wallet_entry(text, text, text, text, text, bigint, text, text, text, jsonb) TO "+quotedRole,
|
|
)
|
|
for _, command := range commands {
|
|
if _, err := tx.Exec(ctx, command); err != nil {
|
|
_ = tx.Rollback(context.Background())
|
|
return safeDBError("provision application role", "", err)
|
|
}
|
|
}
|
|
if err := tx.Commit(ctx); err != nil {
|
|
return safeDBError("commit application role privileges", "", err)
|
|
}
|
|
tableNames := make([]string, len(appTables))
|
|
privileges := make([]string, len(appTables))
|
|
for i, table := range appTables {
|
|
tableNames[i] = "public." + table.name
|
|
privileges[i] = strings.ReplaceAll(table.privileges, " ", "")
|
|
}
|
|
var tablesReady, claimReady, billingReady bool
|
|
err = conn.QueryRow(ctx, `SELECT
|
|
bool_and(has_table_privilege($1, table_name, privileges)) AS tables_ready,
|
|
has_function_privilege($1, 'public.claim_generation_jobs(text,integer,integer)', 'EXECUTE') AS claim_ready,
|
|
has_function_privilege($1,
|
|
'public.billing_post_wallet_entry(text,text,text,text,text,bigint,text,text,text,jsonb)',
|
|
'EXECUTE') AS billing_ready
|
|
FROM unnest($2::text[], $3::text[]) AS required(table_name, privileges)`, role, tableNames, privileges).Scan(&tablesReady, &claimReady, &billingReady)
|
|
if err != nil {
|
|
return safeDBError("verify application role privileges", "", err)
|
|
}
|
|
if !tablesReady || !claimReady || !billingReady {
|
|
return fmt.Errorf("application role privilege verification failed")
|
|
}
|
|
return nil
|
|
}
|