From 4cb7a2a1a431d7a0f83667b4a5c7c1c92e9e9625 Mon Sep 17 00:00:00 2001
From: steden <1470804@qq.com>
Date: Sun, 7 Jul 2024 19:32:46 +0800
Subject: [PATCH] =?UTF-8?q?=E6=96=B0=E5=A2=9E=EF=BC=9A=E5=BC=BA=E5=88=B6?=
 =?UTF-8?q?=E8=B5=B0=E7=B4=A2=E5=BC=95=E6=94=AF=E6=8C=81?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 go.mod      |  1 +
 tableSet.go | 39 +++++++++++++++++++++++++++++----------
 2 files changed, 30 insertions(+), 10 deletions(-)

diff --git a/go.mod b/go.mod
index b90f4cd..fbe8df1 100644
--- a/go.mod
+++ b/go.mod
@@ -12,6 +12,7 @@ require (
 	github.com/stretchr/testify v1.8.4
 	gorm.io/driver/mysql v1.5.6
 	gorm.io/gorm v1.25.10
+	gorm.io/hints v1.1.2
 )
 
 require (
diff --git a/tableSet.go b/tableSet.go
index 1ef1970..dae11c5 100644
--- a/tableSet.go
+++ b/tableSet.go
@@ -9,6 +9,7 @@ import (
 	"gorm.io/gorm"
 	"gorm.io/gorm/clause"
 	"gorm.io/gorm/schema"
+	"gorm.io/hints"
 	"reflect"
 	"strings"
 	"time"
@@ -16,13 +17,15 @@ import (
 
 // TableSet 数据库表操作
 type TableSet[Table any] struct {
-	dbContext    *internalContext  // 上下文（用指针的方式，共享同一个上下文）
-	dbName       string            // 库名
-	tableName    string            // 表名
-	primaryName  []string          // 主键字段名称
-	nameReplacer *strings.Replacer // 替换dbName、tableName
-	ormClient    *gorm.DB          // 最外层的ormClient一定是nil的
-	layer        int               // 链式第几层
+	dbContext      *internalContext  // 上下文（用指针的方式，共享同一个上下文）
+	dbName         string            // 库名
+	tableName      string            // 表名
+	forceIndexName string            // 强制索引名称
+	useIndexName   string            // 推荐使用索引名称
+	primaryName    []string          // 主键字段名称
+	nameReplacer   *strings.Replacer // 替换dbName、tableName
+	ormClient      *gorm.DB          // 最外层的ormClient一定是nil的
+	layer          int               // 链式第几层
 	// 字段筛选（官方再第二次设置时，会覆盖第一次的设置，因此需要暂存）
 	selectList collections.ListAny          // 筛选字段
 	omitList   collections.List[string]     // 过滤字段
@@ -222,6 +225,12 @@ func (receiver *TableSet[Table]) getClient() *gorm.DB {
 		receiver.ormClient.Limit(receiver.limit)
 	}
 
+	// 强制索引
+	if receiver.forceIndexName != "" {
+		receiver.ormClient.Clauses(hints.ForceIndex(receiver.forceIndexName))
+	} else if receiver.useIndexName != "" { // 推荐使用索引
+		receiver.ormClient.Clauses(hints.UseIndex(receiver.useIndexName))
+	}
 	return receiver.ormClient
 }
 
@@ -263,6 +272,16 @@ func (receiver *TableSet[Table]) Omit(columns ...string) *TableSet[Table] {
 	return session
 }
 
+// ForceIndex 强制使用索引
+func (receiver *TableSet[Table]) ForceIndex(idxName string) {
+	receiver.forceIndexName = idxName
+}
+
+// UseIndex 推荐使用索引
+func (receiver *TableSet[Table]) UseIndex(idxName string) {
+	receiver.useIndexName = idxName
+}
+
 // Where 条件
 func (receiver *TableSet[Table]) Where(query any, args ...any) *TableSet[Table] {
 	session := receiver.getOrCreateSession()
@@ -275,7 +294,7 @@ func (receiver *TableSet[Table]) Where(query any, args ...any) *TableSet[Table]
 
 // WhereFindInSet FIND_IN_SET条件
 // FIND_IN_SET (fieldValue,fieldName)
-func (receiver *TableSet[Table]) WhereFindInSet(fieldName string,fieldValue string) *TableSet[Table] {
+func (receiver *TableSet[Table]) WhereFindInSet(fieldName string, fieldValue string) *TableSet[Table] {
 	session := receiver.getOrCreateSession()
 	session.whereList.Add(whereQuery{
 		query: fmt.Sprintf("FIND_IN_SET ('%s' ,%s)", fieldValue, fieldName),
@@ -284,9 +303,9 @@ func (receiver *TableSet[Table]) WhereFindInSet(fieldName string,fieldValue stri
 	return session
 }
 
-// WhereFindInSet FIND_IN_SET条件
+// WhereFindInSetOrEq FIND_IN_SET条件
 // FIND_IN_SET (fieldValue,fieldName) or orFieldName = orFieldValue
-func (receiver *TableSet[Table]) WhereFindInSetOrEq(fieldName ,fieldValue , orFieldName string, orFieldValue any) *TableSet[Table] {
+func (receiver *TableSet[Table]) WhereFindInSetOrEq(fieldName, fieldValue, orFieldName string, orFieldValue any) *TableSet[Table] {
 	session := receiver.getOrCreateSession()
 	session.whereList.Add(whereQuery{
 		query: fmt.Sprintf("(FIND_IN_SET ('%s' ,%s) OR %s = ?)", fieldValue, fieldName, orFieldName),