main.go

package genddl

import (
	"flag"
	"fmt"
	"go/ast"
	"go/types"
	"log"
	"os"
	"path/filepath"
	"sort"
	"strings"

	"golang.org/x/tools/go/packages"
)

type Field struct {
	ColumnDef string
}

type Table struct {
	Name       string
	Fields     []Field
	PrimaryKey string
}

func Run(from string) {
	fromdir := filepath.Dir(from)

	var schemadir, outpath, driverName, tableCollate string
	tableMapOptionArgs := &tableMapOption{}
	flag.StringVar(&schemadir, "schemadir", fromdir, "schema declaretion directory")
	flag.StringVar(&outpath, "outpath", "", "schema target path")
	flag.StringVar(&driverName, "driver", "mysql", "target driver name. support mysql, pg, sqlite3")
	flag.BoolVar(&tableMapOptionArgs.innerIndexDef, "innerindex", false, "Placement of index definition. If this specified, the definition was placement inner of `create table`")
	flag.BoolVar(&tableMapOptionArgs.uniqueWithName, "uniquename", false, "Provides a name for the definition of a unique index.")
	flag.StringVar(&tableCollate, "tablecollate", "", "Provides a collate for the definition of tables.")
	flag.BoolVar(&tableMapOptionArgs.foreignKeyWithName, "foreignkeyname", false, "Provides a name for the definition of a foreign-key.")
	flag.BoolVar(&tableMapOptionArgs.outerForeignKey, "outerforeignkey", false, "Placement of foreign key definition. If this specified, the definition was placement end of DDL file.")
	flag.BoolVar(&tableMapOptionArgs.outerUniqueKey, "outeruniquekey", false, "Placement of unique key definition. If this specified, the definition was placement outer of CREATE TABLE.")
	flag.BoolVar(&tableMapOptionArgs.withoutDropTable, "withoutdroptable", false, "If this specified, the DDL file does not contain DROP TABLE statement.")

	flag.Parse()

	var dialect Dialect
	switch driverName {
	case "mysql":
		dialect = MysqlDialect{Collate: tableCollate}
	case "pg":
		dialect = PostgresqlDialect{Collate: tableCollate}
		// It is not supported by PostgreSQL that foreign key definition with name
		tableMapOptionArgs.foreignKeyWithName = false
		// It is not supported by PostgreSQL that unique index definition with name
		tableMapOptionArgs.uniqueWithName = false
	case "sqlite3":
		dialect = Sqlite3Dialect{}
		// It is not supported by SQLite that placement of index definition inner CREATE TABLE
		tableMapOptionArgs.innerIndexDef = false
		// It is not supported by SQLite that unique index definition with name
		tableMapOptionArgs.uniqueWithName = false
		// It is not supported by SQLite that foreign key definition placement outer of CREATE TABLE
		tableMapOptionArgs.outerForeignKey = false
	default:
		log.Fatalf("undefined driver name: %s", driverName)
	}

	tr, err := retrieveTables(schemadir)
	if err != nil {
		log.Fatalf("parse and retrieve table error: %s", err)
	}

	file, err := os.Create(outpath)
	if err != nil {
		log.Fatal("invalid outpath error:", err)
	}
	tablesMap := map[*ast.StructType]string{}
	var tableNames []string
	for tableName, st := range tr.tables {
		tablesMap[st] = tableName
		tableNames = append(tableNames, tableName)
	}
	sort.Strings(tableNames)
	file.WriteString("-- generated by github.com/mackee/go-genddl. DO NOT EDIT!!!\n")
	endOfDDLFileIndexes := make([]indexer, 0, len(tableNames))
	for _, tableName := range tableNames {
		st := tr.tables[tableName]
		funcs := tr.funcs[st]
		tableMap, err := NewTableMap(tableName, st, funcs, tablesMap, tr.ti, tableMapOptionArgs)
		if err != nil {
			log.Fatalf("failed to create table map: %s", err)
		}
		if tableMap != nil {
			file.WriteString("\n")
			if err := tableMap.WriteDDL(file, dialect, tableMapOptionArgs); err != nil {
				log.Fatalf("failed to write DDL: %s", err)
			}
			endOfDDLFileIndexes = append(endOfDDLFileIndexes, tableMap.EndOfDDLFileIndexes...)
		}
	}
	viewNames := make([]string, 0, len(tr.views))
	for viewName := range tr.views {
		viewNames = append(viewNames, viewName)
	}
	sort.Strings(viewNames)
	for _, viewName := range viewNames {
		v := tr.views[viewName]
		vm, err := NewViewMap(newViewMapInput{
			name:  viewName,
			st:    v,
			funcs: tr.funcs[v],
			ti:    tr.ti,
		})
		if err != nil {
			log.Fatalf("[ERROR] failed to create view map: %s", err)
		}
		if vm == nil {
			log.Println("[ERROR] skip view:", viewName)
			continue
		}
		vm.WriteDDL(file, dialect)
	}
	for _, index := range endOfDDLFileIndexes {
		if _, err := file.WriteString(index.Index(dialect, tablesMap)); err != nil {
			log.Fatalf("failed to write index: %s", err)
		}
		if _, err := file.WriteString(";\n"); err != nil {
			log.Fatalf("failed to write index: %s", err)
		}
	}
}

