pg: recognize streaming prepared tx msg (#1071) (#1073)

Co-authored-by: jchappelow <[email protected]>
kwilteam · Oct 22, 2024 · f11e0f7 · f11e0f7
1 parent f4aa5da
commit f11e0f7
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/internal/sql/pg/repl.go b/internal/sql/pg/repl.go
@@ -501,7 +501,7 @@ func decodeWALData(hasher hash.Hash, walData []byte, relations map[uint32]*pglog
 		hasher.Reset()
 		changesetWriter.fail() // discard changeset
 
-	// v2 Stream control messages.  Not expected for kwil
+	// v2 Stream control messages.  Only expected with large transactions.
 	case *pglogrepl.StreamStartMessageV2:
 		*inStream = true
 		logger.Warnf(" [msg] StreamStartMessageV2: xid %d, first segment? %d", logicalMsg.Xid, logicalMsg.FirstSegment)

diff --git a/internal/sql/pg/repl_msgs.go b/internal/sql/pg/repl_msgs.go
@@ -14,13 +14,16 @@ const (
 	MessageTypeBeginPrepare     pglogrepl.MessageType = 'b'
 	MessageTypeCommitPrepared   pglogrepl.MessageType = 'K'
 	MessageTypeRollbackPrepared pglogrepl.MessageType = 'r'
+	MessageTypeStreamPrepare    pglogrepl.MessageType = 'p'
 )
 
 // msgTypeToString is helpful for debugging, but normally unused.
 func msgTypeToString(t pglogrepl.MessageType) string { //nolint:unused
 	switch t {
 	case MessageTypePrepare:
 		return "Prepare"
+	case MessageTypeStreamPrepare:
+		return "Stream Prepared"
 	case MessageTypeBeginPrepare:
 		return "Begin Prepared"
 	case MessageTypeCommitPrepared:
@@ -39,7 +42,7 @@ func parseV3(data []byte, inStream bool) (m pglogrepl.Message, err error) {
 	var decoder pglogrepl.MessageDecoder // v1 and v3 have same Decode signature (stream not relevant)
 
 	switch msgType {
-	case MessageTypePrepare:
+	case MessageTypePrepare, MessageTypeStreamPrepare: // same encoding
 		decoder = new(PrepareMessageV3)
 	case MessageTypeBeginPrepare:
 		decoder = new(BeginPrepareMessageV3)