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

     1  // Copyright 2020-2021 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  	"os"
    19  
    20  	"github.com/dolthub/go-mysql-server/sql/transform"
    21  
    22  	"github.com/dolthub/go-mysql-server/sql"
    23  	"github.com/dolthub/go-mysql-server/sql/plan"
    24  )
    25  
    26  var updateQueryProgressEachRow bool
    27  
    28  const updateQueryProcessEachRowEnvKey = "DETAILED_QUERY_PROGRESS"
    29  
    30  func init() {
    31  	if v, ok := os.LookupEnv(updateQueryProcessEachRowEnvKey); ok && len(v) > 0 {
    32  		updateQueryProgressEachRow = true
    33  	}
    34  }
    35  
    36  // trackProcess will wrap the query in a process node and add progress items
    37  // to the already existing process.
    38  func trackProcess(ctx *sql.Context, a *Analyzer, n sql.Node, scope *plan.Scope, sel RuleSelector) (sql.Node, transform.TreeIdentity, error) {
    39  	if !n.Resolved() {
    40  		return n, transform.SameTree, nil
    41  	}
    42  
    43  	if _, ok := n.(*plan.QueryProcess); ok {
    44  		return n, transform.SameTree, nil
    45  	}
    46  
    47  	processList := ctx.ProcessList
    48  
    49  	var seen = make(map[string]struct{})
    50  	n, _, err := transform.Node(n, func(n sql.Node) (sql.Node, transform.TreeIdentity, error) {
    51  		switch n := n.(type) {
    52  		case *plan.ResolvedTable:
    53  			switch n.Table.(type) {
    54  			case *plan.ProcessTable, *plan.ProcessIndexableTable:
    55  				return n, transform.SameTree, nil
    56  			}
    57  
    58  			name := n.Table.Name()
    59  			if _, ok := seen[name]; ok {
    60  				return n, transform.SameTree, nil
    61  			}
    62  
    63  			var total int64 = -1
    64  			if counter, ok := n.Table.(sql.PartitionCounter); ok {
    65  				count, err := counter.PartitionCount(ctx)
    66  				if err != nil {
    67  					return nil, transform.SameTree, err
    68  				}
    69  				total = count
    70  			}
    71  			processList.AddTableProgress(ctx.Pid(), name, total)
    72  
    73  			seen[name] = struct{}{}
    74  
    75  			onPartitionDone := func(partitionName string) {
    76  				processList.UpdateTableProgress(ctx.Pid(), name, 1)
    77  				processList.RemovePartitionProgress(ctx.Pid(), name, partitionName)
    78  			}
    79  
    80  			onPartitionStart := func(partitionName string) {
    81  				processList.AddPartitionProgress(ctx.Pid(), name, partitionName, -1)
    82  			}
    83  
    84  			var onRowNext plan.NamedNotifyFunc
    85  			// TODO: coarser default for row updates (like updating every 100 rows) that doesn't kill performance
    86  			if updateQueryProgressEachRow {
    87  				onRowNext = func(partitionName string) {
    88  					processList.UpdatePartitionProgress(ctx.Pid(), name, partitionName, 1)
    89  				}
    90  			}
    91  
    92  			var t sql.Table
    93  			switch table := n.Table.(type) {
    94  			case sql.DriverIndexableTable:
    95  				t = plan.NewProcessIndexableTable(table, onPartitionDone, onPartitionStart, onRowNext)
    96  			default:
    97  				t = plan.NewProcessTable(table, onPartitionDone, onPartitionStart, onRowNext)
    98  			}
    99  
   100  			rt, err := n.ReplaceTable(t)
   101  			if err != nil {
   102  				return nil, false, err
   103  			}
   104  			return rt, transform.NewTree, nil
   105  		default:
   106  			return n, transform.SameTree, nil
   107  		}
   108  	})
   109  	if err != nil {
   110  		return nil, transform.SameTree, err
   111  	}
   112  
   113  	// Don't wrap CreateIndex in a QueryProcess, as it is a CreateIndexProcess.
   114  	// CreateIndex will take care of marking the process as done on its own.
   115  	if _, ok := n.(*plan.CreateIndex); ok {
   116  		return n, transform.SameTree, nil
   117  	}
   118  
   119  	// Remove QueryProcess nodes from the subqueries and trigger bodies. Otherwise, the process
   120  	// will be marked as done as soon as a subquery / trigger finishes.
   121  	node, _, err := transform.Node(n, func(n sql.Node) (sql.Node, transform.TreeIdentity, error) {
   122  		if sq, ok := n.(*plan.SubqueryAlias); ok {
   123  			if qp, ok := sq.Child.(*plan.QueryProcess); ok {
   124  				n, err := sq.WithChildren(qp.Child())
   125  				return n, transform.NewTree, err
   126  			}
   127  		}
   128  		if t, ok := n.(*plan.TriggerExecutor); ok {
   129  			if qp, ok := t.Right().(*plan.QueryProcess); ok {
   130  				n, err := t.WithChildren(t.Left(), qp.Child())
   131  				return n, transform.NewTree, err
   132  			}
   133  		}
   134  		return n, transform.SameTree, nil
   135  	})
   136  	if err != nil {
   137  		return nil, transform.SameTree, err
   138  	}
   139  
   140  	return plan.NewQueryProcess(node, func() {
   141  		processList.EndQuery(ctx)
   142  		if span := ctx.RootSpan(); span != nil {
   143  			span.End()
   144  		}
   145  	}), transform.NewTree, nil
   146  }