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

     1  // Copyright 2024 Block, Inc.
     2  
     3  package data
     4  
     5  import (
     6  	"bytes"
     7  	"database/sql"
     8  
     9  	"github.com/square/finch"
    10  )
    11  
    12  // Column is a special Generator that is used to save (Scan) values from rows
    13  // or insert ID, then return those values (Value) to other statements.
    14  type Column struct {
    15  	quoteValue bool
    16  	val        interface{}
    17  	bytes      *bytes.Buffer
    18  	useBytes   bool
    19  }
    20  
    21  var _ Generator = &Column{}
    22  var _ sql.Scanner = &Column{}
    23  
    24  func NewColumn(params map[string]string) *Column {
    25  	return &Column{
    26  		quoteValue: finch.Bool(params["quote-value"]),
    27  	}
    28  }
    29  
    30  func (g *Column) Name() string { return "column" }
    31  
    32  func (g *Column) Format() (uint, string) {
    33  	if g.quoteValue {
    34  		return 1, "'%v'"
    35  	}
    36  	return 1, "%v"
    37  }
    38  
    39  func (g *Column) Copy() Generator {
    40  	return &Column{
    41  		quoteValue: g.quoteValue,
    42  	}
    43  }
    44  
    45  func (g *Column) Scan(any interface{}) error {
    46  	// @todo column type won't change, so maybe sync.Once to set val or bytes
    47  	// will make this more efficient?
    48  	switch any.(type) {
    49  	case []byte:
    50  		g.useBytes = true // is reference; copy bytes
    51  		if g.bytes == nil {
    52  			g.bytes = bytes.NewBuffer(make([]byte, len(any.([]byte))))
    53  		}
    54  		g.bytes.Reset()
    55  		g.bytes.Write(any.([]byte))
    56  	default:
    57  		g.useBytes = false // not reference; copy value
    58  		g.val = any
    59  	}
    60  	return nil
    61  }
    62  
    63  func (g *Column) Values(_ RunCount) []interface{} {
    64  	if g.useBytes {
    65  		return []interface{}{g.bytes.String()}
    66  	}
    67  	return []interface{}{g.val}
    68  }
    69  
    70  // --------------------------------------------------------------------------
    71  
    72  var Noop = noop{}
    73  
    74  type noop struct{}
    75  
    76  var _ Generator = Noop
    77  var _ sql.Scanner = Noop
    78  
    79  func (g noop) Name() string                    { return "no-op" }
    80  func (g noop) Format() (uint, string)          { return 0, "" }
    81  func (g noop) Copy() Generator                 { return Noop }
    82  func (g noop) Values(_ RunCount) []interface{} { return nil }
    83  func (g noop) Scan(any interface{}) error      { return nil }