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  }