Support SQL formatting inside textproto string fields marked with '# txtpbfmt sql-format'.

PiperOrigin-RevId: 869641916

Author: txtpbfmt-copybara-robot

impl/impl.go

@@ -15,6 +15,7 @@ import (
 	"github.com/protocolbuffers/txtpbfmt/descriptor"
 	"github.com/protocolbuffers/txtpbfmt/quote"
 	"github.com/protocolbuffers/txtpbfmt/sort"
+	"github.com/protocolbuffers/txtpbfmt/sql"
 	"github.com/protocolbuffers/txtpbfmt/wrap"
 )
 
@@ -188,6 +189,9 @@ func ParseWithMetaCommentConfig(in []byte, c config.Config) ([]*ast.Node, error)
 	if p.index < p.length {
 		return nil, fmt.Errorf("parser didn't consume all input. Stopped at %s", p.errorContext())
 	}
+	if err := sql.Format(nodes); err != nil {
+		return nil, err
+	}
 	if err := wrap.Strings(nodes, 0, c); err != nil {
 		return nil, err
 	}
@@ -201,7 +205,7 @@ func ParseWithMetaCommentConfig(in []byte, c config.Config) ([]*ast.Node, error)
 // have the equal sign. Currently there are only two MetaComments that are in the former format:
 //
 //	"sort_repeated_fields_by_subfield": If this appears multiple times, then they will all be added
