to #246 fix: bugs in replication.

catalog/internal_tables.go

@@ -180,8 +180,8 @@ var InternalTables = struct {
 		Schema:       "__sys__",
 		Name:         "pg_subscription",
 		KeyColumns:   []string{"subname"},
-		ValueColumns: []string{"subconninfo", "subpublication", "subenabled", "subskiplsn"},
-		DDL:          "subname TEXT PRIMARY KEY, subconninfo TEXT, subpublication TEXT, subenabled BOOLEAN, subskiplsn TEXT",
+		ValueColumns: []string{"subconninfo", "subpublication", "subskiplsn", "subenabled"},
+		DDL:          "subname TEXT PRIMARY KEY, subconninfo TEXT, subpublication TEXT, subskiplsn TEXT, subenabled BOOLEAN",
 	},
 	GlobalStatus: InternalTable{
 		Schema:       "performance_schema",

pgserver/logrepl/replication.go

@@ -45,6 +45,7 @@ type rcvMsg struct {
 }
 
 type LogicalReplicator struct {
+	subscription  string
 	primaryDns    string
 	flushInterval time.Duration
 
@@ -59,8 +60,9 @@ type LogicalReplicator struct {
 // NewLogicalReplicator creates a new logical replicator instance which connects to the primary and replication
 // databases using the connection strings provided. The connection to the replica is established immediately, and the
 // connection to the primary is established when StartReplication is called.
-func NewLogicalReplicator(primaryDns string) (*LogicalReplicator, error) {
+func NewLogicalReplicator(subscription, primaryDns string) (*LogicalReplicator, error) {
 	return &LogicalReplicator{
+		subscription:  subscription,
 		primaryDns:    primaryDns,
 		flushInterval: 200 * time.Millisecond,
 		mu:            &sync.Mutex{},
@@ -221,7 +223,7 @@ func (r *LogicalReplicator) StartReplication(sqlCtx *sql.Context, slotName strin
 	standbyMessageTimeout := 10 * time.Second
 	nextStandbyMessageDeadline := time.Now().Add(standbyMessageTimeout)
 
-	lastWrittenLsn, err := SelectSubscriptionLsn(sqlCtx, slotName)
+	lastWrittenLsn, err := SelectSubscriptionLsn(sqlCtx, r.subscription)
 	if err != nil {
 		return err
 	}
@@ -979,7 +981,7 @@ func (r *LogicalReplicator) commitOngoingTxn(state *replicationState, flushReaso
 	}
 
 	r.logger.Debugf("Writing LSN %s\n", state.lastCommitLSN)
-	if err = UpdateSubscriptionLsn(state.replicaCtx, state.slotName, state.lastCommitLSN.String()); err != nil {
+	if err = UpdateSubscriptionLsn(state.replicaCtx, state.lastCommitLSN.String(), r.subscription); err != nil {
 		return err
 	}
 

pgserver/logrepl/replication_test.go

@@ -687,7 +687,7 @@ func RunReplicationScript(t *testing.T, dsn string, script ReplicationTest) {
 }
 
 func newReplicator(t *testing.T, primaryDns string) *logrepl.LogicalReplicator {
-	r, err := logrepl.NewLogicalReplicator(primaryDns)
+	r, err := logrepl.NewLogicalReplicator("mySubTest", primaryDns)
 	require.NoError(t, err)
 	return r
 }

pgserver/logrepl/subscription.go

@@ -15,39 +15,18 @@ type Subscription struct {
 	Subscription string
 	Conn         string
 	Publication  string
-	Lsn          pglogrepl.LSN
+	LsnStr       string
 	Enabled      bool
 	Replicator   *LogicalReplicator
 }
 
 var keyColumns = []string{"subname"}
-var statusValueColumns = []string{"substatus"}
+var statusValueColumns = []string{"subenabled"}
 var lsnValueColumns = []string{"subskiplsn"}
 
 var subscriptionMap = sync.Map{}
 var mu sync.Mutex
 
-func GetSubscription(ctx *sql.Context, name string) (*Subscription, error) {
-	if value, ok := subscriptionMap.Load(name); ok {
-		if sub, ok := value.(*Subscription); ok {
-			return sub, nil
-		}
-	}
-
-	// Attempt to reload all subscriptions if not found
-	if err := UpdateSubscriptions(ctx); err != nil {
-		return nil, err
-	}
-
-	if value, ok := subscriptionMap.Load(name); ok {
-		if sub, ok := value.(*Subscription); ok {
-			return sub, nil
-		}
-	}
-
-	return nil, nil
-}
-
 func UpdateSubscriptions(ctx *sql.Context) error {
 	mu.Lock()
 	defer mu.Unlock()
@@ -59,23 +38,24 @@ func UpdateSubscriptions(ctx *sql.Context) error {
 
 	var tempMap = make(map[string]*Subscription)
 	for rows.Next() {
-		var name, conn, pub string
+		var name, conn, pub, lsn string
 		var enabled bool
-		if err := rows.Scan(&name, &conn, &pub, &enabled); err != nil {
+		if err := rows.Scan(&name, &conn, &pub, &lsn, &enabled); err != nil {
 			return err
 		}
 		tempMap[name] = &Subscription{
 			Subscription: name,
 			Conn:         conn,
 			Publication:  pub,
+			LsnStr:       lsn,
 			Enabled:      enabled,
 			Replicator:   nil,
 		}
 	}
 
 	for tempName, tempSub := range tempMap {
 		if _, loaded := subscriptionMap.LoadOrStore(tempName, tempSub); !loaded {
-			replicator, err := NewLogicalReplicator(tempSub.Conn)
+			replicator, err := NewLogicalReplicator(tempName, tempSub.Conn)
 			if err != nil {
 				return fmt.Errorf("failed to create logical replicator: %v", err)
 			}
@@ -90,8 +70,11 @@ func UpdateSubscriptions(ctx *sql.Context) error {
 			if err != nil {
 				return fmt.Errorf("failed to create replication slot: %v", err)
 			}
+			if tempSub.Enabled {
+				go replicator.StartReplication(ctx, tempSub.Publication)
+			}
 		} else {
-			if sub, ok := subscriptionMap.Load(tempSub); ok {
+			if sub, ok := subscriptionMap.Load(tempName); ok {
 				if subscription, ok := sub.(*Subscription); ok {
 					if tempSub.Enabled != subscription.Enabled {
 						subscription.Enabled = tempSub.Enabled
@@ -124,8 +107,8 @@ func CreateSubscription(ctx *sql.Context, name, conn, pub, lsn string, enabled b
 	return err
 }
 
-func UpdateSubscriptionStatus(ctx *sql.Context, name string, enabled bool) error {
-	_, err := adapter.ExecCatalogInTxn(ctx, catalog.InternalTables.PgSubscription.UpdateStmt(keyColumns, statusValueColumns), name, enabled)
+func UpdateSubscriptionStatus(ctx *sql.Context, enabled bool, name string) error {
+	_, err := adapter.ExecCatalogInTxn(ctx, catalog.InternalTables.PgSubscription.UpdateStmt(keyColumns, statusValueColumns), enabled, name)
 	return err
 }
 
@@ -134,8 +117,8 @@ func DeleteSubscription(ctx *sql.Context, name string) error {
 	return err
 }
 
-func UpdateSubscriptionLsn(ctx *sql.Context, name, lsn string) error {
-	_, err := adapter.ExecCatalogInTxn(ctx, catalog.InternalTables.PgSubscription.UpdateStmt(keyColumns, lsnValueColumns), name, lsn)
+func UpdateSubscriptionLsn(ctx *sql.Context, lsn, name string) error {
+	_, err := adapter.ExecCatalogInTxn(ctx, catalog.InternalTables.PgSubscription.UpdateStmt(keyColumns, lsnValueColumns), lsn, name)
 	return err
 }
 

pgserver/subscription_handler.go

@@ -181,7 +181,7 @@ func (h *ConnectionHandler) executeAlterEnable(subscriptionConfig *SubscriptionC
 		return fmt.Errorf("failed to create context for query: %w", err)
 	}
 
-	err = logrepl.UpdateSubscriptionStatus(sqlCtx, subscriptionConfig.SubscriptionName, true)
+	err = logrepl.UpdateSubscriptionStatus(sqlCtx, true, subscriptionConfig.SubscriptionName)
 	if err != nil {
 		return fmt.Errorf("failed to delete subscription: %w", err)
 	}
@@ -195,7 +195,7 @@ func (h *ConnectionHandler) executeAlterDisable(subscriptionConfig *Subscription
 		return fmt.Errorf("failed to create context for query: %w", err)
 	}
 
-	err = logrepl.UpdateSubscriptionStatus(sqlCtx, subscriptionConfig.SubscriptionName, false)
+	err = logrepl.UpdateSubscriptionStatus(sqlCtx, false, subscriptionConfig.SubscriptionName)
 	if err != nil {
 		return fmt.Errorf("failed to delete subscription: %w", err)
 	}