Merge pull request #8 from troyanov/feature/custom-table-name-support

Feature/custom table name support

adapter.go

@@ -8,6 +8,7 @@ import (
 	"github.com/casbin/casbin/v2/persist"
 	"github.com/go-pg/pg/v9"
 	"github.com/go-pg/pg/v9/orm"
+	"github.com/go-pg/pg/v9/types"
 	"github.com/mmcloughlin/meow"
 )
 
@@ -31,9 +32,12 @@ type Filter struct {
 // Adapter represents the github.com/go-pg/pg adapter for policy storage.
 type Adapter struct {
 	db        *pg.DB
+	tableName string
 	filtered  bool
 }
 
+type Option func(a *Adapter)
+
 // NewAdapter is the constructor for Adapter.
 // arg should be a PostgreS URL string or of type *pg.Options
 // The adapter will create a DB named "casbin" if it doesn't exist
@@ -54,14 +58,31 @@ func NewAdapter(arg interface{}) (*Adapter, error) {
 
 // NewAdapterByDB creates new Adapter by using existing DB connection
 // creates table from CasbinRule struct if it doesn't exist
-func NewAdapterByDB(db *pg.DB) (*Adapter, error) {
+func NewAdapterByDB(db *pg.DB, opts ...Option) (*Adapter, error) {
 	a := &Adapter{db: db}
+	for _, opt := range opts {
+		opt(a)
+	}
+
+	if len(a.tableName) > 0 {
+		a.db.Model((*CasbinRule)(nil)).TableModel().Table().Name = a.tableName
+		a.db.Model((*CasbinRule)(nil)).TableModel().Table().FullName = (types.Safe)(a.tableName)
+		a.db.Model((*CasbinRule)(nil)).TableModel().Table().FullNameForSelects = (types.Safe)(a.tableName)
+	}
+
 	if err := a.createTableifNotExists(); err != nil {
 		return nil, fmt.Errorf("pgadapter.NewAdapter: %v", err)
 	}
 	return a, nil
 }
 
+// WithTableName can be used to pass custom table name for Casbin rules
+func WithTableName(tableName string) Option {
+	return func(a *Adapter) {
+		a.tableName = tableName
+	}
+}
+
 func createCasbinDatabase(arg interface{}) (*pg.DB, error) {
 	var opts *pg.Options
 	var err error
@@ -145,7 +166,7 @@ func (r *CasbinRule) String() string {
 func (a *Adapter) LoadPolicy(model model.Model) error {
 	var lines []*CasbinRule
 
-	if _, err := a.db.Query(&lines, `SELECT * FROM casbin_rules`); err != nil {
+	if err := a.db.Model(&lines).Select(); err != nil {
 		return err
 	}
 
@@ -236,7 +257,7 @@ func (a *Adapter) AddPolicy(sec string, ptype string, rule []string) error {
 // AddPolicies adds policy rules to the storage.
 func (a *Adapter) AddPolicies(sec string, ptype string, rules [][]string) error {
 	var lines []*CasbinRule
-	for _,rule := range rules{
+	for _, rule := range rules {
 		line := savePolicyLine(ptype, rule)
 		lines = append(lines, line)
 	}
@@ -261,7 +282,7 @@ func (a *Adapter) RemovePolicy(sec string, ptype string, rule []string) error {
 // RemovePolicies removes policy rules from the storage.
 func (a *Adapter) RemovePolicies(sec string, ptype string, rules [][]string) error {
 	var lines []*CasbinRule
-	for _,rule := range rules{
+	for _, rule := range rules {
 		line := savePolicyLine(ptype, rule)
 		lines = append(lines, line)
 	}