Print more nodes

Author: kyleconroy

internal/engine/postgresql/print_test.go

@@ -5,9 +5,8 @@ import (
 	"strings"
 	"testing"
 
-	"github.com/google/go-cmp/cmp"
+	pg_query "github.com/pganalyze/pg_query_go/v4"
 
-	"github.com/sqlc-dev/sqlc/internal/debug"
 	"github.com/sqlc-dev/sqlc/internal/sql/ast"
 )
 
@@ -16,11 +15,27 @@ func TestPrintAst(t *testing.T) {
 
 	queries := []string{
 		`SELECT * FROM foo;`,
+		`SELECT * 
+FROM foo;`,
 		`SELECT * FROM foo,bar;`,
 		`SELECT * FROM foo WHERE EXISTS (SELECT * FROM foo);`,
 		`WITH bar AS (SELECT * FROM foo), bat AS (SELECT 1) SELECT * FROM foo;`,
 		`SELECT t.* FROM foo t;`,
 		`SELECT *,*,foo.* FROM foo;`,
+		`SELECT 'foo';`,
+		`SELECT true;`,
+		`SELECT 1.2;`,
+		`SELECT "foo";`,
+		`SELECT * FROM foo LIMIT 1;`,
+		`SELECT * FROM foo OFFSET 1;`,
+		`SELECT * FROM foo LIMIT 1 OFFSET 1;`,
+		`SELECT * FROM foo ORDER BY name;`,
+		`SELECT DISTINCT * FROM foo;`,
+		`SELECT DISTINCT ON (location) location, time, report
+		FROM weather_reports
+		ORDER BY location, time DESC;`,
+		`SELECT * FROM (SELECT * FROM mytable FOR SHARE) ss WHERE col1 = 5;`,
+		`INSERT INTO myschema.foo (a, b) VALUES ($1, $2);`,
 	}
 
 	// Use astutils to look for select nodes
@@ -29,17 +44,27 @@ func TestPrintAst(t *testing.T) {
 	for i, q := range queries {
 		q := q
 		t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
+			expected, err := pg_query.Fingerprint(q)
+			if err != nil {
+				t.Fatal(err)
+			}
 			stmts, err := p.Parse(strings.NewReader(q))
 			if err != nil {
 				t.Fatal(err)
 			}
-			for _, stmt := range stmts {
-				out := ast.Format(stmt.Raw)
-				if diff := cmp.Diff(q, out); diff != "" {
-					debug.Dump(stmt)
-					t.Errorf("- %s", q)
-					t.Errorf("+ %s", out)
-				}
+			if len(stmts) != 1 {
+				t.Fatal("expected one statement")
+			}
+			out := ast.Format(stmts[0].Raw)
+			actual, err := pg_query.Fingerprint(out)
+			if err != nil {
+				t.Error(err)
+			}
+			if expected != actual {
+				t.Errorf("- %s", expected)
+				t.Errorf("- %s", q)
+				t.Errorf("+ %s", actual)
+				t.Errorf("+ %s", out)
 			}
 		})
 	}

internal/sql/ast/a_const.go

@@ -13,5 +13,11 @@ func (n *A_Const) Format(buf *TrackedBuffer) {
 	if n == nil {
 		return
 	}
-	buf.astFormat(n.Val)
+	if _, ok := n.Val.(*String); ok {
+		buf.WriteString("'")
+		buf.astFormat(n.Val)
+		buf.WriteString("'")
+	} else {
+		buf.astFormat(n.Val)
+	}
 }

internal/sql/ast/a_expr.go

@@ -11,3 +11,14 @@ type A_Expr struct {
 func (n *A_Expr) Pos() int {
 	return n.Location
 }
+
+func (n *A_Expr) Format(buf *TrackedBuffer) {
+	if n == nil {
+		return
+	}
+	buf.astFormat(n.Lexpr)
+	buf.WriteString(" ")
+	buf.astFormat(n.Name)
+	buf.WriteString(" ")
+	buf.astFormat(n.Rexpr)
+}

internal/sql/ast/boolean.go

@@ -1,9 +1,22 @@
 package ast
 
+import "fmt"
+
 type Boolean struct {
 	Boolval bool
 }
 
 func (n *Boolean) Pos() int {
 	return 0
 }
+
+func (n *Boolean) Format(buf *TrackedBuffer) {
+	if n == nil {
+		return
+	}
+	if n.Boolval {
+		fmt.Fprintf(buf, "true")
+	} else {
+		fmt.Fprintf(buf, "false")
+	}
+}

internal/sql/ast/float.go

@@ -7,3 +7,10 @@ type Float struct {
 func (n *Float) Pos() int {
 	return 0
 }
+
+func (n *Float) Format(buf *TrackedBuffer) {
+	if n == nil {
+		return
+	}
+	buf.WriteString(n.Str)
+}

internal/sql/ast/insert_stmt.go

@@ -1,5 +1,7 @@
 package ast
 
+import "github.com/sqlc-dev/sqlc/internal/debug"
+
 type InsertStmt struct {
 	Relation         *RangeVar
 	Cols             *List
@@ -13,3 +15,29 @@ type InsertStmt struct {
 func (n *InsertStmt) Pos() int {
 	return 0
 }
+
+func (n *InsertStmt) Format(buf *TrackedBuffer) {
+	if n == nil {
+		return
+	}
+
+	if n.WithClause != nil {
+		buf.astFormat(n.WithClause)
+		buf.WriteString(" ")
+	}
+
+	buf.WriteString("INSERT INTO ")
+	if n.Relation != nil {
+		debug.Dump(n.Relation)
+		buf.astFormat(n.Relation)
+	}
+	if items(n.Cols) {
+		buf.WriteString(" (")
+		buf.astFormat(n.Cols)
+		buf.WriteString(") ")
+	}
+
+	if set(n.SelectStmt) {
+		buf.astFormat(n.SelectStmt)
+	}
+}

internal/sql/ast/locking_clause.go

@@ -9,3 +9,16 @@ type LockingClause struct {
 func (n *LockingClause) Pos() int {
 	return 0
 }
+
+func (n *LockingClause) Format(buf *TrackedBuffer) {
+	if n == nil {
+		return
+	}
+	buf.WriteString("FOR ")
+	switch n.Strength {
+	case 3:
+		buf.WriteString("SHARE")
+	case 5:
+		buf.WriteString("UPDATE")
+	}
+}

internal/sql/ast/param_ref.go

@@ -1,5 +1,7 @@
 package ast
 
+import "fmt"
+
 type ParamRef struct {
 	Number   int
 	Location int
@@ -9,3 +11,10 @@ type ParamRef struct {
 func (n *ParamRef) Pos() int {
 	return n.Location
 }
+
+func (n *ParamRef) Format(buf *TrackedBuffer) {
+	if n == nil {
+		return
+	}
+	fmt.Fprintf(buf, "$%d", n.Number)
+}

internal/sql/ast/print.go

@@ -37,3 +37,30 @@ func Format(n Node) string {
 	}
 	return tb.String()
 }
+
+func set(n Node) bool {
+	if n == nil {
+		return false
+	}
+	_, ok := n.(*TODO)
+	if ok {
+		return false
+	}
+	return true
+}
+
+func items(n *List) bool {
+	if n == nil {
+		return false
+	}
+	return len(n.Items) > 0
+}
+
+func todo(n *List) bool {
+	for _, item := range n.Items {
+		if _, ok := item.(*TODO); !ok {
+			return false
+		}
+	}
+	return true
+}

internal/sql/ast/range_subselect.go

@@ -9,3 +9,16 @@ type RangeSubselect struct {
 func (n *RangeSubselect) Pos() int {
 	return 0
 }
+
+func (n *RangeSubselect) Format(buf *TrackedBuffer) {
+	if n == nil {
+		return
+	}
+	buf.WriteString("(")
+	buf.astFormat(n.Subquery)
+	buf.WriteString(")")
+	if n.Alias != nil {
+		buf.WriteString(" ")
+		buf.astFormat(n.Alias)
+	}
+}