github.com/dolthub/go-mysql-server@v0.18.0/sql/processlist.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 sql
    16  
    17  import (
    18  	"context"
    19  	"fmt"
    20  	"time"
    21  )
    22  
    23  type ProcessList interface {
    24  	// Processes returns the list of current running processes
    25  	Processes() []Process
    26  
    27  	// AddConnection adds a new connection to the process list. Must be matched with RemoveConnection.
    28  	AddConnection(connID uint32, addr string)
    29  
    30  	// Transitions a connection from Connect to Sleep.
    31  	ConnectionReady(sess Session)
    32  
    33  	// RemoveConnection removes the connection from the process list.
    34  	RemoveConnection(connID uint32)
    35  
    36  	// BeginQuery transitions an existing connection in the processlist from Command "Sleep" to Command "Query".
    37  	// Returns a new context which will be canceled when this query is done.
    38  	BeginQuery(ctx *Context, query string) (*Context, error)
    39  
    40  	// EndQuery transitions a previously transitioned connection from Command "Query" to Command "Sleep".
    41  	EndQuery(ctx *Context)
    42  
    43  	// Kill terminates all queries for a given connection id
    44  	Kill(connID uint32)
    45  
    46  	// UpdateTableProgress updates the progress of the table with the given name for the
    47  	// process with the given pid.
    48  	UpdateTableProgress(pid uint64, name string, delta int64)
    49  
    50  	// UpdatePartitionProgress updates the progress of the table partition with the
    51  	// given name for the process with the given pid.
    52  	UpdatePartitionProgress(pid uint64, tableName, partitionName string, delta int64)
    53  
    54  	// AddTableProgress adds a new item to track progress from to the process with
    55  	// the given pid. If the pid does not exist, it will do nothing.
    56  	AddTableProgress(pid uint64, name string, total int64)
    57  
    58  	// AddPartitionProgress adds a new item to track progress from to the process with
    59  	// the given pid. If the pid or the table does not exist, it will do nothing.
    60  	AddPartitionProgress(pid uint64, tableName, partitionName string, total int64)
    61  
    62  	// RemoveTableProgress removes an existing item tracking progress from the
    63  	// process with the given pid, if it exists.
    64  	RemoveTableProgress(pid uint64, name string)
    65  
    66  	// RemovePartitionProgress removes an existing partition tracking progress from the
    67  	// process with the given pid, if it exists.
    68  	RemovePartitionProgress(pid uint64, tableName, partitionName string)
    69  }
    70  
    71  type ProcessCommand string
    72  
    73  const (
    74  	// During initial connection and handshake.
    75  	ProcessCommandConnect ProcessCommand = "Connect"
    76  	// Connected, not running a query.
    77  	ProcessCommandSleep ProcessCommand = "Sleep"
    78  	// Currently running a query, possibly streaming the response.
    79  	ProcessCommandQuery ProcessCommand = "Query"
    80  )
    81  
    82  // Process represents a process in the SQL server.
    83  type Process struct {
    84  	Connection uint32
    85  	Host       string
    86  	Database   string
    87  	User       string
    88  	Command    ProcessCommand
    89  
    90  	// The time of the last Command transition...
    91  	StartedAt time.Time
    92  
    93  	QueryPid uint64
    94  	Query    string
    95  	Progress map[string]TableProgress
    96  	Kill     context.CancelFunc
    97  }
    98  
    99  // Done needs to be called when this process has finished.
   100  func (p *Process) Done() { p.Kill() }
   101  
   102  // Seconds returns the number of seconds this process has been running.
   103  func (p *Process) Seconds() uint64 {
   104  	return uint64(time.Since(p.StartedAt) / time.Second)
   105  }
   106  
   107  // Progress between done items and total items
   108  type Progress struct {
   109  	Name  string
   110  	Done  int64
   111  	Total int64
   112  }
   113  
   114  func (p Progress) totalString() string {
   115  	var total = "?"
   116  	if p.Total > 0 {
   117  		total = fmt.Sprint(p.Total)
   118  	}
   119  	return total
   120  }
   121  
   122  // TableProgress keeps track of a table progress, and for each of its partitions
   123  type TableProgress struct {
   124  	Progress
   125  	PartitionsProgress map[string]PartitionProgress
   126  }
   127  
   128  func NewTableProgress(name string, total int64) TableProgress {
   129  	return TableProgress{
   130  		Progress: Progress{
   131  			Name:  name,
   132  			Total: total,
   133  		},
   134  		PartitionsProgress: make(map[string]PartitionProgress),
   135  	}
   136  }
   137  
   138  func (p TableProgress) String() string {
   139  	return fmt.Sprintf("%s (%d/%s partitions)", p.Name, p.Done, p.totalString())
   140  }
   141  
   142  // PartitionProgress keeps track of a partition progress
   143  type PartitionProgress struct {
   144  	Progress
   145  }
   146  
   147  func (p PartitionProgress) String() string {
   148  	return fmt.Sprintf("%s (%d/%s rows)", p.Name, p.Done, p.totalString())
   149  }
   150  
   151  // EmptyProcessList is a no-op implementation of ProcessList suitable for use in tests or other installations that
   152  // don't require a process list
   153  type EmptyProcessList struct{}
   154  
   155  var _ ProcessList = EmptyProcessList{}
   156  
   157  func (e EmptyProcessList) Processes() []Process {
   158  	return nil
   159  }
   160  
   161  func (e EmptyProcessList) AddConnection(id uint32, addr string) {}
   162  func (e EmptyProcessList) ConnectionReady(Session)              {}
   163  func (e EmptyProcessList) RemoveConnection(uint32)              {}
   164  
   165  func (e EmptyProcessList) BeginQuery(ctx *Context, query string) (*Context, error) {
   166  	return ctx, nil
   167  }
   168  func (e EmptyProcessList) EndQuery(ctx *Context) {}
   169  
   170  func (e EmptyProcessList) Kill(connID uint32)                                       {}
   171  func (e EmptyProcessList) Done(pid uint64)                                          {}
   172  func (e EmptyProcessList) UpdateTableProgress(pid uint64, name string, delta int64) {}
   173  func (e EmptyProcessList) UpdatePartitionProgress(pid uint64, tableName, partitionName string, delta int64) {
   174  }
   175  func (e EmptyProcessList) AddTableProgress(pid uint64, name string, total int64) {}
   176  func (e EmptyProcessList) AddPartitionProgress(pid uint64, tableName, partitionName string, total int64) {
   177  }
   178  func (e EmptyProcessList) RemoveTableProgress(pid uint64, name string)                         {}
   179  func (e EmptyProcessList) RemovePartitionProgress(pid uint64, tableName, partitionName string) {}