github.com/hasnat/dolt/go@v0.0.0-20210628190320-9eb5d843fbb7/libraries/doltcore/dtestutils/schema.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  	"context"
    19  	"math"
    20  	"testing"
    21  
    22  	"github.com/google/go-cmp/cmp"
    23  	"github.com/stretchr/testify/require"
    24  
    25  	"github.com/dolthub/dolt/go/libraries/doltcore/doltdb"
    26  	"github.com/dolthub/dolt/go/libraries/doltcore/env"
    27  	"github.com/dolthub/dolt/go/libraries/doltcore/row"
    28  	"github.com/dolthub/dolt/go/libraries/doltcore/schema"
    29  	"github.com/dolthub/dolt/go/libraries/doltcore/schema/encoding"
    30  	"github.com/dolthub/dolt/go/libraries/doltcore/table"
    31  	"github.com/dolthub/dolt/go/libraries/doltcore/table/editor"
    32  	"github.com/dolthub/dolt/go/libraries/doltcore/table/typed/noms"
    33  	"github.com/dolthub/dolt/go/store/types"
    34  )
    35  
    36  // CreateSchema returns a schema from the columns given, panicking on any errors.
    37  func CreateSchema(columns ...schema.Column) schema.Schema {
    38  	colColl := schema.NewColCollection(columns...)
    39  	return schema.MustSchemaFromCols(colColl)
    40  }
    41  
    42  // Creates a row with the schema given, having the values given. Starts at tag 0 and counts up.
    43  func NewRow(sch schema.Schema, values ...types.Value) row.Row {
    44  	taggedVals := make(row.TaggedValues)
    45  	for i := range values {
    46  		taggedVals[uint64(i)] = values[i]
    47  	}
    48  	r, err := row.New(types.Format_Default, sch, taggedVals)
    49  
    50  	if err != nil {
    51  		panic(err)
    52  	}
    53  
    54  	return r
    55  }
    56  
    57  // AddColumnToSchema returns a new schema by adding the given column to the given schema. Will panic on an invalid
    58  // schema, e.g. tag collision.
    59  func AddColumnToSchema(sch schema.Schema, col schema.Column) schema.Schema {
    60  	columns := sch.GetAllCols()
    61  	columns = columns.Append(col)
    62  	return schema.MustSchemaFromCols(columns)
    63  }
    64  
    65  // RemoveColumnFromSchema returns a new schema with the given tag missing, but otherwise identical. At least one
    66  // primary column must remain.
    67  func RemoveColumnFromSchema(sch schema.Schema, tagToRemove uint64) schema.Schema {
    68  	var newCols []schema.Column
    69  	err := sch.GetAllCols().Iter(func(tag uint64, col schema.Column) (stop bool, err error) {
    70  		if tag != tagToRemove {
    71  			newCols = append(newCols, col)
    72  		}
    73  		return false, nil
    74  	})
    75  
    76  	if err != nil {
    77  		panic(err)
    78  	}
    79  
    80  	columns := schema.NewColCollection(newCols...)
    81  	return schema.MustSchemaFromCols(columns)
    82  }
    83  
    84  // Compares two noms Floats for approximate equality
    85  var FloatComparer = cmp.Comparer(func(x, y types.Float) bool {
    86  	return math.Abs(float64(x)-float64(y)) < .001
    87  })
    88  
    89  var TimestampComparer = cmp.Comparer(func(x, y types.Timestamp) bool {
    90  	return x.Equals(y)
    91  })
    92  
    93  // CreateTestTable creates a new test table with the name, schema, and rows given.
    94  func CreateTestTable(t *testing.T, dEnv *env.DoltEnv, tableName string, sch schema.Schema, rs ...row.Row) {
    95  	imt := table.NewInMemTable(sch)
    96  
    97  	for _, r := range rs {
    98  		_ = imt.AppendRow(r)
    99  	}
   100  
   101  	ctx := context.Background()
   102  	vrw := dEnv.DoltDB.ValueReadWriter()
   103  	rd := table.NewInMemTableReader(imt)
   104  	wr := noms.NewNomsMapCreator(ctx, vrw, sch)
   105  
   106  	_, _, err := table.PipeRows(ctx, rd, wr, false)
   107  	require.NoError(t, err)
   108  	err = rd.Close(ctx)
   109  	require.NoError(t, err)
   110  	err = wr.Close(ctx)
   111  	require.NoError(t, err)
   112  
   113  	schVal, err := encoding.MarshalSchemaAsNomsValue(ctx, vrw, sch)
   114  	require.NoError(t, err)
   115  	empty, err := types.NewMap(ctx, vrw)
   116  	require.NoError(t, err)
   117  	tbl, err := doltdb.NewTable(ctx, vrw, schVal, wr.GetMap(), empty, nil)
   118  	require.NoError(t, err)
   119  	tbl, err = editor.RebuildAllIndexes(ctx, tbl)
   120  	require.NoError(t, err)
   121  
   122  	sch, err = tbl.GetSchema(ctx)
   123  	require.NoError(t, err)
   124  	rows, err := tbl.GetRowData(ctx)
   125  	require.NoError(t, err)
   126  	indexes, err := tbl.GetIndexData(ctx)
   127  	require.NoError(t, err)
   128  	err = dEnv.PutTableToWorking(ctx, sch, rows, indexes, tableName, nil)
   129  	require.NoError(t, err)
   130  }
   131  
   132  // MustSchema takes a variable number of columns and returns a schema.
   133  func MustSchema(cols ...schema.Column) schema.Schema {
   134  	hasPKCols := false
   135  	for _, col := range cols {
   136  		if col.IsPartOfPK {
   137  			hasPKCols = true
   138  			break
   139  		}
   140  	}
   141  
   142  	colColl := schema.NewColCollection(cols...)
   143  
   144  	if !hasPKCols {
   145  		return schema.UnkeyedSchemaFromCols(colColl)
   146  	} else {
   147  		return schema.MustSchemaFromCols(colColl)
   148  	}
   149  }