gitee.com/eden-framework/sqlx@v0.0.3/db.go (about)

     1  package sqlx
     2  
     3  import (
     4  	"context"
     5  	"database/sql"
     6  	"errors"
     7  	"time"
     8  
     9  	"gitee.com/eden-framework/sqlx/builder"
    10  )
    11  
    12  var ErrNotTx = errors.New("db is not *sql.Tx")
    13  var ErrNotDB = errors.New("db is not *sql.DB")
    14  
    15  type SqlExecutor interface {
    16  	ExecContext(ctx context.Context, query string, args ...interface{}) (sql.Result, error)
    17  	QueryContext(ctx context.Context, query string, args ...interface{}) (*sql.Rows, error)
    18  }
    19  
    20  type SqlxExecutor interface {
    21  	SqlExecutor
    22  	ExecExpr(expr builder.SqlExpr) (sql.Result, error)
    23  	QueryExpr(expr builder.SqlExpr) (*sql.Rows, error)
    24  
    25  	QueryExprAndScan(expr builder.SqlExpr, v interface{}) error
    26  }
    27  
    28  type Migrator interface {
    29  	Migrate(ctx context.Context, db DBExecutor) error
    30  }
    31  
    32  type DBExecutor interface {
    33  	SqlxExecutor
    34  
    35  	// dialect of databases
    36  	Dialect() builder.Dialect
    37  	// return database which is connecting
    38  	D() *Database
    39  	// switch database schema
    40  	WithSchema(schema string) DBExecutor
    41  	// return table of the connecting database
    42  	T(model builder.Model) *builder.Table
    43  
    44  	Context() context.Context
    45  	WithContext(ctx context.Context) DBExecutor
    46  }
    47  
    48  type MaybeTxExecutor interface {
    49  	IsTx() bool
    50  	BeginTx(*sql.TxOptions) (DBExecutor, error)
    51  	Begin() (DBExecutor, error)
    52  	Commit() error
    53  	Rollback() error
    54  }
    55  
    56  type DB struct {
    57  	dialect builder.Dialect
    58  	*Database
    59  	SqlExecutor
    60  	ctx context.Context
    61  }
    62  
    63  func (d *DB) WithContext(ctx context.Context) DBExecutor {
    64  	dd := new(DB)
    65  	*dd = *d
    66  	dd.ctx = ctx
    67  	return dd
    68  }
    69  
    70  func (d *DB) Context() context.Context {
    71  	if d.ctx != nil {
    72  		return d.ctx
    73  	}
    74  	return context.Background()
    75  }
    76  
    77  func (d DB) WithSchema(schema string) DBExecutor {
    78  	d.Database = d.Database.WithSchema(schema)
    79  	return &d
    80  }
    81  
    82  func (d *DB) Dialect() builder.Dialect {
    83  	return d.dialect
    84  }
    85  
    86  func (d *DB) Migrate(ctx context.Context, db DBExecutor) error {
    87  	if migrator, ok := d.dialect.(Migrator); ok {
    88  		return migrator.Migrate(ctx, db)
    89  	}
    90  	return nil
    91  }
    92  
    93  func (d *DB) D() *Database {
    94  	return d.Database
    95  }
    96  
    97  func (d *DB) ExecExpr(expr builder.SqlExpr) (sql.Result, error) {
    98  	e := builder.ResolveExprContext(d.Context(), expr)
    99  	if builder.IsNilExpr(e) {
   100  		return nil, nil
   101  	}
   102  	if err := e.Err(); err != nil {
   103  		return nil, err
   104  	}
   105  	result, err := d.ExecContext(d.Context(), e.Query(), e.Args()...)
   106  	if err != nil {
   107  		if d.dialect.IsErrorConflict(err) {
   108  			return nil, NewSqlError(sqlErrTypeConflict, err.Error())
   109  		}
   110  		return nil, err
   111  	}
   112  	return result, nil
   113  }
   114  
   115  func (d *DB) QueryExpr(expr builder.SqlExpr) (*sql.Rows, error) {
   116  	e := builder.ResolveExprContext(d.Context(), expr)
   117  	if builder.IsNilExpr(e) {
   118  		return nil, nil
   119  	}
   120  	if err := e.Err(); err != nil {
   121  		return nil, err
   122  	}
   123  	return d.QueryContext(d.Context(), e.Query(), e.Args()...)
   124  }
   125  
   126  func (d *DB) QueryExprAndScan(expr builder.SqlExpr, v interface{}) error {
   127  	rows, err := d.QueryExpr(expr)
   128  	if err != nil {
   129  		return err
   130  	}
   131  	return Scan(rows, v)
   132  }
   133  
   134  func (d *DB) IsTx() bool {
   135  	_, ok := d.SqlExecutor.(*sql.Tx)
   136  	return ok
   137  }
   138  
   139  func (d *DB) Begin() (DBExecutor, error) {
   140  	return d.BeginTx(nil)
   141  }
   142  
   143  func (d *DB) BeginTx(opt *sql.TxOptions) (DBExecutor, error) {
   144  	if d.IsTx() {
   145  		return nil, ErrNotDB
   146  	}
   147  	db, err := d.SqlExecutor.(*sql.DB).BeginTx(d.Context(), opt)
   148  	if err != nil {
   149  		return nil, err
   150  	}
   151  	return &DB{
   152  		Database:    d.Database,
   153  		dialect:     d.dialect,
   154  		SqlExecutor: db,
   155  		ctx:         d.Context(),
   156  	}, nil
   157  }
   158  
   159  func (d *DB) Commit() error {
   160  	if !d.IsTx() {
   161  		return ErrNotTx
   162  	}
   163  	if d.Context().Err() == context.Canceled {
   164  		return context.Canceled
   165  	}
   166  	return d.SqlExecutor.(*sql.Tx).Commit()
   167  }
   168  
   169  func (d *DB) Rollback() error {
   170  	if !d.IsTx() {
   171  		return ErrNotTx
   172  	}
   173  	if d.Context().Err() == context.Canceled {
   174  		return context.Canceled
   175  	}
   176  	return d.SqlExecutor.(*sql.Tx).Rollback()
   177  }
   178  
   179  func (d *DB) SetMaxOpenConns(n int) {
   180  	d.SqlExecutor.(*sql.DB).SetMaxOpenConns(n)
   181  }
   182  
   183  func (d *DB) SetMaxIdleConns(n int) {
   184  	d.SqlExecutor.(*sql.DB).SetMaxIdleConns(n)
   185  }
   186  
   187  func (d *DB) SetConnMaxLifetime(t time.Duration) {
   188  	d.SqlExecutor.(*sql.DB).SetConnMaxLifetime(t)
   189  }