app-cloudep-trade-service/internal/model/wallet_model.go

package model

import (
	"app-cloudep-trade-service/internal/domain/wallet"
	"context"
	"database/sql"
	"errors"
	"fmt"
	"strings"

	"github.com/zeromicro/go-zero/core/stores/sqlc"

	"github.com/zeromicro/go-zero/core/stores/sqlx"
)

var _ WalletModel = (*customWalletModel)(nil)

type (
	// WalletModel is an interface to be customized, add more methods here,
	// and implement the added methods in customWalletModel.
	WalletModel interface {
		walletModel
		InsertMany(ctx context.Context, wallets []*Wallet) (sql.Result, error)
		// Balances 給他 select for update 的鎖先加上去
		Balances(ctx context.Context, req BalanceReq, withLock bool, tx sqlx.Session) ([]Wallet, error)
		// BalancesByIDs 給他 select for update 的鎖先加上去
		BalancesByIDs(ctx context.Context, ids []int64, withLock bool, tx sqlx.Session) ([]Wallet, error)
	}

	customWalletModel struct {
		*defaultWalletModel
	}

	BalanceReq struct {
		UID      []string
		Currency []string
		Kind     []wallet.BalanceType
	}
)

// NewWalletModel returns a model for the database table.
func NewWalletModel(conn sqlx.SqlConn) WalletModel {
	return &customWalletModel{
		defaultWalletModel: newWalletModel(conn),
	}
}

func (m *customWalletModel) InsertMany(ctx context.Context, wallets []*Wallet) (sql.Result, error) {
	if len(wallets) == 0 {
		return nil, fmt.Errorf("no data to insert")
	}

	// 構建多條記錄的佔位符，例如: (?, ?, ?, ?, ?, ?, ?), (?, ?, ?, ?, ?, ?, ?), ...
	valueStrings := make([]string, 0, len(wallets))
	valueArgs := make([]interface{}, 0, len(wallets)*7) // 每條記錄有7個值

	for _, wallet := range wallets {
		valueStrings = append(valueStrings, "(?, ?, ?, ?, ?, ?, ?)")
		valueArgs = append(valueArgs, wallet.Uid, wallet.Brand, wallet.Currency, wallet.Balance, wallet.WalletType, wallet.CreatedAt, wallet.UpdatedAt)
	}

	// 構建批量插入的 SQL 語句
	query := fmt.Sprintf("insert into %s (%s) values %s", m.table, walletRowsExpectAutoSet, strings.Join(valueStrings, ","))

	// 使用單一連線執行批量插入
	return m.conn.ExecCtx(ctx, query, valueArgs...)
}

func (m *customWalletModel) Balances(ctx context.Context, req BalanceReq, isLock bool, tx sqlx.Session) ([]Wallet, error) {
	baseQuery := fmt.Sprintf("SELECT `id`, `currency`, `balance`, `wallet_type` FROM %s", m.table)

	// 構建條件字典
	conditions := map[string][]any{
		"`uid`":         convertSliceToInterface(req.UID),
		"`currency`":    convertSliceToInterface(req.Currency),
		"`wallet_type`": convertSliceToInterface(req.Kind),
	}
	// 使用 queryBuilder 構建完整查詢
	query, args := queryBindINBuilder(baseQuery, conditions)

	if isLock {
		// 加上排他鎖
		query += " FOR UPDATE"
	}

	// 執行查詢
	var wallets []Wallet
	err := tx.QueryRowsPartialCtx(ctx, &wallets, query, args...)
	switch {
	case err == nil:
		return wallets, nil
	case errors.Is(err, sqlc.ErrNotFound):
		return nil, ErrNotFound
	default:
		return nil, err
	}
}

func (m *customWalletModel) BalancesByIDs(ctx context.Context, ids []int64, withLock bool, tx sqlx.Session) ([]Wallet, error) {
	baseQuery := fmt.Sprintf("SELECT `id`, `currency`, `balance`, `wallet_type` FROM %s", m.table)

	// 構建條件字典
	conditions := map[string][]any{
		"`id`": convertSliceToInterface(ids),
	}
	// 使用 queryBuilder 構建完整查詢
	query, args := queryBindINBuilder(baseQuery, conditions)

	if withLock {
		// 加上排他鎖
		query += " FOR UPDATE"
	}

	// 執行查詢
	var wallets []Wallet
	err := tx.QueryRowsPartialCtx(ctx, &wallets, query, args...)
	switch {
	case err == nil:
		return wallets, nil
	case errors.Is(err, sqlc.ErrNotFound):
		return nil, ErrNotFound
	default:
		return nil, err
	}
}