github.com/cloudfoundry/postgres-release/src/acceptance-tests@v0.0.0-20240511030151-872bdd2e0dba/testing/helpers/postgres_load_sets.go (about)

     1  package helpers
     2  
     3  import (
     4  	"fmt"
     5  	"strings"
     6  	"time"
     7  
     8  	pq "github.com/lib/pq"
     9  )
    10  
    11  type PGLoadTable struct {
    12  	Name        string
    13  	ColumnNames []string
    14  	ColumnTypes []string
    15  	SampleRow   []interface{}
    16  	NumRows     int
    17  }
    18  
    19  var RowSamples = [][]interface{}{
    20  	[]interface{}{"character varying not null", "short_string"},
    21  	[]interface{}{"character varying not null", "long_string_xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx"},
    22  	[]interface{}{"integer", 0},
    23  }
    24  
    25  type LoadType struct {
    26  	NumTables  int
    27  	NumColumns int
    28  	NumRows    int
    29  }
    30  
    31  var Test1Load = LoadType{NumTables: 1, NumColumns: 1, NumRows: 1}
    32  var Test2Load = LoadType{NumTables: 2, NumColumns: 4, NumRows: 5}
    33  var SmallLoad = LoadType{NumTables: 2, NumColumns: 10, NumRows: 50}
    34  var MediumLoad = LoadType{NumTables: 10, NumColumns: 10, NumRows: 100}
    35  var LargeLoad = LoadType{NumTables: 100, NumColumns: 10, NumRows: 20000}
    36  
    37  func GetSampleLoadWithPrefix(loadType LoadType, prefix string) []PGLoadTable {
    38  	if loadType.NumTables <= 0 {
    39  		return nil
    40  	}
    41  	var result = make([]PGLoadTable, loadType.NumTables)
    42  	rowTypesNum := len(RowSamples)
    43  	for i := 0; i < loadType.NumTables; i++ {
    44  		var table = PGLoadTable{
    45  			Name:        fmt.Sprintf("%s_%d", prefix, i),
    46  			ColumnNames: make([]string, loadType.NumColumns),
    47  			ColumnTypes: make([]string, loadType.NumColumns),
    48  			SampleRow:   make([]interface{}, loadType.NumColumns),
    49  			NumRows:     loadType.NumRows,
    50  		}
    51  		for j := 0; j < loadType.NumColumns; j++ {
    52  			idx := j % rowTypesNum
    53  			table.ColumnNames[j] = fmt.Sprintf("column%d", j)
    54  			table.ColumnTypes[j] = RowSamples[idx][0].(string)
    55  			table.SampleRow[j] = RowSamples[idx][1]
    56  		}
    57  		result[i] = table
    58  	}
    59  	return result
    60  }
    61  
    62  func (table PGLoadTable) PrepareCreate() string {
    63  	columns := make([]string, len(table.ColumnNames))
    64  	for idx, name := range table.ColumnNames {
    65  		var dataType string
    66  		if idx > len(table.ColumnTypes)-1 {
    67  			dataType = "character varying"
    68  		} else {
    69  			dataType = table.ColumnTypes[idx]
    70  		}
    71  		columns[idx] = fmt.Sprintf("%s %s", name, dataType)
    72  	}
    73  	return fmt.Sprintf("CREATE TABLE %s (%s);", table.Name, strings.Join(columns, ",\n"))
    74  }
    75  
    76  func (table PGLoadTable) PrepareCreateIndex() string {
    77  	return fmt.Sprintf("CREATE INDEX %s_index ON %s USING hash (%s) ;", table.Name, table.Name, table.ColumnNames[0])
    78  }
    79  
    80  func (table PGLoadTable) PrepareStatement() string {
    81  	var result string
    82  	if len(table.ColumnNames) > 0 && table.Name != "" {
    83  		result = pq.CopyIn(table.Name, table.ColumnNames...)
    84  	}
    85  	return result
    86  }
    87  func (table PGLoadTable) PrepareRow(rowIdx int) []interface{} {
    88  	var result []interface{}
    89  	if len(table.ColumnNames) > 0 {
    90  		for idx, value := range table.SampleRow {
    91  			if idx > len(table.ColumnNames)-1 {
    92  				break
    93  			}
    94  			var out interface{}
    95  			switch value.(type) {
    96  			case string:
    97  				out = fmt.Sprintf("%s%d", value, rowIdx)
    98  			case int32, int64, int:
    99  				out = rowIdx
   100  			case time.Time:
   101  				out = time.Now()
   102  			default:
   103  				out = value
   104  			}
   105  			result = append(result, out)
   106  		}
   107  	}
   108  	return result
   109  }