-//	to the config and the order is perserved.
+//	to the config and the order is preserved.
 //	"wrap_strings_at_column": The <val> is expected to be an integer. If it is not, then it will be
 //	ignored. If this appears multiple times, only the last one saved.
 func addToConfig(metaComment string, c *config.Config) error {

sql.go

@@ -0,0 +1,97 @@
+// Package sql provides functions for formatting SQL inside textproto string fields.
+package sql
+
+import (
+	"strconv"
+	"strings"
+
+	"google3/storage/googlesql/public/go/lenientformatter"
+	"github.com/protocolbuffers/txtpbfmt/ast"
+	"github.com/protocolbuffers/txtpbfmt/unquote"
+)
+
+const sqlFormatComment = "# txtpbfmt sql-format"
+
+// Format formats SQL strings in the given nodes.
+func Format(nodes []*ast.Node) error {
+	for _, nd := range nodes {
+		if nd.ChildrenSameLine {
+			continue
+		}
+		if err := formatNode(nd); err != nil {
+			return err
+		}
+		if err := Format(nd.Children); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func formatNode(nd *ast.Node) error {
+	hasComment := false
+	for _, c := range nd.PreComments {
+		if strings.TrimSpace(c) == sqlFormatComment {
+			hasComment = true
+			break
+		}
+	}
+	if !hasComment {
+		return nil
+	}
+	if len(nd.Values) == 0 {
+		return nil
+	}
+
+	// Unquote the string(s).
+	sql, _, err := unquote.Unquote(nd)
+	if err != nil {
+		// If it's not a valid string, we can't format it as SQL.
+		return nil
+	}
+
+	formatted, err := lenientformatter.FormatSQL(sql)
+	if err != nil {
+		// If SQL formatting fails, we keep the original.
+		return nil
+	}
+
+	lines := strings.Split(formatted, "\n")
+	// If there's a trailing newline, Split might produce an extra empty line.
+	if len(lines) > 1 && lines[len(lines)-1] == "" {
+		lines = lines[:len(lines)-1]
+	}
+
+	newValues := make([]*ast.Value, 0, len(lines))
+	for i, line := range lines {
+		val := line
+		if i < len(lines)-1 {
+			// Add a space instead of a newline to separate SQL lines.
+			// This results in a cleaner look in the textproto.
+			if !strings.HasSuffix(val, " ") {
+				val += " "
+			}
+		}
+		newValues = append(newValues, &ast.Value{
+			Value: strconv.Quote(val),
+		})
+	}
+
+	// Preserve comments if possible.
+	if len(nd.Values) > 0 {
+		if len(nd.Values[0].PreComments) > 0 {
+			newValues[0].PreComments = nd.Values[0].PreComments
+		}
+		lastIdx := len(nd.Values) - 1
+		if len(nd.Values[lastIdx].InlineComment) > 0 {
+			newValues[len(newValues)-1].InlineComment = nd.Values[lastIdx].InlineComment
+		}
+	}
+
+	nd.Values = newValues
+	if len(lines) > 1 {
+		nd.PutSingleValueOnNextLine = true
+	}
+
+	return nil
+}

sql_test.go

@@ -0,0 +1,75 @@
+package sql_test
+
+import (
+	"testing"
+
+	"github.com/kylelemons/godebug/diff"
+	"github.com/protocolbuffers/txtpbfmt/parser"
+	"github.com/protocolbuffers/txtpbfmt/sql"
+)
+
+func TestFormat(t *testing.T) {
+	tests := []struct {
+		name string
+		in   string
+		want string
+	}{
+		{
+			name: "simple SQL",
+			in: `# txtpbfmt sql-format
+query: "select * from table where id = 1"
+`,
+			want: `# txtpbfmt sql-format
+query: "SELECT * FROM table WHERE id = 1"
+`,
+		},
+		{
+			name: "multiline SQL",
+			in: `# txtpbfmt sql-format
+query: "select a, b, c from table1 join table2 on table1.id = table2.id where a > 10 group by 1, 2, 3 order by 1"
+`,
+			want: `# txtpbfmt sql-format
+query:
+  "SELECT a, b, c "
+  "FROM table1 "
+  "JOIN table2 "
+  "  ON table1.id = table2.id "
+  "WHERE a > 10 "
+  "GROUP BY 1, 2, 3 "
+  "ORDER BY 1"
+`,
+		},
+		{
+			name: "no comment, no format",
+			in: `query: "select * from table"
+`,
+			want: `query: "select * from table"
+`,
+		},
+		{
+			name: "wrong comment, no format",
+			in: `# some other comment
+query: "select * from table"
+`,
+			want: `# some other comment
+query: "select * from table"
+`,
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			nodes, err := parser.Parse([]byte(tc.in))
+			if err != nil {
+				t.Fatalf("Parse() err = %v", err)
+			}
+			if err := sql.Format(nodes); err != nil {
+				t.Fatalf("Format() err = %v", err)
+			}
+			got := parser.Pretty(nodes, 0)
+			if diff := diff.Diff(tc.want, got); diff != "" {
+				t.Errorf("Format() diff (-want +got):\n%s", diff)
+			}
+		})
+	}
+}

wrap/wrap.go

@@ -14,7 +14,10 @@ import (
 
 var tagRegex = regexp.MustCompile(`<.*>`)
 
-const indentSpaces = "  "
+const (
+	indentSpaces     = "  "
+	sqlFormatComment = "# txtpbfmt sql-format"
+)
 
 // Strings wraps the strings in the given nodes.
 func Strings(nodes []*ast.Node, depth int, c config.Config) error {
@@ -25,6 +28,9 @@ func Strings(nodes []*ast.Node, depth int, c config.Config) error {
 		if nd.ChildrenSameLine {
 			continue
 		}
+		if hasSQLFormatComment(nd) {
+			continue
+		}
 		if err := wrapNodeStrings(nd, depth, c); err != nil {
 			return err
 		}
@@ -35,6 +41,15 @@ func Strings(nodes []*ast.Node, depth int, c config.Config) error {
 	return nil
 }
 
+func hasSQLFormatComment(nd *ast.Node) bool {
+	for _, c := range nd.PreComments {
+		if strings.TrimSpace(c) == sqlFormatComment {
+			return true
+		}
+	}
+	return false
+}
+
 func wrapNodeStrings(nd *ast.Node, depth int, c config.Config) error {
 	if c.WrapStringsAtColumn > 0 && needsWrappingAtColumn(nd, depth, c) {
 		if err := wrapLinesAtColumn(nd, depth, c); err != nil {