gtsocial-umbx

Unnamed repository; edit this file 'description' to name the repository.
Log | Files | Refs | README | LICENSE

hook.go (2083B)

      1 package bun
      2 
      3 import (
      4 	"context"
      5 	"database/sql"
      6 	"strings"
      7 	"sync/atomic"
      8 	"time"
      9 	"unicode"
     10 
     11 	"github.com/uptrace/bun/schema"
     12 )
     13 
     14 type QueryEvent struct {
     15 	DB *DB
     16 
     17 	QueryAppender schema.QueryAppender // DEPRECATED: use IQuery instead
     18 	IQuery        Query
     19 	Query         string
     20 	QueryTemplate string
     21 	QueryArgs     []interface{}
     22 	Model         Model
     23 
     24 	StartTime time.Time
     25 	Result    sql.Result
     26 	Err       error
     27 
     28 	Stash map[interface{}]interface{}
     29 }
     30 
     31 func (e *QueryEvent) Operation() string {
     32 	if e.IQuery != nil {
     33 		return e.IQuery.Operation()
     34 	}
     35 	return queryOperation(e.Query)
     36 }
     37 
     38 func queryOperation(query string) string {
     39 	queryOp := strings.TrimLeftFunc(query, unicode.IsSpace)
     40 
     41 	if idx := strings.IndexByte(queryOp, ' '); idx > 0 {
     42 		queryOp = queryOp[:idx]
     43 	}
     44 	if len(queryOp) > 16 {
     45 		queryOp = queryOp[:16]
     46 	}
     47 	return queryOp
     48 }
     49 
     50 type QueryHook interface {
     51 	BeforeQuery(context.Context, *QueryEvent) context.Context
     52 	AfterQuery(context.Context, *QueryEvent)
     53 }
     54 
     55 func (db *DB) beforeQuery(
     56 	ctx context.Context,
     57 	iquery Query,
     58 	queryTemplate string,
     59 	queryArgs []interface{},
     60 	query string,
     61 	model Model,
     62 ) (context.Context, *QueryEvent) {
     63 	atomic.AddUint32(&db.stats.Queries, 1)
     64 
     65 	if len(db.queryHooks) == 0 {
     66 		return ctx, nil
     67 	}
     68 
     69 	event := &QueryEvent{
     70 		DB: db,
     71 
     72 		Model:         model,
     73 		QueryAppender: iquery,
     74 		IQuery:        iquery,
     75 		Query:         query,
     76 		QueryTemplate: queryTemplate,
     77 		QueryArgs:     queryArgs,
     78 
     79 		StartTime: time.Now(),
     80 	}
     81 
     82 	for _, hook := range db.queryHooks {
     83 		ctx = hook.BeforeQuery(ctx, event)
     84 	}
     85 
     86 	return ctx, event
     87 }
     88 
     89 func (db *DB) afterQuery(
     90 	ctx context.Context,
     91 	event *QueryEvent,
     92 	res sql.Result,
     93 	err error,
     94 ) {
     95 	switch err {
     96 	case nil, sql.ErrNoRows:
     97 		// nothing
     98 	default:
     99 		atomic.AddUint32(&db.stats.Errors, 1)
    100 	}
    101 
    102 	if event == nil {
    103 		return
    104 	}
    105 
    106 	event.Result = res
    107 	event.Err = err
    108 
    109 	db.afterQueryFromIndex(ctx, event, len(db.queryHooks)-1)
    110 }
    111 
    112 func (db *DB) afterQueryFromIndex(ctx context.Context, event *QueryEvent, hookIndex int) {
    113 	for ; hookIndex >= 0; hookIndex-- {
    114 		db.queryHooks[hookIndex].AfterQuery(ctx, event)
    115 	}
    116 }