m2pool_payment/internal/db/sqlite.go

package db

import (
	"database/sql"
	"fmt"

	_ "modernc.org/sqlite" // 导入驱动
)

type SQLite struct {
	DB *sql.DB
	// Ch chan any
}

// 初始化连接
func NewSQLite(path string) (*SQLite, error) {
	db, err := sql.Open("sqlite", path)
	// ch := make(chan any, 1000)
	if err != nil {
		return nil, fmt.Errorf("open sqlite failed: %v", err)
	}
	return &SQLite{
		DB: db,
		// Ch: ch,
	}, nil
}

// 关闭数据库
func (s *SQLite) Close() {
	if s.DB != nil {
		s.DB.Close()
	}
}

// 新建表
func (s *SQLite) CreateTable(sql string) error {
	_, err := s.DB.Exec(sql)
	if err != nil {
		return fmt.Errorf("Exec DB error: %w", err)
	}
	return nil
}

// 插入数据（支持事务和批量）
func (s *SQLite) Insert(sqlStr string, args ...[]any) error {
	tx, err := s.DB.Begin()
	if err != nil {
		return err
	}
	stmt, err := tx.Prepare(sqlStr)
	if err != nil {
		return err
	}
	defer stmt.Close()

	for _, a := range args {
		_, err = stmt.Exec(a...)
		if err != nil {
			tx.Rollback()
			return fmt.Errorf("exec insert error: %v", err)
		}
	}
	return tx.Commit()
}

// 删除数据
func (s *SQLite) Delete(sqlStr string, args ...any) (int64, error) {
	res, err := s.DB.Exec(sqlStr, args...)
	if err != nil {
		return 0, fmt.Errorf("delete error: %v", err)
	}
	rows, _ := res.RowsAffected()
	return rows, nil
}

// 更新数据
func (s *SQLite) Update(sqlStr string, args ...any) (int64, error) {
	res, err := s.DB.Exec(sqlStr, args...)
	if err != nil {
		return 0, fmt.Errorf("update error: %v", err)
	}
	rows, _ := res.RowsAffected()
	return rows, nil
}

// 通用查询方法（返回[]map[string]any）
func (s *SQLite) Query_(query string, args ...any) ([]map[string]any, error) {
	rows, err := s.DB.Query(query, args...)
	if err != nil {
		return nil, fmt.Errorf("query error: %v", err)
	}
	defer rows.Close()

	columns, err := rows.Columns()
	if err != nil {
		return nil, err
	}

	var results []map[string]any
	for rows.Next() {
		// 创建一个与列数相同的 slice 用于 Scan
		values := make([]any, len(columns))
		valuePtrs := make([]any, len(columns))
		for i := range values {
			valuePtrs[i] = &values[i]
		}

		if err := rows.Scan(valuePtrs...); err != nil {
			return nil, err
		}

		// 将行转换为 map
		rowMap := make(map[string]any)
		for i, col := range columns {
			val := values[i]
			if b, ok := val.([]byte); ok {
				rowMap[col] = string(b)
			} else {
				rowMap[col] = val
			}
		}
		results = append(results, rowMap)
	}
	return results, nil
}

// 事务增删改(多条增删改串行，只要有一条增删改失败则全部回滚并返回错误)
func (s *SQLite) ExecuteTransactions(str_sqls []string, params [][]any) error {
	// 检查 SQL 和参数的数量是否匹配
	if len(str_sqls) != len(params) {
		return fmt.Errorf("sql length != params length")
	}

	// 开始事务
	tx, err := s.DB.Begin()
	if err != nil {
		return fmt.Errorf("failed to begin transaction: %v", err)
	}

	// 确保在函数结束时提交或回滚事务
	defer func() {
		if err != nil {
			// 发生错误时回滚事务
			if rollbackErr := tx.Rollback(); rollbackErr != nil {
				err = fmt.Errorf("failed to rollback transaction: %v", rollbackErr)
			}
		} else {
			// 如果没有错误，提交事务
			if commitErr := tx.Commit(); commitErr != nil {
				err = fmt.Errorf("failed to commit transaction: %v", commitErr)
			}
		}
	}()

	// 执行每个 SQL 语句
	for i, sql_str := range str_sqls {
		// 使用事务对象 tx 来执行 SQL
		_, err := tx.Exec(sql_str, params[i]...)
		if err != nil {
			// 如果执行失败，立即返回并且触发回滚
			return fmt.Errorf("failed to execute SQL: %v", err)
		}
	}

	// 如果所有 SQL 执行成功，则返回 nil
	return nil
}