From dded6048788490b2bebb35d61992d21e4f003ee1 Mon Sep 17 00:00:00 2001
From: EpicMo <1982742309@qq.com>
Date: Mon, 31 Jul 2023 10:06:15 +0800
Subject: [PATCH] feat:add before find hook

---
 callbacks/callbacks.go  |  1 +
 callbacks/interfaces.go |  4 ++++
 callbacks/query.go      | 12 ++++++++++++
 schema/schema.go        |  4 ++--
 4 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/callbacks/callbacks.go b/callbacks/callbacks.go
index d681aef3..db5b865e 100644
--- a/callbacks/callbacks.go
+++ b/callbacks/callbacks.go
@@ -48,6 +48,7 @@ func RegisterDefaultCallbacks(db *gorm.DB, config *Config) {
 	createCallback.Clauses = config.CreateClauses
 
 	queryCallback := db.Callback().Query()
+	queryCallback.Register("gorm:before_query", BeforeQuery)
 	queryCallback.Register("gorm:query", Query)
 	queryCallback.Register("gorm:preload", Preload)
 	queryCallback.Register("gorm:after_query", AfterQuery)
diff --git a/callbacks/interfaces.go b/callbacks/interfaces.go
index 2302470f..758fb2c7 100644
--- a/callbacks/interfaces.go
+++ b/callbacks/interfaces.go
@@ -37,3 +37,7 @@ type AfterDeleteInterface interface {
 type AfterFindInterface interface {
 	AfterFind(*gorm.DB) error
 }
+
+type BeforeFindInterface interface {
+	BeforeFind(*gorm.DB) error
+}
diff --git a/callbacks/query.go b/callbacks/query.go
index e89dd199..4be7b928 100644
--- a/callbacks/query.go
+++ b/callbacks/query.go
@@ -314,3 +314,15 @@ func AfterQuery(db *gorm.DB) {
 		})
 	}
 }
+
+func BeforeQuery(db *gorm.DB) {
+	if db.Error == nil && db.Statement.Schema != nil && !db.Statement.SkipHooks && db.Statement.Schema.BeforeFind {
+		callMethod(db, func(value interface{}, tx *gorm.DB) bool {
+			if i, ok := value.(BeforeFindInterface); ok {
+				db.AddError(i.BeforeFind(tx))
+				return true
+			}
+			return false
+		})
+	}
+}
diff --git a/schema/schema.go b/schema/schema.go
index e13a5ed1..35184ad1 100644
--- a/schema/schema.go
+++ b/schema/schema.go
@@ -38,7 +38,7 @@ type Schema struct {
 	BeforeUpdate, AfterUpdate bool
 	BeforeDelete, AfterDelete bool
 	BeforeSave, AfterSave     bool
-	AfterFind                 bool
+	BeforeFind, AfterFind     bool
 	err                       error
 	initialized               chan struct{}
 	namer                     Namer
@@ -288,7 +288,7 @@ func ParseWithSpecialTableName(dest interface{}, cacheStore *sync.Map, namer Nam
 		}
 	}
 
-	callbacks := []string{"BeforeCreate", "AfterCreate", "BeforeUpdate", "AfterUpdate", "BeforeSave", "AfterSave", "BeforeDelete", "AfterDelete", "AfterFind"}
+	callbacks := []string{"BeforeCreate", "AfterCreate", "BeforeUpdate", "AfterUpdate", "BeforeSave", "AfterSave", "BeforeDelete", "AfterDelete", "AfterFind", "BeforeFind"}
 	for _, name := range callbacks {
 		if methodValue := modelValue.MethodByName(name); methodValue.IsValid() {
 			switch methodValue.Type().String() {