strip pg_dump 17+ meta-commands

Author: sofuture

pkg/dbutil/dbutil.go

@@ -57,22 +57,30 @@ func RunCommand(name string, args ...string) ([]byte, error) {
 	return stdout.Bytes(), nil
 }
 
-// TrimLeadingSQLComments removes sql comments and blank lines from the beginning of text
-// generally when performing sql dumps these contain host-specific information such as
-// client/server version numbers
+// TrimLeadingSQLComments removes sql comments and blank lines from the beginning of text,
+// and psql meta-commands (lines starting with \) from anywhere in the text.
+// It also collapses 3+ consecutive blank lines to 2 for consistent formatting.
+// Generally when performing sql dumps these contain host-specific information such as
+// client/server version numbers.
 func TrimLeadingSQLComments(data []byte) ([]byte, error) {
 	// create decent size buffer
 	out := bytes.NewBuffer(make([]byte, 0, len(data)))
 
 	// iterate over sql lines
 	preamble := true
+	consecutiveBlankLines := 0
 	scanner := bufio.NewScanner(bytes.NewReader(data))
 	for scanner.Scan() {
 		// we read bytes directly for premature performance optimization
 		line := scanner.Bytes()
 
-		if preamble && (len(line) == 0 || bytes.Equal(line[0:2], []byte("--"))) {
-			// header section, skip this line in output buffer
+		// always skip psql meta-commands (e.g., \restrict, \unrestrict from pg_dump 17+)
+		if len(line) >= 1 && line[0] == '\\' {
+			continue
+		}
+
+		if preamble && (len(line) == 0 || (len(line) >= 2 && bytes.Equal(line[0:2], []byte("--")))) {
+			// header section, skip blank lines and SQL comments (--)
 			continue
 		}
 
@@ -82,6 +90,16 @@ func TrimLeadingSQLComments(data []byte) ([]byte, error) {
 		// trim trailing whitespace
 		line = bytes.TrimRightFunc(line, unicode.IsSpace)
 
+		// track consecutive blank lines and collapse 3+ to 2
+		if len(line) == 0 {
+			consecutiveBlankLines++
+			if consecutiveBlankLines > 2 {
+				continue
+			}
+		} else {
+			consecutiveBlankLines = 0
+		}
+
 		// copy bytes to output buffer
 		if _, err := out.Write(line); err != nil {
 			return nil, err