From b571d918586f137d4bc5819909de4f13befdbf1c Mon Sep 17 00:00:00 2001
From: Lunny Xiao <xiaolunwen@gmail.com>
Date: Fri, 15 Dec 2023 02:17:13 +0000
Subject: [PATCH] Add index hint support (#2375)

Fix #1456

Reviewed-on: https://gitea.com/xorm/xorm/pulls/2375
---
 dialects/sqlite3_test.go         |  4 +--
 engine.go                        |  7 +++++
 internal/statements/index.go     | 54 ++++++++++++++++++++++++++++++++
 internal/statements/query.go     |  1 +
 internal/statements/statement.go |  6 ++++
 session_schema.go                |  5 +++
 tests/session_test.go            | 11 +++++++
 7 files changed, 86 insertions(+), 2 deletions(-)
 create mode 100644 internal/statements/index.go

diff --git a/dialects/sqlite3_test.go b/dialects/sqlite3_test.go
index ec973ae6..75115c6c 100644
--- a/dialects/sqlite3_test.go
+++ b/dialects/sqlite3_test.go
@@ -11,7 +11,7 @@ import (
 )
 
 func TestSplitColStr(t *testing.T) {
-	var kases = []struct {
+	kases := []struct {
 		colStr string
 		fields []string
 	}{
@@ -31,7 +31,7 @@ func TestSplitColStr(t *testing.T) {
 			colStr: `    id               INTEGER           not null
 primary key autoincrement`,
 			fields: []string{
-				"id", "INTEGER", "not null", "primary", "key", "autoincrement",
+				"id", "INTEGER", "not", "null", "primary", "key", "autoincrement",
 			},
 		},
 	}
diff --git a/engine.go b/engine.go
index 0cbfdede..31a6b688 100644
--- a/engine.go
+++ b/engine.go
@@ -1433,3 +1433,10 @@ func (engine *Engine) Transaction(f func(*Session) (interface{}, error)) (interf
 
 	return result, nil
 }
+
+func (engine *Engine) IndexHint(op, indexerOrColName string) *Session {
+	session := engine.NewSession()
+	session.isAutoClose = true
+	session.statement.LastError = session.statement.IndexHint(op, indexerOrColName)
+	return session
+}
diff --git a/internal/statements/index.go b/internal/statements/index.go
new file mode 100644
index 00000000..8f54d242
--- /dev/null
+++ b/internal/statements/index.go
@@ -0,0 +1,54 @@
+// Copyright 2023 The Xorm Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package statements
+
+import (
+	"strings"
+
+	"xorm.io/builder"
+	"xorm.io/xorm/schemas"
+)
+
+type ErrInvalidIndexHintOperator struct {
+	Op string
+}
+
+func (e ErrInvalidIndexHintOperator) Error() string {
+	return "invalid index hint operator: " + e.Op
+}
+
+func (statement *Statement) IndexHint(op, indexName string) error {
+	op = strings.ToUpper(op)
+	statement.indexHints = append(statement.indexHints, indexHint{
+		op:        op,
+		indexName: indexName,
+	})
+	return nil
+}
+
+func (statement *Statement) writeIndexHints(w *builder.BytesWriter) error {
+	if len(statement.indexHints) == 0 {
+		return nil
+	}
+
+	switch statement.dialect.URI().DBType {
+	case schemas.MYSQL:
+		return statement.writeIndexHintsMySQL(w)
+	default:
+		return ErrNotImplemented
+	}
+}
+
+func (statement *Statement) writeIndexHintsMySQL(w *builder.BytesWriter) error {
+	for _, hint := range statement.indexHints {
+		if hint.op != "USE" && hint.op != "FORCE" && hint.op != "IGNORE" {
+			return ErrInvalidIndexHintOperator{Op: hint.op}
+		}
+		if err := statement.writeStrings(" ", hint.op, " INDEX(", hint.indexName, ")")(w); err != nil {
+			return err
+		}
+	}
+	return nil
+}
diff --git a/internal/statements/query.go b/internal/statements/query.go
index 8a9e59e4..a36d8ad3 100644
--- a/internal/statements/query.go
+++ b/internal/statements/query.go
@@ -254,6 +254,7 @@ func (statement *Statement) writeSelect(buf *builder.BytesWriter, columnStr stri
 	return statement.writeMultiple(buf,
 		statement.writeSelectColumns(columnStr),
 		statement.writeFrom,
+		statement.writeIndexHints,
 		statement.writeWhere,
 		statement.writeGroupBy,
 		statement.writeHaving,
diff --git a/internal/statements/statement.go b/internal/statements/statement.go
index 55a3d89e..31073ad1 100644
--- a/internal/statements/statement.go
+++ b/internal/statements/statement.go
@@ -41,6 +41,11 @@ type join struct {
 	args      []interface{}
 }
 
+type indexHint struct {
+	op        string
+	indexName string
+}
+
 // Statement save all the sql info for executing SQL
 type Statement struct {
 	RefTable        *schemas.Table
@@ -84,6 +89,7 @@ type Statement struct {
 	BufferSize      int
 	Context         contexts.ContextCache
 	LastError       error
+	indexHints      []indexHint
 }
 
 // NewStatement creates a new statement
diff --git a/session_schema.go b/session_schema.go
index 830ba08a..f7ab2657 100644
--- a/session_schema.go
+++ b/session_schema.go
@@ -311,3 +311,8 @@ func (session *Session) Import(r io.Reader) ([]sql.Result, error) {
 
 	return results, lastError
 }
+
+func (session *Session) IndexHint(op, indexerOrColName string) *Session {
+	session.statement.IndexHint(op, indexerOrColName)
+	return session
+}
diff --git a/tests/session_test.go b/tests/session_test.go
index 261f2c48..aeba1141 100644
--- a/tests/session_test.go
+++ b/tests/session_test.go
@@ -62,3 +62,14 @@ func TestEnableSessionId(t *testing.T) {
 	_, err := testEngine.Table("userinfo").MustLogSQL(true).Get(new(Userinfo))
 	assert.NoError(t, err)
 }
+
+func TestIndexHint(t *testing.T) {
+	assert.NoError(t, PrepareEngine())
+	assertSync(t, new(Userinfo))
+	if testEngine.Dialect().URI().DBType != "mysql" {
+		return
+	}
+
+	_, err := testEngine.Table("userinfo").IndexHint("USE", "UQE_userinfo_username").Get(new(Userinfo))
+	assert.NoError(t, err)
+}