u

code/go-project/base-farmework/pkg/utils/db/db.go

@@ -2,37 +2,67 @@ package db
 
 import (
 	"base-framework/pkg/config"
-	"base-framework/pkg/reponse"
 	"base-framework/pkg/router"
 	"database/sql"
+	"errors"
 	"fmt"
 )
 
-// Client 封装数据库连接对象
 type Client struct {
-	Conn *sql.DB
+	conn *sql.DB
+	tx   *sql.Tx
 }
 
-// NewClient 根据 context 获取 orgId 并返回 Client 对象
-func NewClient(c *router.Context) *Client {
+// JdbcTemplate 创建 JdbcTemplate
+func JdbcTemplate(c *router.Context) (*Client, error) {
 	val, ok := c.Get("orgId")
-	fmt.Println(val)
 	if !ok {
-		reponse.Error(c).Code(reponse.CodeInvalidOrgCode).Send()
-		return nil
+		return nil, errors.New("missing orgId in context")
 	}
-
 	orgId, ok := val.(string)
 	if !ok || orgId == "" {
-		reponse.Error(c).Code(reponse.CodeInvalidOrgCode).Send()
-		return nil
+		return nil, errors.New("invalid orgId in context")
 	}
-
 	conn, ok := config.GetDB(orgId)
 	if !ok || conn == nil {
-		reponse.Error(c).Code(reponse.CodeInvalidOrgCode).Send()
-		return 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)
 	}
 
-	return &Client{Conn: conn}
+	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()
 }