vitess.io/vitess@v0.16.2/go/vt/vtgate/planbuilder/union.go (about)

     1  /*
     2  Copyright 2019 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 planbuilder
    18  
    19  import (
    20  	"fmt"
    21  
    22  	"vitess.io/vitess/go/vt/vtgate/planbuilder/plancontext"
    23  
    24  	"vitess.io/vitess/go/vt/vterrors"
    25  
    26  	"vitess.io/vitess/go/mysql"
    27  
    28  	"vitess.io/vitess/go/vt/sqlparser"
    29  )
    30  
    31  func buildUnionPlan(string) stmtPlanner {
    32  	return func(stmt sqlparser.Statement, reservedVars *sqlparser.ReservedVars, vschema plancontext.VSchema) (*planResult, error) {
    33  		union := stmt.(*sqlparser.Union)
    34  		if union.With != nil {
    35  			return nil, vterrors.VT12001("WITH expression in UNION statement")
    36  		}
    37  		// For unions, create a pb with anonymous scope.
    38  		pb := newPrimitiveBuilder(vschema, newJointab(reservedVars))
    39  		if err := pb.processUnion(union, reservedVars, nil); err != nil {
    40  			return nil, err
    41  		}
    42  		if err := pb.plan.Wireup(pb.plan, pb.jt); err != nil {
    43  			return nil, err
    44  		}
    45  		return newPlanResult(pb.plan.Primitive()), nil
    46  	}
    47  }
    48  
    49  func (pb *primitiveBuilder) processUnion(union *sqlparser.Union, reservedVars *sqlparser.ReservedVars, outer *symtab) error {
    50  	if err := pb.processPart(union.Left, reservedVars, outer); err != nil {
    51  		return err
    52  	}
    53  
    54  	rpb := newPrimitiveBuilder(pb.vschema, pb.jt)
    55  	if err := rpb.processPart(union.Right, reservedVars, outer); err != nil {
    56  		return err
    57  	}
    58  	err := unionRouteMerge(pb.plan, rpb.plan, union)
    59  	if err != nil {
    60  		// we are merging between two routes - let's check if we can see so that we have the same amount of columns on both sides of the union
    61  		lhsCols := len(pb.plan.ResultColumns())
    62  		rhsCols := len(rpb.plan.ResultColumns())
    63  		if lhsCols != rhsCols {
    64  			return &mysql.SQLError{
    65  				Num:     mysql.ERWrongNumberOfColumnsInSelect,
    66  				State:   "21000",
    67  				Message: "The used SELECT statements have a different number of columns",
    68  				Query:   sqlparser.String(union),
    69  			}
    70  		}
    71  
    72  		pb.plan = &concatenate{
    73  			lhs: pb.plan,
    74  			rhs: rpb.plan,
    75  		}
    76  
    77  		if union.Distinct {
    78  			pb.plan = newDistinctV3(pb.plan)
    79  		}
    80  	}
    81  	pb.st.Outer = outer
    82  
    83  	if err := setLock(pb.plan, union.Lock); err != nil {
    84  		return err
    85  	}
    86  
    87  	if err := pb.pushOrderBy(union.OrderBy); err != nil {
    88  		return err
    89  	}
    90  	return pb.pushLimit(union.Limit)
    91  }
    92  
    93  func (pb *primitiveBuilder) processPart(part sqlparser.SelectStatement, reservedVars *sqlparser.ReservedVars, outer *symtab) error {
    94  	switch part := part.(type) {
    95  	case *sqlparser.Union:
    96  		return pb.processUnion(part, reservedVars, outer)
    97  	case *sqlparser.Select:
    98  		if part.SQLCalcFoundRows {
    99  			return vterrors.VT12001("SQL_CALC_FOUND_ROWS not supported with UNION")
   100  		}
   101  		return pb.processSelect(part, reservedVars, outer, "")
   102  	}
   103  	return vterrors.VT13001(fmt.Sprintf("unexpected SELECT type: %T", part))
   104  }
   105  
   106  // TODO (systay) we never use this as an actual error. we should rethink the return type
   107  func unionRouteMerge(left, right logicalPlan, us *sqlparser.Union) error {
   108  	lroute, ok := left.(*route)
   109  	if !ok {
   110  		return vterrors.VT12001("SELECT of UNION is non-trivial")
   111  	}
   112  	rroute, ok := right.(*route)
   113  	if !ok {
   114  		return vterrors.VT12001("SELECT of UNION is non-trivial")
   115  	}
   116  	mergeSuccess := lroute.MergeUnion(rroute, us.Distinct)
   117  	if !mergeSuccess {
   118  		return vterrors.VT12001("execute UNION as a single route")
   119  	}
   120  
   121  	lroute.Select = &sqlparser.Union{Left: lroute.Select, Right: us.Right, Distinct: us.Distinct}
   122  
   123  	return nil
   124  }
   125  
   126  // planLock pushes "FOR UPDATE", "LOCK IN SHARE MODE" down to all routes
   127  func setLock(in logicalPlan, lock sqlparser.Lock) error {
   128  	_, err := visit(in, func(plan logicalPlan) (bool, logicalPlan, error) {
   129  		switch node := in.(type) {
   130  		case *route:
   131  			node.Select.SetLock(lock)
   132  			return false, node, nil
   133  		case *sqlCalcFoundRows, *vindexFunc:
   134  			return false, nil, vterrors.VT13001(fmt.Sprintf("unreachable %T.locking", in))
   135  		}
   136  		return true, plan, nil
   137  	})
   138  	if err != nil {
   139  		return err
   140  	}
   141  	return nil
   142  }