github.com/Azareal/Gosora@v0.0.0-20210729070923-553e66b59003/cmd/query_gen/spitter.go (about)

     1  package main
     2  
     3  import "strings"
     4  import "github.com/Azareal/Gosora/query_gen"
     5  
     6  type PrimaryKeySpitter struct {
     7  	keys map[string]string
     8  }
     9  
    10  func NewPrimaryKeySpitter() *PrimaryKeySpitter {
    11  	return &PrimaryKeySpitter{make(map[string]string)}
    12  }
    13  
    14  func (spit *PrimaryKeySpitter) Hook(name string, args ...interface{}) error {
    15  	if name == "CreateTableStart" {
    16  		var found string
    17  		var table = args[0].(*qgen.DBInstallTable)
    18  		for _, key := range table.Keys {
    19  			if key.Type == "primary" {
    20  				expl := strings.Split(key.Columns, ",")
    21  				if len(expl) > 1 {
    22  					continue
    23  				}
    24  				found = key.Columns
    25  			}
    26  			if found != "" {
    27  				table := table.Name
    28  				spit.keys[table] = found
    29  			}
    30  		}
    31  	}
    32  	return nil
    33  }
    34  
    35  func (spit *PrimaryKeySpitter) Write() error {
    36  	out := `// Generated by Gosora's Query Generator. DO NOT EDIT.
    37  package main
    38  
    39  var dbTablePrimaryKeys = map[string]string{
    40  `
    41  	for table, key := range spit.keys {
    42  		out += "\t\"" + table + "\":\"" + key + "\",\n"
    43  	}
    44  	return writeFile("./gen_tables.go", out+"}\n")
    45  }