gosora/cmd/query_gen/spitter.go

package main

import "strings"
import "github.com/Azareal/Gosora/query_gen"

type PrimaryKeySpitter struct {
    keys map[string]string
}

func NewPrimaryKeySpitter() *PrimaryKeySpitter {
    return &PrimaryKeySpitter{make(map[string]string)}
}

func (spit *PrimaryKeySpitter) Hook(name string, args ...interface{}) error {
    if name == "CreateTableStart" {
        var found string
        var table = args[0].(*qgen.DBInstallTable)
        for _, key := range table.Keys {
            if key.Type == "primary" {
                expl := strings.Split(key.Columns, ",")
                if len(expl) > 1 {
                    continue
                }
                found = key.Columns
            }
            if found != "" {
                table := table.Name
                spit.keys[table] = found
            }
        }
    }
    return nil
}

func (spit *PrimaryKeySpitter) Write() error {
    out := `// Generated by Gosora's Query Generator. DO NOT EDIT.
package main

var dbTablePrimaryKeys = map[string]string{
`
    for table, key := range spit.keys {
        out += "\t\"" + table + "\":\"" + key + "\",\n"
    }
    return writeFile("./gen_tables.go", out+"}\n")
}