vitess.io/vitess@v0.16.2/go/vt/vtgate/engine/replace_variables.go (about)

     1  /*
     2  Copyright 2021 The Vitess Authors.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8      http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  package engine
    18  
    19  import (
    20  	"context"
    21  
    22  	"vitess.io/vitess/go/sqltypes"
    23  	querypb "vitess.io/vitess/go/vt/proto/query"
    24  )
    25  
    26  var _ Primitive = (*ReplaceVariables)(nil)
    27  
    28  // ReplaceVariables is used in SHOW VARIABLES statements so that it replaces the values for vitess-aware variables
    29  type ReplaceVariables struct {
    30  	Input Primitive
    31  	noTxNeeded
    32  }
    33  
    34  // NewReplaceVariables is used to create a new ReplaceVariables primitive
    35  func NewReplaceVariables(input Primitive) *ReplaceVariables {
    36  	return &ReplaceVariables{Input: input}
    37  }
    38  
    39  // RouteType implements the Primitive interface
    40  func (r *ReplaceVariables) RouteType() string {
    41  	return r.Input.RouteType()
    42  }
    43  
    44  // GetKeyspaceName implements the Primitive interface
    45  func (r *ReplaceVariables) GetKeyspaceName() string {
    46  	return r.Input.GetKeyspaceName()
    47  }
    48  
    49  // GetTableName implements the Primitive interface
    50  func (r *ReplaceVariables) GetTableName() string {
    51  	return r.Input.GetTableName()
    52  }
    53  
    54  // TryExecute implements the Primitive interface
    55  func (r *ReplaceVariables) TryExecute(ctx context.Context, vcursor VCursor, bindVars map[string]*querypb.BindVariable, wantfields bool) (*sqltypes.Result, error) {
    56  	qr, err := vcursor.ExecutePrimitive(ctx, r.Input, bindVars, wantfields)
    57  	if err != nil {
    58  		return nil, err
    59  	}
    60  	replaceVariables(qr, bindVars)
    61  	return qr, nil
    62  }
    63  
    64  // TryStreamExecute implements the Primitive interface
    65  func (r *ReplaceVariables) TryStreamExecute(ctx context.Context, vcursor VCursor, bindVars map[string]*querypb.BindVariable, wantfields bool, callback func(*sqltypes.Result) error) error {
    66  	innerCallback := callback
    67  	callback = func(result *sqltypes.Result) error {
    68  		replaceVariables(result, bindVars)
    69  		return innerCallback(result)
    70  	}
    71  	return vcursor.StreamExecutePrimitive(ctx, r.Input, bindVars, wantfields, callback)
    72  }
    73  
    74  // GetFields implements the Primitive interface
    75  func (r *ReplaceVariables) GetFields(ctx context.Context, vcursor VCursor, bindVars map[string]*querypb.BindVariable) (*sqltypes.Result, error) {
    76  	return r.Input.GetFields(ctx, vcursor, bindVars)
    77  }
    78  
    79  // Inputs implements the Primitive interface
    80  func (r *ReplaceVariables) Inputs() []Primitive {
    81  	return []Primitive{r.Input}
    82  }
    83  
    84  // description implements the Primitive interface
    85  func (r *ReplaceVariables) description() PrimitiveDescription {
    86  	return PrimitiveDescription{
    87  		OperatorType: "ReplaceVariables",
    88  	}
    89  }
    90  
    91  func replaceVariables(qr *sqltypes.Result, bindVars map[string]*querypb.BindVariable) {
    92  	for i, row := range qr.Rows {
    93  		variableName := row[0].ToString()
    94  		res, found := bindVars["__vt"+variableName]
    95  		if found {
    96  			qr.Rows[i][1] = sqltypes.NewVarChar(string(res.GetValue()))
    97  		}
    98  	}
    99  }