diff --git a/backend/internal/services/archive/partitions.go b/backend/internal/services/archive/partitions.go new file mode 100644 index 00000000..90a44157 --- /dev/null +++ b/backend/internal/services/archive/partitions.go @@ -0,0 +1,432 @@ +package archive + +import ( + "context" + "encoding/json" + "fmt" + "io" + "os" + "reflect" + "strings" + "time" + + "gorm.io/gorm" + "gorm.io/gorm/schema" + + "github.com/kurodakayn/mpp-backend/internal/models" + "github.com/kurodakayn/mpp-backend/internal/pkg/objectstorage" +) + +type monthlyPartitionArchiveSpec struct { + Table string + BeforeDrop func(context.Context, *gorm.DB, coldMonthlyPartition) error + CanDropPartition func(context.Context, *gorm.DB, coldMonthlyPartition) (bool, error) +} + +type coldMonthlyPartition struct { + ParentSchema string + ParentName string + Schema string + Name string + Start time.Time + End time.Time +} + +type archivePartitionLine struct { + SchemaVersion int `json:"schema_version"` + Table string `json:"table"` + Partition string `json:"partition"` + ArchivedAt time.Time `json:"archived_at"` + RetentionCutoff time.Time `json:"retention_cutoff"` + PartitionStart time.Time `json:"partition_start"` + PartitionEnd time.Time `json:"partition_end"` + Row json.RawMessage `json:"row"` +} + +var publishEventsPartitionSpec = monthlyPartitionArchiveSpec{ + Table: "publish_events", +} + +var extensionExecutionEventsPartitionSpec = monthlyPartitionArchiveSpec{ + Table: "extension_execution_events", + BeforeDrop: deleteExtensionExecutionEventClaimsForPartition, +} + +var projectActivitiesPartitionSpec = monthlyPartitionArchiveSpec{ + Table: "project_activities", +} + +var workspaceActivitiesPartitionSpec = monthlyPartitionArchiveSpec{ + Table: "workspace_activities", +} + +var remoteBrowserSessionsPartitionSpec = monthlyPartitionArchiveSpec{ + Table: "remote_browser_sessions", + CanDropPartition: remoteBrowserSessionPartitionIsTerminal, +} + +func archivePartitionedModel[T any]( + ctx context.Context, + db *gorm.DB, + storage objectstorage.Client, + config Config, + spec monthlyPartitionArchiveSpec, + retention time.Duration, + now time.Time, + scope func(*gorm.DB, time.Time) *gorm.DB, + beforeDelete archiveBeforeDeleteHook[T], +) (TableResult, error) { + cutoff := now.Add(-retention) + result := TableResult{Table: spec.Table, Cutoff: cutoff} + + if db.Name() == "postgres" { + partitionResult, err := archiveColdMonthlyPartitions[T](ctx, db, storage, config, spec, retention, now) + if err != nil { + return result, err + } + mergeTableResult(&result, partitionResult) + } + + rowResult, err := archiveModelWithDeleteHook[T](ctx, db, storage, config, spec.Table, retention, now, scope, beforeDelete) + if err != nil { + return result, err + } + mergeTableResult(&result, rowResult) + return result, nil +} + +func archiveColdMonthlyPartitions[T any]( + ctx context.Context, + db *gorm.DB, + storage objectstorage.Client, + config Config, + spec monthlyPartitionArchiveSpec, + retention time.Duration, + now time.Time, +) (TableResult, error) { + cutoff := now.Add(-retention) + result := TableResult{Table: spec.Table, Cutoff: cutoff} + partitions, err := listColdMonthlyPartitions(ctx, db, spec.Table, cutoff) + if err != nil { + return result, err + } + + for _, partition := range partitions { + if spec.CanDropPartition != nil { + canDrop, err := spec.CanDropPartition(ctx, db, partition) + if err != nil { + return result, err + } + if !canDrop { + continue + } + } + + partitionResult, err := archiveColdMonthlyPartition[T](ctx, db, storage, config, spec, cutoff, now, partition) + if err != nil { + return result, err + } + mergeTableResult(&result, partitionResult) + } + return result, nil +} + +func listColdMonthlyPartitions(ctx context.Context, db *gorm.DB, table string, cutoff time.Time) ([]coldMonthlyPartition, error) { + rows, err := db.WithContext(ctx).Raw(` + SELECT + parent_namespace.nspname AS parent_schema, + parent.relname AS parent_name, + child_namespace.nspname AS partition_schema, + child.relname AS partition_name + FROM pg_inherits inheritance + JOIN pg_class parent ON parent.oid = inheritance.inhparent + JOIN pg_namespace parent_namespace ON parent_namespace.oid = parent.relnamespace + JOIN pg_class child ON child.oid = inheritance.inhrelid + JOIN pg_namespace child_namespace ON child_namespace.oid = child.relnamespace + WHERE parent.oid = to_regclass(?) + ORDER BY child.relname ASC + `, table).Rows() + if err != nil { + return nil, fmt.Errorf("list %s monthly partitions: %w", table, err) + } + defer func() { + _ = rows.Close() + }() + + var partitions []coldMonthlyPartition + for rows.Next() { + var partition coldMonthlyPartition + if err := rows.Scan(&partition.ParentSchema, &partition.ParentName, &partition.Schema, &partition.Name); err != nil { + return nil, fmt.Errorf("scan %s monthly partition: %w", table, err) + } + start, end, ok := monthlyPartitionBounds(partition.ParentName, partition.Name) + if !ok || end.After(cutoff) { + continue + } + partition.Start = start + partition.End = end + partitions = append(partitions, partition) + } + if err := rows.Err(); err != nil { + return nil, fmt.Errorf("iterate %s monthly partitions: %w", table, err) + } + return partitions, nil +} + +func archiveColdMonthlyPartition[T any]( + ctx context.Context, + db *gorm.DB, + storage objectstorage.Client, + config Config, + spec monthlyPartitionArchiveSpec, + cutoff time.Time, + now time.Time, + partition coldMonthlyPartition, +) (TableResult, error) { + result := TableResult{Table: spec.Table, Cutoff: cutoff} + tempFile, err := os.CreateTemp("", fmt.Sprintf("%s-%s-*.jsonl", spec.Table, partition.Name)) + if err != nil { + return result, fmt.Errorf("create %s partition %s temp archive file: %w", spec.Table, partition.Name, err) + } + tempName := tempFile.Name() + defer func() { + _ = tempFile.Close() + _ = os.Remove(tempName) + }() + + rowsArchived, err := encodePartitionJSONLines[T](ctx, db, spec, cutoff, now, partition, tempFile) + if err != nil { + return result, err + } + key := "" + if rowsArchived > 0 { + if err := tempFile.Sync(); err != nil { + return result, fmt.Errorf("sync %s partition %s temp archive file: %w", spec.Table, partition.Name, err) + } + if _, err := tempFile.Seek(0, io.SeekStart); err != nil { + return result, fmt.Errorf("rewind %s partition %s temp archive file: %w", spec.Table, partition.Name, err) + } + + key = archivePartitionObjectKey(config.ObjectKeyPrefix, spec.Table, partition.Name, partition.Start, partition.End, now) + if _, err := storage.PutObject(ctx, objectstorage.UploadObjectInput{ + Key: key, + ContentType: archiveContentType, + Body: tempFile, + }); err != nil { + return result, fmt.Errorf("upload %s partition %s archive object: %w", spec.Table, partition.Name, err) + } + } + + if err := dropArchivedMonthlyPartition(ctx, db, spec, partition); err != nil { + return result, err + } + + result.RowsArchived = rowsArchived + if key != "" { + result.ObjectKey = key + result.ObjectKeys = []string{key} + } + result.PartitionsArchived = []string{partition.Name} + return result, nil +} + +func dropArchivedMonthlyPartition(ctx context.Context, db *gorm.DB, spec monthlyPartitionArchiveSpec, partition coldMonthlyPartition) error { + return db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { + if spec.BeforeDrop != nil { + if err := spec.BeforeDrop(ctx, tx, partition); err != nil { + return err + } + } + if err := tx.Exec(fmt.Sprintf( + "ALTER TABLE %s DETACH PARTITION %s", + partition.qualifiedParentName(), + partition.qualifiedName(), + )).Error; err != nil { + return fmt.Errorf("detach %s partition %s: %w", spec.Table, partition.Name, err) + } + if err := tx.Exec(fmt.Sprintf( + "DROP TABLE %s", + partition.qualifiedName(), + )).Error; err != nil { + return fmt.Errorf("drop %s partition %s: %w", spec.Table, partition.Name, err) + } + return nil + }) +} + +func encodePartitionJSONLines[T any]( + ctx context.Context, + db *gorm.DB, + spec monthlyPartitionArchiveSpec, + cutoff time.Time, + archivedAt time.Time, + partition coldMonthlyPartition, + writer io.Writer, +) (int, error) { + stmt := &gorm.Statement{DB: db} + if err := stmt.Parse(new(T)); err != nil { + return 0, fmt.Errorf("parse %s partition %s row schema: %w", spec.Table, partition.Name, err) + } + columns := make([]string, 0, len(stmt.Schema.Fields)) + for _, field := range stmt.Schema.Fields { + if field.DBName == "" { + continue + } + columns = append(columns, field.DBName) + } + query := fmt.Sprintf( + "SELECT %s FROM %s ORDER BY %s ASC, %s ASC", + archiveQuotedColumnList(columns), + partition.qualifiedName(), + quoteArchivePostgresIdentifier("created_at"), + quoteArchivePostgresIdentifier("id"), + ) + rows, err := db.WithContext(ctx).Raw(query).Rows() + if err != nil { + return 0, fmt.Errorf("query %s partition %s rows: %w", spec.Table, partition.Name, err) + } + defer func() { + _ = rows.Close() + }() + + encoder := json.NewEncoder(writer) + rowCount := 0 + for rows.Next() { + var record T + if err := db.WithContext(ctx).ScanRows(rows, &record); err != nil { + return 0, fmt.Errorf("scan %s partition %s row: %w", spec.Table, partition.Name, err) + } + rowJSON, err := archivePartitionRowJSON(record, stmt.Schema.Fields) + if err != nil { + return 0, fmt.Errorf("marshal %s partition %s row: %w", spec.Table, partition.Name, err) + } + if err := encoder.Encode(archivePartitionLine{ + SchemaVersion: 1, + Table: spec.Table, + Partition: partition.Name, + ArchivedAt: archivedAt, + RetentionCutoff: cutoff, + PartitionStart: partition.Start, + PartitionEnd: partition.End, + Row: json.RawMessage(rowJSON), + }); err != nil { + return 0, fmt.Errorf("encode %s partition %s rows: %w", spec.Table, partition.Name, err) + } + rowCount++ + } + if err := rows.Err(); err != nil { + return 0, fmt.Errorf("iterate %s partition %s rows: %w", spec.Table, partition.Name, err) + } + return rowCount, nil +} + +func archivePartitionRowJSON(record any, fields []*schema.Field) ([]byte, error) { + value := reflect.ValueOf(record) + if value.Kind() == reflect.Pointer { + if value.IsNil() { + return nil, fmt.Errorf("archive record is nil") + } + value = value.Elem() + } + if value.Kind() != reflect.Struct { + return nil, fmt.Errorf("archive record must be a struct, got %s", value.Kind()) + } + + payload := make(map[string]any, len(fields)) + for _, field := range fields { + if field.DBName == "" { + continue + } + fieldValue := value.FieldByName(field.Name) + if !fieldValue.IsValid() { + return nil, fmt.Errorf("archive record is missing field %q", field.Name) + } + payload[field.Name] = fieldValue.Interface() + } + return json.Marshal(payload) +} + +func archiveQuotedColumnList(columns []string) string { + quoted := make([]string, 0, len(columns)) + for _, column := range columns { + quoted = append(quoted, quoteArchivePostgresIdentifier(column)) + } + return strings.Join(quoted, ", ") +} + +func deleteExtensionExecutionEventClaimsForPartition(ctx context.Context, tx *gorm.DB, partition coldMonthlyPartition) error { + if !tx.Migrator().HasTable(&models.ExtensionExecutionEventClaim{}) { + return nil + } + return tx.WithContext(ctx).Exec(fmt.Sprintf(` + DELETE FROM extension_execution_event_claims + WHERE record_id IN (SELECT id FROM %s) + `, partition.qualifiedName())).Error +} + +func remoteBrowserSessionPartitionIsTerminal(ctx context.Context, db *gorm.DB, partition coldMonthlyPartition) (bool, error) { + var hasNonTerminal bool + err := db.WithContext(ctx).Raw(fmt.Sprintf(` + SELECT EXISTS ( + SELECT 1 + FROM %s + WHERE status NOT IN (?, ?, ?) + ) + `, partition.qualifiedName()), + models.BrowserSessionStatusConnected, + models.BrowserSessionStatusExpired, + models.BrowserSessionStatusFailed, + ).Scan(&hasNonTerminal).Error + if err != nil { + return false, fmt.Errorf("check remote browser session partition %s terminal status: %w", partition.Name, err) + } + return !hasNonTerminal, nil +} + +func monthlyPartitionBounds(parentTable string, partitionName string) (time.Time, time.Time, bool) { + suffix := strings.TrimPrefix(partitionName, parentTable+"_") + if suffix == partitionName || suffix == "default" { + return time.Time{}, time.Time{}, false + } + start, err := time.Parse("2006_01", suffix) + if err != nil { + return time.Time{}, time.Time{}, false + } + start = time.Date(start.Year(), start.Month(), 1, 0, 0, 0, 0, time.UTC) + return start, start.AddDate(0, 1, 0), true +} + +func archivePartitionObjectKey(prefix string, table string, partitionName string, partitionStart time.Time, partitionEnd time.Time, archivedAt time.Time) string { + normalizedPrefix := strings.Trim(strings.TrimSpace(prefix), "/") + if normalizedPrefix == "" { + normalizedPrefix = defaultArchiveObjectPrefix + } + return fmt.Sprintf( + "%s/%s/partitions/partition_start=%s/partition_end=%s/%s-%d.jsonl", + normalizedPrefix, + table, + partitionStart.UTC().Format("2006-01-02"), + partitionEnd.UTC().Format("2006-01-02"), + partitionName, + archivedAt.UTC().UnixNano(), + ) +} + +func quoteArchivePostgresIdentifier(identifier string) string { + return `"` + strings.ReplaceAll(identifier, `"`, `""`) + `"` +} + +func quoteArchivePostgresQualifiedIdentifier(schema string, name string) string { + if strings.TrimSpace(schema) == "" { + return quoteArchivePostgresIdentifier(name) + } + return quoteArchivePostgresIdentifier(schema) + "." + quoteArchivePostgresIdentifier(name) +} + +func (p coldMonthlyPartition) qualifiedName() string { + return quoteArchivePostgresQualifiedIdentifier(p.Schema, p.Name) +} + +func (p coldMonthlyPartition) qualifiedParentName() string { + return quoteArchivePostgresQualifiedIdentifier(p.ParentSchema, p.ParentName) +} diff --git a/backend/internal/services/archive/partitions_test.go b/backend/internal/services/archive/partitions_test.go new file mode 100644 index 00000000..d1206cf6 --- /dev/null +++ b/backend/internal/services/archive/partitions_test.go @@ -0,0 +1,307 @@ +package archive + +import ( + "bytes" + "context" + "database/sql" + "database/sql/driver" + "fmt" + "io" + "strings" + "sync" + "testing" + "time" + + "github.com/google/uuid" + "gorm.io/datatypes" + "gorm.io/driver/postgres" + "gorm.io/gorm" + + "github.com/kurodakayn/mpp-backend/internal/models" + "github.com/kurodakayn/mpp-backend/internal/pkg/objectstorage" +) + +func TestMonthlyPartitionBoundsParsesManagedPartitionName(t *testing.T) { + start, end, ok := monthlyPartitionBounds("publish_events", "publish_events_2026_01") + + if !ok { + t.Fatalf("expected managed monthly partition to parse") + } + if !start.Equal(time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC)) { + t.Fatalf("unexpected partition start %s", start) + } + if !end.Equal(time.Date(2026, 2, 1, 0, 0, 0, 0, time.UTC)) { + t.Fatalf("unexpected partition end %s", end) + } +} + +func TestMonthlyPartitionBoundsRejectsDefaultOrUnknownNames(t *testing.T) { + for _, partitionName := range []string{ + "publish_events_default", + "publish_events_legacy_20260101000000", + "extension_execution_events_2026_01", + } { + if _, _, ok := monthlyPartitionBounds("publish_events", partitionName); ok { + t.Fatalf("expected %q to be ignored", partitionName) + } + } +} + +func TestArchivePartitionObjectKeyUsesPartitionRange(t *testing.T) { + key := archivePartitionObjectKey( + " /cold/database/ ", + "publish_events", + "publish_events_2026_01", + time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC), + time.Date(2026, 2, 1, 0, 0, 0, 0, time.UTC), + time.Unix(123, 456).UTC(), + ) + + if !strings.HasPrefix(key, "cold/database/publish_events/partitions/partition_start=2026-01-01/partition_end=2026-02-01/") { + t.Fatalf("unexpected archive key prefix %q", key) + } + if !strings.HasSuffix(key, "/publish_events_2026_01-123000000456.jsonl") { + t.Fatalf("unexpected archive key suffix %q", key) + } +} + +func TestArchiveQuotedColumnListQuotesIdentifiers(t *testing.T) { + columns := archiveQuotedColumnList([]string{"id", `strange"column`}) + + if columns != `"id", "strange""column"` { + t.Fatalf("unexpected column list %q", columns) + } +} + +func TestEncodePartitionJSONLinesPreservesStructFieldNames(t *testing.T) { + db := setupArchiveTestDB(t) + if err := db.Exec(`CREATE TABLE publish_events_2026_01 ( + id TEXT PRIMARY KEY, + publication_id TEXT NOT NULL, + project_id TEXT NOT NULL, + user_id TEXT NOT NULL, + platform TEXT NOT NULL, + job_id TEXT NOT NULL, + idempotency_key TEXT NOT NULL, + event_type TEXT NOT NULL, + status TEXT NOT NULL, + message TEXT, + remote_id TEXT, + publish_url TEXT, + error_message TEXT, + metadata TEXT NOT NULL DEFAULT '{}', + created_at DATETIME + )`).Error; err != nil { + t.Fatalf("create partition table: %v", err) + } + + archivedAt := time.Date(2026, 6, 11, 10, 0, 0, 0, time.UTC) + partitionStart := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC) + event := models.PublishEvent{ + ID: uuid.New(), + PublicationID: uuid.New(), + ProjectID: uuid.New(), + UserID: uuid.New(), + Platform: "wechat", + JobID: uuid.New(), + IdempotencyKey: "cold-a", + EventType: "publish.completed", + Status: models.PublicationStatusSucceeded, + Metadata: datatypes.JSON(`{"source":"test"}`), + CreatedAt: partitionStart.Add(time.Hour), + } + if err := db.Table("publish_events_2026_01").Create(&event).Error; err != nil { + t.Fatalf("insert partition row: %v", err) + } + + var body bytes.Buffer + rowsArchived, err := encodePartitionJSONLines[models.PublishEvent]( + context.Background(), + db, + monthlyPartitionArchiveSpec{Table: "publish_events"}, + archivedAt.Add(-180*24*time.Hour), + archivedAt, + coldMonthlyPartition{ + Name: "publish_events_2026_01", + Start: partitionStart, + End: partitionStart.AddDate(0, 1, 0), + }, + &body, + ) + if err != nil { + t.Fatalf("encode partition JSONL: %v", err) + } + if rowsArchived != 1 { + t.Fatalf("expected one archived row, got %d", rowsArchived) + } + + lines := jsonLines(t, body.String()) + if len(lines) != 1 { + t.Fatalf("expected one JSONL line, got %d", len(lines)) + } + row := lines[0]["row"].(map[string]any) + if _, ok := row["ID"]; !ok { + t.Fatalf("expected exported row to use Go field names, got %#v", row) + } + if _, ok := row["id"]; ok { + t.Fatalf("expected exported row to avoid snake_case database keys, got %#v", row) + } +} + +func TestArchiveColdMonthlyPartitionDropsEmptyPartitionWithoutUpload(t *testing.T) { + state := &archivePartitionTestState{} + db := openArchivePartitionTestDB(t, state) + storage := &rejectPutPartitionStorage{} + now := time.Date(2026, 6, 11, 10, 0, 0, 0, time.UTC) + partitionStart := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC) + partition := coldMonthlyPartition{ + ParentName: "publish_events", + Name: "publish_events_2026_01", + Start: partitionStart, + End: partitionStart.AddDate(0, 1, 0), + } + + result, err := archiveColdMonthlyPartition[models.PublishEvent]( + context.Background(), + db, + storage, + Config{ObjectKeyPrefix: "cold"}, + monthlyPartitionArchiveSpec{Table: "publish_events"}, + now.Add(-180*24*time.Hour), + now, + partition, + ) + if err != nil { + t.Fatalf("archive empty partition: %v", err) + } + + if storage.putObjects != 0 { + t.Fatalf("expected empty partition to skip upload, got %d uploads", storage.putObjects) + } + if result.RowsArchived != 0 { + t.Fatalf("expected zero archived rows, got %d", result.RowsArchived) + } + if result.ObjectKey != "" || len(result.ObjectKeys) != 0 { + t.Fatalf("expected no archive object for empty partition, got %q %#v", result.ObjectKey, result.ObjectKeys) + } + if len(result.PartitionsArchived) != 1 || result.PartitionsArchived[0] != partition.Name { + t.Fatalf("expected partition to be marked archived, got %#v", result.PartitionsArchived) + } + + expectedExecs := []string{ + `ALTER TABLE "publish_events" DETACH PARTITION "publish_events_2026_01"`, + `DROP TABLE "publish_events_2026_01"`, + } + if len(state.execs) != len(expectedExecs) { + t.Fatalf("expected detach and drop statements, got %#v", state.execs) + } + for i, expected := range expectedExecs { + if state.execs[i] != expected { + t.Fatalf("unexpected exec %d: got %q want %q", i, state.execs[i], expected) + } + } +} + +type rejectPutPartitionStorage struct { + objectstorage.Client + putObjects int +} + +func (s *rejectPutPartitionStorage) PutObject(context.Context, objectstorage.UploadObjectInput) (objectstorage.ObjectInfo, error) { + s.putObjects++ + return objectstorage.ObjectInfo{}, fmt.Errorf("empty partition should not upload an archive object") +} + +var registerArchivePartitionDriverOnce sync.Once +var archivePartitionDriverState *archivePartitionTestState + +type archivePartitionTestState struct { + queries []string + execs []string +} + +type archivePartitionDriver struct{} + +func (archivePartitionDriver) Open(_ string) (driver.Conn, error) { + return &archivePartitionConn{state: archivePartitionDriverState}, nil +} + +type archivePartitionConn struct { + state *archivePartitionTestState +} + +func (c *archivePartitionConn) Prepare(_ string) (driver.Stmt, error) { + return nil, driver.ErrSkip +} + +func (c *archivePartitionConn) Close() error { + return nil +} + +func (c *archivePartitionConn) Begin() (driver.Tx, error) { + return archivePartitionTx{}, nil +} + +func (c *archivePartitionConn) QueryContext(_ context.Context, query string, _ []driver.NamedValue) (driver.Rows, error) { + c.state.queries = append(c.state.queries, query) + if strings.Contains(query, `FROM "publish_events_2026_01"`) { + return archivePartitionRows{}, nil + } + return nil, fmt.Errorf("unexpected archive partition query: %s", query) +} + +func (c *archivePartitionConn) ExecContext(_ context.Context, query string, _ []driver.NamedValue) (driver.Result, error) { + c.state.execs = append(c.state.execs, query) + if query == `ALTER TABLE "publish_events" DETACH PARTITION "publish_events_2026_01"` || + query == `DROP TABLE "publish_events_2026_01"` { + return driver.RowsAffected(0), nil + } + return nil, fmt.Errorf("unexpected archive partition exec: %s", query) +} + +type archivePartitionTx struct{} + +func (archivePartitionTx) Commit() error { + return nil +} + +func (archivePartitionTx) Rollback() error { + return nil +} + +type archivePartitionRows struct{} + +func (archivePartitionRows) Columns() []string { + return []string{"id", "created_at"} +} + +func (archivePartitionRows) Close() error { + return nil +} + +func (archivePartitionRows) Next(_ []driver.Value) error { + return io.EOF +} + +func openArchivePartitionTestDB(t *testing.T, state *archivePartitionTestState) *gorm.DB { + t.Helper() + registerArchivePartitionDriverOnce.Do(func() { + sql.Register("archive-partition-test", archivePartitionDriver{}) + }) + archivePartitionDriverState = state + sqlDB, err := sql.Open("archive-partition-test", "") + if err != nil { + t.Fatalf("open archive partition sql db: %v", err) + } + t.Cleanup(func() { + _ = sqlDB.Close() + }) + db, err := gorm.Open(postgres.New(postgres.Config{ + Conn: sqlDB, + PreferSimpleProtocol: true, + }), &gorm.Config{}) + if err != nil { + t.Fatalf("open archive partition gorm db: %v", err) + } + return db +} diff --git a/backend/internal/services/archive/worker.go b/backend/internal/services/archive/worker.go index ee7615d4..f3f60cdf 100644 --- a/backend/internal/services/archive/worker.go +++ b/backend/internal/services/archive/worker.go @@ -30,11 +30,12 @@ type RunResult struct { } type TableResult struct { - Table string - RowsArchived int - ObjectKey string - ObjectKeys []string - Cutoff time.Time + Table string + RowsArchived int + ObjectKey string + ObjectKeys []string + Cutoff time.Time + PartitionsArchived []string } type archiveLine[T any] struct { @@ -128,27 +129,27 @@ func (w *Worker) runArchiveTables(ctx context.Context, db *gorm.DB, now time.Tim } func (w *Worker) archivePublishEvents(ctx context.Context, db *gorm.DB, now time.Time) (TableResult, error) { - return archiveModel[models.PublishEvent](ctx, db, w.storage, w.config, "publish_events", w.config.PublishEventRetention, now, func(query *gorm.DB, cutoff time.Time) *gorm.DB { + return archivePartitionedModel[models.PublishEvent](ctx, db, w.storage, w.config, publishEventsPartitionSpec, w.config.PublishEventRetention, now, func(query *gorm.DB, cutoff time.Time) *gorm.DB { return query.Where("created_at < ?", cutoff) - }) + }, nil) } func (w *Worker) archiveExtensionExecutionEvents(ctx context.Context, db *gorm.DB, now time.Time) (TableResult, error) { - return archiveModelWithDeleteHook(ctx, db, w.storage, w.config, "extension_execution_events", w.config.ExtensionExecutionEventRetention, now, func(query *gorm.DB, cutoff time.Time) *gorm.DB { + return archivePartitionedModel[models.ExtensionExecutionEvent](ctx, db, w.storage, w.config, extensionExecutionEventsPartitionSpec, w.config.ExtensionExecutionEventRetention, now, func(query *gorm.DB, cutoff time.Time) *gorm.DB { return query.Where("created_at < ?", cutoff) }, deleteExtensionExecutionEventClaims) } func (w *Worker) archiveProjectActivities(ctx context.Context, db *gorm.DB, now time.Time) (TableResult, error) { - return archiveModel[models.ProjectActivity](ctx, db, w.storage, w.config, "project_activities", w.config.ProjectActivityRetention, now, func(query *gorm.DB, cutoff time.Time) *gorm.DB { + return archivePartitionedModel[models.ProjectActivity](ctx, db, w.storage, w.config, projectActivitiesPartitionSpec, w.config.ProjectActivityRetention, now, func(query *gorm.DB, cutoff time.Time) *gorm.DB { return query.Where("created_at < ?", cutoff) - }) + }, nil) } func (w *Worker) archiveWorkspaceActivities(ctx context.Context, db *gorm.DB, now time.Time) (TableResult, error) { - return archiveModel[models.WorkspaceActivity](ctx, db, w.storage, w.config, "workspace_activities", w.config.WorkspaceActivityRetention, now, func(query *gorm.DB, cutoff time.Time) *gorm.DB { + return archivePartitionedModel[models.WorkspaceActivity](ctx, db, w.storage, w.config, workspaceActivitiesPartitionSpec, w.config.WorkspaceActivityRetention, now, func(query *gorm.DB, cutoff time.Time) *gorm.DB { return query.Where("created_at < ?", cutoff) - }) + }, nil) } func (w *Worker) archiveRemoteBrowserSessions(ctx context.Context, db *gorm.DB, now time.Time) (TableResult, error) { @@ -157,22 +158,9 @@ func (w *Worker) archiveRemoteBrowserSessions(ctx context.Context, db *gorm.DB, models.BrowserSessionStatusExpired, models.BrowserSessionStatusFailed, } - return archiveModel[models.RemoteBrowserSession](ctx, db, w.storage, w.config, "remote_browser_sessions", w.config.BrowserSessionHistoryRetention, now, func(query *gorm.DB, cutoff time.Time) *gorm.DB { + return archivePartitionedModel[models.RemoteBrowserSession](ctx, db, w.storage, w.config, remoteBrowserSessionsPartitionSpec, w.config.BrowserSessionHistoryRetention, now, func(query *gorm.DB, cutoff time.Time) *gorm.DB { return query.Where("created_at < ? AND status IN ?", cutoff, terminalStatuses) - }) -} - -func archiveModel[T any]( - ctx context.Context, - db *gorm.DB, - storage objectstorage.Client, - config Config, - table string, - retention time.Duration, - now time.Time, - scope func(*gorm.DB, time.Time) *gorm.DB, -) (TableResult, error) { - return archiveModelWithDeleteHook[T](ctx, db, storage, config, table, retention, now, scope, nil) + }, nil) } func archiveModelWithDeleteHook[T any]( @@ -267,6 +255,21 @@ func archiveModelBatch[T any]( return result, nil } +func mergeTableResult(result *TableResult, next TableResult) { + if result.Table == "" { + result.Table = next.Table + } + if result.Cutoff.IsZero() && !next.Cutoff.IsZero() { + result.Cutoff = next.Cutoff + } + result.RowsArchived += next.RowsArchived + result.ObjectKeys = append(result.ObjectKeys, next.ObjectKeys...) + if result.ObjectKey == "" { + result.ObjectKey = next.ObjectKey + } + result.PartitionsArchived = append(result.PartitionsArchived, next.PartitionsArchived...) +} + func deleteExtensionExecutionEventClaims(tx *gorm.DB, records []models.ExtensionExecutionEvent) error { if len(records) == 0 || !tx.Migrator().HasTable(&models.ExtensionExecutionEventClaim{}) { return nil diff --git a/doc/plan/database-optimization.md b/doc/plan/database-optimization.md index 5e903703..08b2885b 100644 --- a/doc/plan/database-optimization.md +++ b/doc/plan/database-optimization.md @@ -11,7 +11,7 @@ Status definitions: - `Not started`: no clear implementation has been found yet. - `Deferred`: not recommended for the current business stage; only trigger conditions are retained. -Current overall progress: about `67%`. This number is manually estimated by phase weight and can be adjusted later according to actual completed items. +Current overall progress: about `69%`. This number is manually estimated by phase weight and can be adjusted later according to actual completed items. | Phase | Weight | Current completion | Status | Completed | Not done / next steps | | ----- | ------ | ------------------ | ------ | --------- | --------------------- | @@ -19,7 +19,7 @@ Current overall progress: about `67%`. This number is manually estimated by phas | Phase 1: Single-database connection pool, indexing, pagination, and lifecycle governance | 15% | 100% | Done | backend/publish-worker/collab-service application connection pools, Redis client connection pool, PgBouncer writer pool, composite indexes, keyset list pagination, list queries avoiding the large `source_content` field, event/session history retention periods, and R2/S3 cold-event archive worker | None; Phase 2 is complete, and later work moves into Phase 3/4 read replicas, partitioning, and recovery flows | | Phase 2: Read models and cache first | 15% | 100% | Done | Redis and Asynq dependencies are reusable; admin dashboard stats, admin project list, and dashboard account summary have short-TTL Redis cache; stats/project list/account cache misses are merged with singleflight; project, prepublish, publish, and account write paths invalidate the related dashboard cache; `workspace_dashboard_stats` and `project_list_summaries` read models are in place, and APIs prefer read models when coverage is complete; async refresh triggers after project save, platform sync, publish completion, and member changes; admin rebuild API, Asynq queue, and worker support full read-model rebuild | None | | Phase 3: Read/write splitting | 15% | 100% | Done | Optional `DB_READER_*` connection, application-level DB Router, signed sticky writer, consistency routing for project/stats/workspace/platform_account/publish/prepublish/mediaasset/browser_session/extension, consistency-level inventories for dashboard/publish/collab-service, self-hosted PostgreSQL read replica, managed `postgres-reader` entry point, PgBouncer reader pool, replica lag monitoring and automatic fallback to writer when over threshold | None; Phase 4 continues partitioning, archiving, and recovery flows | -| Phase 4: Single-database partitioning, archiving, and hot/cold tiering | 15% | 70% | In progress | Collaborative editing already has state + update batch + compaction foundation; `collab_document_update_batches` has PostgreSQL `document_id` hash partition target schema; event and terminal-session history already have row-level R2/S3 archive worker; `publish_events`, `extension_execution_events`, `project_activities`, `workspace_activities`, and `remote_browser_sessions` have PostgreSQL monthly partition target schema | Cold partition export and recovery flow | +| Phase 4: Single-database partitioning, archiving, and hot/cold tiering | 15% | 85% | In progress | Collaborative editing already has state + update batch + compaction foundation; `collab_document_update_batches` has PostgreSQL `document_id` hash partition target schema; event and terminal-session history already have row-level R2/S3 archive worker; `publish_events`, `extension_execution_events`, `project_activities`, `workspace_activities`, and `remote_browser_sessions` have PostgreSQL monthly partition target schema; the archive worker exports whole cold monthly partitions to R2/S3 before detaching and dropping them | Archive recovery flow | | Phase 5: Citus preparation | 20% | 5% | Not started | Workspace model, `projects.workspace_id`, and personal workspace ID already exist | Global `workspace_id`, Citus distribution column/colocation design, unique constraint and foreign-key review | | Phase 6: Citus distributed PostgreSQL operation | 10% | 0% | Deferred | None | Future Citus cluster design, worker/coordinator monitoring and backup, large-tenant isolation strategy | @@ -64,7 +64,7 @@ Atomic commit guidance: | Dashboard read models | Done | Added `workspace_dashboard_stats` and `project_list_summaries` read models, idempotently recomputed from fact tables by a centralized readmodel service; async refresh is triggered after project save, platform sync, publish completion, and member changes; admin stats and admin project list prefer read models when coverage is complete; admin rebuild API enqueues through Asynq, and API/worker processes can start readmodel workers for full rebuild from fact tables | None | `backend/internal/models/models.go`, `backend/internal/services/readmodel/service.go`, `backend/internal/services/readmodel/queue.go`, `backend/internal/services/readmodel/service_test.go`, `backend/internal/services/readmodel/queue_test.go`, `backend/internal/services/stats/overview.go`, `backend/internal/services/project/lifecycle.go`, `backend/internal/handlers/dashboard.go`, `backend/cmd/api/main.go`, `backend/cmd/publish-worker/main.go` | | Redis read cache | Done | Redis is already used for queues, locks, OAuth, browser sessions, and short-term coordination; admin dashboard stats, admin project list, and dashboard account summary use 15s TTL cache and bypass scoped/sticky-writer strong-consistency paths; stats/project list/account cache misses use singleflight to prevent process-local stampede; stats and account caches use versioned payloads and semantic validation, and Redis read-error fallback is also merged into one DB computation per key; project create/edit/platform save, prepublish sync/draft update, publish queue/execute/fail, and platform account write paths invalidate the related dashboard cache; full read-model rebuild reuses the Redis/Asynq queue | None | `backend/internal/services/stats/overview.go`, `backend/internal/services/stats/overview_test.go`, `backend/internal/services/project/list_cache.go`, `backend/internal/services/project/list_cache_test.go`, `backend/internal/services/prepublish/drafts.go`, `backend/internal/services/publish/service.go`, `backend/internal/services/publish/queue.go`, `backend/internal/services/publish/publication_flow_test.go`, `backend/internal/services/publish/queue_test.go`, `backend/internal/services/platform_account/account_cache.go`, `backend/internal/services/platform_account/account_cache_test.go`, `backend/internal/services/browser_session/complete.go`, `backend/internal/services/browser_session/service_test.go`, `backend/internal/services/readmodel/queue.go` | | Read/write splitting | Done | Supports optional `DB_READER_*` read-replica connection, `DefaultRouter`, and signed sticky writer; project/stats/workspace/platform_account/publish/prepublish/mediaasset/browser_session/extension are wired to strong/eventual/writer routing; dashboard, publish, and collab-service consistency-level inventories are complete, with collab-service online path kept writer-only; writer/reader pools are in self-hosted Kubernetes, and managed overlay provides a `postgres-reader` ExternalName entry point; `DB_READER_MAX_REPLICA_LAG` configures the replica lag threshold, eventual/analytics reads automatically fall back to writer when over threshold or lag is unknown, and `mpp_db_replica_lag_seconds` and `mpp_db_replica_healthy` metrics are exposed | None | `backend/internal/db/db.go`, `backend/internal/db/router.go`, `backend/internal/db/replica_lag.go`, `backend/internal/services/publish/service.go`, `backend/internal/services/prepublish/service.go`, `backend/internal/services/mediaasset/service.go`, `backend/internal/services/browser_session/service.go`, `backend/internal/services/extension/service.go`, `backend/internal/app/runtime.go`, `deploy/kubernetes/data-services/self-hosted/postgres.yaml`, `deploy/kubernetes/data-services/self-hosted/pgbouncer.yaml`, `deploy/kubernetes/data-services/managed/services.yaml`, `script/kubernetes/validation/data_services.rb` | -| Event-table partitioning and archiving | In progress | `publish_events`, `extension_execution_events`, `project_activities`, `workspace_activities`, and terminal `remote_browser_sessions` have default retention periods; the `archive` worker can batch-export JSONL to R2/S3 and delete old hot-table rows after successful upload; PostgreSQL schema initialization now creates monthly `created_at` partitions for `publish_events`, `extension_execution_events`, `project_activities`, `workspace_activities`, and `remote_browser_sessions`, with partition-compatible `(id, created_at)` primary keys and rolling partition creation; PostgreSQL browser-session active-row fallback uses a scoped advisory transaction lock because partitioned unique constraints must include the partition key | Cold partition export and recovery flow are not implemented | `backend/internal/db/monthly_partitions.go`, `backend/internal/db/db.go`, `backend/internal/models/models.go`, `backend/internal/db/db_test.go`, `backend/internal/services/browser_session/start.go`, `backend/internal/services/browser_session/cleanup.go`, `backend/internal/services/archive/worker.go`, `backend/internal/services/archive/worker_test.go` | +| Event-table partitioning and archiving | In progress | `publish_events`, `extension_execution_events`, `project_activities`, `workspace_activities`, and terminal `remote_browser_sessions` have default retention periods; the `archive` worker can batch-export JSONL to R2/S3 and delete old hot-table rows after successful upload; PostgreSQL schema initialization now creates monthly `created_at` partitions for `publish_events`, `extension_execution_events`, `project_activities`, `workspace_activities`, and `remote_browser_sessions`, with partition-compatible `(id, created_at)` primary keys and rolling partition creation; the archive worker exports whole cold monthly partitions as JSONL to R2/S3, then detaches and drops the partition after successful upload; PostgreSQL browser-session active-row fallback uses a scoped advisory transaction lock because partitioned unique constraints must include the partition key | Archive recovery flow is not implemented | `backend/internal/db/monthly_partitions.go`, `backend/internal/db/db.go`, `backend/internal/models/models.go`, `backend/internal/db/db_test.go`, `backend/internal/services/browser_session/start.go`, `backend/internal/services/browser_session/cleanup.go`, `backend/internal/services/archive/worker.go`, `backend/internal/services/archive/partitions.go`, `backend/internal/services/archive/worker_test.go`, `backend/internal/services/archive/partitions_test.go` | | Collaboration batch governance | In progress | `collab_document_states`, `collab_document_update_batches`, and compaction/retention foundations exist; PostgreSQL schema initialization creates a 16-way `document_id` hash-partitioned `collab_document_update_batches` target table and migrates existing regular-table rows into it | Cold archiving is not implemented | `backend/internal/db/hash_partitions.go`, `backend/internal/db/db.go`, `backend/internal/models/collab.go`, `backend/internal/db/db_test.go`, `collab-service/src/persistence/document-persistence.ts` | | Outbox/CDC/event stream | In progress | The publishing queue path has a transactional Outbox: `EnqueuePublishProject` writes `outbox_events` in the same transaction and dispatches immediately after commit; publish worker starts an outbox dispatcher and supports retries for failed/stale processing records; Asynq continues to serve as the task-execution queue, and `PublishEvent` continues to serve as publishing audit | Currently covers only `publish.job_requested`; general business-event outbox, Debezium, and Redpanda/Kafka CDC are not implemented | `backend/internal/services/publish/queue.go`, `backend/internal/services/publish/outbox.go`, `backend/internal/models/models.go` | | Citus target state | Not started | Confirmed `workspace_id` as the most suitable distribution-column direction | Citus distributed tables, reference tables, and colocation are not implemented | Phase 5/6 in this document | @@ -120,7 +120,7 @@ Atomic commit guidance: - [x] Partition `publish_events`, `extension_execution_events`, and activity tables by month. Verification entry point: `backend/internal/db/monthly_partitions.go`, `backend/internal/models/models.go`, `backend/internal/db/db_test.go`. - [x] Partition `remote_browser_sessions` by time or expiration time. Verification entry point: `backend/internal/db/monthly_partitions.go`, `backend/internal/models/models.go`, `backend/internal/services/browser_session/start.go`, `backend/internal/db/db_test.go`. - [x] Hash partition `collab_document_update_batches` by `document_id`. Verification entry point: `backend/internal/db/hash_partitions.go`, `backend/internal/db/db.go`, `backend/internal/models/collab.go`, `backend/internal/db/db_test.go`. -- [ ] Export cold partitions to R2/S3. +- [x] Export cold partitions to R2/S3. Verification entry point: `backend/internal/services/archive/partitions.go`, `backend/internal/services/archive/partitions_test.go`, `backend/internal/services/archive/worker.go`. - [ ] Write archive recovery procedure. #### Phase 5: Citus Preparation