69 lines
1.5 KiB
Go
69 lines
1.5 KiB
Go
package db
|
||
|
||
import (
|
||
"base-framework/pkg/config"
|
||
"base-framework/pkg/router"
|
||
"database/sql"
|
||
"errors"
|
||
"fmt"
|
||
)
|
||
|
||
type Client struct {
|
||
conn *sql.DB
|
||
tx *sql.Tx
|
||
}
|
||
|
||
// JdbcTemplate 创建 JdbcTemplate
|
||
func JdbcTemplate(c *router.Context) (*Client, error) {
|
||
val, ok := c.Get("orgId")
|
||
if !ok {
|
||
return nil, errors.New("missing orgId in context")
|
||
}
|
||
orgId, ok := val.(string)
|
||
if !ok || orgId == "" {
|
||
return nil, errors.New("invalid orgId in context")
|
||
}
|
||
conn, ok := config.GetDB(orgId)
|
||
if !ok || conn == nil {
|
||
return nil, fmt.Errorf("no db connection found for orgId=%s", orgId)
|
||
}
|
||
return &Client{conn: conn}, nil
|
||
}
|
||
|
||
// ------------------------ 内部方法 ------------------------
|
||
|
||
// 执行查询，返回 *sql.Rows
|
||
func (c *Client) query(query string, args ...any) (*sql.Rows, error) {
|
||
if c.tx != nil {
|
||
return c.tx.Query(query, args...)
|
||
}
|
||
return c.conn.Query(query, args...)
|
||
}
|
||
|
||
// 执行执行类语句（insert/update/delete）
|
||
func (c *Client) exec(query string, args ...any) (sql.Result, error) {
|
||
if c.tx != nil {
|
||
return c.tx.Exec(query, args...)
|
||
}
|
||
return c.conn.Exec(query, args...)
|
||
}
|
||
|
||
// WithTransaction 自动处理事务提交或回滚
|
||
func (c *Client) WithTransaction(fn func(txClient *Client) error) error {
|
||
if c.tx != nil {
|
||
// 已经在事务中，直接执行
|
||
return fn(c)
|
||
}
|
||
|
||
tx, err := c.conn.Begin()
|
||
if err != nil {
|
||
return err
|
||
}
|
||
txClient := &Client{conn: c.conn, tx: tx}
|
||
if err := fn(txClient); err != nil {
|
||
_ = tx.Rollback()
|
||
return err
|
||
}
|
||
return tx.Commit()
|
||
}
|