github.com/erda-project/erda-infra@v1.0.9/pkg/trace/inject/sql/conn.go (about)

     1  // Copyright (c) 2021 Terminus, Inc.
     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 sql
    16  
    17  import (
    18  	"context"
    19  	"database/sql/driver"
    20  
    21  	injectcontext "github.com/erda-project/erda-infra/pkg/trace/inject/context"
    22  )
    23  
    24  type wrappedConn struct {
    25  	driver.Conn
    26  }
    27  
    28  var (
    29  	_ driver.Pinger             = (*wrappedConn)(nil)
    30  	_ driver.Execer             = (*wrappedConn)(nil) // nolint
    31  	_ driver.ExecerContext      = (*wrappedConn)(nil)
    32  	_ driver.Queryer            = (*wrappedConn)(nil) // nolint
    33  	_ driver.QueryerContext     = (*wrappedConn)(nil)
    34  	_ driver.Conn               = (*wrappedConn)(nil)
    35  	_ driver.ConnPrepareContext = (*wrappedConn)(nil)
    36  	_ driver.ConnBeginTx        = (*wrappedConn)(nil)
    37  	_ driver.SessionResetter    = (*wrappedConn)(nil)
    38  	_ driver.NamedValueChecker  = (*wrappedConn)(nil)
    39  )
    40  
    41  func (c *wrappedConn) Ping(ctx context.Context) error {
    42  	pinger, ok := c.Conn.(driver.Pinger)
    43  	if !ok {
    44  		return driver.ErrSkip
    45  	}
    46  	return pinger.Ping(injectcontext.ContextWithSpan(ctx))
    47  }
    48  
    49  func (c *wrappedConn) Exec(query string, args []driver.Value) (driver.Result, error) {
    50  	exec, ok := c.Conn.(driver.ExecerContext)
    51  	if ok {
    52  		return exec.ExecContext(injectcontext.ContextWithSpan(context.Background()), query, convertToNamedArgs(args))
    53  	}
    54  	exe, ok := c.Conn.(driver.Execer)
    55  	if ok {
    56  		return exe.Exec(query, args)
    57  	}
    58  	return nil, driver.ErrSkip
    59  }
    60  
    61  func (c *wrappedConn) ExecContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Result, error) {
    62  	execer, ok := c.Conn.(driver.ExecerContext)
    63  	if !ok {
    64  		return nil, driver.ErrSkip
    65  	}
    66  	return execer.ExecContext(injectcontext.ContextWithSpan(ctx), query, args)
    67  }
    68  
    69  func (c *wrappedConn) Query(query string, args []driver.Value) (driver.Rows, error) {
    70  	queryc, ok := c.Conn.(driver.QueryerContext)
    71  	if ok {
    72  		return queryc.QueryContext(injectcontext.ContextWithSpan(context.Background()), query, convertToNamedArgs(args))
    73  	}
    74  	queryer, ok := c.Conn.(driver.Queryer)
    75  	if ok {
    76  		return queryer.Query(query, args)
    77  	}
    78  	return nil, driver.ErrSkip
    79  }
    80  
    81  func (c *wrappedConn) QueryContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Rows, error) {
    82  	queryer, ok := c.Conn.(driver.QueryerContext)
    83  	if !ok {
    84  		return nil, driver.ErrSkip
    85  	}
    86  	return queryer.QueryContext(injectcontext.ContextWithSpan(ctx), query, args)
    87  }
    88  
    89  func (c *wrappedConn) PrepareContext(ctx context.Context, query string) (driver.Stmt, error) {
    90  	conn, ok := c.Conn.(driver.ConnPrepareContext)
    91  	if !ok {
    92  		return nil, driver.ErrSkip
    93  	}
    94  	return conn.PrepareContext(injectcontext.ContextWithSpan(ctx), query)
    95  }
    96  
    97  func (c *wrappedConn) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx, error) {
    98  	tx, ok := c.Conn.(driver.ConnBeginTx)
    99  	if !ok {
   100  		return nil, driver.ErrSkip
   101  	}
   102  	return tx.BeginTx(injectcontext.ContextWithSpan(ctx), opts)
   103  }
   104  
   105  func (c *wrappedConn) ResetSession(ctx context.Context) error {
   106  	r, ok := c.Conn.(driver.SessionResetter)
   107  	if !ok {
   108  		return driver.ErrSkip
   109  	}
   110  	return r.ResetSession(injectcontext.ContextWithSpan(ctx))
   111  }
   112  
   113  func (c *wrappedConn) CheckNamedValue(nv *driver.NamedValue) error {
   114  	nvc, ok := c.Conn.(driver.NamedValueChecker)
   115  	if !ok {
   116  		return driver.ErrSkip
   117  	}
   118  	return nvc.CheckNamedValue(nv)
   119  }
   120  
   121  func convertToNamedArgs(args []driver.Value) []driver.NamedValue {
   122  	nargs := make([]driver.NamedValue, len(args))
   123  	for i, arg := range args {
   124  		nargs[i] = driver.NamedValue{
   125  			Ordinal: i + 1,
   126  			Value:   arg,
   127  		}
   128  	}
   129  	return nargs
   130  }