github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/sql/opt/testutils/testcat/set_zone_config.go (about)

     1  // Copyright 2019 The Cockroach Authors.
     2  //
     3  // Use of this software is governed by the Business Source License
     4  // included in the file licenses/BSL.txt.
     5  //
     6  // As of the Change Date specified in that file, in accordance with
     7  // the Business Source License, use of this software will be governed
     8  // by the Apache License, Version 2.0, included in the file
     9  // licenses/APL.txt.
    10  
    11  package testcat
    12  
    13  import (
    14  	"fmt"
    15  
    16  	"github.com/cockroachdb/cockroach/pkg/config/zonepb"
    17  	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
    18  	"gopkg.in/yaml.v2"
    19  )
    20  
    21  // SetZoneConfig is a partial implementation of the ALTER TABLE ... CONFIGURE
    22  // ZONE USING statement.
    23  func (tc *Catalog) SetZoneConfig(stmt *tree.SetZoneConfig) *zonepb.ZoneConfig {
    24  	// Update the table name to include catalog and schema if not provided.
    25  	tabName := stmt.TableOrIndex.Table
    26  	tc.qualifyTableName(&tabName)
    27  	tab := tc.Table(&tabName)
    28  
    29  	// Handle special case of primary index.
    30  	if stmt.TableOrIndex.Index == "" {
    31  		tab.Indexes[0].IdxZone = makeZoneConfig(stmt.Options)
    32  		return tab.Indexes[0].IdxZone
    33  	}
    34  
    35  	for _, idx := range tab.Indexes {
    36  		if idx.IdxName == string(stmt.TableOrIndex.Index) {
    37  			idx.IdxZone = makeZoneConfig(stmt.Options)
    38  			return idx.IdxZone
    39  		}
    40  	}
    41  	panic(fmt.Errorf("\"%q\" is not an index", stmt.TableOrIndex.Index))
    42  }
    43  
    44  // makeZoneConfig constructs a ZoneConfig from options provided to the CONFIGURE
    45  // ZONE USING statement.
    46  func makeZoneConfig(options tree.KVOptions) *zonepb.ZoneConfig {
    47  	zone := &zonepb.ZoneConfig{}
    48  	for i := range options {
    49  		switch options[i].Key {
    50  		case "constraints":
    51  			constraintsList := &zonepb.ConstraintsList{}
    52  			value := options[i].Value.(*tree.StrVal).RawString()
    53  			if err := yaml.UnmarshalStrict([]byte(value), constraintsList); err != nil {
    54  				panic(err)
    55  			}
    56  			zone.Constraints = constraintsList.Constraints
    57  
    58  		case "lease_preferences":
    59  			value := options[i].Value.(*tree.StrVal).RawString()
    60  			if err := yaml.UnmarshalStrict([]byte(value), &zone.LeasePreferences); err != nil {
    61  				panic(err)
    62  			}
    63  		}
    64  	}
    65  	return zone
    66  }