internal/dinosql: Add COMMENT ON support

internal/catalog/build.go

@@ -370,6 +370,105 @@ func Update(c *pg.Catalog, stmt nodes.Node) error {
 			Arguments:  args,
 			ReturnType: join(n.ReturnType.Names, "."),
 		})
+
+	case nodes.CommentStmt:
+		switch n.Objtype {
+
+		case nodes.OBJECT_SCHEMA:
+			name := n.Object.(nodes.String).Str
+			schema, exists := c.Schemas[name]
+			if !exists {
+				return wrap(pg.ErrorSchemaDoesNotExist(name), raw.StmtLocation)
+			}
+			if n.Comment != nil {
+				schema.Comment = *n.Comment
+			} else {
+				schema.Comment = ""
+			}
+			c.Schemas[name] = schema
+
+		case nodes.OBJECT_TABLE:
+			fqn, err := ParseList(n.Object.(nodes.List))
+			if err != nil {
+				return err
+			}
+			schema, exists := c.Schemas[fqn.Schema]
+			if !exists {
+				return wrap(pg.ErrorSchemaDoesNotExist(fqn.Schema), raw.StmtLocation)
+			}
+			table, exists := schema.Tables[fqn.Rel]
+			if !exists {
+				return wrap(pg.ErrorRelationDoesNotExist(fqn.Rel), raw.StmtLocation)
+			}
+			if n.Comment != nil {
+				table.Comment = *n.Comment
+			} else {
+				table.Comment = ""
+			}
+			schema.Tables[fqn.Rel] = table
+
+		case nodes.OBJECT_COLUMN:
+			colParts := stringSlice(n.Object.(nodes.List))
+			var fqn pg.FQN
+			var col string
+			switch len(colParts) {
+			case 2:
+				col = colParts[1]
+				fqn = pg.FQN{Schema: "public", Rel: colParts[0]}
+			case 3:
+				col = colParts[2]
+				fqn = pg.FQN{Schema: colParts[0], Rel: colParts[1]}
+			case 4:
+				col = colParts[3]
+				fqn = pg.FQN{Catalog: colParts[0], Schema: colParts[1], Rel: colParts[2]}
+			default:
+				return fmt.Errorf("column specifier %q is not the proper format, expected '[catalog.][schema.]colname.tablename'", strings.Join(colParts, "."))
+			}
+			schema, exists := c.Schemas[fqn.Schema]
+			if !exists {
+				return wrap(pg.ErrorSchemaDoesNotExist(fqn.Schema), raw.StmtLocation)
+			}
+			table, exists := schema.Tables[fqn.Rel]
+			if !exists {
+				return wrap(pg.ErrorRelationDoesNotExist(fqn.Rel), raw.StmtLocation)
+			}
+			idx := -1
+			for i, c := range table.Columns {
+				if c.Name == col {
+					idx = i
+				}
+			}
+			if idx < 0 {
+				return wrap(pg.ErrorColumnDoesNotExist(table.Name, col), raw.StmtLocation)
+			}
+			if n.Comment != nil {
+				table.Columns[idx].Comment = *n.Comment
+			} else {
+				table.Columns[idx].Comment = ""
+			}
+
+		case nodes.OBJECT_TYPE:
+			fqn, err := ParseList(n.Object.(nodes.TypeName).Names)
+			if err != nil {
+				return err
+			}
+			schema, exists := c.Schemas[fqn.Schema]
+			if !exists {
+				return wrap(pg.ErrorSchemaDoesNotExist(fqn.Schema), raw.StmtLocation)
+			}
+			enum, exists := schema.Enums[fqn.Rel]
+			if !exists {
+				return wrap(pg.ErrorRelationDoesNotExist(fqn.Rel), raw.StmtLocation)
+			}
+			if n.Comment != nil {
+				enum.Comment = *n.Comment
+			} else {
+				enum.Comment = ""
+			}
+			schema.Enums[fqn.Rel] = enum
+
+		}
+
 	}
 	return nil
 }

internal/catalog/build_test.go

@@ -502,6 +502,71 @@ func TestUpdate(t *testing.T) {
 				},
 			},
 		},
