sqlingo/sqlingo-gen/generator.go

package main

import (
	"database/sql"
	"errors"
	"fmt"
	"go/format"
	"os"
	"strconv"
	"strings"
	"unicode"
)

const (
	sqlingoGeneratorVersion = 2
)

type schemaFetcher interface {
	GetDatabaseName() (dbName string, err error)
	GetTableNames() (tableNames []string, err error)
	GetFieldDescriptors(tableName string) ([]fieldDescriptor, error)
	QuoteIdentifier(identifier string) string
}

type fieldDescriptor struct {
	Name      string
	Type      string
	Size      int
	Unsigned  bool
	AllowNull bool
	Comment   string
}

func convertCase(s string) (result string) {
	nextCharShouldBeUpperCase := true
	for _, ch := range s {
		if ch == '_' {
			nextCharShouldBeUpperCase = true
		} else {
			if nextCharShouldBeUpperCase {
				result += string(unicode.ToUpper(ch))
				nextCharShouldBeUpperCase = false
			} else {
				result += string(ch)
			}
		}
	}
	return
}

func getType(fieldDescriptor fieldDescriptor) (goType string, fieldClass string, err error) {
	switch strings.ToLower(fieldDescriptor.Type) {
	case "tinyint":
		goType = "int8"
		fieldClass = "NumberField"
	case "smallint":
		goType = "int16"
		fieldClass = "NumberField"
	case "int", "mediumint":
		goType = "int32"
		fieldClass = "NumberField"
	case "bigint", "integer":
		goType = "int64"
		fieldClass = "NumberField"
	case "float", "double", "decimal", "real":
		goType = "float64"
		fieldClass = "NumberField"
	case "char", "varchar", "text", "tinytext", "mediumtext", "longtext", "enum", "datetime", "date", "time", "timestamp", "json", "numeric", "character varying":
		goType = "string"
		fieldClass = "StringField"
	case "binary", "varbinary", "blob", "tinyblob", "mediumblob", "longblob":
		// TODO: use []byte ?
		goType = "string"
		fieldClass = "StringField"
	case "bit":
		if fieldDescriptor.Size == 1 {
			goType = "bool"
			fieldClass = "BooleanField"
		} else {
			goType = "string"
			fieldClass = "StringField"
		}
	default:
		err = fmt.Errorf("unknown field type %s", fieldDescriptor.Type)
		return
	}
	if fieldDescriptor.Unsigned {
		goType = "u" + goType
	}
	if fieldDescriptor.AllowNull {
		goType = "*" + goType
	}
	return
}

func getSchemaFetcherFactory(driverName string) func(db *sql.DB) schemaFetcher {
	switch driverName {
	case "mysql":
		return newMySQLSchemaFetcher
	case "sqlite3":
		return newSQLite3SchemaFetcher
	case "postgres":
		return newPostgresSchemaFetcher
	default:
		_, _ = fmt.Fprintln(os.Stderr, "unsupported driver "+driverName)
		os.Exit(2)
		return nil
	}
}

func generate(driverName string, dataSourceName string, tableNames []string) (string, error) {
	schemaFetcherFactory := getSchemaFetcherFactory(driverName)

	db, err := sql.Open(driverName, dataSourceName)
	if err != nil {
		return "", err
	}

	schemaFetcher := schemaFetcherFactory(db)

	dbName, err := schemaFetcher.GetDatabaseName()
	if err != nil {
		return "", err
	}

	if dbName == "" {
		return "", errors.New("no database selected")
	}

	code := "// This file is generated by sqlingo (https://github.com/lqs/sqlingo)\n"
	code += "// DO NOT EDIT.\n\n"
	code += "package " + dbName + "_dsl\n"
	code += "import . \"github.com/lqs/sqlingo\"\n\n"

	code += "type sqlingoRuntimeAndGeneratorVersionsShouldBeTheSame uint32\n\n"

	sqlingoGeneratorVersionString := strconv.Itoa(sqlingoGeneratorVersion)
	code += "const _ = sqlingoRuntimeAndGeneratorVersionsShouldBeTheSame(SqlingoRuntimeVersion - " + sqlingoGeneratorVersionString + ")\n"
	code += "const _ = sqlingoRuntimeAndGeneratorVersionsShouldBeTheSame(" + sqlingoGeneratorVersionString + " - SqlingoRuntimeVersion)\n\n"

	code += "type table interface {\n"
	code += "\tTable\n"
	code += "}\n\n"

	code += "type numberField interface {\n"
	code += "\tNumberField\n"
	code += "}\n\n"

	code += "type stringField interface {\n"
	code += "\tStringField\n"
	code += "}\n\n"

	code += "type booleanField interface {\n"
	code += "\tBooleanField\n"
	code += "}\n\n"

	if len(tableNames) == 0 {
		tableNames, err = schemaFetcher.GetTableNames()
		if err != nil {
			return "", err
		}
	}

	for _, tableName := range tableNames {
		println("Generating", tableName)
		tableCode, err := generateTable(schemaFetcher, tableName)
		if err != nil {
			return "", err
		}
		code += tableCode
	}
	code += generateGetTable(tableNames)
	codeOut, err := format.Source([]byte(code))
	if err != nil {
		return "", err
	}
	return string(codeOut), nil
}

