RenewCMS/internal/infrastructure/persistence/database/postgresBackup.go

package database

import (
	"fmt"
	"strings"
	"time"

	"gorm.io/gorm"
)

const (
	postgresSchema = "public"

	postgresTablesQuery = "SELECT table_name FROM information_schema.tables WHERE table_schema = ? AND table_type = 'BASE TABLE' ORDER BY table_name"

	postgresColumnsQuery = `SELECT a.attname AS name, format_type(a.atttypid, a.atttypmod) AS data_type, a.attnotnull AS not_null, COALESCE(pg_get_expr(d.adbin, d.adrelid), '') AS default_value, a.attidentity::text AS identity
FROM pg_attribute a
JOIN pg_class c ON c.oid = a.attrelid
JOIN pg_namespace n ON n.oid = c.relnamespace
LEFT JOIN pg_attrdef d ON d.adrelid = a.attrelid AND d.adnum = a.attnum
WHERE n.nspname = ? AND c.relname = ? AND a.attnum > 0 AND NOT a.attisdropped
ORDER BY a.attnum`

	postgresConstraintsQuery = `SELECT con.conname AS name, pg_get_constraintdef(con.oid) AS definition
FROM pg_constraint con
JOIN pg_class c ON c.oid = con.conrelid
JOIN pg_namespace n ON n.oid = c.relnamespace
WHERE n.nspname = ? AND c.relname = ? AND con.contype IN ('p', 'u', 'c')
ORDER BY CASE con.contype WHEN 'p' THEN 0 ELSE 1 END, con.conname`

	postgresIndexesQuery = `SELECT i.indexdef
FROM pg_indexes i
WHERE i.schemaname = ? AND i.tablename = ?
AND i.indexname NOT IN (
	SELECT con.conname FROM pg_constraint con
	JOIN pg_class c ON c.oid = con.conrelid
	JOIN pg_namespace n ON n.oid = c.relnamespace
	WHERE n.nspname = ? AND c.relname = ?
)
ORDER BY i.indexname`

	postgresIdentityColumn = "id"
	postgresTimeLayout     = "2006-01-02 15:04:05.999999999Z07:00"
	postgresByteaType      = "BYTEA"
)

var postgresNumericTypes = map[string]bool{
	"INT2": true, "INT4": true, "INT8": true, "FLOAT4": true, "FLOAT8": true, "NUMERIC": true, "OID": true,
}

var postgresSerialTypes = map[string]string{
	"smallint": "smallserial",
	"integer":  "serial",
	"bigint":   "bigserial",
}

var postgresDumpDialect = dumpDialect{
	preamble: []string{
		"SET client_encoding = 'UTF8';",
		"SET standard_conforming_strings = on;",
		"SET TimeZone = 'UTC';",
	},
	quote:        quotePostgresIdentifier,
	literal:      postgresLiteral,
	tableTrailer: postgresSequenceReset,
}

func NewPostgresBackup(db *gorm.DB, directory string, now func() time.Time) *DumpBackup {
	backup := NewDumpBackup(func() (DumpSource, error) { return openPostgresDumpSource(db) }, directory, now)
	backup.dialect = postgresDumpDialect
	return backup
}

type postgresDumpSource struct {
	*txDumpSource
}

var _ DumpSource = &postgresDumpSource{}

func openPostgresDumpSource(db *gorm.DB) (*postgresDumpSource, error) {
	source, err := openTxDumpSource(db, quotePostgresIdentifier, postgresDumpBytes)
	if err != nil {
		return nil, err
	}
	return &postgresDumpSource{source}, nil
}

func (s *postgresDumpSource) Tables() ([]string, error) {
	var tables []string
	err := s.tx.Raw(postgresTablesQuery, postgresSchema).Scan(&tables).Error
	return tables, err
}

type postgresColumn struct {
	Name         string
	DataType     string
	NotNull      bool
	DefaultValue string
	Identity     string
}

type postgresConstraint struct {
	Name       string
	Definition string
}

func (s *postgresDumpSource) CreateStatement(table string) (string, error) {
	var columns []postgresColumn
	if err := s.tx.Raw(postgresColumnsQuery, postgresSchema, table).Scan(&columns).Error; err != nil {
		return "", err
	}
	var constraints []postgresConstraint
	if err := s.tx.Raw(postgresConstraintsQuery, postgresSchema, table).Scan(&constraints).Error; err != nil {
		return "", err
	}
	var indexes []string
	if err := s.tx.Raw(postgresIndexesQuery, postgresSchema, table, postgresSchema, table).Scan(&indexes).Error; err != nil {
		return "", err
	}
	var definitions []string
	for _, column := range columns {
		definitions = append(definitions, "  "+postgresColumnDefinition(column))
	}
	for _, constraint := range constraints {
		definitions = append(definitions, fmt.Sprintf("  CONSTRAINT %s %s", quotePostgresIdentifier(constraint.Name), constraint.Definition))
	}
	statements := []string{fmt.Sprintf("CREATE TABLE %s (\n%s\n)", quotePostgresIdentifier(table), strings.Join(definitions, ",\n"))}
	statements = append(statements, indexes...)
	return strings.Join(statements, ";\n"), nil
}

func postgresColumnDefinition(column postgresColumn) string {
	quoted := quotePostgresIdentifier(column.Name)
	serial, isSerial := postgresSerialTypes[column.DataType]
	if isSerial && strings.HasPrefix(column.DefaultValue, "nextval(") {
		return quoted + " " + serial + " NOT NULL"
	}
	if column.Identity != "" {
		return quoted + " " + column.DataType + " GENERATED BY DEFAULT AS IDENTITY NOT NULL"
	}
	definition := quoted + " " + column.DataType
	if column.DefaultValue != "" {
		definition += " DEFAULT " + column.DefaultValue
	}
	if column.NotNull {
		definition += " NOT NULL"
	}
	return definition
}

func postgresDumpBytes(value any) []byte {
	if moment, isTime := value.(time.Time); isTime {
		return []byte(moment.UTC().Format(postgresTimeLayout))
	}
	return dumpBytes(value)
}

func quotePostgresIdentifier(name string) string {
	return `"` + strings.ReplaceAll(name, `"`, `""`) + `"`
}

func postgresLiteral(databaseType string, value []byte) string {
	if value == nil {
		return "NULL"
	}
	upper := strings.ToUpper(databaseType)
	if postgresNumericTypes[upper] && plainNumber.Match(value) {
		return string(value)
	}
	if upper == postgresByteaType {
		return fmt.Sprintf(`'\x%x'`, value)
	}
	return "'" + strings.ReplaceAll(string(value), "'", "''") + "'"
}

func postgresSequenceReset(table string, columns []DumpColumn) []string {
	for _, column := range columns {
		if column.Name != postgresIdentityColumn {
			continue
		}
		quotedTable, quotedColumn := quotePostgresIdentifier(table), quotePostgresIdentifier(column.Name)
		sequenceTable := strings.ReplaceAll(quotedTable, "'", "''")
		return []string{fmt.Sprintf(
			"SELECT setval(pg_get_serial_sequence('%s', '%s'), COALESCE((SELECT MAX(%s) FROM %s), 1), (SELECT COUNT(*) > 0 FROM %s));",
			sequenceTable, column.Name, quotedColumn, quotedTable, quotedTable,
		)}
	}
	return nil
}