fallback to full table partitions for Timescale compressed hypertables (#3318)

Author: heavycrystal

flow/activities/snapshot_activity.go

@@ -23,8 +23,7 @@ type SlotSnapshotState struct {
 }
 
 type TxSnapshotState struct {
-	SnapshotName     string
-	SupportsTIDScans bool
+	SnapshotName string
 }
 
 type SnapshotActivity struct {
@@ -93,9 +92,8 @@ func (a *SnapshotActivity) SetupReplication(
 	a.Alerter.LogFlowInfo(ctx, config.FlowJobName, "Replication slot and publication setup complete")
 
 	return &protos.SetupReplicationOutput{
-		SlotName:         slotInfo.SlotName,
-		SnapshotName:     slotInfo.SnapshotName,
-		SupportsTidScans: slotInfo.SupportsTIDScans,
+		SlotName:     slotInfo.SlotName,
+		SnapshotName: slotInfo.SnapshotName,
 	}, nil
 }
 
@@ -118,8 +116,7 @@ func (a *SnapshotActivity) MaintainTx(ctx context.Context, sessionID string, pee
 	a.SnapshotStatesMutex.Lock()
 	if exportSnapshotOutput != nil {
 		a.TxSnapshotStates[sessionID] = TxSnapshotState{
-			SnapshotName:     exportSnapshotOutput.SnapshotName,
-			SupportsTIDScans: exportSnapshotOutput.SupportsTidScans,
+			SnapshotName: exportSnapshotOutput.SnapshotName,
 		}
 	} else {
 		a.TxSnapshotStates[sessionID] = TxSnapshotState{}
@@ -169,3 +166,22 @@ func (a *SnapshotActivity) LoadTableSchema(
 ) (*protos.TableSchema, error) {
 	return internal.LoadTableSchemaFromCatalog(ctx, a.CatalogPool, flowName, tableName)
 }
+
+func (a *SnapshotActivity) GetDefaultPartitionKeyForTables(
+	ctx context.Context,
+	input *protos.FlowConnectionConfigs,
+) (*protos.GetDefaultPartitionKeyForTablesOutput, error) {
+	connector, err := connectors.GetByNameAs[connectors.QRepPullConnectorCore](ctx, nil, a.CatalogPool, input.SourceName)
+	if err != nil {
+		return nil, a.Alerter.LogFlowError(ctx, input.FlowJobName, fmt.Errorf("failed to get connector: %w", err))
+	}
+
+	output, err := connector.GetDefaultPartitionKeyForTables(ctx, &protos.GetDefaultPartitionKeyForTablesInput{
+		TableMappings: input.TableMappings,
+	})
+	if err != nil {
+		return nil, a.Alerter.LogFlowError(ctx, input.FlowJobName, fmt.Errorf("failed to check if tables can parallel load: %w", err))
+	}
+
+	return output, nil
+}

flow/connectors/core.go

@@ -225,6 +225,9 @@ type QRepPullConnectorCore interface {
 
 	// GetQRepPartitions returns the partitions for a given table that haven't been synced yet.
 	GetQRepPartitions(ctx context.Context, config *protos.QRepConfig, last *protos.QRepPartition) ([]*protos.QRepPartition, error)
+
+	GetDefaultPartitionKeyForTables(ctx context.Context,
+		input *protos.GetDefaultPartitionKeyForTablesInput) (*protos.GetDefaultPartitionKeyForTablesOutput, error)
 }
 
 type QRepPullConnector interface {

flow/connectors/mongo/qrep.go

@@ -135,6 +135,15 @@ func (c *MongoConnector) GetQRepPartitions(
 	return partitions, nil
 }
 
+func (c *MongoConnector) GetDefaultPartitionKeyForTables(
+	ctx context.Context,
+	input *protos.GetDefaultPartitionKeyForTablesInput,
+) (*protos.GetDefaultPartitionKeyForTablesOutput, error) {
+	return &protos.GetDefaultPartitionKeyForTablesOutput{
+		TableDefaultPartitionKeyMapping: nil,
+	}, nil
+}
+
 func (c *MongoConnector) PullQRepRecords(
 	ctx context.Context,
 	otelManager *otel_metrics.OtelManager,

flow/connectors/mysql/qrep.go

@@ -22,29 +22,6 @@ import (
 	"github.com/PeerDB-io/peerdb/flow/shared/types"
 )
 
-func (c *MySqlConnector) GetDataTypeOfWatermarkColumn(
-	ctx context.Context,
-	watermarkTableName string,
-	watermarkColumn string,
-) (types.QValueKind, byte, error) {
-	if watermarkColumn == "" {
-		return "", 0, errors.New("watermark column is not specified in the config")
-	}
-
-	query := fmt.Sprintf("SELECT `%s` FROM %s LIMIT 0", watermarkColumn, watermarkTableName)
-	rs, err := c.Execute(ctx, query)
-	if err != nil {
-		return "", 0, fmt.Errorf("failed to execute query for watermark column type: %w", err)
-	}
-
-	if len(rs.Fields) == 0 {
-		return "", 0, fmt.Errorf("no fields returned from select query: %s", query)
-	}
-
-	qk, err := qkindFromMysql(rs.Fields[0])
-	return qk, rs.Fields[0].Type, err
-}
-
 func (c *MySqlConnector) GetQRepPartitions(
 	ctx context.Context,
 	config *protos.QRepConfig,
@@ -153,6 +130,15 @@ func (c *MySqlConnector) GetQRepPartitions(
 	return partitionHelper.GetPartitions(), nil
 }
 
+func (c *MySqlConnector) GetDefaultPartitionKeyForTables(
+	ctx context.Context,
+	input *protos.GetDefaultPartitionKeyForTablesInput,
+) (*protos.GetDefaultPartitionKeyForTablesOutput, error) {
+	return &protos.GetDefaultPartitionKeyForTablesOutput{
+		TableDefaultPartitionKeyMapping: nil,
+	}, nil
+}
+
 func (c *MySqlConnector) PullQRepRecords(
 	ctx context.Context,
 	otelManager *otel_metrics.OtelManager,

flow/connectors/postgres/postgres.go

@@ -1200,18 +1200,12 @@ func (c *PostgresConnector) EnsurePullability(
 }
 
 func (c *PostgresConnector) ExportTxSnapshot(ctx context.Context, env map[string]string) (*protos.ExportTxSnapshotOutput, any, error) {
-	pgversion, err := c.MajorVersion(ctx)
-	if err != nil {
-		return nil, nil, fmt.Errorf("[export-snapshot] error getting PG version: %w", err)
-	}
-
 	skipSnapshotExport, err := internal.PeerDBSkipSnapshotExport(ctx, env)
 	if err != nil {
 		c.logger.Error("failed to check PEERDB_SKIP_SNAPSHOT_EXPORT, proceeding with export snapshot", slog.Any("error", err))
 	} else if skipSnapshotExport {
 		return &protos.ExportTxSnapshotOutput{
-			SnapshotName:     "",
-			SupportsTidScans: pgversion >= shared.POSTGRES_13,
+			SnapshotName: "",
 		}, nil, err
 	}
 
@@ -1242,8 +1236,7 @@ func (c *PostgresConnector) ExportTxSnapshot(ctx context.Context, env map[string
 	needRollback = false
 
 	return &protos.ExportTxSnapshotOutput{
-		SnapshotName:     snapshotName,
-		SupportsTidScans: pgversion >= shared.POSTGRES_13,
+		SnapshotName: snapshotName,
 	}, tx, err
 }
 

flow/connectors/postgres/qrep.go

@@ -69,6 +69,67 @@ func (c *PostgresConnector) GetQRepPartitions(
 	return c.getNumRowsPartitions(ctx, getPartitionsTx, config, last)
 }
 
+func (c *PostgresConnector) GetDefaultPartitionKeyForTables(
+	ctx context.Context,
+	input *protos.GetDefaultPartitionKeyForTablesInput,
+) (*protos.GetDefaultPartitionKeyForTablesOutput, error) {
+	c.logger.Info("Evaluating if tables can perform parallel load")
+
+	output := &protos.GetDefaultPartitionKeyForTablesOutput{
+		TableDefaultPartitionKeyMapping: make(map[string]string, len(input.TableMappings)),
+	}
+
+	pgVersion, err := shared.GetMajorVersion(ctx, c.conn)
+	if err != nil {
+		return nil, fmt.Errorf("failed to determine server version: %w", err)
+	}
+	supportsTidScans := pgVersion >= shared.POSTGRES_14
+
+	if supportsTidScans {
+		for _, tm := range input.TableMappings {
+			output.TableDefaultPartitionKeyMapping[tm.SourceTableIdentifier] = "ctid"
+		}
+	}
+
+	if !supportsTidScans {
+		// older versions fall back to full table partitions anyway; nothing more to do
+		c.logger.Warn("Postgres version does not support TID scans, falling back to full table partitions")
+		return output, nil
+	}
+
+	var hasTimescale bool
+	if err := c.conn.QueryRow(ctx,
+		"SELECT EXISTS (SELECT 1 FROM pg_extension WHERE extname = 'timescaledb')").Scan(&hasTimescale); err != nil {
+		return nil, fmt.Errorf("failed to check for timescaledb extension: %w", err)
+	}
+	if !hasTimescale {
+		return output, nil
+	}
+
+	// compressed hypertables cannot do ctid scans, so disable for them
+	// NOTE: it appears that the hypercore "TAM" may give us access to ctid scans, but that's to be removed in Timescale 2.22
+	rows, err := c.conn.Query(ctx, `SELECT DISTINCT hypertable_schema, hypertable_name
+		FROM timescaledb_information.chunks
+		WHERE is_compressed='t';`)
+	if err != nil {
+		return nil, fmt.Errorf("query compressed hypertables: %w", err)
+	}
+	var schema, name string
+	if _, err := pgx.ForEachRow(rows, []any{&schema, &name}, func() error {
+		if _, ok := output.TableDefaultPartitionKeyMapping[fmt.Sprintf("%s.%s", schema, name)]; ok {
+			table := fmt.Sprintf("%s.%s", schema, name)
+			delete(output.TableDefaultPartitionKeyMapping, table)
+			c.logger.Warn("table is a compressed hypertable, falling back to full table partition",
+				slog.String("table", table))
+		}
+		return nil
+	}); err != nil {
+		return nil, fmt.Errorf("failed to get compressed hypertables: %w", err)
+	}
+
+	return output, nil
+}
+
 func (c *PostgresConnector) setTransactionSnapshot(ctx context.Context, tx pgx.Tx, snapshot string) error {
 	if snapshot != "" {
 		if _, err := tx.Exec(ctx, "SET TRANSACTION SNAPSHOT "+utils.QuoteLiteral(snapshot)); err != nil {

flow/workflows/snapshot_flow.go

@@ -244,7 +244,6 @@ func (s *SnapshotFlowExecution) cloneTables(
 	snapshotType snapshotType,
 	slotName string,
 	snapshotName string,
-	supportsTIDScans bool,
 	maxParallelClones int,
 ) error {
 	if snapshotType == SNAPSHOT_TYPE_SLOT {
@@ -253,14 +252,21 @@ func (s *SnapshotFlowExecution) cloneTables(
 		s.logger.Info("cloning tables in tx snapshot mode", slog.String("snapshot", snapshotName))
 	}
 
-	boundSelector := shared.NewBoundSelector(ctx, "CloneTablesSelector", maxParallelClones)
+	getParallelLoadKeyForTablesCtx := workflow.WithActivityOptions(ctx, workflow.ActivityOptions{
+		StartToCloseTimeout: 10 * time.Minute,
+		RetryPolicy: &temporal.RetryPolicy{
+			InitialInterval: 1 * time.Minute,
+		},
+	})
 
-	defaultPartitionCol := "ctid"
-	if !supportsTIDScans {
-		s.logger.Info("Postgres version too old for TID scans, might use full table partitions!")
-		defaultPartitionCol = ""
+	var res *protos.GetDefaultPartitionKeyForTablesOutput
+	if err := workflow.ExecuteActivity(getParallelLoadKeyForTablesCtx,
+		snapshot.GetDefaultPartitionKeyForTables, s.config).Get(ctx, &res); err != nil {
+		return fmt.Errorf("failed to get default partition keys for tables: %w", err)
 	}
 
+	boundSelector := shared.NewBoundSelector(ctx, "CloneTablesSelector", maxParallelClones)
+
 	for _, v := range s.config.TableMappings {
 		source := v.SourceTableIdentifier
 		destination := v.DestinationTableIdentifier
@@ -269,7 +275,7 @@ func (s *SnapshotFlowExecution) cloneTables(
 			slog.String("snapshotName", snapshotName),
 		)
 		if v.PartitionKey == "" {
-			v.PartitionKey = defaultPartitionCol
+			v.PartitionKey = res.TableDefaultPartitionKeyMapping[source]
 		}
 		if err := s.cloneTable(ctx, boundSelector, snapshotName, v); err != nil {
 			s.logger.Error("failed to start clone child workflow", slog.Any("error", err))
@@ -304,19 +310,16 @@ func (s *SnapshotFlowExecution) cloneTablesWithSlot(
 	}()
 	var slotName string
 	var snapshotName string
-	var supportsTidScans bool
 	if slotInfo != nil {
 		slotName = slotInfo.SlotName
 		snapshotName = slotInfo.SnapshotName
-		supportsTidScans = slotInfo.SupportsTidScans
 	}
 
 	s.logger.Info("cloning tables in parallel", slog.Int("parallelism", numTablesInParallel))
 	if err := s.cloneTables(ctx,
 		SNAPSHOT_TYPE_SLOT,
 		slotName,
 		snapshotName,
-		supportsTidScans,
 		numTablesInParallel,
 	); err != nil {
 		s.logger.Error("failed to clone tables", slog.Any("error", err))
@@ -411,7 +414,6 @@ func SnapshotFlowWorkflow(
 			SNAPSHOT_TYPE_TX,
 			"",
 			txnSnapshotState.SnapshotName,
-			txnSnapshotState.SupportsTIDScans,
 			numTablesInParallel,
 		); err != nil {
 			return fmt.Errorf("failed to clone tables: %w", err)

protos/flow.proto

@@ -156,9 +156,9 @@ message SetupReplicationInput {
 }
 
 message SetupReplicationOutput {
+  reserved 3;
   string slot_name = 1;
   string snapshot_name = 2;
-  bool supports_tid_scans = 3;
 }
 
 message CreateRawTableInput {
@@ -460,8 +460,8 @@ message IsQRepPartitionSyncedInput {
 }
 
 message ExportTxSnapshotOutput {
+  reserved 2;
   string snapshot_name = 1;
-  bool supports_tid_scans = 2;
 }
 
 enum DynconfValueType {
@@ -560,3 +560,11 @@ message FlowContextMetadata{
 message AdditionalContextMetadata{
   FlowOperation operation = 1;
 }
+
+message GetDefaultPartitionKeyForTablesInput {
+  repeated TableMapping table_mappings = 1;
+}
+
+message GetDefaultPartitionKeyForTablesOutput {
+  map<string, string> table_default_partition_key_mapping = 1;
+}