+		{
+			`
+			CREATE SCHEMA foo;
+			CREATE TABLE foo.bar (baz text);
+			CREATE TYPE foo.bat AS ENUM ('bat');
+			COMMENT ON SCHEMA foo IS 'Schema comment';
+			COMMENT ON TABLE foo.bar IS 'Table comment';
+			COMMENT ON COLUMN foo.bar.baz IS 'Column comment';
+			COMMENT ON TYPE foo.bat IS 'Enum comment';
+			`,
+			pg.Catalog{
+				Schemas: map[string]pg.Schema{
+					"foo": {
+						Comment: "Schema comment",
+						Tables: map[string]pg.Table{
+							"bar": {
+								Comment: "Table comment",
+								Name:    "bar",
+								Columns: []pg.Column{
+									{
+										Name:     "baz",
+										DataType: "text",
+										Table:    pg.FQN{Schema: "foo", Rel: "bar"},
+										Comment:  "Column comment",
+									},
+								},
+							},
+						},
+						Enums: map[string]pg.Enum{"bat": {Comment: "Enum comment", Name: "bat", Vals: []string{"bat"}}},
+						Funcs: map[string][]pg.Function{},
+					},
+				},
+			},
+		},
+		{
+			`
+			CREATE TABLE bar (baz text);
+			CREATE TYPE bat AS ENUM ('bat');
+			COMMENT ON TABLE bar IS 'Table comment';
+			COMMENT ON COLUMN bar.baz IS 'Column comment';
+			COMMENT ON TYPE bat IS 'Enum comment';
+			`,
+			pg.Catalog{
+				Schemas: map[string]pg.Schema{
+					"public": {
+						Tables: map[string]pg.Table{
+							"bar": {
+								Comment: "Table comment",
+								Name:    "bar",
+								Columns: []pg.Column{
+									{
+										Name:     "baz",
+										DataType: "text",
+										Table:    pg.FQN{Schema: "public", Rel: "bar"},
+										Comment:  "Column comment",
+									},
+								},
+							},
+						},
+						Enums: map[string]pg.Enum{"bat": {Comment: "Enum comment", Name: "bat", Vals: []string{"bat"}}},
+						Funcs: map[string][]pg.Function{},
+					},
+				},
+			},
+		},
 	} {
 		test := tc
 		t.Run(strconv.Itoa(i), func(t *testing.T) {

internal/dinosql/gen.go

@@ -25,13 +25,15 @@ type GoConstant struct {
 
 type GoEnum struct {
 	Name      string
+	Comment   string
 	Constants []GoConstant
 }
 
 type GoField struct {
-	Name string
-	Type string
-	Tags map[string]string
+	Name    string
+	Type    string
+	Tags    map[string]string
+	Comment string
 }
 
 func (gf GoField) Tag() string {
@@ -47,9 +49,10 @@ func (gf GoField) Tag() string {
 }
 
 type GoStruct struct {
-	Table  *core.FQN
-	Name   string
-	Fields []GoField
+	Table   *core.FQN
+	Name    string
+	Fields  []GoField
+	Comment string
 }
 
 // TODO: Terrible name
@@ -396,7 +399,8 @@ func (r Result) Enums() []GoEnum {
 				enumName = name + "_" + enum.Name
 			}
 			e := GoEnum{
-				Name: r.structName(enumName),
+				Name:    r.structName(enumName),
+				Comment: enum.Comment,
 			}
 			for _, v := range enum.Vals {
 				name := ""
@@ -447,14 +451,16 @@ func (r Result) Structs() []GoStruct {
 				tableName = name + "_" + table.Name
 			}
 			s := GoStruct{
-				Table: &core.FQN{Schema: name, Rel: table.Name},
-				Name:  inflection.Singular(r.structName(tableName)),
+				Table:   &core.FQN{Schema: name, Rel: table.Name},
+				Name:    inflection.Singular(r.structName(tableName)),
+				Comment: table.Comment,
 			}
 			for _, column := range table.Columns {
 				s.Fields = append(s.Fields, GoField{
-					Name: r.structName(column.Name),
-					Type: r.goType(column),
-					Tags: map[string]string{"json:": column.Name},
+					Name:    r.structName(column.Name),
+					Type:    r.goType(column),
+					Tags:    map[string]string{"json:": column.Name},
+					Comment: column.Comment,
 				})
 			}
 			structs = append(structs, s)
@@ -860,6 +866,7 @@ import (
 )
 
 {{range .Enums}}
+{{if .Comment}}// {{.Comment}}{{end}}
 type {{.Name}} string
 
 const (
@@ -875,7 +882,11 @@ func (e *{{.Name}}) Scan(src interface{}) error {
 {{end}}
 
 {{range .Structs}}
+{{if .Comment}}// {{.Comment}}{{end}}
 type {{.Name}} struct { {{- range .Fields}}
+  {{- if .Comment}}
+  // {{.Comment}}{{else}}
+  {{- end}}
   {{.Name}} {{.Type}} {{if $.EmitJSONTags}}{{$.Q}}{{.Tag}}{{$.Q}}{{end}}
   {{- end}}
 }

internal/dinosql/testdata/ondeck/schema/0002_venue.sql

@@ -1,4 +1,5 @@
 CREATE TYPE status AS ENUM ('open', 'closed');
+COMMENT ON TYPE status IS 'Venues can be either open or closed';
 
 CREATE TABLE venues (
     id               SERIAL primary key,
@@ -11,4 +12,7 @@ CREATE TABLE venues (
     spotify_playlist varchar      not null,
     songkick_id      text,
     tags             text[]
-)
+);
+COMMENT ON TABLE venues IS 'Venues are places where muisc happens';
+COMMENT ON COLUMN venues.slug IS 'This value appears in public URLs';
+

internal/pg/catalog.go

@@ -91,39 +91,44 @@ func (c Catalog) LookupFunctionN(fqn FQN, argn int) (Function, error) {
 }
 
 type Schema struct {
-	Name   string
-	Tables map[string]Table
-	Enums  map[string]Enum
-	Funcs  map[string][]Function
+	Name    string
+	Tables  map[string]Table
+	Enums   map[string]Enum
+	Funcs   map[string][]Function
+	Comment string
 }
 
 type Table struct {
 	ID      FQN
 	Name    string
 	Columns []Column
+	Comment string
 }
 
 type Column struct {
 	Name     string
 	DataType string
 	NotNull  bool
 	IsArray  bool
+	Comment  string
 
 	// XXX: Figure out what PostgreSQL calls `foo.id`
 	Scope string
 	Table FQN
 }
 
 type Enum struct {
-	Name string
-	Vals []string
+	Name    string
+	Vals    []string
+	Comment string
 }
 
 type Function struct {
 	Name       string
 	ArgN       int
 	Arguments  []Argument // not recorded for builtins
 	ReturnType string
+	Comment    string
 }
 
 type Argument struct {