github.com/dolthub/go-mysql-server@v0.18.0/sql/analyzer/inline_subquery_aliases.go (about)

     1  // Copyright 2023 Dolthub, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package analyzer
    16  
    17  import (
    18  	"strings"
    19  
    20  	"github.com/dolthub/go-mysql-server/sql"
    21  	"github.com/dolthub/go-mysql-server/sql/expression"
    22  	"github.com/dolthub/go-mysql-server/sql/plan"
    23  	"github.com/dolthub/go-mysql-server/sql/transform"
    24  )
    25  
    26  type aliasScope struct {
    27  	aliases map[string]sql.Expression
    28  	parent  *aliasScope
    29  }
    30  
    31  func (a *aliasScope) push() *aliasScope {
    32  	return &aliasScope{
    33  		parent: a,
    34  	}
    35  }
    36  
    37  func (a *aliasScope) addRef(alias *expression.Alias) {
    38  	if a.aliases == nil {
    39  		a.aliases = make(map[string]sql.Expression)
    40  	}
    41  	a.aliases[alias.Name()] = alias.Child
    42  }
    43  
    44  func (a *aliasScope) isOuterRef(name string) (sql.Expression, bool) {
    45  	if a.aliases != nil {
    46  		if a, ok := a.aliases[name]; ok {
    47  			return a, false
    48  		}
    49  	}
    50  	if a.parent == nil {
    51  		return nil, false
    52  	}
    53  	found, _ := a.parent.isOuterRef(name)
    54  	if found != nil {
    55  		return found, true
    56  	}
    57  	return nil, false
    58  }
    59  
    60  // inlineSubqueryAliasRefs matches the pattern:
    61  // SELECT expr as <alias>, (SELECT <alias> ...) ...
    62  // and performs a variable replacement:
    63  // SELECT expr as <alias>, (SELECT expr ...) ...
    64  // Outer alias references can occur anywhere in subquery expressions,
    65  // as written this is a fairly unflexible rule.
    66  // TODO: extend subquery search to WHERE filters and other scalar expressions
    67  // TODO: convert subquery expressions to lateral joins to avoid this hack
    68  func inlineSubqueryAliasRefs(ctx *sql.Context, a *Analyzer, n sql.Node, scope *plan.Scope, sel RuleSelector) (sql.Node, transform.TreeIdentity, error) {
    69  	ret, err := inlineSubqueryAliasRefsHelper(&aliasScope{}, n)
    70  	return ret, transform.NewTree, err
    71  }
    72  
    73  func inlineSubqueryAliasRefsHelper(scope *aliasScope, n sql.Node) (sql.Node, error) {
    74  	ret := n
    75  	switch n := n.(type) {
    76  	case *plan.Project:
    77  		var newProj []sql.Expression
    78  		for i, e := range n.Projections {
    79  			e, same, err := transform.Expr(e, func(e sql.Expression) (sql.Expression, transform.TreeIdentity, error) {
    80  				switch e := e.(type) {
    81  				case *expression.AliasReference:
    82  				case *expression.Alias:
    83  					// new def
    84  					if !e.Unreferencable() {
    85  						scope.addRef(e)
    86  					}
    87  				case *expression.GetField:
    88  					// is an alias ref?
    89  					// check if in parent scope
    90  					if alias, inOuter := scope.isOuterRef(strings.ToLower(e.Name())); e.Table() == "" && alias != nil && inOuter {
    91  						return alias, transform.NewTree, nil
    92  					}
    93  				case *plan.Subquery:
    94  					subqScope := scope.push()
    95  					newQ, err := inlineSubqueryAliasRefsHelper(subqScope, e.Query)
    96  					if err != nil {
    97  						return e, transform.SameTree, err
    98  					}
    99  					ret := *e
   100  					ret.Query = newQ
   101  					return &ret, transform.NewTree, nil
   102  				default:
   103  				}
   104  				return e, transform.SameTree, nil
   105  			})
   106  			if err != nil {
   107  				return nil, err
   108  			}
   109  			if !same {
   110  				if newProj == nil {
   111  					newProj = make([]sql.Expression, len(n.Projections))
   112  					copy(newProj, n.Projections)
   113  				}
   114  				newProj[i] = e
   115  			}
   116  		}
   117  		if newProj != nil {
   118  			ret = plan.NewProject(newProj, n.Child)
   119  		}
   120  	default:
   121  	}
   122  
   123  	newChildren := make([]sql.Node, len(n.Children()))
   124  	var err error
   125  	for i, c := range ret.Children() {
   126  		newChildren[i], err = inlineSubqueryAliasRefsHelper(scope, c)
   127  		if err != nil {
   128  			return nil, err
   129  		}
   130  	}
   131  	ret, err = ret.WithChildren(newChildren...)
   132  	if err != nil {
   133  		return nil, err
   134  	}
   135  	if err != nil {
   136  		panic(err)
   137  	}
   138  	return ret, nil
   139  }