enhance utilisation of postgres features

2026-04-20 10:19:27 +10:00
parent 98e92a8264
commit 8ccf5a7009
28 changed files with 2836 additions and 422 deletions
@@ -10,6 +10,7 @@ import (
 	"strconv"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"time"

 	"vctp/db/queries"
@@ -49,6 +50,21 @@ type loggerContextKey struct{}

 var ensureOnceRegistry sync.Map

+var vmHourlyStatsPostgresPartitioningEnabled atomic.Bool
+
+func init() {
+	vmHourlyStatsPostgresPartitioningEnabled.Store(false)
+}
+
+// SetVmHourlyStatsPostgresPartitioningEnabled toggles postgres monthly partitioning for vm_hourly_stats.
+func SetVmHourlyStatsPostgresPartitioningEnabled(enabled bool) {
+	vmHourlyStatsPostgresPartitioningEnabled.Store(enabled)
+}
+
+func vmHourlyStatsPartitioningEnabled() bool {
+	return vmHourlyStatsPostgresPartitioningEnabled.Load()
+}
+
 // WithLoggerContext stores a logger in context for downstream DB helper logging.
 func WithLoggerContext(ctx context.Context, logger *slog.Logger) context.Context {
 	if ctx == nil {
@@ -700,7 +716,18 @@ func CheckpointDatabase(ctx context.Context, dbConn *sqlx.DB) (string, error) {

 // EnsureVmHourlyStats creates the shared per-snapshot cache table used by Go aggregations.
 func EnsureVmHourlyStats(ctx context.Context, dbConn *sqlx.DB) error {
-	ddl := `
+	driver := strings.ToLower(dbConn.DriverName())
+	if (driver == "pgx" || driver == "postgres") && vmHourlyStatsPartitioningEnabled() {
+		return ensureOncePerDB(dbConn, "vm_hourly_stats_partitioned", func() error {
+			if err := ensureVmHourlyStatsPartitionedPostgres(ctx, dbConn); err != nil {
+				return err
+			}
+			return ensureVmHourlyStatsIndexes(ctx, dbConn, driver)
+		})
+	}
+
+	return ensureOncePerDB(dbConn, "vm_hourly_stats_unpartitioned", func() error {
+		ddl := `
 CREATE TABLE IF NOT EXISTS vm_hourly_stats (
  "SnapshotTime" BIGINT NOT NULL,
  "Vcenter" TEXT NOT NULL,
@@ -721,27 +748,241 @@ CREATE TABLE IF NOT EXISTS vm_hourly_stats (
  "SrmPlaceholder" TEXT,
  PRIMARY KEY ("Vcenter","VmId","SnapshotTime")
 );`
-	return ensureOncePerDB(dbConn, "vm_hourly_stats", func() error {
 		if _, err := execLog(ctx, dbConn, ddl); err != nil {
 			return err
 		}
-		indexQueries := []string{
-			`CREATE INDEX IF NOT EXISTS vm_hourly_stats_vmuuid_time_idx ON vm_hourly_stats ("VmUuid","SnapshotTime")`,
-			`CREATE INDEX IF NOT EXISTS vm_hourly_stats_vmid_time_idx ON vm_hourly_stats ("VmId","SnapshotTime")`,
-			`CREATE INDEX IF NOT EXISTS vm_hourly_stats_name_time_idx ON vm_hourly_stats (lower("Name"),"SnapshotTime")`,
-			`CREATE INDEX IF NOT EXISTS vm_hourly_stats_snapshottime_idx ON vm_hourly_stats ("SnapshotTime")`,
+		return ensureVmHourlyStatsIndexes(ctx, dbConn, driver)
+	})
+}
+
+// EnsureVmHourlyStatsPartitionForSnapshot creates the month partition for snapshotUnix when postgres partitioning is enabled.
+func EnsureVmHourlyStatsPartitionForSnapshot(ctx context.Context, dbConn *sqlx.DB, snapshotUnix int64) error {
+	driver := strings.ToLower(dbConn.DriverName())
+	if driver != "pgx" && driver != "postgres" {
+		return nil
+	}
+	if !vmHourlyStatsPartitioningEnabled() || snapshotUnix <= 0 {
+		return nil
+	}
+	snapshotMonth := time.Unix(snapshotUnix, 0).UTC().Format("200601")
+	key := "vm_hourly_stats_partition_month_" + snapshotMonth
+	return ensureOncePerDB(dbConn, key, func() error {
+		partitioned, err := isVmHourlyStatsPartitioned(ctx, dbConn)
+		if err != nil {
+			return err
 		}
-		failedIndexes := 0
-		for _, q := range indexQueries {
-			if _, err := execLog(ctx, dbConn, q); err != nil {
-				failedIndexes++
+		if !partitioned {
+			return nil
+		}
+		return ensureVmHourlyStatsMonthPartition(ctx, dbConn, time.Unix(snapshotUnix, 0).UTC())
+	})
+}
+
+func ensureVmHourlyStatsIndexes(ctx context.Context, dbConn *sqlx.DB, driver string) error {
+	indexQueries := []string{
+		`CREATE INDEX IF NOT EXISTS vm_hourly_stats_snapshottime_idx ON vm_hourly_stats ("SnapshotTime")`,
+		`CREATE INDEX IF NOT EXISTS vm_hourly_stats_vmuuid_time_idx ON vm_hourly_stats ("VmUuid","SnapshotTime")`,
+		`CREATE INDEX IF NOT EXISTS vm_hourly_stats_vmid_time_idx ON vm_hourly_stats ("VmId","SnapshotTime")`,
+		`CREATE INDEX IF NOT EXISTS vm_hourly_stats_name_time_idx ON vm_hourly_stats (lower("Name"),"SnapshotTime")`,
+	}
+
+	if driver == "pgx" || driver == "postgres" {
+		indexQueries = append(indexQueries,
+			`CREATE INDEX IF NOT EXISTS vm_hourly_stats_vcenter_time_idx ON vm_hourly_stats ("Vcenter","SnapshotTime")`,
+			`CREATE INDEX IF NOT EXISTS vm_hourly_stats_vcenter_vmid_time_idx ON vm_hourly_stats ("Vcenter","VmId","SnapshotTime")`,
+			`CREATE INDEX IF NOT EXISTS vm_hourly_stats_vcenter_vmuuid_time_idx ON vm_hourly_stats ("Vcenter","VmUuid","SnapshotTime")`,
+			`CREATE INDEX IF NOT EXISTS vm_hourly_stats_vcenter_name_time_idx ON vm_hourly_stats ("Vcenter",lower("Name"),"SnapshotTime")`,
+		)
+	}
+
+	failedIndexes := 0
+	for _, q := range indexQueries {
+		if _, err := execLog(ctx, dbConn, q); err != nil {
+			failedIndexes++
+		}
+	}
+	if failedIndexes > 0 {
+		slog.Warn("vm_hourly_stats index ensure incomplete; continuing without retries until restart", "failed_indexes", failedIndexes)
+	}
+	return nil
+}
+
+func ensureVmHourlyStatsPartitionedPostgres(ctx context.Context, dbConn *sqlx.DB) error {
+	partitioned, err := isVmHourlyStatsPartitioned(ctx, dbConn)
+	if err != nil {
+		return err
+	}
+	if !partitioned {
+		exists := TableExists(ctx, dbConn, "vm_hourly_stats")
+		if exists {
+			if err := migrateVmHourlyStatsToPartitioned(ctx, dbConn); err != nil {
+				return err
+			}
+		} else {
+			if _, err := execLog(ctx, dbConn, vmHourlyStatsPartitionedDDL()); err != nil {
+				return err
 			}
 		}
-		if failedIndexes > 0 {
-			slog.Warn("vm_hourly_stats index ensure incomplete; continuing without retries until restart", "failed_indexes", failedIndexes)
-		}
+	}
+
+	if err := ensureVmHourlyStatsDefaultPartition(ctx, dbConn); err != nil {
+		return err
+	}
+	if err := ensureVmHourlyStatsPartitionsForExistingData(ctx, dbConn); err != nil {
+		return err
+	}
+	nowUTC := time.Now().UTC()
+	if err := ensureVmHourlyStatsMonthPartition(ctx, dbConn, nowUTC); err != nil {
+		return err
+	}
+	if err := ensureVmHourlyStatsMonthPartition(ctx, dbConn, nowUTC.AddDate(0, 1, 0)); err != nil {
+		return err
+	}
+	return nil
+}
+
+func vmHourlyStatsPartitionedDDL() string {
+	return `
+CREATE TABLE IF NOT EXISTS vm_hourly_stats (
+  "SnapshotTime" BIGINT NOT NULL,
+  "Vcenter" TEXT NOT NULL,
+  "VmId" TEXT,
+  "VmUuid" TEXT,
+  "Name" TEXT,
+  "CreationTime" BIGINT,
+  "DeletionTime" BIGINT,
+  "ResourcePool" TEXT,
+  "Datacenter" TEXT,
+  "Cluster" TEXT,
+  "Folder" TEXT,
+  "ProvisionedDisk" REAL,
+  "VcpuCount" BIGINT,
+  "RamGB" BIGINT,
+  "IsTemplate" TEXT,
+  "PoweredOn" TEXT,
+  "SrmPlaceholder" TEXT,
+  CONSTRAINT vm_hourly_stats_partitioned_pk PRIMARY KEY ("Vcenter","VmId","SnapshotTime")
+) PARTITION BY RANGE ("SnapshotTime");
+`
+}
+
+func isVmHourlyStatsPartitioned(ctx context.Context, dbConn *sqlx.DB) (bool, error) {
+	var count int
+	err := getLog(ctx, dbConn, &count, `
+SELECT COUNT(1)
+FROM pg_partitioned_table pt
+JOIN pg_class c ON c.oid = pt.partrelid
+JOIN pg_namespace n ON n.oid = c.relnamespace
+WHERE n.nspname = 'public'
+  AND c.relname = 'vm_hourly_stats'
+`)
+	if err != nil {
+		return false, err
+	}
+	return count > 0, nil
+}
+
+func migrateVmHourlyStatsToPartitioned(ctx context.Context, dbConn *sqlx.DB) error {
+	tx, err := dbConn.BeginTxx(ctx, nil)
+	if err != nil {
+		return err
+	}
+	defer tx.Rollback()
+
+	if _, err := tx.ExecContext(ctx, `LOCK TABLE vm_hourly_stats IN ACCESS EXCLUSIVE MODE`); err != nil {
+		return err
+	}
+
+	var partitionedCount int
+	if err := tx.GetContext(ctx, &partitionedCount, `
+SELECT COUNT(1)
+FROM pg_partitioned_table pt
+JOIN pg_class c ON c.oid = pt.partrelid
+JOIN pg_namespace n ON n.oid = c.relnamespace
+WHERE n.nspname = 'public'
+  AND c.relname = 'vm_hourly_stats'
+`); err != nil {
+		return err
+	}
+	if partitionedCount > 0 {
+		return tx.Commit()
+	}
+
+	backupTable := fmt.Sprintf("vm_hourly_stats_unpartitioned_%d", time.Now().UTC().UnixNano())
+	if _, err := SafeTableName(backupTable); err != nil {
+		return err
+	}
+	if _, err := tx.ExecContext(ctx, fmt.Sprintf(`ALTER TABLE vm_hourly_stats RENAME TO %s`, backupTable)); err != nil {
+		return err
+	}
+	if _, err := tx.ExecContext(ctx, vmHourlyStatsPartitionedDDL()); err != nil {
+		return err
+	}
+	if _, err := tx.ExecContext(ctx, `CREATE TABLE IF NOT EXISTS vm_hourly_stats_default PARTITION OF vm_hourly_stats DEFAULT`); err != nil {
+		return err
+	}
+	if _, err := tx.ExecContext(ctx, fmt.Sprintf(`INSERT INTO vm_hourly_stats SELECT * FROM %s`, backupTable)); err != nil {
+		return err
+	}
+	if _, err := tx.ExecContext(ctx, fmt.Sprintf(`DROP TABLE %s`, backupTable)); err != nil {
+		return err
+	}
+	return tx.Commit()
+}
+
+func ensureVmHourlyStatsDefaultPartition(ctx context.Context, dbConn *sqlx.DB) error {
+	_, err := execLog(ctx, dbConn, `CREATE TABLE IF NOT EXISTS vm_hourly_stats_default PARTITION OF vm_hourly_stats DEFAULT`)
+	return err
+}
+
+func ensureVmHourlyStatsPartitionsForExistingData(ctx context.Context, dbConn *sqlx.DB) error {
+	var bounds struct {
+		Min sql.NullInt64 `db:"min_snapshot"`
+		Max sql.NullInt64 `db:"max_snapshot"`
+	}
+	if err := getLog(ctx, dbConn, &bounds, `
+SELECT MIN("SnapshotTime") AS min_snapshot, MAX("SnapshotTime") AS max_snapshot
+FROM vm_hourly_stats
+`); err != nil {
+		return err
+	}
+	if !bounds.Min.Valid || !bounds.Max.Valid {
 		return nil
-	})
+	}
+	start := monthStartUTC(time.Unix(bounds.Min.Int64, 0).UTC())
+	end := monthStartUTC(time.Unix(bounds.Max.Int64, 0).UTC())
+	guard := 0
+	for m := start; !m.After(end); m = m.AddDate(0, 1, 0) {
+		if err := ensureVmHourlyStatsMonthPartition(ctx, dbConn, m); err != nil {
+			return err
+		}
+		guard++
+		if guard > 240 {
+			return fmt.Errorf("vm_hourly_stats partition range guard exceeded while creating existing-data partitions")
+		}
+	}
+	return nil
+}
+
+func ensureVmHourlyStatsMonthPartition(ctx context.Context, dbConn *sqlx.DB, month time.Time) error {
+	start := monthStartUTC(month.UTC())
+	end := start.AddDate(0, 1, 0)
+	partitionName := fmt.Sprintf("vm_hourly_stats_%s", start.Format("200601"))
+	if _, err := SafeTableName(partitionName); err != nil {
+		return err
+	}
+	query := fmt.Sprintf(
+		`CREATE TABLE IF NOT EXISTS %s PARTITION OF vm_hourly_stats FOR VALUES FROM (%d) TO (%d)`,
+		partitionName,
+		start.Unix(),
+		end.Unix(),
+	)
+	_, err := execLog(ctx, dbConn, query)
+	return err
+}
+
+func monthStartUTC(t time.Time) time.Time {
+	return time.Date(t.Year(), t.Month(), 1, 0, 0, 0, 0, time.UTC)
 }

 // EnsureVmLifecycleCache creates an upsert cache for first/last seen VM info.