github.com/dolthub/dolt/go@v0.40.5-0.20240520175717-68db7794bea6/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  	"github.com/dolthub/dolt/go/libraries/doltcore/row"
    19  	"github.com/dolthub/dolt/go/libraries/doltcore/schema"
    20  	"github.com/dolthub/dolt/go/store/types"
    21  )
    22  
    23  // CreateSchema returns a schema from the columns given, panicking on any errors.
    24  func CreateSchema(columns ...schema.Column) schema.Schema {
    25  	colColl := schema.NewColCollection(columns...)
    26  	sch := schema.MustSchemaFromCols(colColl)
    27  	sch.SetCollation(schema.Collation_Default)
    28  	return sch
    29  }
    30  
    31  // NewRow creates a row with the schema given, having the values given. Starts at tag 0 and counts up.
    32  func NewRow(sch schema.Schema, values ...types.Value) row.Row {
    33  	taggedVals := make(row.TaggedValues)
    34  	for i := range values {
    35  		taggedVals[uint64(i)] = values[i]
    36  	}
    37  	r, err := row.New(types.Format_Default, sch, taggedVals)
    38  
    39  	if err != nil {
    40  		panic(err)
    41  	}
    42  
    43  	return r
    44  }
    45  
    46  // AddColumnToSchema returns a new schema by adding the given column to the given schema. Will panic on an invalid
    47  // schema, e.g. tag collision.
    48  // Note the AddColumnToSchema relies on being called from the engine (GMS) to correctly update defaults. Directly calling
    49  // this method in Dolt only adds a new column to schema but does not apply the default.
    50  func AddColumnToSchema(sch schema.Schema, col schema.Column) schema.Schema {
    51  	columns := sch.GetAllCols()
    52  	columns = columns.Append(col)
    53  	newSch := schema.MustSchemaFromCols(columns)
    54  	newSch.SetCollation(sch.GetCollation())
    55  	return newSch
    56  }
    57  
    58  // MustSchema takes a variable number of columns and returns a schema.
    59  func MustSchema(cols ...schema.Column) schema.Schema {
    60  	hasPKCols := false
    61  	for _, col := range cols {
    62  		if col.IsPartOfPK {
    63  			hasPKCols = true
    64  			break
    65  		}
    66  	}
    67  
    68  	colColl := schema.NewColCollection(cols...)
    69  
    70  	if !hasPKCols {
    71  		return schema.UnkeyedSchemaFromCols(colColl)
    72  	} else {
    73  		return schema.MustSchemaFromCols(colColl)
    74  	}
    75  }