Remove slot signal

flow/activities/flowable.go

@@ -31,8 +31,7 @@ import (
 	"github.com/PeerDB-io/peerdb/flow/shared"
 )
 
-// CheckConnectionResult is the result of a CheckConnection call.
-type CheckConnectionResult struct {
+type CheckMetadataTablesResult struct {
 	NeedsSetupMetadataTables bool
 }
 
@@ -54,18 +53,36 @@ type StreamCloser interface {
 func (a *FlowableActivity) CheckConnection(
 	ctx context.Context,
 	config *protos.SetupInput,
-) (*CheckConnectionResult, error) {
+) error {
 	ctx = context.WithValue(ctx, shared.FlowNameKey, config.FlowName)
-	dstConn, err := connectors.GetByNameAs[connectors.CDCSyncConnector](ctx, config.Env, a.CatalogPool, config.PeerName)
+	conn, err := connectors.GetByNameAs[connectors.CDCSyncConnector](ctx, config.Env, a.CatalogPool, config.PeerName)
+	if err != nil {
+		if errors.Is(err, errors.ErrUnsupported) {
+			return nil
+		}
+		a.Alerter.LogFlowError(ctx, config.FlowName, err)
+		return fmt.Errorf("failed to get connector: %w", err)
+	}
+	defer connectors.CloseConnector(ctx, conn)
+
+	return conn.ConnectionActive(ctx)
+}
+
+func (a *FlowableActivity) CheckMetadataTables(
+	ctx context.Context,
+	config *protos.SetupInput,
+) (*CheckMetadataTablesResult, error) {
+	ctx = context.WithValue(ctx, shared.FlowNameKey, config.FlowName)
+	conn, err := connectors.GetByNameAs[connectors.CDCSyncConnector](ctx, config.Env, a.CatalogPool, config.PeerName)
 	if err != nil {
 		a.Alerter.LogFlowError(ctx, config.FlowName, err)
 		return nil, fmt.Errorf("failed to get connector: %w", err)
 	}
-	defer connectors.CloseConnector(ctx, dstConn)
+	defer connectors.CloseConnector(ctx, conn)
 
-	needsSetup := dstConn.NeedsSetupMetadataTables(ctx)
+	needsSetup := conn.NeedsSetupMetadataTables(ctx)
 
-	return &CheckConnectionResult{
+	return &CheckMetadataTablesResult{
 		NeedsSetupMetadataTables: needsSetup,
 	}, nil
 }
@@ -939,15 +956,15 @@ func (a *FlowableActivity) AddTablesToPublication(ctx context.Context, cfg *prot
 	}
 	defer connectors.CloseConnector(ctx, srcConn)
 
-	err = srcConn.AddTablesToPublication(ctx, &protos.AddTablesToPublicationInput{
+	if err := srcConn.AddTablesToPublication(ctx, &protos.AddTablesToPublicationInput{
 		FlowJobName:      cfg.FlowJobName,
 		PublicationName:  cfg.PublicationName,
 		AdditionalTables: additionalTableMappings,
-	})
-	if err != nil {
+	}); err != nil {
 		a.Alerter.LogFlowError(ctx, cfg.FlowJobName, err)
+		return err
 	}
-	return err
+	return nil
 }
 
 func (a *FlowableActivity) RemoveTablesFromPublication(
@@ -962,15 +979,15 @@ func (a *FlowableActivity) RemoveTablesFromPublication(
 	}
 	defer connectors.CloseConnector(ctx, srcConn)
 
-	err = srcConn.RemoveTablesFromPublication(ctx, &protos.RemoveTablesFromPublicationInput{
+	if err := srcConn.RemoveTablesFromPublication(ctx, &protos.RemoveTablesFromPublicationInput{
 		FlowJobName:     cfg.FlowJobName,
 		PublicationName: cfg.PublicationName,
 		TablesToRemove:  removedTablesMapping,
-	})
-	if err != nil {
+	}); err != nil {
 		a.Alerter.LogFlowError(ctx, cfg.FlowJobName, err)
+		return err
 	}
-	return err
+	return nil
 }
 
 func (a *FlowableActivity) RemoveTablesFromRawTable(

flow/activities/snapshot_activity.go

@@ -2,7 +2,6 @@ package activities
 
 import (
 	"context"
-	"errors"
 	"fmt"
 	"log/slog"
 	"sync"
@@ -13,14 +12,13 @@ import (
 
 	"github.com/PeerDB-io/peerdb/flow/alerting"
 	"github.com/PeerDB-io/peerdb/flow/connectors"
-	connpostgres "github.com/PeerDB-io/peerdb/flow/connectors/postgres"
 	"github.com/PeerDB-io/peerdb/flow/generated/protos"
 	"github.com/PeerDB-io/peerdb/flow/shared"
 )
 
 type SlotSnapshotState struct {
-	connector    connectors.CDCPullConnector
-	signal       connpostgres.SlotSignal
+	connector    connectors.CDCPullConnectorCore
+	slotConn     interface{ Close(context.Context) error }
 	snapshotName string
 }
 
@@ -43,7 +41,9 @@ func (a *SnapshotActivity) CloseSlotKeepAlive(ctx context.Context, flowJobName s
 	defer a.SnapshotStatesMutex.Unlock()
 
 	if s, ok := a.SlotSnapshotStates[flowJobName]; ok {
-		close(s.signal.CloneComplete)
+		if s.slotConn != nil {
+			s.slotConn.Close(ctx)
+		}
 		connectors.CloseConnector(ctx, s.connector)
 		delete(a.SlotSnapshotStates, flowJobName)
 	}
@@ -61,30 +61,22 @@ func (a *SnapshotActivity) SetupReplication(
 
 	a.Alerter.LogFlowEvent(ctx, config.FlowJobName, "Started Snapshot Flow Job")
 
-	conn, err := connectors.GetByNameAs[*connpostgres.PostgresConnector](ctx, nil, a.CatalogPool, config.PeerName)
+	conn, err := connectors.GetByNameAs[connectors.CDCPullConnectorCore](ctx, nil, a.CatalogPool, config.PeerName)
 	if err != nil {
-		if errors.Is(err, errors.ErrUnsupported) {
-			logger.Info("setup replication is no-op for non-postgres source")
-			return nil, nil
-		}
 		return nil, fmt.Errorf("failed to get connector: %w", err)
 	}
 
-	closeConnectionForError := func(err error) {
+	logger.Info("waiting for slot to be created...")
+	slotInfo, err := conn.SetupReplication(ctx, config)
+
+	if err != nil {
 		a.Alerter.LogFlowError(ctx, config.FlowJobName, err)
 		// it is important to close the connection here as it is not closed in CloseSlotKeepAlive
 		connectors.CloseConnector(ctx, conn)
-	}
-
-	slotSignal := connpostgres.NewSlotSignal()
-	go conn.SetupReplication(ctx, slotSignal, config)
-
-	logger.Info("waiting for slot to be created...")
-	slotInfo := <-slotSignal.SlotCreated
-
-	if slotInfo.Err != nil {
-		closeConnectionForError(slotInfo.Err)
-		return nil, fmt.Errorf("slot error: %w", slotInfo.Err)
+		return nil, fmt.Errorf("slot error: %w", err)
+	} else if slotInfo.Conn == nil && slotInfo.SlotName == "" {
+		logger.Info("replication setup without slot")
+		return nil, nil
 	} else {
 		logger.Info("slot created", slog.String("SlotName", slotInfo.SlotName))
 	}
@@ -93,7 +85,7 @@ func (a *SnapshotActivity) SetupReplication(
 	defer a.SnapshotStatesMutex.Unlock()
 
 	a.SlotSnapshotStates[config.FlowJobName] = SlotSnapshotState{
-		signal:       slotSignal,
+		slotConn:     slotInfo.Conn,
 		snapshotName: slotInfo.SnapshotName,
 		connector:    conn,
 	}

flow/connectors/core.go

@@ -44,7 +44,7 @@ type ValidationConnector interface {
 type GetTableSchemaConnector interface {
 	Connector
 
-	// GetTableSchema returns the schema of a table in terms of QValueKind.
+	// GetTableSchema returns the schema of a table in terms of type system.
 	GetTableSchema(
 		ctx context.Context,
 		env map[string]string,
@@ -67,6 +67,9 @@ type CDCPullConnectorCore interface {
 	// `any` from ExportSnapshot passed here when done, allowing transaction to commit
 	FinishExport(any) error
 
+	// Setup replication in prep for initial copy
+	SetupReplication(context.Context, *protos.SetupReplicationInput) (model.SetupReplicationResult, error)
+
 	// Methods related to retrieving and pushing records for this connector as a source and destination.
 	SetupReplConn(context.Context) error
 

flow/connectors/postgres/client.go

@@ -221,16 +221,15 @@ func (c *PostgresConnector) getNullableColumns(ctx context.Context, relID uint32
 
 func (c *PostgresConnector) tableExists(ctx context.Context, schemaTable *utils.SchemaTable) (bool, error) {
 	var exists pgtype.Bool
-	err := c.conn.QueryRow(ctx,
+	if err := c.conn.QueryRow(ctx,
 		`SELECT EXISTS (
 			SELECT FROM pg_tables
 			WHERE schemaname = $1
 			AND tablename = $2
 		)`,
 		schemaTable.Schema,
 		schemaTable.Table,
-	).Scan(&exists)
-	if err != nil {
+	).Scan(&exists); err != nil {
 		return false, fmt.Errorf("error checking if table exists: %w", err)
 	}
 
@@ -362,53 +361,46 @@ func (c *PostgresConnector) CreatePublication(
 // createSlotAndPublication creates the replication slot and publication.
 func (c *PostgresConnector) createSlotAndPublication(
 	ctx context.Context,
-	signal SlotSignal,
 	s SlotCheckResult,
 	slot string,
 	publication string,
 	tableNameMapping map[string]model.NameAndExclude,
 	doInitialCopy bool,
-) {
+) (model.SetupReplicationResult, error) {
 	// iterate through source tables and create publication,
 	// expecting tablenames to be schema qualified
 	if !s.PublicationExists {
 		srcTableNames := make([]string, 0, len(tableNameMapping))
 		for srcTableName := range tableNameMapping {
 			parsedSrcTableName, err := utils.ParseSchemaTable(srcTableName)
 			if err != nil {
-				signal.SlotCreated <- SlotCreationResult{
-					Err: fmt.Errorf("[publication-creation] source table identifier %s is invalid", srcTableName),
-				}
-				return
+				return model.SetupReplicationResult{}, fmt.Errorf("[publication-creation] source table identifier %s is invalid", srcTableName)
 			}
 			srcTableNames = append(srcTableNames, parsedSrcTableName.String())
 		}
 		if err := c.CreatePublication(ctx, srcTableNames, publication); err != nil {
-			signal.SlotCreated <- SlotCreationResult{Err: err}
-			return
+			return model.SetupReplicationResult{}, err
 		}
 	}
 
 	// create slot only after we succeeded in creating publication.
 	if !s.SlotExists {
 		conn, err := c.CreateReplConn(ctx)
 		if err != nil {
-			signal.SlotCreated <- SlotCreationResult{Err: fmt.Errorf("[slot] error acquiring connection: %w", err)}
-			return
+			return model.SetupReplicationResult{}, fmt.Errorf("[slot] error acquiring connection: %w", err)
 		}
-		defer conn.Close(ctx)
 
 		c.logger.Warn(fmt.Sprintf("Creating replication slot '%s'", slot))
 
 		// THIS IS NOT IN A TX!
 		if _, err := conn.Exec(ctx, "SET idle_in_transaction_session_timeout=0"); err != nil {
-			signal.SlotCreated <- SlotCreationResult{Err: fmt.Errorf("[slot] error setting idle_in_transaction_session_timeout: %w", err)}
-			return
+			conn.Close(ctx)
+			return model.SetupReplicationResult{}, fmt.Errorf("[slot] error setting idle_in_transaction_session_timeout: %w", err)
 		}
 
 		if _, err := conn.Exec(ctx, "SET lock_timeout=0"); err != nil {
-			signal.SlotCreated <- SlotCreationResult{Err: fmt.Errorf("[slot] error setting lock_timeout: %w", err)}
-			return
+			conn.Close(ctx)
+			return model.SetupReplicationResult{}, fmt.Errorf("[slot] error setting lock_timeout: %w", err)
 		}
 
 		opts := pglogrepl.CreateReplicationSlotOptions{
@@ -417,39 +409,30 @@ func (c *PostgresConnector) createSlotAndPublication(
 		}
 		res, err := pglogrepl.CreateReplicationSlot(ctx, conn.PgConn(), slot, "pgoutput", opts)
 		if err != nil {
-			signal.SlotCreated <- SlotCreationResult{Err: fmt.Errorf("[slot] error creating replication slot: %w", err)}
-			return
+			conn.Close(ctx)
+			return model.SetupReplicationResult{}, fmt.Errorf("[slot] error creating replication slot: %w", err)
 		}
 
 		pgversion, err := c.MajorVersion(ctx)
 		if err != nil {
-			signal.SlotCreated <- SlotCreationResult{Err: fmt.Errorf("[slot] error getting PG version: %w", err)}
-			return
+			conn.Close(ctx)
+			return model.SetupReplicationResult{}, fmt.Errorf("[slot] error getting PG version: %w", err)
 		}
 
 		c.logger.Info(fmt.Sprintf("Created replication slot '%s'", slot))
-		slotDetails := SlotCreationResult{
+		return model.SetupReplicationResult{
+			Conn:             conn,
 			SlotName:         res.SlotName,
 			SnapshotName:     res.SnapshotName,
-			Err:              nil,
 			SupportsTIDScans: pgversion >= shared.POSTGRES_13,
-		}
-		signal.SlotCreated <- slotDetails
-		c.logger.Info("Waiting for clone to complete")
-		<-signal.CloneComplete
-		c.logger.Info("Clone complete")
+		}, nil
 	} else {
 		c.logger.Info(fmt.Sprintf("Replication slot '%s' already exists", slot))
-		slotDetails := SlotCreationResult{
-			SlotName:         slot,
-			SnapshotName:     "",
-			Err:              nil,
-			SupportsTIDScans: false,
-		}
+		var err error
 		if doInitialCopy {
-			slotDetails.Err = ErrSlotAlreadyExists
+			err = ErrSlotAlreadyExists
 		}
-		signal.SlotCreated <- slotDetails
+		return model.SetupReplicationResult{SlotName: slot}, err
 	}
 }
 

flow/connectors/postgres/postgres.go

@@ -125,7 +125,7 @@ func (c *PostgresConnector) fetchCustomTypeMapping(ctx context.Context) (map[uin
 }
 
 func (c *PostgresConnector) CreateReplConn(ctx context.Context) (*pgx.Conn, error) {
-	// create a separate connection pool for non-replication queries as replication connections cannot
+	// create a separate connection for non-replication queries as replication connections cannot
 	// be used for extended query protocol, i.e. prepared statements
 	replConfig, err := pgx.ParseConfig(c.connStr)
 	if err != nil {
@@ -1090,13 +1090,13 @@ func (c *PostgresConnector) FinishExport(tx any) error {
 	return pgtx.Commit(timeout)
 }
 
-// SetupReplication sets up replication for the source connector.
-func (c *PostgresConnector) SetupReplication(ctx context.Context, signal SlotSignal, req *protos.SetupReplicationInput) {
+// SetupReplication sets up replication for the source connector
+func (c *PostgresConnector) SetupReplication(
+	ctx context.Context,
+	req *protos.SetupReplicationInput,
+) (model.SetupReplicationResult, error) {
 	if !shared.IsValidReplicationName(req.FlowJobName) {
-		signal.SlotCreated <- SlotCreationResult{
-			Err: fmt.Errorf("invalid flow job name: `%s`, it should be ^[a-z_][a-z0-9_]*$", req.FlowJobName),
-		}
-		return
+		return model.SetupReplicationResult{}, fmt.Errorf("invalid flow job name: `%s`, it should be ^[a-z_][a-z0-9_]*$", req.FlowJobName)
 	}
 
 	// Slotname would be the job name prefixed with "peerflow_slot_"
@@ -1113,8 +1113,7 @@ func (c *PostgresConnector) SetupReplication(ctx context.Context, signal SlotSig
 	// Check if the replication slot and publication exist
 	exists, err := c.checkSlotAndPublication(ctx, slotName, publicationName)
 	if err != nil {
-		signal.SlotCreated <- SlotCreationResult{Err: err}
-		return
+		return model.SetupReplicationResult{}, err
 	}
 
 	tableNameMapping := make(map[string]model.NameAndExclude, len(req.TableNameMapping))
@@ -1125,7 +1124,7 @@ func (c *PostgresConnector) SetupReplication(ctx context.Context, signal SlotSig
 		}
 	}
 	// Create the replication slot and publication
-	c.createSlotAndPublication(ctx, signal, exists, slotName, publicationName, tableNameMapping, req.DoInitialSnapshot)
+	return c.createSlotAndPublication(ctx, exists, slotName, publicationName, tableNameMapping, req.DoInitialSnapshot)
 }
 
 func (c *PostgresConnector) PullFlowCleanup(ctx context.Context, jobName string) error {