github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/sql/sqltestutils/sql_test_utils.go (about)

     1  // Copyright 2020 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 sqltestutils provides helper methods for testing sql packages.
    12  package sqltestutils
    13  
    14  import (
    15  	gosql "database/sql"
    16  	"testing"
    17  
    18  	"github.com/cockroachdb/cockroach/pkg/config/zonepb"
    19  	"github.com/cockroachdb/cockroach/pkg/sql/sqlbase"
    20  	"github.com/cockroachdb/cockroach/pkg/util/protoutil"
    21  	"github.com/stretchr/testify/require"
    22  )
    23  
    24  // AddImmediateGCZoneConfig set the GC TTL to 0 for the given table ID. One must
    25  // make sure to disable strict GC TTL enforcement when using this.
    26  func AddImmediateGCZoneConfig(sqlDB *gosql.DB, id sqlbase.ID) (zonepb.ZoneConfig, error) {
    27  	cfg := zonepb.DefaultZoneConfig()
    28  	cfg.GC.TTLSeconds = 0
    29  	buf, err := protoutil.Marshal(&cfg)
    30  	if err != nil {
    31  		return cfg, err
    32  	}
    33  	_, err = sqlDB.Exec(`UPSERT INTO system.zones VALUES ($1, $2)`, id, buf)
    34  	return cfg, err
    35  }
    36  
    37  // DisableGCTTLStrictEnforcement sets the cluster setting to disable strict
    38  // GC TTL enforcement.
    39  func DisableGCTTLStrictEnforcement(t *testing.T, db *gosql.DB) (cleanup func()) {
    40  	_, err := db.Exec(`SET CLUSTER SETTING kv.gc_ttl.strict_enforcement.enabled = false`)
    41  	require.NoError(t, err)
    42  	return func() {
    43  		_, err := db.Exec(`SET CLUSTER SETTING kv.gc_ttl.strict_enforcement.enabled = DEFAULT`)
    44  		require.NoError(t, err)
    45  	}
    46  }
    47  
    48  // AddDefaultZoneConfig adds an entry for the given id into system.zones.
    49  func AddDefaultZoneConfig(sqlDB *gosql.DB, id sqlbase.ID) (zonepb.ZoneConfig, error) {
    50  	cfg := zonepb.DefaultZoneConfig()
    51  	buf, err := protoutil.Marshal(&cfg)
    52  	if err != nil {
    53  		return cfg, err
    54  	}
    55  	_, err = sqlDB.Exec(`UPSERT INTO system.zones VALUES ($1, $2)`, id, buf)
    56  	return cfg, err
    57  }