From 9d2b65f8c9604651197b9d864500d05ddce2cc99 Mon Sep 17 00:00:00 2001
From: Dozer <mail@dozer.cc>
Date: Fri, 6 Dec 2019 09:16:51 +0800
Subject: [PATCH] add query hint support (#2351)

* add query hint support

* remove add extra space

* add test and fix bug

* fix ut

* fix ut
---
 callback_query.go     |  5 +++++
 callback_row_query.go |  5 +++++
 main_test.go          | 24 ++++++++++++++++++++++++
 3 files changed, 34 insertions(+)

diff --git a/callback_query.go b/callback_query.go
index 7facc42b..544afd63 100644
--- a/callback_query.go
+++ b/callback_query.go
@@ -60,6 +60,11 @@ func queryCallback(scope *Scope) {
 
 	if !scope.HasError() {
 		scope.db.RowsAffected = 0
+
+		if str, ok := scope.Get("gorm:query_hint"); ok {
+			scope.SQL = fmt.Sprint(str) + scope.SQL
+		}
+
 		if str, ok := scope.Get("gorm:query_option"); ok {
 			scope.SQL += addExtraSpaceIfExist(fmt.Sprint(str))
 		}
diff --git a/callback_row_query.go b/callback_row_query.go
index 687b0039..323b1605 100644
--- a/callback_row_query.go
+++ b/callback_row_query.go
@@ -23,6 +23,11 @@ type RowsQueryResult struct {
 func rowQueryCallback(scope *Scope) {
 	if result, ok := scope.InstanceGet("row_query_result"); ok {
 		scope.prepareQuerySQL()
+
+		if str, ok := scope.Get("gorm:query_hint"); ok {
+			scope.SQL = fmt.Sprint(str) + scope.SQL
+		}
+
 		if str, ok := scope.Get("gorm:query_option"); ok {
 			scope.SQL += addExtraSpaceIfExist(fmt.Sprint(str))
 		}
diff --git a/main_test.go b/main_test.go
index 98ea4694..b51fe413 100644
--- a/main_test.go
+++ b/main_test.go
@@ -1333,6 +1333,30 @@ func TestCountWithQueryOption(t *testing.T) {
 	}
 }
 
+func TestQueryHint1(t *testing.T) {
+	db := DB.New()
+
+	_, err := db.Model(User{}).Raw("select 1").Rows()
+
+	if err != nil {
+		t.Error("Unexpected error on query count with query_option")
+	}
+}
+
+func TestQueryHint2(t *testing.T) {
+	type TestStruct struct {
+		ID   string `gorm:"primary_key"`
+		Name string
+	}
+	DB.DropTable(&TestStruct{})
+	DB.AutoMigrate(&TestStruct{})
+
+	data := TestStruct{ID: "uuid", Name: "hello"}
+	if err := DB.Set("gorm:query_hint", "/*master*/").Save(&data).Error; err != nil {
+		t.Error("Unexpected error on query count with query_option")
+	}
+}
+
 func TestFloatColumnPrecision(t *testing.T) {
 	if dialect := os.Getenv("GORM_DIALECT"); dialect != "mysql" && dialect != "sqlite" {
 		t.Skip()