github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/testutils/sqlutils/zone.go (about)

     1  // Copyright 2017 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 sqlutils
    12  
    13  import (
    14  	gosql "database/sql"
    15  	"fmt"
    16  	"testing"
    17  
    18  	"github.com/cockroachdb/cockroach/pkg/config/zonepb"
    19  	"github.com/cockroachdb/cockroach/pkg/sql/lex"
    20  	"github.com/cockroachdb/cockroach/pkg/util/protoutil"
    21  )
    22  
    23  // ZoneRow represents a row returned by SHOW ZONE CONFIGURATION.
    24  type ZoneRow struct {
    25  	ID     uint32
    26  	Config zonepb.ZoneConfig
    27  }
    28  
    29  func (row ZoneRow) sqlRowString() ([]string, error) {
    30  	configProto, err := protoutil.Marshal(&row.Config)
    31  	if err != nil {
    32  		return nil, err
    33  	}
    34  	return []string{
    35  		fmt.Sprintf("%d", row.ID),
    36  		string(configProto),
    37  	}, nil
    38  }
    39  
    40  // RemoveAllZoneConfigs removes all installed zone configs.
    41  func RemoveAllZoneConfigs(t testing.TB, sqlDB *SQLRunner) {
    42  	t.Helper()
    43  	for _, row := range sqlDB.QueryStr(t, "SHOW ALL ZONE CONFIGURATIONS") {
    44  		target := row[0]
    45  		if target == fmt.Sprintf("RANGE %s", zonepb.DefaultZoneName) {
    46  			// The default zone cannot be removed.
    47  			continue
    48  		}
    49  		DeleteZoneConfig(t, sqlDB, target)
    50  	}
    51  }
    52  
    53  // DeleteZoneConfig deletes the specified zone config through the SQL interface.
    54  func DeleteZoneConfig(t testing.TB, sqlDB *SQLRunner, target string) {
    55  	t.Helper()
    56  	sqlDB.Exec(t, fmt.Sprintf("ALTER %s CONFIGURE ZONE DISCARD", target))
    57  }
    58  
    59  // SetZoneConfig updates the specified zone config through the SQL interface.
    60  func SetZoneConfig(t testing.TB, sqlDB *SQLRunner, target string, config string) {
    61  	t.Helper()
    62  	sqlDB.Exec(t, fmt.Sprintf("ALTER %s CONFIGURE ZONE = %s",
    63  		target, lex.EscapeSQLString(config)))
    64  }
    65  
    66  // TxnSetZoneConfig updates the specified zone config through the SQL interface
    67  // using the provided transaction.
    68  func TxnSetZoneConfig(t testing.TB, sqlDB *SQLRunner, txn *gosql.Tx, target string, config string) {
    69  	t.Helper()
    70  	_, err := txn.Exec(fmt.Sprintf("ALTER %s CONFIGURE ZONE = %s",
    71  		target, lex.EscapeSQLString(config)))
    72  	if err != nil {
    73  		t.Fatal(err)
    74  	}
    75  }
    76  
    77  // VerifyZoneConfigForTarget verifies that the specified zone matches the specified
    78  // ZoneRow.
    79  func VerifyZoneConfigForTarget(t testing.TB, sqlDB *SQLRunner, target string, row ZoneRow) {
    80  	t.Helper()
    81  	sqlRow, err := row.sqlRowString()
    82  	if err != nil {
    83  		t.Fatal(err)
    84  	}
    85  	sqlDB.CheckQueryResults(t, fmt.Sprintf(`
    86  SELECT zone_id, raw_config_protobuf
    87    FROM [SHOW ZONE CONFIGURATION FOR %s]`, target),
    88  		[][]string{sqlRow})
    89  }
    90  
    91  // VerifyAllZoneConfigs verifies that the specified ZoneRows exactly match the
    92  // list of active zone configs.
    93  func VerifyAllZoneConfigs(t testing.TB, sqlDB *SQLRunner, rows ...ZoneRow) {
    94  	t.Helper()
    95  	expected := make([][]string, len(rows))
    96  	for i, row := range rows {
    97  		var err error
    98  		expected[i], err = row.sqlRowString()
    99  		if err != nil {
   100  			t.Fatal(err)
   101  		}
   102  	}
   103  	sqlDB.CheckQueryResults(t, `SELECT zone_id, raw_config_protobuf FROM crdb_internal.zones`, expected)
   104  }
   105  
   106  // ZoneConfigExists returns whether a zone config with the provided name exists.
   107  func ZoneConfigExists(t testing.TB, sqlDB *SQLRunner, name string) bool {
   108  	t.Helper()
   109  	var exists bool
   110  	sqlDB.QueryRow(
   111  		t, "SELECT EXISTS (SELECT 1 FROM crdb_internal.zones WHERE target = $1)", name,
   112  	).Scan(&exists)
   113  	return exists
   114  }