Something went wrong. Try again.
This repository has no description
Something went wrong. Try again.
7.1 kB · 198 lines
Go
at next
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199package statedb
import ( "context" "encoding/json" "fmt" "reflect" "strings" "sync" "time"
"gorm.io/driver/postgres" "gorm.io/driver/sqlite" "gorm.io/gorm" "gorm.io/gorm/clause" "gorm.io/gorm/schema" "stream.place/streamplace/pkg/config" "stream.place/streamplace/pkg/log")
// Copying a state database from one engine to another (the sqlite → Postgres// move when a node grows up). Rows travel through the Go models, so every// difference between the two schemas — bytea vs blob, jsonb vs text,// timestamptz vs datetime text, bool vs int — is GORM's type mapping doing// what it already does on every read and write; nothing is hand-mapped.//// The copy is idempotent (insert on conflict do nothing, keyed by each// table's primary key), so it can be run once while the old node is still up// to prove the target out, and again after stopping it to catch the delta.
// CopyReport is one table's outcome.type CopyReport struct { Table string Source int64 // rows in the source (soft-deleted included) Inserted int64 // rows this run wrote (inserted, or reconciled on a later pass) Target int64 // rows in the target afterwards}
// dialectorFor parses a state database URL the way MakeDB does.func dialectorFor(dbURL string) (gorm.Dialector, DBType, error) { switch { case dbURL == ":memory:": return sqlite.Open(":memory:"), DBTypeSQLite, nil case strings.HasPrefix(dbURL, "sqlite://"): return sqlite.Open(dbURL[len("sqlite://"):]), DBTypeSQLite, nil case strings.HasPrefix(dbURL, "postgres://") || strings.HasPrefix(dbURL, "postgresql://"): return postgres.Open(dbURL), DBTypePostgres, nil } return nil, "", fmt.Errorf("unsupported database URL (most start with sqlite:// or postgresql://): %s", redactDBURL(dbURL))}
// openSource opens a state database to read from: no AutoMigrate, nothing// written.func openSource(dbURL string) (*gorm.DB, DBType, error) { dial, dbType, err := dialectorFor(dbURL) if err != nil { return nil, "", err } db, err := openDB(dial) if err != nil { return nil, "", fmt.Errorf("error opening source database: %w", err) } if dbType == DBTypeSQLite { if err := sqlitePragmas(db); err != nil { return nil, "", err } } return db, dbType, nil}
// CopyState copies every table in StatefulDBModels from fromURL into toURL.// The target is opened exactly as a node would open it (created if missing,// AutoMigrated to the current schema). batch is rows per INSERT.func CopyState(ctx context.Context, fromURL, toURL string, batch int) ([]CopyReport, error) { if batch <= 0 { batch = 500 } src, srcType, err := openSource(fromURL) if err != nil { return nil, err } log.Log(ctx, "copy: source open", "url", redactDBURL(fromURL), "type", srcType) dst, err := MakeDB(ctx, &config.CLI{DBURL: toURL}, nil, nil) if err != nil { return nil, fmt.Errorf("error opening target database: %w", err) } log.Log(ctx, "copy: target open", "url", redactDBURL(toURL), "type", dst.Type)
var reports []CopyReport var problems []string for _, m := range StatefulDBModels { r, err := copyTable(ctx, src, dst, m, batch) if err != nil { return reports, err } reports = append(reports, r) if r.Target < r.Source { problems = append(problems, fmt.Sprintf("%s: source has %d rows, target has %d", r.Table, r.Source, r.Target)) } } if len(problems) > 0 { return reports, fmt.Errorf("copy finished but some tables are short:\n %s", strings.Join(problems, "\n ")) } return reports, nil}
func copyTable(ctx context.Context, src *gorm.DB, dst *StatefulDB, m any, batch int) (CopyReport, error) { s, err := schema.Parse(m, &sync.Map{}, dst.DB.NamingStrategy) if err != nil { return CopyReport{}, err } r := CopyReport{Table: s.Table} if err := src.Unscoped().Model(m).Count(&r.Source).Error; err != nil { return r, fmt.Errorf("%s: counting source rows: %w", s.Table, err) } started := time.Now() // A pointer to a []Model, built by reflection since the model list is // []any; FindInBatches walks the source in primary-key order. rows := reflect.New(reflect.SliceOf(reflect.TypeOf(m))) // Hooks off: rows are copied as they are, not re-created. Unscoped on // the target too, so a soft-deleted row is inserted with its deleted_at. writer := dst.DB.Session(&gorm.Session{SkipHooks: true, Context: ctx}).Unscoped() res := src.WithContext(ctx).Unscoped().Model(m).FindInBatches(rows.Interface(), batch, func(tx *gorm.DB, n int) error { normalizeJSON(rows.Elem()) // Upsert, not insert-or-skip: the documented second (delta) pass // after the cutover has to carry rows that changed since the first // pass (tasks, configs, branding, soft deletions), which keep their // primary keys and would otherwise be left at their old values. ins := writer.Clauses(clause.OnConflict{UpdateAll: true}).Create(rows.Interface()) if ins.Error != nil { return ins.Error } r.Inserted += ins.RowsAffected return nil }) if res.Error != nil { return r, fmt.Errorf("%s: %w", s.Table, res.Error) } if dst.Type == DBTypePostgres { if err := bumpSequence(ctx, dst.DB, s); err != nil { return r, err } } if err := dst.DB.Unscoped().Model(m).Count(&r.Target).Error; err != nil { return r, fmt.Errorf("%s: counting target rows: %w", s.Table, err) } log.Log(ctx, "copy: table done", "table", s.Table, "source", r.Source, "inserted", r.Inserted, "target", r.Target, "took", time.Since(started).Round(time.Millisecond)) return r, nil}
// bumpSequence moves an autoincrement table's sequence past the ids that were// copied in with explicit values, so the node's first insert doesn't collide.// Tables whose key isn't a serial have no sequence (pg_get_serial_sequence is// NULL) and setval of NULL is a no-op.func bumpSequence(ctx context.Context, db *gorm.DB, s *schema.Schema) error { pk := s.PrioritizedPrimaryField if pk == nil || !pk.AutoIncrement || (pk.DataType != schema.Int && pk.DataType != schema.Uint) { return nil } q := fmt.Sprintf( `SELECT setval(pg_get_serial_sequence('%s', '%s'), COALESCE((SELECT MAX(%s) FROM %s), 0) + 1, false)`, s.Table, pk.DBName, pk.DBName, s.Table, ) if err := db.WithContext(ctx).Exec(q).Error; err != nil { return fmt.Errorf("%s: resetting sequence: %w", s.Table, err) } return nil}
var rawMessageType = reflect.TypeOf(json.RawMessage{})
// normalizeJSON turns an empty or invalid json.RawMessage into NULL in every// row of a batch: sqlite stored whatever it was handed, but a jsonb column// rejects "" on the way in.func normalizeJSON(rows reflect.Value) { if rows.Len() == 0 { return } var fields []int t := rows.Index(0).Type() for i := 0; i < t.NumField(); i++ { if t.Field(i).Type == rawMessageType { fields = append(fields, i) } } if len(fields) == 0 { return } for i := 0; i < rows.Len(); i++ { row := rows.Index(i) for _, f := range fields { v := row.Field(f) if b := v.Bytes(); len(b) > 0 && !json.Valid(b) || len(b) == 0 && !v.IsNil() { v.Set(reflect.Zero(rawMessageType)) } } }}