ast: Add support for CREATE / DROP SCHEMA (#389)

Author: kyleconroy

Makefile

@@ -1,5 +1,5 @@
 build:
-	go build ./...
+	go build --tags=exp ./...
 
 test:
 	go test --tags=exp ./...

internal/endtoend/testdata/experimental_elephant/query.sql

@@ -14,10 +14,13 @@ DROP TABLE bar;
 DROP TABLE IF EXISTS bar;
 DROP TABLE IF EXISTS baz;
 
-
 CREATE TABLE baz (name text);
 ALTER TABLE baz ADD COLUMN email text;
 ALTER TABLE baz ADD COLUMN bar text;
 ALTER TABLE baz DROP COLUMN IF EXISTS foo;
 ALTER TABLE baz DROP COLUMN bar;
 ALTER TABLE baz DROP COLUMN IF EXISTS bar;
+
+CREATE SCHEMA baz;
+CREATE TABLE baz.foo ();
+DROP SCHEMA baz;

internal/postgresql/parse.go

@@ -223,20 +223,43 @@ func translate(node nodes.Node) (ast.Node, error) {
 		}
 		return stmt, nil
 
+	case nodes.CreateSchemaStmt:
+		return &ast.CreateSchemaStmt{
+			Name:        n.Schemaname,
+			IfNotExists: n.IfNotExists,
+		}, nil
+
 	case nodes.DropStmt:
-		drop := &ast.DropTableStmt{
-			IfExists: n.MissingOk,
-		}
-		for _, obj := range n.Objects.Items {
-			if n.RemoveType == nodes.OBJECT_TABLE {
+		switch n.RemoveType {
+
+		case nodes.OBJECT_SCHEMA:
+			drop := &ast.DropSchemaStmt{
+				MissingOk: n.MissingOk,
+			}
+			for _, obj := range n.Objects.Items {
+				val, ok := obj.(nodes.String)
+				if !ok {
+					return nil, fmt.Errorf("nodes.DropStmt: unknown type in objects list: %T", obj)
+				}
+				drop.Schemas = append(drop.Schemas, &ast.String{Str: val.Str})
+			}
+			return drop, nil
+
+		case nodes.OBJECT_TABLE:
+			drop := &ast.DropTableStmt{
+				IfExists: n.MissingOk,
+			}
+			for _, obj := range n.Objects.Items {
 				name, err := parseTableName(obj)
 				if err != nil {
 					return nil, err
 				}
 				drop.Tables = append(drop.Tables, name)
 			}
+			return drop, nil
+
 		}
-		return drop, nil
+		return nil, nil
 
 	default:
 		return nil, nil

internal/sql/ast/ast.go

@@ -66,6 +66,15 @@ func (n *CreateEnumStmt) Pos() int {
 	return 0
 }
 
+type CreateSchemaStmt struct {
+	Name        *string
+	IfNotExists bool
+}
+
+func (n *CreateSchemaStmt) Pos() int {
+	return 0
+}
+
 type CreateTableStmt struct {
 	IfNotExists bool
 	Name        *TableName
@@ -76,6 +85,15 @@ func (n *CreateTableStmt) Pos() int {
 	return 0
 }
 
+type DropSchemaStmt struct {
+	Schemas   []*String
+	MissingOk bool
+}
+
+func (n *DropSchemaStmt) Pos() int {
+	return 0
+}
+
 type DropTableStmt struct {
 	IfExists bool
 	Tables   []*TableName

internal/sql/catalog/catalog.go

@@ -2,6 +2,7 @@ package catalog
 
 import (
 	"errors"
+	"fmt"
 
 	"github.com/kyleconroy/sqlc/internal/sql/ast"
 )
@@ -23,8 +24,12 @@ func Build(stmts []ast.Statement) (*Catalog, error) {
 			err = c.alterTable(n)
 		case *ast.CreateEnumStmt:
 			err = c.createEnum(n)
+		case *ast.CreateSchemaStmt:
+			err = c.createSchema(n)
 		case *ast.CreateTableStmt:
 			err = c.createTable(n)
+		case *ast.DropSchemaStmt:
+			err = c.dropSchema(n)
 		case *ast.DropTableStmt:
 			err = c.dropTable(n)
 		}
@@ -203,6 +208,20 @@ func (c *Catalog) createEnum(stmt *ast.CreateEnumStmt) error {
 	return nil
 }
 
+func (c *Catalog) createSchema(stmt *ast.CreateSchemaStmt) error {
+	if stmt.Name == nil {
+		return fmt.Errorf("create schema: empty name")
+	}
+	if _, err := c.getSchema(*stmt.Name); err == nil {
+		if !stmt.IfNotExists {
+			// return wrap(pg.ErrorSchemaAlreadyExists(name), raw.StmtLocation)
+			return ErrRelationAlreadyExists
+		}
+	}
+	c.Schemas = append(c.Schemas, &Schema{Name: *stmt.Name})
+	return nil
+}
+
 func (c *Catalog) createTable(stmt *ast.CreateTableStmt) error {
 	ns := stmt.Name.Schema
 	if ns == "" {
@@ -231,6 +250,26 @@ func (c *Catalog) createTable(stmt *ast.CreateTableStmt) error {
 	return nil
 }
 
+func (c *Catalog) dropSchema(stmt *ast.DropSchemaStmt) error {
+	// TODO: n^2 in the worst-case
+	for _, name := range stmt.Schemas {
+		idx := -1
+		for i := range c.Schemas {
+			if c.Schemas[i].Name == name.Str {
+				idx = i
+			}
+		}
+		if idx == -1 {
+			if stmt.MissingOk {
+				continue
+			}
+			return ErrSchemaNotFound
+		}
+		c.Schemas = append(c.Schemas[:idx], c.Schemas[idx+1:]...)
+	}
+	return nil
+}
+
 func (c *Catalog) dropTable(stmt *ast.DropTableStmt) error {
 	for _, name := range stmt.Tables {
 		ns := name.Schema