package stardb import ( "context" "database/sql" "strings" "time" ) // StarDB is a simple wrapper around sql.DB providing enhanced functionality type StarDB struct { db *sql.DB ManualScan bool // If true, rows won't be automatically parsed StrictORM bool // If true, Orm requires all tagged columns to exist in query results // batchInsertMaxRows controls batch split size for BatchInsert/BatchInsertStructs. // <= 0 means no split (single SQL statement). batchInsertMaxRows int64 batchInsertMaxParams int64 runtime sqlRuntime } // NewStarDB creates a new StarDB instance func NewStarDB() *StarDB { return &StarDB{} } // NewStarDBWithDB creates a new StarDB instance with an existing *sql.DB func NewStarDBWithDB(db *sql.DB) *StarDB { return &StarDB{db: db} } // DB returns the underlying *sql.DB func (s *StarDB) DB() *sql.DB { return s.db } // SetDB sets the underlying *sql.DB func (s *StarDB) SetDB(db *sql.DB) { s.db = db } // SetStrictORM enables or disables strict column validation for Orm mapping. func (s *StarDB) SetStrictORM(strict bool) { if s == nil { return } s.StrictORM = strict } func (s *StarDB) ensureDB() error { if s == nil || s.db == nil { return ErrDBNotInitialized } return nil } // Open opens a new database connection func (s *StarDB) Open(driver, connStr string) error { var err error s.db, err = sql.Open(driver, connStr) return err } // Close closes the database connection func (s *StarDB) Close() error { if err := s.ensureDB(); err != nil { return err } return s.db.Close() } // Ping verifies the database connection is alive func (s *StarDB) Ping() error { if err := s.ensureDB(); err != nil { return err } return s.db.Ping() } // PingContext verifies the database connection with context func (s *StarDB) PingContext(ctx context.Context) error { if err := s.ensureDB(); err != nil { return err } return s.db.PingContext(ctx) } // Stats returns database statistics func (s *StarDB) Stats() sql.DBStats { if s == nil || s.db == nil { return sql.DBStats{} } return s.db.Stats() } // SetMaxOpenConns sets the maximum number of open connections func (s *StarDB) SetMaxOpenConns(n int) { if s == nil || s.db == nil { return } s.db.SetMaxOpenConns(n) } // SetMaxIdleConns sets the maximum number of idle connections func (s *StarDB) SetMaxIdleConns(n int) { if s == nil || s.db == nil { return } s.db.SetMaxIdleConns(n) } // Conn returns a single connection from the pool func (s *StarDB) Conn(ctx context.Context) (*sql.Conn, error) { if err := s.ensureDB(); err != nil { return nil, err } return s.db.Conn(ctx) } // Query executes a query that returns rows // Usage: Query("SELECT * FROM users WHERE id = ?", 1) func (s *StarDB) Query(query string, args ...interface{}) (*StarRows, error) { return s.query(nil, query, args...) } // QueryContext executes a query with context func (s *StarDB) QueryContext(ctx context.Context, query string, args ...interface{}) (*StarRows, error) { return s.query(ctx, query, args...) } // QueryRaw executes a query and returns *sql.Rows without automatic parsing. func (s *StarDB) QueryRaw(query string, args ...interface{}) (*sql.Rows, error) { return s.queryRaw(nil, query, args...) } // QueryRawContext executes a query with context and returns *sql.Rows without automatic parsing. func (s *StarDB) QueryRawContext(ctx context.Context, query string, args ...interface{}) (*sql.Rows, error) { return s.queryRaw(ctx, query, args...) } func (s *StarDB) queryRaw(ctx context.Context, query string, args ...interface{}) (*sql.Rows, error) { if err := s.ensureDB(); err != nil { return nil, err } if strings.TrimSpace(query) == "" { return nil, ErrQueryEmpty } query, beforeHook, afterHook, hookArgs, slowThreshold := s.prepareSQLCall(query, args) hookCtx := s.hookContext(ctx, query, beforeHook, afterHook) if beforeHook != nil { beforeHook(hookCtx, query, hookArgs) } start := time.Now() var ( rows *sql.Rows err error ) if ctx == nil { rows, err = s.db.Query(query, args...) } else { rows, err = s.db.QueryContext(ctx, query, args...) } duration := time.Since(start) if shouldRunAfterHook(afterHook, slowThreshold, duration, err) { afterHook(hookCtx, query, hookArgs, duration, err) } if err != nil { return nil, err } return rows, nil } // query is the internal query implementation func (s *StarDB) query(ctx context.Context, query string, args ...interface{}) (*StarRows, error) { rows, err := s.queryRaw(ctx, query, args...) if err != nil { return nil, err } starRows := &StarRows{ rows: rows, db: s, } if !s.ManualScan { if err := starRows.parse(); err != nil { _ = rows.Close() return nil, err } } return starRows, nil } // Exec executes a query that doesn't return rows // Usage: Exec("INSERT INTO users (name) VALUES (?)", "John") func (s *StarDB) Exec(query string, args ...interface{}) (sql.Result, error) { return s.exec(nil, query, args...) } // ExecContext executes a query with context func (s *StarDB) ExecContext(ctx context.Context, query string, args ...interface{}) (sql.Result, error) { return s.exec(ctx, query, args...) } // exec is the internal exec implementation func (s *StarDB) exec(ctx context.Context, query string, args ...interface{}) (sql.Result, error) { if err := s.ensureDB(); err != nil { return nil, err } if strings.TrimSpace(query) == "" { return nil, ErrQueryEmpty } query, beforeHook, afterHook, hookArgs, slowThreshold := s.prepareSQLCall(query, args) hookCtx := s.hookContext(ctx, query, beforeHook, afterHook) if beforeHook != nil { beforeHook(hookCtx, query, hookArgs) } start := time.Now() var ( result sql.Result err error ) if ctx == nil { result, err = s.db.Exec(query, args...) } else { result, err = s.db.ExecContext(ctx, query, args...) } duration := time.Since(start) if shouldRunAfterHook(afterHook, slowThreshold, duration, err) { afterHook(hookCtx, query, hookArgs, duration, err) } if err != nil { return nil, err } return result, nil } // Prepare creates a prepared statement func (s *StarDB) Prepare(query string) (*StarStmt, error) { if err := s.ensureDB(); err != nil { return nil, err } if strings.TrimSpace(query) == "" { return nil, ErrQueryEmpty } query, beforeHook, afterHook, _, slowThreshold := s.prepareSQLCall(query, nil) hookCtx := s.hookContext(nil, query, beforeHook, afterHook) if beforeHook != nil { beforeHook(hookCtx, query, nil) } start := time.Now() stmt, err := s.db.Prepare(query) duration := time.Since(start) if shouldRunAfterHook(afterHook, slowThreshold, duration, err) { afterHook(hookCtx, query, nil, duration, err) } if err != nil { return nil, err } return &StarStmt{stmt: stmt, db: s, sqlText: query}, nil } // PrepareContext creates a prepared statement with context func (s *StarDB) PrepareContext(ctx context.Context, query string) (*StarStmt, error) { if err := s.ensureDB(); err != nil { return nil, err } if strings.TrimSpace(query) == "" { return nil, ErrQueryEmpty } query, beforeHook, afterHook, _, slowThreshold := s.prepareSQLCall(query, nil) hookCtx := s.hookContext(ctx, query, beforeHook, afterHook) if beforeHook != nil { beforeHook(hookCtx, query, nil) } start := time.Now() stmt, err := s.db.PrepareContext(ctx, query) duration := time.Since(start) if shouldRunAfterHook(afterHook, slowThreshold, duration, err) { afterHook(hookCtx, query, nil, duration, err) } if err != nil { return nil, err } return &StarStmt{stmt: stmt, db: s, sqlText: query}, nil } // QueryStmt executes a prepared statement query // Usage: QueryStmt("SELECT * FROM users WHERE id = ?", 1) func (s *StarDB) QueryStmt(query string, args ...interface{}) (*StarRows, error) { if strings.TrimSpace(query) == "" { return nil, ErrQueryEmpty } stmt, err := s.Prepare(query) if err != nil { return nil, err } defer stmt.Close() return stmt.Query(args...) } // QueryStmtContext executes a prepared statement query with context func (s *StarDB) QueryStmtContext(ctx context.Context, query string, args ...interface{}) (*StarRows, error) { if strings.TrimSpace(query) == "" { return nil, ErrQueryEmpty } stmt, err := s.PrepareContext(ctx, query) if err != nil { return nil, err } defer stmt.Close() return stmt.QueryContext(ctx, args...) } // ExecStmt executes a prepared statement // Usage: ExecStmt("INSERT INTO users (name) VALUES (?)", "John") func (s *StarDB) ExecStmt(query string, args ...interface{}) (sql.Result, error) { if strings.TrimSpace(query) == "" { return nil, ErrQueryEmpty } stmt, err := s.Prepare(query) if err != nil { return nil, err } defer stmt.Close() return stmt.Exec(args...) } // ExecStmtContext executes a prepared statement with context func (s *StarDB) ExecStmtContext(ctx context.Context, query string, args ...interface{}) (sql.Result, error) { if strings.TrimSpace(query) == "" { return nil, ErrQueryEmpty } stmt, err := s.PrepareContext(ctx, query) if err != nil { return nil, err } defer stmt.Close() return stmt.ExecContext(ctx, args...) } // Begin starts a transaction func (s *StarDB) Begin() (*StarTx, error) { if err := s.ensureDB(); err != nil { return nil, err } tx, err := s.db.Begin() if err != nil { return nil, err } return &StarTx{tx: tx, db: s}, nil } // BeginTx starts a transaction with options func (s *StarDB) BeginTx(ctx context.Context, opts *sql.TxOptions) (*StarTx, error) { if err := s.ensureDB(); err != nil { return nil, err } tx, err := s.db.BeginTx(ctx, opts) if err != nil { return nil, err } return &StarTx{tx: tx, db: s}, nil } // StarStmt represents a prepared statement type StarStmt struct { stmt *sql.Stmt db *StarDB sqlText string } func (s *StarStmt) ensureStmt() error { if s == nil || s.stmt == nil { return ErrStmtNotInitialized } return nil } func (s *StarStmt) ensureStmtWithDB() error { if err := s.ensureStmt(); err != nil { return err } if s.db == nil { return ErrStmtDBNotInitialized } return nil } // Query executes a prepared statement query func (s *StarStmt) Query(args ...interface{}) (*StarRows, error) { return s.query(nil, args...) } // QueryContext executes a prepared statement query with context func (s *StarStmt) QueryContext(ctx context.Context, args ...interface{}) (*StarRows, error) { return s.query(ctx, args...) } // QueryRaw executes a prepared statement query and returns *sql.Rows without automatic parsing. func (s *StarStmt) QueryRaw(args ...interface{}) (*sql.Rows, error) { return s.queryRaw(nil, args...) } // QueryRawContext executes a prepared statement query with context and returns *sql.Rows without automatic parsing. func (s *StarStmt) QueryRawContext(ctx context.Context, args ...interface{}) (*sql.Rows, error) { return s.queryRaw(ctx, args...) } func (s *StarStmt) queryRaw(ctx context.Context, args ...interface{}) (*sql.Rows, error) { if err := s.ensureStmt(); err != nil { return nil, err } var beforeHook SQLBeforeHook var afterHook SQLAfterHook var slowThreshold time.Duration if s.db != nil { beforeHook, afterHook, slowThreshold = s.db.sqlHooks() } var hookArgs []interface{} if beforeHook != nil || afterHook != nil { hookArgs = cloneHookArgs(args) } hookCtx := ctx if s.db != nil { hookCtx = s.db.hookContext(ctx, s.sqlText, beforeHook, afterHook) } if beforeHook != nil { beforeHook(hookCtx, s.sqlText, hookArgs) } start := time.Now() var ( rows *sql.Rows err error ) if ctx == nil { rows, err = s.stmt.Query(args...) } else { rows, err = s.stmt.QueryContext(ctx, args...) } duration := time.Since(start) if shouldRunAfterHook(afterHook, slowThreshold, duration, err) { afterHook(hookCtx, s.sqlText, hookArgs, duration, err) } if err != nil { return nil, err } return rows, nil } // query is the internal query implementation func (s *StarStmt) query(ctx context.Context, args ...interface{}) (*StarRows, error) { if err := s.ensureStmtWithDB(); err != nil { return nil, err } rows, err := s.queryRaw(ctx, args...) if err != nil { return nil, err } starRows := &StarRows{ rows: rows, db: s.db, } if !s.db.ManualScan { if err := starRows.parse(); err != nil { _ = rows.Close() return nil, err } } return starRows, nil } // Exec executes a prepared statement func (s *StarStmt) Exec(args ...interface{}) (sql.Result, error) { return s.exec(nil, args...) } // ExecContext executes a prepared statement with context func (s *StarStmt) ExecContext(ctx context.Context, args ...interface{}) (sql.Result, error) { return s.exec(ctx, args...) } // exec is the internal exec implementation func (s *StarStmt) exec(ctx context.Context, args ...interface{}) (sql.Result, error) { if err := s.ensureStmt(); err != nil { return nil, err } var beforeHook SQLBeforeHook var afterHook SQLAfterHook var slowThreshold time.Duration if s.db != nil { beforeHook, afterHook, slowThreshold = s.db.sqlHooks() } var hookArgs []interface{} if beforeHook != nil || afterHook != nil { hookArgs = cloneHookArgs(args) } hookCtx := ctx if s.db != nil { hookCtx = s.db.hookContext(ctx, s.sqlText, beforeHook, afterHook) } if beforeHook != nil { beforeHook(hookCtx, s.sqlText, hookArgs) } start := time.Now() var ( result sql.Result err error ) if ctx == nil { result, err = s.stmt.Exec(args...) } else { result, err = s.stmt.ExecContext(ctx, args...) } duration := time.Since(start) if shouldRunAfterHook(afterHook, slowThreshold, duration, err) { afterHook(hookCtx, s.sqlText, hookArgs, duration, err) } if err != nil { return nil, err } return result, nil } // Close closes the prepared statement func (s *StarStmt) Close() error { if err := s.ensureStmt(); err != nil { return err } return s.stmt.Close() }