github.com/blend/go-sdk@v1.20220411.3/db/query_event.go (about)

     1  /*
     2  
     3  Copyright (c) 2022 - Present. Blend Labs, Inc. All rights reserved
     4  Use of this source code is governed by a MIT license that can be found in the LICENSE file.
     5  
     6  */
     7  
     8  package db
     9  
    10  import (
    11  	"context"
    12  	"fmt"
    13  	"io"
    14  	"time"
    15  
    16  	"github.com/blend/go-sdk/ansi"
    17  	"github.com/blend/go-sdk/logger"
    18  	"github.com/blend/go-sdk/stringutil"
    19  	"github.com/blend/go-sdk/timeutil"
    20  )
    21  
    22  // Logger flags
    23  const (
    24  	QueryFlag = "db.query"
    25  )
    26  
    27  // these are compile time assertions
    28  var (
    29  	_ logger.Event        = (*QueryEvent)(nil)
    30  	_ logger.TextWritable = (*QueryEvent)(nil)
    31  	_ logger.JSONWritable = (*QueryEvent)(nil)
    32  )
    33  
    34  // NewQueryEvent creates a new query event.
    35  func NewQueryEvent(body string, elapsed time.Duration, options ...QueryEventOption) QueryEvent {
    36  	qe := QueryEvent{
    37  		Body:    body,
    38  		Elapsed: elapsed,
    39  	}
    40  	for _, opt := range options {
    41  		opt(&qe)
    42  	}
    43  	return qe
    44  }
    45  
    46  // NewQueryEventListener returns a new listener for query events.
    47  func NewQueryEventListener(listener func(context.Context, QueryEvent)) logger.Listener {
    48  	return func(ctx context.Context, e logger.Event) {
    49  		if typed, isTyped := e.(QueryEvent); isTyped {
    50  			listener(ctx, typed)
    51  		}
    52  	}
    53  }
    54  
    55  // NewQueryEventFilter returns a new query event filter.
    56  func NewQueryEventFilter(filter func(context.Context, QueryEvent) (QueryEvent, bool)) logger.Filter {
    57  	return func(ctx context.Context, e logger.Event) (logger.Event, bool) {
    58  		if typed, isTyped := e.(QueryEvent); isTyped {
    59  			return filter(ctx, typed)
    60  		}
    61  		return e, false
    62  	}
    63  }
    64  
    65  // QueryEventOption mutates a query event.
    66  type QueryEventOption func(*QueryEvent)
    67  
    68  // OptQueryEventBody sets a field on the query event.
    69  func OptQueryEventBody(value string) QueryEventOption {
    70  	return func(e *QueryEvent) { e.Body = value }
    71  }
    72  
    73  // OptQueryEventDatabase sets a field on the query event.
    74  func OptQueryEventDatabase(value string) QueryEventOption {
    75  	return func(e *QueryEvent) { e.Database = value }
    76  }
    77  
    78  // OptQueryEventEngine sets a field on the query event.
    79  func OptQueryEventEngine(value string) QueryEventOption {
    80  	return func(e *QueryEvent) { e.Engine = value }
    81  }
    82  
    83  // OptQueryEventUsername sets a field on the query event.
    84  func OptQueryEventUsername(value string) QueryEventOption {
    85  	return func(e *QueryEvent) { e.Username = value }
    86  }
    87  
    88  // OptQueryEventLabel sets a field on the query event.
    89  func OptQueryEventLabel(label string) QueryEventOption {
    90  	return func(e *QueryEvent) { e.Label = label }
    91  }
    92  
    93  // OptQueryEventElapsed sets a field on the query event.
    94  func OptQueryEventElapsed(value time.Duration) QueryEventOption {
    95  	return func(e *QueryEvent) { e.Elapsed = value }
    96  }
    97  
    98  // OptQueryEventErr sets a field on the query event.
    99  func OptQueryEventErr(value error) QueryEventOption {
   100  	return func(e *QueryEvent) { e.Err = value }
   101  }
   102  
   103  // QueryEvent represents a database query.
   104  type QueryEvent struct {
   105  	Database string
   106  	Engine   string
   107  	Username string
   108  	Label    string
   109  	Body     string
   110  	Elapsed  time.Duration
   111  	Err      error
   112  }
   113  
   114  // GetFlag implements Event.
   115  func (e QueryEvent) GetFlag() string { return QueryFlag }
   116  
   117  // WriteText writes the event text to the output.
   118  func (e QueryEvent) WriteText(tf logger.TextFormatter, wr io.Writer) {
   119  	fmt.Fprint(wr, "[")
   120  	if len(e.Engine) > 0 {
   121  		fmt.Fprint(wr, tf.Colorize(e.Engine, ansi.ColorLightWhite))
   122  		fmt.Fprint(wr, logger.Space)
   123  	}
   124  	if len(e.Username) > 0 {
   125  		fmt.Fprint(wr, tf.Colorize(e.Username, ansi.ColorLightWhite))
   126  		fmt.Fprint(wr, "@")
   127  	}
   128  	fmt.Fprint(wr, tf.Colorize(e.Database, ansi.ColorLightWhite))
   129  	fmt.Fprint(wr, "]")
   130  
   131  	if len(e.Label) > 0 {
   132  		fmt.Fprint(wr, logger.Space)
   133  		fmt.Fprintf(wr, "[%s]", tf.Colorize(e.Label, ansi.ColorLightWhite))
   134  	}
   135  
   136  	if len(e.Body) > 0 {
   137  		fmt.Fprint(wr, logger.Space)
   138  		fmt.Fprint(wr, stringutil.CompressSpace(e.Body))
   139  	}
   140  
   141  	fmt.Fprint(wr, logger.Space)
   142  	fmt.Fprint(wr, e.Elapsed.String())
   143  
   144  	if e.Err != nil {
   145  		fmt.Fprint(wr, logger.Space)
   146  		fmt.Fprint(wr, tf.Colorize("failed", ansi.ColorRed))
   147  	}
   148  }
   149  
   150  // Decompose implements JSONWritable.
   151  func (e QueryEvent) Decompose() map[string]interface{} {
   152  	return map[string]interface{}{
   153  		"engine":   e.Engine,
   154  		"database": e.Database,
   155  		"username": e.Username,
   156  		"label":    e.Label,
   157  		"body":     e.Body,
   158  		"err":      e.Err,
   159  		"elapsed":  timeutil.Milliseconds(e.Elapsed),
   160  	}
   161  }