github.com/vescale/zgraph@v0.0.0-20230410094002-959c02d50f95/stmtctx/context.go (about)

     1  // Copyright 2022 zGraph Authors. All rights reserved.
     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 stmtctx
    16  
    17  import (
    18  	"strings"
    19  	"sync"
    20  	"sync/atomic"
    21  
    22  	"github.com/vescale/zgraph/catalog"
    23  	"github.com/vescale/zgraph/storage/kv"
    24  )
    25  
    26  // Context represent the intermediate state of a query execution and will be
    27  // reset after a query finished.
    28  type Context struct {
    29  	store   kv.Storage
    30  	catalog *catalog.Catalog
    31  
    32  	mu struct {
    33  		sync.RWMutex
    34  
    35  		currentGraph string
    36  		txn          *LazyTxn
    37  
    38  		affectedRows uint64
    39  		foundRows    uint64
    40  		records      uint64
    41  		deleted      uint64
    42  		updated      uint64
    43  		copied       uint64
    44  		touched      uint64
    45  
    46  		warnings   []SQLWarn
    47  		errorCount uint16
    48  	}
    49  
    50  	// TODO: perhaps we can move these to a separate struct.
    51  	planID       atomic.Int64
    52  	planColumnID atomic.Int64
    53  }
    54  
    55  // New returns a session statement context instance.
    56  func New(store kv.Storage, catalog *catalog.Catalog) *Context {
    57  	return &Context{
    58  		store:   store,
    59  		catalog: catalog,
    60  	}
    61  }
    62  
    63  // Reset resets all variables associated to execute a query.
    64  func (sc *Context) Reset() {
    65  	sc.mu.Lock()
    66  	defer sc.mu.Unlock()
    67  
    68  	sc.mu.affectedRows = 0
    69  	sc.mu.foundRows = 0
    70  	sc.mu.records = 0
    71  	sc.mu.deleted = 0
    72  	sc.mu.updated = 0
    73  	sc.mu.copied = 0
    74  	sc.mu.touched = 0
    75  	sc.mu.warnings = sc.mu.warnings[:0]
    76  	sc.mu.errorCount = 0
    77  }
    78  
    79  // Store returns the storage instance.
    80  func (sc *Context) Store() kv.Storage {
    81  	return sc.store
    82  }
    83  
    84  // Catalog returns the catalog object.
    85  func (sc *Context) Catalog() *catalog.Catalog {
    86  	return sc.catalog
    87  }
    88  
    89  // CurrentGraph returns the current chosen catalog graph
    90  func (sc *Context) CurrentGraph() *catalog.Graph {
    91  	sc.mu.RLock()
    92  	defer sc.mu.RUnlock()
    93  
    94  	return sc.Catalog().Graph(sc.mu.currentGraph)
    95  }
    96  
    97  // SetCurrentGraphName changes the current graph name.
    98  func (sc *Context) SetCurrentGraphName(graphName string) {
    99  	sc.mu.Lock()
   100  	defer sc.mu.Unlock()
   101  
   102  	sc.mu.currentGraph = strings.ToLower(graphName)
   103  }
   104  
   105  // CurrentGraphName returns the current chosen graph name.
   106  func (sc *Context) CurrentGraphName() string {
   107  	sc.mu.RLock()
   108  	defer sc.mu.RUnlock()
   109  
   110  	return sc.mu.currentGraph
   111  }
   112  
   113  func (sc *Context) AllocPlanID() int {
   114  	return int(sc.planID.Add(1))
   115  }
   116  
   117  func (sc *Context) AllocPlanColumnID() int64 {
   118  	return sc.planColumnID.Add(1)
   119  }