move src/dst getPeerType out of cloneTable loop (#3559)

Author: heavycrystal

flow/connectors/core.go

@@ -317,9 +317,9 @@ func LoadPeerType(ctx context.Context, catalogPool shared.CatalogPool, peerName
 	err := row.Scan(&dbtype)
 	if err != nil {
 		if errors.Is(err, pgx.ErrNoRows) {
-			return 0, fmt.Errorf("peer not found: %s", peerName)
+			return 0, exceptions.NewNotFoundError(fmt.Errorf("peer not found: %s", peerName))
 		}
-		return 0, fmt.Errorf("failed to load peer type for %s: %w", peerName, err)
+		return 0, exceptions.NewCatalogError(fmt.Errorf("failed to load peer type for %s: %w", peerName, err))
 	}
 	return dbtype, nil
 }

flow/workflows/snapshot_flow.go

@@ -106,6 +106,8 @@ func (s *SnapshotFlowExecution) cloneTable(
 	boundSelector *shared.BoundSelector,
 	snapshotName string,
 	mapping *protos.TableMapping,
+	sourcePeerType protos.DBType,
+	destinationPeerType protos.DBType,
 ) error {
 	flowName := s.config.FlowJobName
 	cloneLog := slog.Group("clone-log",
@@ -173,9 +175,7 @@ func (s *SnapshotFlowExecution) cloneTable(
 	// usually MySQL supports double quotes with ANSI_QUOTES, but Vitess doesn't
 	// Vitess currently only supports initial load so change here is enough
 	srcTableEscaped := parsedSrcTable.String()
-	if dbtype, err := getPeerType(ctx, s.config.SourceName); err != nil {
-		return err
-	} else if dbtype == protos.DBType_MYSQL {
+	if sourcePeerType == protos.DBType_MYSQL {
 		srcTableEscaped = parsedSrcTable.MySQL()
 	}
 
@@ -208,9 +208,7 @@ func (s *SnapshotFlowExecution) cloneTable(
 
 	// ensure document IDs are synchronized across initial load and CDC
 	// for the same document
-	if dbtype, err := getPeerType(ctx, s.config.DestinationName); err != nil {
-		return err
-	} else if dbtype == protos.DBType_ELASTICSEARCH {
+	if destinationPeerType == protos.DBType_ELASTICSEARCH {
 		if err := initTableSchema(); err != nil {
 			return err
 		}
@@ -277,6 +275,15 @@ func (s *SnapshotFlowExecution) cloneTables(
 
 	boundSelector := shared.NewBoundSelector(ctx, "CloneTablesSelector", maxParallelClones)
 
+	sourcePeerType, err := getPeerType(ctx, s.config.SourceName)
+	if err != nil {
+		return err
+	}
+	destinationPeerType, err := getPeerType(ctx, s.config.DestinationName)
+	if err != nil {
+		return err
+	}
+
 	for _, v := range s.config.TableMappings {
 		source := v.SourceTableIdentifier
 		destination := v.DestinationTableIdentifier
@@ -287,7 +294,7 @@ func (s *SnapshotFlowExecution) cloneTables(
 		if v.PartitionKey == "" {
 			v.PartitionKey = res.TableDefaultPartitionKeyMapping[source]
 		}
-		if err := s.cloneTable(ctx, boundSelector, snapshotName, v); err != nil {
+		if err := s.cloneTable(ctx, boundSelector, snapshotName, v, sourcePeerType, destinationPeerType); err != nil {
 			s.logger.Error("failed to start clone child workflow", slog.Any("error", err))
 			return err
 		}