func generateGetTable(tableNames []string) string {
	code := "func GetTable(name string) Table {\n"
	code += "\tswitch name {\n"
	for _, tableName := range tableNames {
		code += "\tcase " + strconv.Quote(tableName) + ": return " + convertCase(tableName) + "\n"
	}
	code += "\tdefault: return nil\n"
	code += "\t}\n"
	code += "}\n\n"
	return code
}

func generateTable(schemaFetcher schemaFetcher, tableName string) (string, error) {
	fieldDescriptors, err := schemaFetcher.GetFieldDescriptors(tableName)
	if err != nil {
		return "", err
	}

	className := convertCase(tableName)
	tableStructName := "t" + className
	tableObjectName := "o" + className

	modelClassName := className + "Model"

	tableLines := ""
	modelLines := ""
	objectLines := "\ttable: " + tableObjectName + ",\n\n"
	fieldCaseLines := ""
	classLines := ""

	fields := ""
	fieldsSQL := ""
	fullFieldsSQL := ""
	values := ""

	for _, fieldDescriptor := range fieldDescriptors {

		goName := convertCase(fieldDescriptor.Name)
		goType, fieldClass, err := getType(fieldDescriptor)
		if err != nil {
			return "", err
		}

		privateFieldClass := string(fieldClass[0]+'a'-'A') + fieldClass[1:]

		commentLine := ""
		if fieldDescriptor.Comment != "" {
			commentLine = "\t// " + strings.ReplaceAll(fieldDescriptor.Comment, "\n", " ") + "\n"
		}

		fieldStructName := "f" + className + goName

		tableLines += commentLine
		tableLines += "\t" + goName + " " + fieldStructName + "\n"

		modelLines += commentLine
		modelLines += "\t" + goName + " " + goType + "\n"

		objectLines += commentLine
		objectLines += "\t" + goName + ": " + fieldStructName + "{"
		objectLines += "New" + fieldClass + "(" + tableObjectName + ", " + strconv.Quote(fieldDescriptor.Name) + ")},\n"

		fieldCaseLines += "\tcase " + strconv.Quote(fieldDescriptor.Name) + ": return t." + goName + "\n"

		classLines += "type " + fieldStructName + " struct{ " + privateFieldClass + " }\n"

		fields += "t." + goName + ", "

		if fieldsSQL != "" {
			fieldsSQL += ", "
		}
		fieldsSQL += schemaFetcher.QuoteIdentifier(fieldDescriptor.Name)

		if fullFieldsSQL != "" {
			fullFieldsSQL += ", "
		}
		fullFieldsSQL += schemaFetcher.QuoteIdentifier(tableName) + "." + schemaFetcher.QuoteIdentifier(fieldDescriptor.Name)

		values += "m." + goName + ", "
	}
	code := ""
	code += "type " + tableStructName + " struct {\n\ttable\n\n"
	code += tableLines
	code += "}\n\n"

	code += classLines

	code += "var " + tableObjectName + " = NewTable(" + strconv.Quote(tableName) + ")\n"
	code += "var " + className + " = " + tableStructName + "{\n"
	code += objectLines
	code += "}\n\n"

	code += "func (t t" + className + ") GetFields() []Field {\n"
	code += "\treturn []Field{" + fields + "}\n"
	code += "}\n\n"

	code += "func (t t" + className + ") GetFieldByName(name string) Field {\n"
	code += "\tswitch name {\n"
	code += fieldCaseLines
	code += "\tdefault: return nil\n"
	code += "\t}\n"
	code += "}\n\n"

	code += "func (t t" + className + ") GetFieldsSQL() string {\n"
	code += "\treturn " + strconv.Quote(fieldsSQL) + "\n"
	code += "}\n\n"

	code += "func (t t" + className + ") GetFullFieldsSQL() string {\n"
	code += "\treturn " + strconv.Quote(fullFieldsSQL) + "\n"
	code += "}\n\n"

	code += "type " + modelClassName + " struct {\n"
	code += modelLines
	code += "}\n\n"

	code += "func (m " + modelClassName + ") GetTable() Table {\n"
	code += "\treturn " + className + "\n"
	code += "}\n\n"

	code += "func (m " + modelClassName + ") GetValues() []interface{} {\n"
	code += "\treturn []interface{}{" + values + "}\n"
	code += "}\n\n"
	return code, nil
}