github.com/TeaOSLab/EdgeNode@v1.3.8/internal/utils/dbs/stmt.go (about)

     1  // Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
     2  
     3  package dbs
     4  
     5  import (
     6  	"context"
     7  	"database/sql"
     8  	fsutils "github.com/TeaOSLab/EdgeNode/internal/utils/fs"
     9  )
    10  
    11  type Stmt struct {
    12  	db      *DB
    13  	rawStmt *sql.Stmt
    14  	query   string
    15  
    16  	enableStat bool
    17  }
    18  
    19  func NewStmt(db *DB, rawStmt *sql.Stmt, query string) *Stmt {
    20  	return &Stmt{
    21  		db:      db,
    22  		rawStmt: rawStmt,
    23  		query:   query,
    24  	}
    25  }
    26  
    27  func (this *Stmt) EnableStat() {
    28  	this.enableStat = true
    29  }
    30  
    31  func (this *Stmt) ExecContext(ctx context.Context, args ...any) (result sql.Result, err error) {
    32  	// check database status
    33  	if this.db.BeginUpdating() {
    34  		defer this.db.EndUpdating()
    35  	} else {
    36  		return nil, errDBIsClosed
    37  	}
    38  
    39  	if this.enableStat {
    40  		defer SharedQueryStatManager.AddQuery(this.query).End()
    41  	}
    42  	fsutils.WriterLimiter.Ack()
    43  	result, err = this.rawStmt.ExecContext(ctx, args...)
    44  	fsutils.WriterLimiter.Release()
    45  	return
    46  }
    47  
    48  func (this *Stmt) Exec(args ...any) (result sql.Result, err error) {
    49  	// check database status
    50  	if this.db.BeginUpdating() {
    51  		defer this.db.EndUpdating()
    52  	} else {
    53  		return nil, errDBIsClosed
    54  	}
    55  
    56  	if this.enableStat {
    57  		defer SharedQueryStatManager.AddQuery(this.query).End()
    58  	}
    59  
    60  	fsutils.WriterLimiter.Ack()
    61  	result, err = this.rawStmt.Exec(args...)
    62  	fsutils.WriterLimiter.Release()
    63  	return
    64  }
    65  
    66  func (this *Stmt) QueryContext(ctx context.Context, args ...any) (*sql.Rows, error) {
    67  	if this.enableStat {
    68  		defer SharedQueryStatManager.AddQuery(this.query).End()
    69  	}
    70  	return this.rawStmt.QueryContext(ctx, args...)
    71  }
    72  
    73  func (this *Stmt) Query(args ...any) (*sql.Rows, error) {
    74  	if this.enableStat {
    75  		defer SharedQueryStatManager.AddQuery(this.query).End()
    76  	}
    77  	rows, err := this.rawStmt.Query(args...)
    78  	if err != nil {
    79  		return nil, err
    80  	}
    81  	var rowsErr = rows.Err()
    82  	if rowsErr != nil {
    83  		_ = rows.Close()
    84  		return nil, rowsErr
    85  	}
    86  	return rows, nil
    87  }
    88  
    89  func (this *Stmt) QueryRowContext(ctx context.Context, args ...any) *sql.Row {
    90  	if this.enableStat {
    91  		defer SharedQueryStatManager.AddQuery(this.query).End()
    92  	}
    93  	return this.rawStmt.QueryRowContext(ctx, args...)
    94  }
    95  
    96  func (this *Stmt) QueryRow(args ...any) *sql.Row {
    97  	if this.enableStat {
    98  		defer SharedQueryStatManager.AddQuery(this.query).End()
    99  	}
   100  	return this.rawStmt.QueryRow(args...)
   101  }
   102  
   103  func (this *Stmt) Close() error {
   104  	return this.rawStmt.Close()
   105  }