github.com/square/finch@v0.0.0-20240412205204-6530c03e2b96/data/id.go (about)

     1  // Copyright 2024 Block, Inc.
     2  
     3  package data
     4  
     5  import (
     6  	"fmt"
     7  	"strings"
     8  
     9  	"github.com/rs/xid"
    10  
    11  	"github.com/square/finch"
    12  )
    13  
    14  // Xid implments the xid data generator.
    15  type Xid struct {
    16  	val string
    17  }
    18  
    19  var _ Generator = &Xid{}
    20  
    21  func NewXid() *Xid {
    22  	return &Xid{}
    23  }
    24  
    25  func (g *Xid) Name() string               { return "xid" }
    26  func (g *Xid) Format() (uint, string)     { return 1, "'%s'" }
    27  func (g *Xid) Scan(any interface{}) error { return nil }
    28  
    29  func (g *Xid) Copy() Generator {
    30  	return NewXid()
    31  }
    32  
    33  func (g *Xid) Values(c RunCount) []interface{} {
    34  	return []interface{}{xid.New().String()}
    35  }
    36  
    37  // --------------------------------------------------------------------------
    38  
    39  // ClientId implments the client-id data generator.
    40  type ClientId struct {
    41  	ids []byte
    42  }
    43  
    44  var _ Generator = &ClientId{}
    45  
    46  func NewClientId(params map[string]string) (*ClientId, error) {
    47  	ids := []byte{}
    48  	if len(params) == 0 {
    49  		ids = append(ids, CLIENT) // just client ID by default
    50  	} else {
    51  		csvIds := params["ids"]
    52  		if csvIds != "" {
    53  			runlevels := strings.Split(csvIds, ",")
    54  			for _, r := range runlevels {
    55  				switch r {
    56  				case finch.SCOPE_STATEMENT:
    57  					ids = append(ids, STATEMENT)
    58  				case finch.SCOPE_TRX:
    59  					ids = append(ids, TRX)
    60  				case finch.SCOPE_ITER:
    61  					ids = append(ids, ITER)
    62  				case "conn":
    63  					ids = append(ids, CONN)
    64  				case finch.SCOPE_CLIENT:
    65  					ids = append(ids, CLIENT)
    66  				case finch.SCOPE_CLIENT_GROUP:
    67  					ids = append(ids, CLIENT_GROUP)
    68  				case finch.SCOPE_EXEC_GROUP:
    69  					ids = append(ids, EXEC_GROUP)
    70  				default:
    71  					return nil, fmt.Errorf("invalid scope: %s", r)
    72  				}
    73  			}
    74  		}
    75  	}
    76  	return &ClientId{ids: ids}, nil
    77  }
    78  
    79  func (g *ClientId) Name() string               { return "client-id" }
    80  func (g *ClientId) Format() (uint, string)     { return uint(len(g.ids)), "%d" }
    81  func (g *ClientId) Scan(any interface{}) error { return nil }
    82  func (g *ClientId) Copy() Generator            { return &ClientId{ids: g.ids} }
    83  
    84  func (g *ClientId) Values(rc RunCount) []interface{} {
    85  	// This data generator can be shared (e.g. client-group scope), so the
    86  	// caller can be a different client each time, so don't save the last value
    87  	// and don't return something like g.val (a struct field) because Go slices
    88  	// are pointers, so each call would overwrite the return value of the previous.
    89  	val := make([]interface{}, len(g.ids)) // a new slice each call
    90  	for i := range g.ids {
    91  		val[i] = rc[g.ids[i]]
    92  	}
    93  	return val
    94  }