github.com/dolthub/dolt/go@v0.40.5-0.20240520175717-68db7794bea6/libraries/doltcore/dtestutils/data.go (about)

     1  // Copyright 2019 Dolthub, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package dtestutils
    16  
    17  import (
    18  	"github.com/google/uuid"
    19  
    20  	"github.com/dolthub/dolt/go/libraries/doltcore/row"
    21  	"github.com/dolthub/dolt/go/libraries/doltcore/schema"
    22  	"github.com/dolthub/dolt/go/store/types"
    23  )
    24  
    25  var uuids = []uuid.UUID{
    26  	uuid.Must(uuid.Parse("00000000-0000-0000-0000-000000000000")),
    27  	uuid.Must(uuid.Parse("00000000-0000-0000-0000-000000000001")),
    28  	uuid.Must(uuid.Parse("00000000-0000-0000-0000-000000000002"))}
    29  var names = []string{"Bill Billerson", "John Johnson", "Rob Robertson"}
    30  var ages = []uint64{32, 25, 21}
    31  var titles = []string{"Senior Dufus", "Dufus", ""}
    32  var maritalStatus = []bool{true, false, false}
    33  
    34  const (
    35  	IdTag uint64 = iota
    36  	NameTag
    37  	AgeTag
    38  	IsMarriedTag
    39  	TitleTag
    40  	NextTag // leave last
    41  )
    42  
    43  const (
    44  	IndexName = "idx_name"
    45  )
    46  
    47  // Schema returns the schema for the `people` test table.
    48  func Schema() (schema.Schema, error) {
    49  	var typedColColl = schema.NewColCollection(
    50  		schema.NewColumn("id", IdTag, types.StringKind, true, schema.NotNullConstraint{}),
    51  		schema.NewColumn("name", NameTag, types.StringKind, false, schema.NotNullConstraint{}),
    52  		schema.NewColumn("age", AgeTag, types.UintKind, false, schema.NotNullConstraint{}),
    53  		schema.NewColumn("is_married", IsMarriedTag, types.IntKind, false, schema.NotNullConstraint{}),
    54  		schema.NewColumn("title", TitleTag, types.StringKind, false),
    55  	)
    56  	sch := schema.MustSchemaFromCols(typedColColl)
    57  
    58  	_, err := sch.Indexes().AddIndexByColTags(IndexName, []uint64{NameTag}, nil, schema.IndexProperties{IsUnique: false, Comment: ""})
    59  	if err != nil {
    60  		return nil, err
    61  	}
    62  
    63  	_, err = sch.Checks().AddCheck("test-check", "age < 123", true)
    64  	if err != nil {
    65  		return nil, err
    66  	}
    67  
    68  	return sch, err
    69  }
    70  
    71  // RowsAndSchema returns the schema and rows for the `people` test table.
    72  func RowsAndSchema() ([]row.Row, schema.Schema, error) {
    73  	sch, err := Schema()
    74  	if err != nil {
    75  		return nil, nil, err
    76  	}
    77  
    78  	rows := make([]row.Row, len(uuids))
    79  	for i := 0; i < len(uuids); i++ {
    80  		married := types.Int(0)
    81  		if maritalStatus[i] {
    82  			married = types.Int(1)
    83  		}
    84  		taggedVals := row.TaggedValues{
    85  			IdTag:        types.String(uuids[i].String()),
    86  			NameTag:      types.String(names[i]),
    87  			AgeTag:       types.Uint(ages[i]),
    88  			TitleTag:     types.String(titles[i]),
    89  			IsMarriedTag: married,
    90  		}
    91  
    92  		r, err := row.New(types.Format_Default, sch, taggedVals)
    93  
    94  		if err != nil {
    95  			panic(err)
    96  		}
    97  
    98  		rows[i] = r
    99  	}
   100  
   101  	return rows, sch, err
   102  }
   103  
   104  // MustTuple contructs a types.Tuple for a slice of types.Values.
   105  func MustTuple(vals ...types.Value) types.Tuple {
   106  	tup, err := types.NewTuple(types.Format_Default, vals...)
   107  	if err != nil {
   108  		panic(err)
   109  	}
   110  	return tup
   111  }