diff --git a/engine.go b/engine.go index 459c2f44..e16bb3e8 100644 --- a/engine.go +++ b/engine.go @@ -51,7 +51,7 @@ type Engine struct { // NewEngine new a db manager according to the parameter. Currently support four // drivers -func NewEngine(driverName string, dataSourceName string) (*Engine, error) { +func NewEngine(driverName string, dataSourceName string, driverOptions ...func(db *sql.DB) error) (*Engine, error) { dialect, err := dialects.OpenDialect(driverName, dataSourceName) if err != nil { return nil, err @@ -62,6 +62,12 @@ func NewEngine(driverName string, dataSourceName string) (*Engine, error) { return nil, err } + for _, driverOption := range driverOptions { + if err := driverOption(db.DB); err != nil { + return nil, err + } + } + return newEngine(driverName, dataSourceName, dialect, db) }