stardb/stardb.go

package stardb

import (
	"context"
	"database/sql"
	"errors"
)

// 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
}

// 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
}

// 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 {
	return s.db.Close()
}

// Ping verifies the database connection is alive
func (s *StarDB) Ping() error {
	return s.db.Ping()
}

// PingContext verifies the database connection with context
func (s *StarDB) PingContext(ctx context.Context) error {
	return s.db.PingContext(ctx)
}

// Stats returns database statistics
func (s *StarDB) Stats() sql.DBStats {
	return s.db.Stats()
}

// SetMaxOpenConns sets the maximum number of open connections
func (s *StarDB) SetMaxOpenConns(n int) {
	s.db.SetMaxOpenConns(n)
}

// SetMaxIdleConns sets the maximum number of idle connections
func (s *StarDB) SetMaxIdleConns(n int) {
	s.db.SetMaxIdleConns(n)
}

// Conn returns a single connection from the pool
func (s *StarDB) Conn(ctx context.Context) (*sql.Conn, error) {
	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...)
}

// query is the internal query implementation
func (s *StarDB) query(ctx context.Context, query string, args ...interface{}) (*StarRows, error) {
	if err := s.db.Ping(); err != nil {
		return nil, err
	}

	var rows *sql.Rows
	var err error

	if ctx == nil {
		rows, err = s.db.Query(query, args...)
	} else {
		rows, err = s.db.QueryContext(ctx, query, args...)
	}

	if err != nil {
		return nil, err
	}

	starRows := &StarRows{
		rows: rows,
		db:   s,
	}

	if !s.ManualScan {
		err = starRows.parse()
	}

	return starRows, err
}

// 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.db.Ping(); err != nil {
		return nil, err
	}

	if ctx == nil {
		return s.db.Exec(query, args...)
	}
	return s.db.ExecContext(ctx, query, args...)
}

// Prepare creates a prepared statement
func (s *StarDB) Prepare(query string) (*StarStmt, error) {
	stmt, err := s.db.Prepare(query)
	if err != nil {
		return nil, err
	}
	return &StarStmt{stmt: stmt, db: s}, nil
}

// PrepareContext creates a prepared statement with context
func (s *StarDB) PrepareContext(ctx context.Context, query string) (*StarStmt, error) {
	stmt, err := s.db.PrepareContext(ctx, query)
	if err != nil {
		return nil, err
	}
	return &StarStmt{stmt: stmt, db: s}, 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 query == "" {
		return nil, errors.New("query string cannot be empty")
	}
	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 query == "" {
		return nil, errors.New("query string cannot be empty")
	}
	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 query == "" {
		return nil, errors.New("query string cannot be empty")
	}
	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 query == "" {
		return nil, errors.New("query string cannot be empty")
	}
	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) {
	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) {
	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
}

// 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...)
}

// query is the internal query implementation
func (s *StarStmt) query(ctx context.Context, args ...interface{}) (*StarRows, error) {
	var rows *sql.Rows
	var err error

	if ctx == nil {
		rows, err = s.stmt.Query(args...)
	} else {
		rows, err = s.stmt.QueryContext(ctx, args...)
	}

	if err != nil {
		return nil, err
	}

	starRows := &StarRows{
		rows: rows,
		db:   s.db,
	}

	if !s.db.ManualScan {
		err = starRows.parse()
	}

	return starRows, err
}

// 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 ctx == nil {
		return s.stmt.Exec(args...)
	}
	return s.stmt.ExecContext(ctx, args...)
}

// Close closes the prepared statement
func (s *StarStmt) Close() error {
	return s.stmt.Close()
}