github.com/RevenueMonster/sqlike@v1.0.6/sql/instrumented/conn.go (about)

     1  package instrumented
     2  
     3  import (
     4  	"context"
     5  	"database/sql/driver"
     6  )
     7  
     8  // Conn :
     9  type Conn interface {
    10  	driver.Conn
    11  	driver.Pinger
    12  	driver.ConnBeginTx
    13  	driver.ConnPrepareContext
    14  	driver.ExecerContext
    15  	driver.QueryerContext
    16  }
    17  
    18  type wrappedConn struct {
    19  	itpr Interceptor
    20  	conn Conn
    21  }
    22  
    23  var _ Conn = (*wrappedConn)(nil)
    24  
    25  // Begin :
    26  func (w wrappedConn) Begin() (driver.Tx, error) {
    27  	tx, err := w.conn.Begin()
    28  	if err != nil {
    29  		return nil, err
    30  	}
    31  	return wrappedTx{ctx: context.Background(), itpr: w.itpr, tx: tx}, nil
    32  }
    33  
    34  // BeginTx :
    35  func (w wrappedConn) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx, error) {
    36  	tx, err := w.itpr.ConnBeginTx(ctx, w.conn, opts)
    37  	if err != nil {
    38  		return nil, err
    39  	}
    40  	return wrappedTx{ctx: ctx, itpr: w.itpr, tx: tx}, nil
    41  }
    42  
    43  // Ping :
    44  func (w wrappedConn) Ping(ctx context.Context) error {
    45  	return w.itpr.ConnPing(ctx, w.conn)
    46  }
    47  
    48  // Prepare :
    49  func (w wrappedConn) Prepare(query string) (driver.Stmt, error) {
    50  	ctx := context.Background()
    51  	stmt, err := w.itpr.ConnPrepareContext(ctx, w.conn, query)
    52  	if err != nil {
    53  		return nil, err
    54  	}
    55  	x, ok := stmt.(Stmt)
    56  	if !ok {
    57  		return nil, driver.ErrBadConn
    58  	}
    59  	return wrappedStmt{ctx: ctx, itpr: w.itpr, query: query, stmt: x}, nil
    60  }
    61  
    62  // PrepareContext :
    63  func (w wrappedConn) PrepareContext(ctx context.Context, query string) (driver.Stmt, error) {
    64  	stmt, err := w.itpr.ConnPrepareContext(ctx, w.conn, query)
    65  	if err != nil {
    66  		return nil, err
    67  	}
    68  	x, ok := stmt.(Stmt)
    69  	if !ok {
    70  		return nil, driver.ErrBadConn
    71  	}
    72  	return wrappedStmt{ctx: ctx, itpr: w.itpr, query: query, stmt: x}, nil
    73  }
    74  
    75  // ExecContext :
    76  func (w wrappedConn) ExecContext(ctx context.Context, query string, args []driver.NamedValue) (res driver.Result, err error) {
    77  	result, err := w.itpr.ConnExecContext(ctx, w.conn, query, args)
    78  	if err != nil {
    79  		return nil, err
    80  	}
    81  	return wrappedResult{ctx: ctx, itpr: w.itpr, result: result}, nil
    82  }
    83  
    84  // QueryContext :
    85  func (w wrappedConn) QueryContext(ctx context.Context, query string, args []driver.NamedValue) (res driver.Rows, err error) {
    86  	rows, err := w.itpr.ConnQueryContext(ctx, w.conn, query, args)
    87  	if err != nil {
    88  		return nil, err
    89  	}
    90  	x, ok := rows.(Rows)
    91  	if !ok {
    92  		return nil, driver.ErrSkip
    93  	}
    94  	return wrappedRows{ctx: ctx, itpr: w.itpr, rows: x}, nil
    95  }
    96  
    97  // Close :
    98  func (w wrappedConn) Close() error {
    99  	return w.conn.Close()
   100  }