github.com/vescale/zgraph@v0.0.0-20230410094002-959c02d50f95/session/session.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 session
    16  
    17  import (
    18  	"context"
    19  	"sync"
    20  	"sync/atomic"
    21  
    22  	"github.com/pingcap/errors"
    23  	"github.com/vescale/zgraph/catalog"
    24  	"github.com/vescale/zgraph/compiler"
    25  	"github.com/vescale/zgraph/parser"
    26  	"github.com/vescale/zgraph/parser/ast"
    27  	"github.com/vescale/zgraph/stmtctx"
    28  	"github.com/vescale/zgraph/storage/kv"
    29  )
    30  
    31  var (
    32  	idGenerator atomic.Int64
    33  	parserPool  = &sync.Pool{New: func() interface{} { return parser.New() }}
    34  )
    35  
    36  // Session represents the session to interact with zGraph database instance.
    37  // Typically, the number of session will be same as the concurrent thread
    38  // count of the application.
    39  // All execution intermediate variables should be placed in the Context.
    40  type Session struct {
    41  	// Protect the current session will not be used concurrently.
    42  	mu       sync.Mutex
    43  	id       int64
    44  	sc       *stmtctx.Context
    45  	wg       sync.WaitGroup
    46  	store    kv.Storage
    47  	catalog  *catalog.Catalog
    48  	closed   atomic.Bool
    49  	cancelFn context.CancelFunc
    50  
    51  	// Callback function while session closing.
    52  	closeCallback func(s *Session)
    53  }
    54  
    55  // New returns a new session instance.
    56  func New(store kv.Storage, catalog *catalog.Catalog) *Session {
    57  	return &Session{
    58  		id:      idGenerator.Add(1),
    59  		sc:      stmtctx.New(store, catalog),
    60  		store:   store,
    61  		catalog: catalog,
    62  	}
    63  }
    64  
    65  // ID returns a integer identifier of the current session.
    66  func (s *Session) ID() int64 {
    67  	return s.id
    68  }
    69  
    70  // StmtContext returns the statement context object.
    71  func (s *Session) StmtContext() *stmtctx.Context {
    72  	return s.sc
    73  }
    74  
    75  // Execute executes a query and reports whether the query executed successfully or not.
    76  // A result set will be non-empty if execute successfully.
    77  func (s *Session) Execute(ctx context.Context, query string) (ResultSet, error) {
    78  	s.mu.Lock()
    79  	defer s.mu.Unlock()
    80  
    81  	ctx, cancelFn := context.WithCancel(ctx)
    82  	s.cancelFn = cancelFn
    83  	s.wg.Add(1)
    84  	defer s.wg.Done()
    85  
    86  	p := parserPool.Get().(*parser.Parser)
    87  	defer parserPool.Put(p)
    88  
    89  	stmts, warns, err := p.Parse(query)
    90  	if err != nil {
    91  		return nil, err
    92  	}
    93  	for _, warn := range warns {
    94  		s.sc.AppendWarning(errors.Annotate(warn, "parse warning"))
    95  	}
    96  	if len(stmts) == 0 {
    97  		return emptyResultSet{}, nil
    98  	}
    99  	if len(stmts) > 1 {
   100  		return nil, ErrMultipleStatementsNotSupported
   101  	}
   102  
   103  	return s.executeStmt(ctx, stmts[0])
   104  }
   105  
   106  func (s *Session) executeStmt(ctx context.Context, node ast.StmtNode) (ResultSet, error) {
   107  	// TODO: support transaction
   108  
   109  	// Reset the current statement context and prepare for executing the next statement.
   110  	s.sc.Reset()
   111  
   112  	exec, err := compiler.Compile(s.sc, node)
   113  	if err != nil {
   114  		return nil, err
   115  	}
   116  	err = exec.Open(ctx)
   117  	if err != nil {
   118  		return nil, err
   119  	}
   120  
   121  	return newQueryResultSet(exec), nil
   122  }
   123  
   124  // Close terminates the current session.
   125  func (s *Session) Close() {
   126  	if s.closed.Swap(true) {
   127  		return
   128  	}
   129  
   130  	// Wait the current execution finished.
   131  	if s.cancelFn != nil {
   132  		s.cancelFn()
   133  	}
   134  	s.wg.Wait()
   135  
   136  	if s.closeCallback != nil {
   137  		s.closeCallback(s)
   138  	}
   139  }
   140  
   141  // OnClosed sets the closed callback which will invoke after session closed.
   142  func (s *Session) OnClosed(cb func(session *Session)) {
   143  	s.mu.Lock()
   144  	s.mu.Unlock()
   145  	s.closeCallback = cb
   146  }