github.com/square/finch@v0.0.0-20240412205204-6530c03e2b96/data/generator.go (about) 1 // Copyright 2024 Block, Inc. 2 3 package data 4 5 import ( 6 "fmt" 7 "math/rand" 8 "strconv" 9 "time" 10 11 "github.com/square/finch" 12 ) 13 14 type ValueFunc func(RunCount) []interface{} 15 16 // Generator generates data values for a data key (@d). 17 type Generator interface { 18 Format() (uint, string) 19 Copy() Generator 20 Values(RunCount) []interface{} 21 Scan(any interface{}) error 22 Name() string 23 } 24 25 func init() { 26 rand.Seed(time.Now().UnixNano()) 27 /* 28 Generator names here must match factory.Make switch cases below 29 */ 30 // Integer 31 Register("int", f) 32 Register("int-gaps", f) 33 Register("int-range", f) 34 Register("int-range-seq", f) 35 Register("auto-inc", f) 36 // String 37 Register("str-fill-az", f) 38 // ID 39 Register("xid", f) 40 Register("client-id", f) 41 // Column 42 Register("column", f) 43 } 44 45 // Factory makes data generators from day keys (@d). 46 type Factory interface { 47 Make(name, dataKey string, params map[string]string) (Generator, error) 48 } 49 50 type factory struct{} 51 52 var f = &factory{} 53 54 func (f factory) Make(name, dataKey string, params map[string]string) (Generator, error) { 55 finch.Debug("Make %s %s", name, dataKey) 56 var g Generator 57 var err error 58 switch name { 59 // Integer 60 case "int": 61 g, err = NewInt(params) 62 case "int-gaps": 63 g, err = NewIntGaps(params) 64 case "int-range": 65 g, err = NewIntRange(params) 66 case "int-range-seq": 67 g, err = NewIntRangeSeq(params) 68 case "auto-inc": 69 g, err = NewAutoInc(params) 70 // String 71 case "str-fill-az": 72 g, err = NewStrFillAz(params) 73 // ID 74 case "xid": 75 g = NewXid() 76 case "client-id": 77 g, err = NewClientId(params) 78 // Column 79 case "column": 80 g = NewColumn(params) 81 default: 82 err = fmt.Errorf("built-in data factory cannot make %s data generator", name) 83 } 84 if err != nil { 85 return nil, err 86 } 87 return g, nil 88 } 89 90 var r = &repo{ 91 f: map[string]Factory{}, 92 } 93 94 type repo struct { 95 f map[string]Factory 96 } 97 98 func Register(name string, f Factory) error { 99 _, have := r.f[name] 100 if have { 101 return fmt.Errorf("data.Generator %s already registered", name) 102 } 103 r.f[name] = f 104 return nil 105 } 106 107 // Make makes a data generator by name with the given generator-specific params. 108 func Make(name, dataKey string, params map[string]string) (Generator, error) { 109 f, have := r.f[name] 110 if !have { 111 return nil, fmt.Errorf("data.Generator %s not registered", name) 112 } 113 return f.Make(name, dataKey, params) 114 } 115 116 func int64From(params map[string]string, key string, n *int64, required bool) error { 117 s, ok := params[key] 118 if !ok { 119 if required { 120 return fmt.Errorf("%s required", key) 121 } 122 return nil 123 } 124 i, err := strconv.ParseInt(s, 10, 64) 125 if err != nil { 126 return fmt.Errorf("invalid %s=%s: %s", key, s, err) 127 } 128 *n = i 129 return nil 130 }