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 }