gosora/query_gen/spitter.go

package main

import "strings"
import "./lib"

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
		for _, key := range args[4].([]qgen.DBTableKey) {
			if key.Type == "primary" {
				expl := strings.Split(key.Columns, ",")
				if len(expl) > 1 {
					continue
				}
				found = key.Columns
			}
			if found != "" {
				table := args[0].(string)
				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")
}