github.com/kunlun-qilian/sqlx/v3@v3.0.0/generator/utils.go (about)

     1  package generator
     2  
     3  import (
     4  	"go/types"
     5  	"reflect"
     6  	"regexp"
     7  	"strings"
     8  
     9  	"github.com/go-courier/codegen"
    10  	"github.com/kunlun-qilian/sqlx/v3/builder"
    11  )
    12  
    13  var (
    14  	defRegexp = regexp.MustCompile(`@def ([^\n]+)`)
    15  	relRegexp = regexp.MustCompile(`@rel ([^\n]+)`)
    16  )
    17  
    18  type Keys struct {
    19  	Primary       []string
    20  	Indexes       builder.Indexes
    21  	UniqueIndexes builder.Indexes
    22  	Partition     []string
    23  }
    24  
    25  func (ks *Keys) PatchUniqueIndexesWithSoftDelete(softDeleteField string) {
    26  	if len(ks.UniqueIndexes) > 0 {
    27  		for name, fieldNames := range ks.UniqueIndexes {
    28  			ks.UniqueIndexes[name] = stringUniq(append(fieldNames, softDeleteField))
    29  		}
    30  	}
    31  }
    32  
    33  func (ks *Keys) Bind(table *builder.Table) {
    34  	if len(ks.Primary) > 0 {
    35  		key := &builder.Key{
    36  			Name:     "primary",
    37  			IsUnique: true,
    38  			Def:      *builder.ParseIndexDef(ks.Primary...),
    39  		}
    40  
    41  		_ = key.Def.TableExpr(table)
    42  
    43  		table.AddKey(key)
    44  	}
    45  
    46  	if len(ks.UniqueIndexes) > 0 {
    47  		for indexNameAndMethod, fieldNames := range ks.UniqueIndexes {
    48  			indexName, method := builder.ResolveIndexNameAndMethod(indexNameAndMethod)
    49  
    50  			key := &builder.Key{
    51  				Name:     indexName,
    52  				Method:   method,
    53  				IsUnique: true,
    54  				Def:      *builder.ParseIndexDef(fieldNames...),
    55  			}
    56  
    57  			_ = key.Def.TableExpr(table)
    58  
    59  			table.AddKey(key)
    60  		}
    61  	}
    62  
    63  	if len(ks.Indexes) > 0 {
    64  		for indexNameAndMethod, fieldNames := range ks.Indexes {
    65  			indexName, method := builder.ResolveIndexNameAndMethod(indexNameAndMethod)
    66  
    67  			key := &builder.Key{
    68  				Name:   indexName,
    69  				Method: method,
    70  				Def:    *builder.ParseIndexDef(fieldNames...),
    71  			}
    72  
    73  			_ = key.Def.TableExpr(table)
    74  
    75  			table.AddKey(key)
    76  		}
    77  	}
    78  }
    79  
    80  func parseColRelFromComment(doc string) (string, []string) {
    81  	others := make([]string, 0)
    82  
    83  	rel := ""
    84  
    85  	for _, line := range strings.Split(doc, "\n") {
    86  		if len(line) == 0 {
    87  			continue
    88  		}
    89  
    90  		matches := relRegexp.FindAllStringSubmatch(line, 1)
    91  
    92  		if matches == nil {
    93  			others = append(others, line)
    94  			continue
    95  		}
    96  
    97  		if len(matches) == 1 {
    98  			rel = matches[0][1]
    99  		}
   100  	}
   101  
   102  	return rel, others
   103  }
   104  
   105  func parseKeysFromDoc(doc string) (*Keys, []string) {
   106  	ks := &Keys{}
   107  
   108  	others := make([]string, 0)
   109  
   110  	for _, line := range strings.Split(doc, "\n") {
   111  		if len(line) == 0 {
   112  			continue
   113  		}
   114  
   115  		matches := defRegexp.FindAllStringSubmatch(line, -1)
   116  
   117  		if matches == nil {
   118  			others = append(others, line)
   119  			continue
   120  		}
   121  
   122  		for _, subMatch := range matches {
   123  			if len(subMatch) == 2 {
   124  				def := builder.ParseIndexDefine(subMatch[1])
   125  
   126  				switch def.Kind {
   127  				case "primary":
   128  					ks.Primary = def.ToDefs()
   129  				case "unique_index":
   130  					if ks.UniqueIndexes == nil {
   131  						ks.UniqueIndexes = builder.Indexes{}
   132  					}
   133  					ks.UniqueIndexes[def.ID()] = def.ToDefs()
   134  				case "index":
   135  					if ks.Indexes == nil {
   136  						ks.Indexes = builder.Indexes{}
   137  					}
   138  					ks.Indexes[def.ID()] = def.ToDefs()
   139  				case "partition":
   140  					ks.Partition = append([]string{def.Name}, def.ToDefs()...)
   141  				}
   142  			}
   143  		}
   144  	}
   145  
   146  	return ks, others
   147  }
   148  
   149  func toDefaultTableName(name string) string {
   150  	return codegen.LowerSnakeCase("t_" + name)
   151  }
   152  
   153  func forEachStructField(structType *types.Struct, fn func(fieldVar *types.Var, columnName string, tagValue string)) {
   154  	for i := 0; i < structType.NumFields(); i++ {
   155  		field := structType.Field(i)
   156  		tag := structType.Tag(i)
   157  		if field.Exported() {
   158  			structTag := reflect.StructTag(tag)
   159  			tagValue, exists := structTag.Lookup("db")
   160  			if exists {
   161  				if tagValue != "-" {
   162  					fn(field, builder.GetColumnName(field.Name(), tagValue), tagValue)
   163  				}
   164  			} else if field.Anonymous() {
   165  				if nextStructType, ok := field.Type().Underlying().(*types.Struct); ok {
   166  					forEachStructField(nextStructType, fn)
   167  				}
   168  				continue
   169  			}
   170  		}
   171  	}
   172  }
   173  
   174  func stringPartition(list []string, checker func(item string, i int) bool) ([]string, []string) {
   175  	newLeftList := make([]string, 0)
   176  	newRightList := make([]string, 0)
   177  	for i, item := range list {
   178  		if checker(item, i) {
   179  			newLeftList = append(newLeftList, item)
   180  		} else {
   181  			newRightList = append(newRightList, item)
   182  		}
   183  	}
   184  	return newLeftList, newRightList
   185  }
   186  
   187  func stringFilter(list []string, checker func(item string, i int) bool) []string {
   188  	newList, _ := stringPartition(list, checker)
   189  	return newList
   190  }
   191  
   192  func stringUniq(list []string) (result []string) {
   193  	strMap := make(map[string]bool)
   194  	for _, str := range list {
   195  		strMap[str] = true
   196  	}
   197  
   198  	for i := range list {
   199  		str := list[i]
   200  		if _, ok := strMap[str]; ok {
   201  			delete(strMap, str)
   202  			result = append(result, str)
   203  		}
   204  	}
   205  	return
   206  }