From f4bfc435cc84824e0ca3a9c4e21458996bce67d3 Mon Sep 17 00:00:00 2001
From: Jinzhu <wosmvp@gmail.com>
Date: Tue, 23 Jun 2020 09:38:51 +0800
Subject: [PATCH] Add register plugin API

---
 errors.go     |  2 ++
 gorm.go       | 15 +++++++++++++++
 interfaces.go |  7 +++++++
 3 files changed, 24 insertions(+)

diff --git a/errors.go b/errors.go
index 2506ecc5..b41eefae 100644
--- a/errors.go
+++ b/errors.go
@@ -27,4 +27,6 @@ var (
 	ErrorModelValueRequired = errors.New("model value required")
 	// ErrUnsupportedDriver unsupported driver
 	ErrUnsupportedDriver = errors.New("unsupported driver")
+	// ErrRegistered registered
+	ErrRegistered = errors.New("registered")
 )
diff --git a/gorm.go b/gorm.go
index 47a209ab..c506c6f3 100644
--- a/gorm.go
+++ b/gorm.go
@@ -39,6 +39,8 @@ type Config struct {
 	ConnPool ConnPool
 	// Dialector database dialector
 	Dialector
+	// Plugins registered plugins
+	Plugins map[string]Plugin
 
 	callbacks  *callbacks
 	cacheStore *sync.Map
@@ -309,3 +311,16 @@ func (db *DB) SetupJoinTable(model interface{}, field string, joinTable interfac
 
 	return nil
 }
+
+func (db *DB) Use(plugin Plugin) (err error) {
+	name := plugin.Name()
+	if _, ok := db.Plugins[name]; !ok {
+		if err = plugin.Initialize(db); err == nil {
+			db.Plugins[name] = plugin
+		}
+	} else {
+		return ErrRegistered
+	}
+
+	return err
+}
diff --git a/interfaces.go b/interfaces.go
index 96289a90..b2ce59b3 100644
--- a/interfaces.go
+++ b/interfaces.go
@@ -20,6 +20,12 @@ type Dialector interface {
 	Explain(sql string, vars ...interface{}) string
 }
 
+// Plugin GORM plugin interface
+type Plugin interface {
+	Name() string
+	Initialize(*DB) error
+}
+
 // ConnPool db conns pool interface
 type ConnPool interface {
 	PrepareContext(ctx context.Context, query string) (*sql.Stmt, error)
@@ -28,6 +34,7 @@ type ConnPool interface {
 	QueryRowContext(ctx context.Context, query string, args ...interface{}) *sql.Row
 }
 
+// SavePointerDialectorInterface save pointer interface
 type SavePointerDialectorInterface interface {
 	SavePoint(tx *DB, name string) error
 	RollbackTo(tx *DB, name string) error