stardb/stardb.go

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()
}