Add support for insert stmts

Author: kyleconroy

go.mod

@@ -4,5 +4,6 @@ go 1.12
 
 require (
 	github.com/davecgh/go-spew v1.1.1
+	github.com/google/go-cmp v0.3.0
 	github.com/lfittl/pg_query_go v1.0.0
 )

go.sum

@@ -4,6 +4,8 @@ github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/elgris/sqrl v0.0.0-20181124135704-90ecf730640a h1:VRAv/FIe+jL6t/kHB43hrwhUIyP2/cq+vC/9bdS8v4o=
 github.com/elgris/sqrl v0.0.0-20181124135704-90ecf730640a/go.mod h1:hQPgqeM4LmbfKCaBkcedRq5y1yfb8Qb8iYdbuNjE4FU=
+github.com/google/go-cmp v0.3.0 h1:crn/baboCvb5fXaQ0IJ1SGTsTVrWpDsCWC8EGETZijY=
+github.com/google/go-cmp v0.3.0/go.mod h1:8QqcDgzrUqlUb/G2PQTWiueGozuR1884gddMywk6iLU=
 github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo=
 github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
 github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=

parser.go

@@ -59,10 +59,10 @@ func parse(s *postgres.Schema, tree pg.ParsetreeList) {
 			for _, elt := range n.TableElts.Items {
 				switch n := elt.(type) {
 				case nodes.ColumnDef:
-					// spew.Dump(n)
 					// log.Printf("not null: %t", n.IsNotNull)
 					table.Columns = append(table.Columns, postgres.Column{
 						Name:    *n.Colname,
+						Type:    join(n.TypeName.Names, "."),
 						GoName:  structName(*n.Colname),
 						NotNull: isNotNull(n),
 					})
@@ -75,6 +75,16 @@ func parse(s *postgres.Schema, tree pg.ParsetreeList) {
 	}
 }
 
+func join(list nodes.List, sep string) string {
+	items := []string{}
+	for _, item := range list.Items {
+		if n, ok := item.(nodes.String); ok {
+			items = append(items, n.Str)
+		}
+	}
+	return strings.Join(items, sep)
+}
+
 func isNotNull(n nodes.ColumnDef) bool {
 	if n.IsNotNull {
 		return true
@@ -101,6 +111,7 @@ type Query struct {
 	SQL        string
 	Args       []Arg
 	Table      postgres.Table
+	ReturnType string
 }
 
 type Result struct {
@@ -179,6 +190,7 @@ func parseFuncs(s *postgres.Schema, r *Result, source string, tree pg.ParsetreeL
 
 			tab := getTable(s, t)
 			r.Queries[i].Table = tab
+			r.Queries[i].ReturnType = tab.GoName
 			r.Queries[i].Args = parseArgs(tab, refs)
 			r.Queries[i].SQL = strings.Replace(rawSQL, "*", strings.Join(c, ", "), 1)
 		case nodes.DeleteStmt:
@@ -189,8 +201,20 @@ func parseFuncs(s *postgres.Schema, r *Result, source string, tree pg.ParsetreeL
 
 			tab := getTable(s, t)
 			r.Queries[i].Table = tab
+			r.Queries[i].ReturnType = tab.GoName
 			r.Queries[i].Args = parseArgs(tab, refs)
 			r.Queries[i].SQL = rawSQL
+		case nodes.InsertStmt:
+			t := tableName(n)
+			c := columnNames(s, t)
+			rawSQL, _ := pluckQuery(source, raw)
+			refs := extractArgs(n)
+
+			tab := getTable(s, t)
+			r.Queries[i].Table = tab
+			r.Queries[i].ReturnType = tab.GoName
+			r.Queries[i].Args = parseArgs(tab, refs)
+			r.Queries[i].SQL = strings.Replace(rawSQL, "*", strings.Join(c, ", "), 1)
 		default:
 			log.Printf("%T\n", n)
 		}
@@ -221,6 +245,16 @@ func findRefs(r []paramRef, parent, n nodes.Node) []paramRef {
 		r = findRefs(r, n.WhereClause, nil)
 		r = findRefs(r, n.LimitCount, nil)
 		r = findRefs(r, n.LimitOffset, nil)
+	case nodes.InsertStmt:
+		switch s := n.SelectStmt.(type) {
+		case nodes.SelectStmt:
+			for _, vl := range s.ValuesLists {
+				for i, v := range vl {
+					// TODO: Index error
+					r = findRefs(r, n.Cols.Items[i], v)
+				}
+			}
+		}
 	case nodes.BoolExpr:
 		for _, item := range n.Args.Items {
 			r = findRefs(r, item, nil)
@@ -265,6 +299,15 @@ func parseArgs(t postgres.Table, args []paramRef) []Arg {
 					panic("unknown column: " + key)
 				}
 			}
+		case nodes.ResTarget:
+			key := *n.Name
+			if typ, ok := typeMap[key]; ok {
+				a = append(a, Arg{Name: key, Type: typ})
+			} else {
+				panic("unknown column: " + key)
+			}
+		case nodes.ParamRef:
+			a = append(a, Arg{Name: "_", Type: "interface{}"})
 		default:
 			panic(fmt.Sprintf("unsupported type: %T", n))
 		}
@@ -296,6 +339,8 @@ func tableName(n nodes.Node) string {
 		}
 	case nodes.DeleteStmt:
 		return *n.Relation.Relname
+	case nodes.InsertStmt:
+		return *n.Relation.Relname
 	}
 	return ""
 }
@@ -304,11 +349,12 @@ var hh = `package {{.Package}}
 import (
 	"context"
 	"database/sql"
+	"time"
 )
 
 {{range .Schema.Tables}}
 type {{.GoName}} struct { {{- range .Columns}}
-  {{.GoName}} {{if .NotNull }}string{{else}}sql.NullString{{end}}
+  {{.GoName}} {{.GoType}}
   {{- end}}
 }
 {{end}}
@@ -358,7 +404,7 @@ const {{.QueryName}} = {{$.Q}}{{.SQL}}
 {{$.Q}}
 
 {{if eq .Type ":one"}}
-func (q *Queries) {{.MethodName}}(ctx context.Context, {{range .Args}}{{.Name}} {{.Type}},{{end}}) ({{.Table.GoName}}, error) {
+func (q *Queries) {{.MethodName}}(ctx context.Context, {{range .Args}}{{.Name}} {{.Type}},{{end}}) ({{.ReturnType}}, error) {
 	var row *sql.Row
 	switch {
 	case q.{{.StmtName}} != nil && q.tx != nil:
@@ -375,7 +421,7 @@ func (q *Queries) {{.MethodName}}(ctx context.Context, {{range .Args}}{{.Name}}
 {{end}}
 
 {{if eq .Type ":many"}}
-func (q *Queries) {{.MethodName}}(ctx context.Context, {{range .Args}}{{.Name}} {{.Type}},{{end}}) ([]{{.Table.GoName}}, error) {
+func (q *Queries) {{.MethodName}}(ctx context.Context, {{range .Args}}{{.Name}} {{.Type}},{{end}}) ([]{{.ReturnType}}, error) {
 	var rows *sql.Rows
 	var err error
 	switch {

parser_test.go

@@ -1,12 +1,12 @@
 package strongdb
 
 import (
-	"fmt"
 	"io/ioutil"
 	"log"
 	"path/filepath"
 	"testing"
 
+	"github.com/google/go-cmp/cmp"
 	pg "github.com/lfittl/pg_query_go"
 	nodes "github.com/lfittl/pg_query_go/nodes"
 )
@@ -88,8 +88,8 @@ func TestParseSchema(t *testing.T) {
 		log.Fatal(err)
 	}
 
-	if source != string(blob) {
-		t.Errorf("output differs")
-		fmt.Println(source)
+	if diff := cmp.Diff(source, string(blob)); diff != "" {
+		t.Errorf("genreated code differed (-want +got):\n%s", diff)
+		t.Log(source)
 	}
 }

postgres/schema.go

@@ -16,3 +16,21 @@ type Column struct {
 	Type    string
 	NotNull bool
 }
+
+func (c Column) GoType() string {
+	// {{.GoName}} {{if .NotNull }}string{{else}}sql.NullString{{end}}
+	switch c.Type {
+	case "text":
+		if c.NotNull {
+			return "string"
+		} else {
+			return "sql.NullString"
+		}
+	case "serial":
+		return "int"
+	case "pg_catalog.timestamp":
+		return "time.Time"
+	default:
+		return "interface{}"
+	}
+}

testdata/ondeck/db.go

@@ -3,6 +3,7 @@ package ondeck
 import (
 	"context"
 	"database/sql"
+	"time"
 )
 
 type City struct {
@@ -11,6 +12,8 @@ type City struct {
 }
 
 type Venue struct {
+	ID              int
+	CreatedAt       time.Time
 	Slug            string
 	Name            string
 	City            sql.NullString
@@ -32,12 +35,18 @@ func New(db dbtx) *Queries {
 func Prepare(ctx context.Context, db dbtx) (*Queries, error) {
 	q := Queries{db: db}
 	var err error
+	if q.createCity, err = db.PrepareContext(ctx, createCity); err != nil {
+		return nil, err
+	}
 	if q.deleteVenue, err = db.PrepareContext(ctx, deleteVenue); err != nil {
 		return nil, err
 	}
 	if q.getCity, err = db.PrepareContext(ctx, getCity); err != nil {
 		return nil, err
 	}
+	if q.getVenue, err = db.PrepareContext(ctx, getVenue); err != nil {
+		return nil, err
+	}
 	if q.listCities, err = db.PrepareContext(ctx, listCities); err != nil {
 		return nil, err
 	}
@@ -51,8 +60,10 @@ type Queries struct {
 	db dbtx
 
 	tx          *sql.Tx
+	createCity  *sql.Stmt
 	deleteVenue *sql.Stmt
 	getCity     *sql.Stmt
+	getVenue    *sql.Stmt
 	listCities  *sql.Stmt
 	listVenues  *sql.Stmt
 }
@@ -61,13 +72,40 @@ func (q *Queries) WithTx(tx *sql.Tx) *Queries {
 	return &Queries{
 		tx:          tx,
 		db:          tx,
+		createCity:  q.createCity,
 		deleteVenue: q.deleteVenue,
 		getCity:     q.getCity,
+		getVenue:    q.getVenue,
 		listCities:  q.listCities,
 		listVenues:  q.listVenues,
 	}
 }
 
+const createCity = `-- name: CreateCity :one
+INSERT INTO city (
+    name,
+    slug
+) VALUES (
+    $1,
+    $2
+) RETURNING slug, name
+`
+
+func (q *Queries) CreateCity(ctx context.Context, name string, slug string) (City, error) {
+	var row *sql.Row
+	switch {
+	case q.createCity != nil && q.tx != nil:
+		row = q.tx.StmtContext(ctx, q.createCity).QueryRowContext(ctx, name, slug)
+	case q.createCity != nil:
+		row = q.createCity.QueryRowContext(ctx, name, slug)
+	default:
+		row = q.db.QueryRowContext(ctx, createCity, name, slug)
+	}
+	i := City{}
+	err := row.Scan(&i.Slug, &i.Name)
+	return i, err
+}
+
 const deleteVenue = `-- name: DeleteVenue :exec
 DELETE FROM venue
 WHERE slug = $1
@@ -107,6 +145,27 @@ func (q *Queries) GetCity(ctx context.Context, slug string) (City, error) {
 	return i, err
 }
 
+const getVenue = `-- name: GetVenue :one
+SELECT id, created_at, slug, name, city, spotify_playlist, songkick_id
+FROM venue
+WHERE slug = $1 AND city = $2
+`
+
+func (q *Queries) GetVenue(ctx context.Context, slug string, city string) (Venue, error) {
+	var row *sql.Row
+	switch {
+	case q.getVenue != nil && q.tx != nil:
+		row = q.tx.StmtContext(ctx, q.getVenue).QueryRowContext(ctx, slug, city)
+	case q.getVenue != nil:
+		row = q.getVenue.QueryRowContext(ctx, slug, city)
+	default:
+		row = q.db.QueryRowContext(ctx, getVenue, slug, city)
+	}
+	i := Venue{}
+	err := row.Scan(&i.ID, &i.CreatedAt, &i.Slug, &i.Name, &i.City, &i.SpotifyPlaylist, &i.SongkickID)
+	return i, err
+}
+
 const listCities = `-- name: ListCities :many
 SELECT slug, name
 FROM city
@@ -146,7 +205,7 @@ func (q *Queries) ListCities(ctx context.Context) ([]City, error) {
 }
 
 const listVenues = `-- name: ListVenues :many
-SELECT slug, name, city, spotify_playlist, songkick_id
+SELECT id, created_at, slug, name, city, spotify_playlist, songkick_id
 FROM venue
 WHERE city = $1
 ORDER BY name
@@ -170,7 +229,7 @@ func (q *Queries) ListVenues(ctx context.Context, city string) ([]Venue, error)
 	items := []Venue{}
 	for rows.Next() {
 		i := Venue{}
-		if err := rows.Scan(&i.Slug, &i.Name, &i.City, &i.SpotifyPlaylist, &i.SongkickID); err != nil {
+		if err := rows.Scan(&i.ID, &i.CreatedAt, &i.Slug, &i.Name, &i.City, &i.SpotifyPlaylist, &i.SongkickID); err != nil {
 			return nil, err
 		}
 		items = append(items, i)

testdata/ondeck/query/queries.sql

@@ -17,3 +17,17 @@ ORDER BY name;
 -- name: DeleteVenue :exec
 DELETE FROM venue
 WHERE slug = $1;
+
+-- name: GetVenue :one
+SELECT *
+FROM venue
+WHERE slug = $1 AND city = $2;
+
+-- name: CreateCity :one
+INSERT INTO city (
+    name,
+    slug
+) VALUES (
+    $1,
+    $2
+) RETURNING *;

testdata/ondeck/schema/0002_venue.sql

@@ -1,5 +1,7 @@
 CREATE TABLE venue (
-    slug             text primary key,
+    id               SERIAL primary key,
+    created_at       TIMESTAMP NOT NULL DEFAULT NOW(),
+    slug             text not null,
     name             text not null,
     city             text references city(slug),
     spotify_playlist text not null,