Merge pull request #73 from Despire/custom_table

add custom table support

Author: hsluoyz

README.md

@@ -63,6 +63,52 @@ func main() {
 }
 ```
 
+## Customize table columns example
+You can change the gorm struct tags, but the table structure must stay the same.
+```go
+package main
+
+import (
+	"github.com/casbin/casbin/v2"
+	gormadapter "github.com/casbin/gorm-adapter/v3"
+	"gorm.io/gorm"
+)
+
+func main() {
+	// Increase the column size to 512.
+	type CasbinRule struct {
+		ID    uint   `gorm:"primaryKey;autoIncrement"`
+		Ptype string `gorm:"size:512;uniqueIndex:unique_index"`
+		V0    string `gorm:"size:512;uniqueIndex:unique_index"`
+		V1    string `gorm:"size:512;uniqueIndex:unique_index"`
+		V2    string `gorm:"size:512;uniqueIndex:unique_index"`
+		V3    string `gorm:"size:512;uniqueIndex:unique_index"`
+		V4    string `gorm:"size:512;uniqueIndex:unique_index"`
+		V5    string `gorm:"size:512;uniqueIndex:unique_index"`
+	}
+
+	db, _ := gorm.Open(...)
+
+	// Initialize a Gorm adapter and use it in a Casbin enforcer:
+	// The adapter will use an existing gorm.DB instnace.
+	a, _ := gormadapter.NewAdapterByDBWithCustomTable(db, &CasbinRule{}) 
+	e, _ := casbin.NewEnforcer("examples/rbac_model.conf", a)
+	
+	// Load the policy from DB.
+	e.LoadPolicy()
+	
+	// Check the permission.
+	e.Enforce("alice", "data1", "read")
+	
+	// Modify the policy.
+	// e.AddPolicy(...)
+	// e.RemovePolicy(...)
+	
+	// Save the policy back to DB.
+	e.SavePolicy()
+}
+```
+
 ## Getting Help
 
 - [Casbin](https://github.com/casbin/casbin)

adapter.go

@@ -15,6 +15,7 @@
 package gormadapter
 
 import (
+	"context"
 	"errors"
 	"runtime"
 	"strings"
@@ -33,6 +34,8 @@ const (
 	defaultTableName    = "casbin_rule"
 )
 
+type customTableKey struct{}
+
 type CasbinRule struct {
 	ID    uint   `gorm:"primaryKey;autoIncrement"`
 	Ptype string `gorm:"size:100;uniqueIndex:unique_index"`
@@ -164,7 +167,7 @@ func NewAdapterByDBUseTableName(db *gorm.DB, prefix string, tableName string) (*
 		tableName:   tableName,
 	}
 
-	a.db = db.Scopes(a.casbinRuleTable()).Session(&gorm.Session{})
+	a.db = db.Scopes(a.casbinRuleTable()).Session(&gorm.Session{Context: db.Statement.Context})
 	err := a.createTable()
 	if err != nil {
 		return nil, err
@@ -189,6 +192,17 @@ func NewAdapterByDB(db *gorm.DB) (*Adapter, error) {
 	return NewAdapterByDBUseTableName(db, "", defaultTableName)
 }
 
+func NewAdapterByDBWithCustomTable(db *gorm.DB, t interface{}) (*Adapter, error) {
+	ctx := db.Statement.Context
+	if ctx == nil {
+		ctx = context.Background()
+	}
+
+	ctx = context.WithValue(ctx, customTableKey{}, t)
+
+	return NewAdapterByDBUseTableName(db.WithContext(ctx), "", defaultTableName)
+}
+
 func openDBConnection(driverName, dataSourceName string) (*gorm.DB, error) {
 	var err error
 	var db *gorm.DB
@@ -280,11 +294,21 @@ func (a *Adapter) casbinRuleTable() func(db *gorm.DB) *gorm.DB {
 }
 
 func (a *Adapter) createTable() error {
-	return a.db.AutoMigrate(a.getTableInstance())
+	t := a.db.Statement.Context.Value(customTableKey{})
+	if t == nil {
+		return a.db.AutoMigrate(a.getTableInstance())
+	}
+
+	return a.db.AutoMigrate(t)
 }
 
 func (a *Adapter) dropTable() error {
-	return a.db.Migrator().DropTable(a.getTableInstance())
+	t := a.db.Statement.Context.Value(customTableKey{})
+	if t == nil {
+		return a.db.Migrator().DropTable(a.getTableInstance())
+	}
+
+	return a.db.Migrator().DropTable(t)
 }
 
 func loadPolicyLine(line CasbinRule, model model.Model) {

adapter_test.go

@@ -15,6 +15,7 @@
 package gormadapter
 
 import (
+	"github.com/jackc/pgconn"
 	"log"
 	"testing"
 
@@ -107,6 +108,29 @@ func initAdapterWithGormInstance(t *testing.T, db *gorm.DB) *Adapter {
 	return a
 }
 
+func initAdapterWithGormInstanceAndCustomTable(t *testing.T, db *gorm.DB) *Adapter {
+	type CasbinRule struct {
+		ID    uint   `gorm:"primaryKey;autoIncrement"`
+		Ptype string `gorm:"size:128;uniqueIndex:unique_index"`
+		V0    string `gorm:"size:128;uniqueIndex:unique_index"`
+		V1    string `gorm:"size:128;uniqueIndex:unique_index"`
+		V2    string `gorm:"size:128;uniqueIndex:unique_index"`
+		V3    string `gorm:"size:128;uniqueIndex:unique_index"`
+		V4    string `gorm:"size:128;uniqueIndex:unique_index"`
+		V5    string `gorm:"size:128;uniqueIndex:unique_index"`
+	}
+
+	// Create an adapter
+	a, _ := NewAdapterByDBWithCustomTable(db, &CasbinRule{})
+	// Initialize some policy in DB.
+	initPolicy(t, a)
+	// Now the DB has policy, so we can provide a normal use case.
+	// Note: you don't need to look at the above code
+	// if you already have a working DB with policy inside.
+
+	return a
+}
+
 func initAdapterWithGormInstanceByName(t *testing.T, db *gorm.DB, name string) *Adapter {
 	//Create an Adapter
 	a, _ := NewAdapterByDBUseTableName(db, "", name)
@@ -210,6 +234,32 @@ func testUpdatePolicy(t *testing.T, a *Adapter) {
 	testGetPolicy(t, e, [][]string{{"alice", "data1", "write"}, {"bob", "data2", "write"}, {"data2_admin", "data2", "read"}, {"data2_admin", "data2", "write"}})
 }
 
+func TestAdapterWithCustomTable(t *testing.T) {
+	db, err := gorm.Open(postgres.Open("user=postgres host=127.0.0.1 port=5432 sslmode=disable"), &gorm.Config{})
+	if err != nil {
+		panic(err)
+	}
+
+	if err = db.Exec("CREATE DATABASE casbin_custom_table").Error; err != nil {
+		// 42P04 is	duplicate_database
+		if err.(*pgconn.PgError).Code != "42P04" {
+			panic(err)
+		}
+	}
+
+	db, err = gorm.Open(postgres.Open("user=postgres host=127.0.0.1 port=5432 sslmode=disable dbname=casbin_custom_table"), &gorm.Config{})
+	if err != nil {
+		panic(err)
+	}
+
+	a := initAdapterWithGormInstanceAndCustomTable(t, db)
+	testAutoSave(t, a)
+	testSaveLoad(t, a)
+
+	a = initAdapterWithGormInstanceAndCustomTable(t, db)
+	testFilteredPolicy(t, a)
+}
+
 func TestAdapters(t *testing.T) {
 	a := initAdapter(t, "mysql", "root:@tcp(127.0.0.1:3306)/", "casbin", "casbin_rule")
 	testAutoSave(t, a)