var typeNameStructMap = map[string]*ast.StructType{}

type retrieveTablesResult struct {
	tables map[string]*ast.StructType
	views  map[string]*ast.StructType
	funcs  map[*ast.StructType][]*ast.FuncDecl
	ti     *types.Info
}

func retrieveTables(schemadir string) (*retrieveTablesResult, error) {
	path, err := filepath.Abs(schemadir)
	if err != nil {
		return nil, err
	}

	conf := &packages.Config{
		Mode: packages.NeedCompiledGoFiles |
			packages.NeedSyntax |
			packages.NeedTypes |
			packages.NeedTypesInfo,
		Dir: path,
	}
	pkgs, err := packages.Load(conf)
	if err != nil {
		return nil, fmt.Errorf("retrieveTables: fail to parse from dir: dir=%s, error=%w", schemadir, err)
	}

	var decls []ast.Decl
	var ti *types.Info
	for _, pkg := range pkgs {
		ti = pkg.TypesInfo
		for _, file := range pkg.Syntax {
			decls = append(decls, file.Decls...)
		}
	}

	tables := map[string]*ast.StructType{}
	views := map[string]*ast.StructType{}
	funcs := []*ast.FuncDecl{}
	funcMap := map[*ast.StructType][]*ast.FuncDecl{}
	for _, decl := range decls {
		if funcDecl, ok := decl.(*ast.FuncDecl); ok {
			funcs = append(funcs, funcDecl)
		}
		if genDecl, ok := decl.(*ast.GenDecl); ok {
			if genDecl.Doc == nil {
				continue
			}

			for _, comment := range genDecl.Doc.List {
				if a := trimAnnotation(comment.Text); a.annotationType != annotationTypeNone {
					tableName := a.tableName
					spec := genDecl.Specs[0]
					ts, ok := spec.(*ast.TypeSpec)
					if !ok {
						continue
					}
					st, ok := ts.Type.(*ast.StructType)
					if !ok {
						continue
					}

					switch a.annotationType {
					case annotationTypeTable:
						tables[tableName] = st
					case annotationTypeView:
						views[tableName] = st
					}
					funcMap[st] = make([]*ast.FuncDecl, 0)
					typeNameStructMap[ts.Name.Name] = st
					break
				}
			}
		}
	}

	for _, funcDecl := range funcs {
		if funcDecl.Recv.NumFields() < 0 {
			continue
		}
		if funcDecl.Recv == nil || len(funcDecl.Recv.List) == 0 {
			continue
		}
		recv := funcDecl.Recv.List[0]
		ident, ok := recv.Type.(*ast.Ident)
		if !ok {
			continue
		}
		if ident.Obj == nil || ident.Obj.Decl == nil {
			continue
		}
		ts, ok := ident.Obj.Decl.(*ast.TypeSpec)
		if !ok {
			continue
		}
		st, ok := ts.Type.(*ast.StructType)
		if !ok {
			continue
		}
		if funcs, ok := funcMap[st]; ok {
			funcMap[st] = append(funcs, funcDecl)
		}
	}

	return &retrieveTablesResult{
		tables: tables,
		views:  views,
		funcs:  funcMap,
		ti:     ti,
	}, nil
}

type annotationType string

const (
	annotationTypeNone  annotationType = "none"
	annotationTypeTable annotationType = "table"
	annotationTypeView  annotationType = "view"
)

type annotation struct {
	tableName      string
	annotationType annotationType
}

func trimAnnotation(comment string) annotation {
	prefixes := map[string]annotationType{
		"//+table:":       annotationTypeTable,
		"// +table:":      annotationTypeTable,
		"//genddl:table ": annotationTypeTable,
		"//genddl:view ":  annotationTypeView,
	}
	for prefix, at := range prefixes {
		if trimmed := strings.TrimPrefix(comment, prefix); trimmed != comment {
			trimmed = strings.TrimSpace(trimmed)
			return annotation{
				tableName:      trimmed,
				annotationType: at,
			}
		}
	}
	return annotation{
		tableName:      comment,
		annotationType: annotationTypeNone,
	}
}