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

     1  // Copyright 2018 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  	"context"
    15  	gojson "encoding/json"
    16  	"fmt"
    17  	"sort"
    18  
    19  	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
    20  	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
    21  	"github.com/cockroachdb/cockroach/pkg/sql/stats"
    22  	"github.com/cockroachdb/cockroach/pkg/sql/types"
    23  )
    24  
    25  // AlterTable is a partial implementation of the ALTER TABLE statement.
    26  //
    27  // Supported commands:
    28  //  - INJECT STATISTICS: imports table statistics from a JSON object.
    29  //  - ADD CONSTRAINT FOREIGN KEY: add a foreign key reference.
    30  //
    31  func (tc *Catalog) AlterTable(stmt *tree.AlterTable) {
    32  	tn := stmt.Table.ToTableName()
    33  	// Update the table name to include catalog and schema if not provided.
    34  	tc.qualifyTableName(&tn)
    35  	tab := tc.Table(&tn)
    36  
    37  	for _, cmd := range stmt.Cmds {
    38  		switch t := cmd.(type) {
    39  		case *tree.AlterTableInjectStats:
    40  			injectTableStats(tab, t.Stats)
    41  
    42  		case *tree.AlterTableAddConstraint:
    43  			switch d := t.ConstraintDef.(type) {
    44  			case *tree.ForeignKeyConstraintTableDef:
    45  				tc.resolveFK(tab, d)
    46  
    47  			default:
    48  				panic(fmt.Sprintf("unsupported constraint type %v", d))
    49  			}
    50  
    51  		default:
    52  			panic(fmt.Sprintf("unsupported ALTER TABLE command %T", t))
    53  		}
    54  	}
    55  }
    56  
    57  // injectTableStats sets the table statistics as specified by a JSON object.
    58  func injectTableStats(tt *Table, statsExpr tree.Expr) {
    59  	ctx := context.Background()
    60  	semaCtx := tree.MakeSemaContext()
    61  	evalCtx := tree.MakeTestingEvalContext(cluster.MakeTestingClusterSettings())
    62  	typedExpr, err := tree.TypeCheckAndRequire(
    63  		ctx, statsExpr, &semaCtx, types.Jsonb, "INJECT STATISTICS",
    64  	)
    65  	if err != nil {
    66  		panic(err)
    67  	}
    68  	val, err := typedExpr.Eval(&evalCtx)
    69  	if err != nil {
    70  		panic(err)
    71  	}
    72  
    73  	if val == tree.DNull {
    74  		panic("statistics cannot be NULL")
    75  	}
    76  	jsonStr := val.(*tree.DJSON).JSON.String()
    77  	var stats []stats.JSONStatistic
    78  	if err := gojson.Unmarshal([]byte(jsonStr), &stats); err != nil {
    79  		panic(err)
    80  	}
    81  	tt.Stats = make([]*TableStat, len(stats))
    82  	for i := range stats {
    83  		tt.Stats[i] = &TableStat{js: stats[i], tt: tt}
    84  	}
    85  	// Call ColumnOrdinal on all possible columns to assert that
    86  	// the column names are valid.
    87  	for _, ts := range tt.Stats {
    88  		for i := 0; i < ts.ColumnCount(); i++ {
    89  			ts.ColumnOrdinal(i)
    90  		}
    91  	}
    92  
    93  	// Finally, sort the stats with most recent first.
    94  	sort.Sort(tt.Stats)
    95  }