to #119 feat: make CREATE SUBSCRIPTION work

Author: TianyuZhang1214

pgserver/logrepl/replication.go

@@ -533,76 +533,130 @@ func DropPublication(primaryDns, slotName string) error {
 	return err
 }
 
-// CreatePublication creates a publication with the given name if it does not already exist. Mostly useful for testing.
+// CreatePublicationIfNotExists creates a publication with the given name if it does not already exist. Mostly useful for testing.
 // Customers should run the CREATE PUBLICATION command on their primary server manually, specifying whichever tables
 // they want to replicate.
-func CreatePublication(primaryDns, slotName string) error {
-	conn, err := pgconn.Connect(context.Background(), primaryDns)
+func CreatePublicationIfNotExists(primaryDns, publicationName string) error {
+	// Connect to the primary PostgreSQL server
+	conn, err := pgx.Connect(context.Background(), primaryDns)
 	if err != nil {
-		return err
+		return fmt.Errorf("failed to connect to primary database: %w", err)
 	}
 	defer conn.Close(context.Background())
 
-	result := conn.Exec(context.Background(), fmt.Sprintf("CREATE PUBLICATION %s FOR ALL TABLES;", slotName))
-	_, err = result.ReadAll()
-	return err
+	// Check if the publication exists
+	query := `SELECT EXISTS(SELECT 1 FROM pg_publication WHERE pubname = $1)`
+	var exists bool
+	if err := conn.QueryRow(context.Background(), query, publicationName).Scan(&exists); err != nil {
+		return fmt.Errorf("failed to check publication existence: %w", err)
+	}
+
+	// Create the publication if it does not exist
+	if !exists {
+		createQuery := fmt.Sprintf("CREATE PUBLICATION %s FOR ALL TABLES", publicationName)
+		if _, err := conn.Exec(context.Background(), createQuery); err != nil {
+			if pgErr, ok := err.(*pgconn.PgError); ok && pgErr.Code == "42710" {
+				// Ignore "publication already exists" error
+				return nil
+			}
+			return fmt.Errorf("failed to create publication: %w", err)
+		}
+	}
+
+	return nil
 }
 
-// DropReplicationSlot drops the replication slot with the given name. Any error from the slot not existing is ignored.
-func (r *LogicalReplicator) DropReplicationSlot(slotName string) error {
+// DropReplicationSlotIfExists drops the replication slot with the given name. Any error from the slot not existing is ignored.
+func (r *LogicalReplicator) DropReplicationSlotIfExists(slotName string) error {
+	// Check if the replication slot exists
+	exists, err := r.replicationSlotExists(slotName)
+	if err != nil {
+		return fmt.Errorf("failed to check replication slot existence: %w", err)
+	}
+
+	if !exists {
+		r.logger.Infof("Replication slot '%s' does not exist.", slotName)
+		return nil
+	}
+
+	// Connect to the replication database
 	conn, err := pgconn.Connect(context.Background(), r.ReplicationDns())
 	if err != nil {
-		return err
+		return fmt.Errorf("failed to connect to replication database: %w", err)
+	}
+	defer conn.Close(context.Background())
+
+	if err := pglogrepl.DropReplicationSlot(context.Background(), conn, slotName, pglogrepl.DropReplicationSlotOptions{}); err != nil {
+		return fmt.Errorf("failed to drop replication slot '%s': %w", slotName, err)
 	}
 
-	_ = pglogrepl.DropReplicationSlot(context.Background(), conn, slotName, pglogrepl.DropReplicationSlotOptions{})
+	r.logger.Infof("Replication slot '%s' successfully dropped.", slotName)
 	return nil
 }
 
-// CreateReplicationSlotIfNecessary creates the replication slot named if it doesn't already exist.
-func (r *LogicalReplicator) CreateReplicationSlotIfNecessary(slotName string) error {
-	conn, err := pgx.Connect(context.Background(), r.PrimaryDns())
+// CreateReplicationSlotIfNotExists creates the replication slot named if it doesn't already exist.
+func (r *LogicalReplicator) CreateReplicationSlotIfNotExists(slotName string) error {
+	// Check if the replication slot already exists
+	exists, err := r.replicationSlotExists(slotName)
 	if err != nil {
-		return err
+		return fmt.Errorf("error checking replication slot existence: %w", err)
 	}
 
-	rows, err := conn.Query(context.Background(), "select * from pg_replication_slots where slot_name = $1", slotName)
+	// If the slot already exists, no further action is needed
+	if exists {
+		r.logger.Infof("Replication slot '%s' already exists.", slotName)
+		return nil
+	}
+
+	// Create the replication slot
+	err = r.createReplicationSlot(slotName)
 	if err != nil {
-		return err
+		return fmt.Errorf("error creating replication slot '%s': %w", slotName, err)
 	}
 
-	slotExists := false
-	defer rows.Close()
-	for rows.Next() {
-		_, err := rows.Values()
-		if err != nil {
-			return err
-		}
-		slotExists = true
+	r.logger.Infof("Replication slot '%s' created successfully.", slotName)
+	return nil
+}
+
+// Helper method to check if a replication slot exists
+func (r *LogicalReplicator) replicationSlotExists(slotName string) (bool, error) {
+	conn, err := pgx.Connect(context.Background(), r.PrimaryDns())
+	if err != nil {
+		return false, fmt.Errorf("failed to connect to primary database: %w", err)
 	}
+	defer conn.Close(context.Background())
 
-	if rows.Err() != nil {
-		return rows.Err()
+	var exists bool
+	query := `SELECT EXISTS(SELECT 1 FROM pg_replication_slots WHERE slot_name = $1)`
+	err = conn.QueryRow(context.Background(), query, slotName).Scan(&exists)
+	if err != nil {
+		return false, fmt.Errorf("error querying replication slots: %w", err)
 	}
 
-	// We need a different connection to create the replication slot
-	conn, err = pgx.Connect(context.Background(), r.ReplicationDns())
+	return exists, nil
+}
+
+// Helper method to create a replication slot
+func (r *LogicalReplicator) createReplicationSlot(slotName string) error {
+	conn, err := pgx.Connect(context.Background(), r.ReplicationDns())
 	if err != nil {
-		return err
+		return fmt.Errorf("failed to connect to replication database: %w", err)
 	}
+	defer conn.Close(context.Background())
 
-	if !slotExists {
-		_, err = pglogrepl.CreateReplicationSlot(context.Background(), conn.PgConn(), slotName, outputPlugin, pglogrepl.CreateReplicationSlotOptions{})
-		if err != nil {
-			pgErr, ok := err.(*pgconn.PgError)
-			if ok && pgErr.Code == "42710" {
-				// replication slot already exists, we can ignore this error
-			} else {
-				return err
-			}
+	_, err = pglogrepl.CreateReplicationSlot(
+		context.Background(),
+		conn.PgConn(),
+		slotName,
+		outputPlugin,
+		pglogrepl.CreateReplicationSlotOptions{},
+	)
+	if err != nil {
+		if pgErr, ok := err.(*pgconn.PgError); ok && pgErr.Code == "42710" {
+			// Replication slot already exists; ignore this error
+			return nil
 		}
-
-		r.logger.Infoln("Created replication slot:", slotName)
+		return fmt.Errorf("error creating replication slot: %w", err)
 	}
 
 	return nil

pgserver/logrepl/replication_test.go

@@ -565,7 +565,7 @@ func RunReplicationScripts(t *testing.T, scripts []ReplicationTest) {
 	// We drop and recreate the replication slot once at the beginning of the test suite. Postgres seems to do a little
 	// work in the background with a publication, so we need to wait a little bit before running any test scripts.
 	require.NoError(t, logrepl.DropPublication(primaryDns, slotName))
-	require.NoError(t, logrepl.CreatePublication(primaryDns, slotName))
+	require.NoError(t, logrepl.CreatePublicationIfNotExists(primaryDns, slotName))
 	time.Sleep(500 * time.Millisecond)
 
 	// for i, script := range scripts {
@@ -731,10 +731,10 @@ func connectionForQuery(t *testing.T, query string, connections map[string]*pgx.
 func handlePseudoQuery(t *testing.T, server *pgserver.Server, query string, r *logrepl.LogicalReplicator) bool {
 	switch query {
 	case createReplicationSlot:
-		require.NoError(t, r.CreateReplicationSlotIfNecessary(slotName))
+		require.NoError(t, r.CreateReplicationSlotIfNotExists(slotName))
 		return true
 	case dropReplicationSlot:
-		require.NoError(t, r.DropReplicationSlot(slotName))
+		require.NoError(t, r.DropReplicationSlotIfExists(slotName))
 		return true
 	case startReplication:
 		go func() {

pgserver/subscription_handler.go

@@ -163,7 +163,12 @@ func doCreateSubscription(h *ConnectionHandler, subscriptionConfig *Subscription
 		return fmt.Errorf("failed to create logical replicator: %w", err)
 	}
 
-	err = replicator.CreateReplicationSlotIfNecessary(subscriptionConfig.SubscriptionName)
+	err = logrepl.CreatePublicationIfNotExists(subscriptionConfig.ToDNS(), subscriptionConfig.PublicationName)
+	if err != nil {
+		return fmt.Errorf("failed to create publication: %w", err)
+	}
+
+	err = replicator.CreateReplicationSlotIfNotExists(subscriptionConfig.PublicationName)
 	if err != nil {
 		return fmt.Errorf("failed to create replication slot: %w", err)
 	}
@@ -173,7 +178,7 @@ func doCreateSubscription(h *ConnectionHandler, subscriptionConfig *Subscription
 		return fmt.Errorf("failed to create context for query: %w", err)
 	}
 
-	err = replicator.WriteWALPosition(sqlCtx, subscriptionConfig.SubscriptionName, subscriptionConfig.LSN)
+	err = replicator.WriteWALPosition(sqlCtx, subscriptionConfig.PublicationName, subscriptionConfig.LSN)
 	if err != nil {
 		return fmt.Errorf("failed to write WAL position: %w", err)
 	}
@@ -183,10 +188,7 @@ func doCreateSubscription(h *ConnectionHandler, subscriptionConfig *Subscription
 		return fmt.Errorf("failed to create context for query: %w", err)
 	}
 
-	// TODO(neo.zty): This process has error message with:
-	//  Received 'conn busy' error, waiting and retrying  component=replicator protocol=pg
-	//  Deal with it.
-	go replicator.StartReplication(sqlCtx, subscriptionConfig.SubscriptionName)
+	go replicator.StartReplication(sqlCtx, subscriptionConfig.PublicationName)
 
 	return nil
 }