github.com/authzed/spicedb@v1.32.1-0.20240520085336-ebda56537386/internal/datastore/postgres/tracer.go (about)

     1  package postgres
     2  
     3  import (
     4  	"context"
     5  	"database/sql/driver"
     6  
     7  	"github.com/ngrok/sqlmw"
     8  )
     9  
    10  type traceInterceptor struct {
    11  	sqlmw.NullInterceptor
    12  }
    13  
    14  func (ti *traceInterceptor) ConnBeginTx(ctx context.Context, conn driver.ConnBeginTx, opts driver.TxOptions) (context.Context, driver.Tx, error) {
    15  	ctx, span := tracer.Start(ctx, "ConnBeginTx")
    16  	defer span.End()
    17  
    18  	tx, err := conn.BeginTx(ctx, opts)
    19  	return ctx, tx, err
    20  }
    21  
    22  func (ti *traceInterceptor) ConnPrepareContext(ctx context.Context, conn driver.ConnPrepareContext, query string) (context.Context, driver.Stmt, error) {
    23  	ctx, span := tracer.Start(ctx, "ConnPrepareContext")
    24  	defer span.End()
    25  	stmt, err := conn.PrepareContext(ctx, query)
    26  	return ctx, stmt, err
    27  }
    28  
    29  func (ti *traceInterceptor) ConnPing(ctx context.Context, conn driver.Pinger) error {
    30  	ctx, span := tracer.Start(ctx, "ConnPing")
    31  	defer span.End()
    32  
    33  	return conn.Ping(ctx)
    34  }
    35  
    36  func (ti *traceInterceptor) ConnExecContext(ctx context.Context, conn driver.ExecerContext, query string, args []driver.NamedValue) (driver.Result, error) {
    37  	ctx, span := tracer.Start(ctx, "ConnExecContext")
    38  	defer span.End()
    39  
    40  	return conn.ExecContext(ctx, query, args)
    41  }
    42  
    43  func (ti *traceInterceptor) ConnQueryContext(ctx context.Context, conn driver.QueryerContext, query string, args []driver.NamedValue) (context.Context, driver.Rows, error) {
    44  	ctx, span := tracer.Start(ctx, "ConnQueryContext")
    45  	defer span.End()
    46  
    47  	rows, err := conn.QueryContext(ctx, query, args)
    48  	return ctx, rows, err
    49  }
    50  
    51  // Connector interceptors
    52  func (ti *traceInterceptor) ConnectorConnect(ctx context.Context, conn driver.Connector) (driver.Conn, error) {
    53  	ctx, span := tracer.Start(ctx, "ConnectorConnect")
    54  	defer span.End()
    55  
    56  	return conn.Connect(ctx)
    57  }
    58  
    59  // Rows interceptors
    60  func (ti *traceInterceptor) RowsNext(ctx context.Context, conn driver.Rows, dest []driver.Value) error {
    61  	_, span := tracer.Start(ctx, "RowsNext")
    62  	defer span.End()
    63  
    64  	return conn.Next(dest)
    65  }
    66  
    67  // Stmt interceptors
    68  func (ti *traceInterceptor) StmtExecContext(ctx context.Context, conn driver.StmtExecContext, _ string, args []driver.NamedValue) (driver.Result, error) {
    69  	ctx, span := tracer.Start(ctx, "StmtExecContext")
    70  	defer span.End()
    71  
    72  	return conn.ExecContext(ctx, args)
    73  }
    74  
    75  func (ti *traceInterceptor) StmtQueryContext(ctx context.Context, conn driver.StmtQueryContext, _ string, args []driver.NamedValue) (context.Context, driver.Rows, error) {
    76  	ctx, span := tracer.Start(ctx, "StmtQueryContext")
    77  	defer span.End()
    78  
    79  	rows, err := conn.QueryContext(ctx, args)
    80  	return ctx, rows, err
    81  }
    82  
    83  func (ti *traceInterceptor) StmtClose(ctx context.Context, conn driver.Stmt) error {
    84  	_, span := tracer.Start(ctx, "StmtClose")
    85  	defer span.End()
    86  
    87  	return conn.Close()
    88  }
    89  
    90  // Tx interceptors
    91  func (ti *traceInterceptor) TxCommit(ctx context.Context, conn driver.Tx) error {
    92  	_, span := tracer.Start(ctx, "TxCommit")
    93  	defer span.End()
    94  
    95  	return conn.Commit()
    96  }
    97  
    98  func (ti *traceInterceptor) TxRollback(ctx context.Context, conn driver.Tx) error {
    99  	_, span := tracer.Start(ctx, "TxRollback")
   100  	defer span.End()
   101  
   102  	return conn.Rollback()
   103  }