diff --git a/codebase/app/postgres_worker/connection.go b/codebase/app/postgres_worker/connection.go index 4cb7417..07ea484 100644 --- a/codebase/app/postgres_worker/connection.go +++ b/codebase/app/postgres_worker/connection.go @@ -21,6 +21,10 @@ func getListener(source string, opts *option) (*sql.DB, *pq.Listener) { panic(fmt.Errorf(`[POSTGRES-LISTENER] ERROR: %v, ping: %s`, err, candihelper.MaskingPasswordURL(dsn))) } + if opts.dbOption != nil { + opts.dbOption(db) + } + ec := &eventCallback{onErrorFunc: opts.onErrorConnectionFunc} listener := pq.NewListener(dsn, opts.minReconnectInterval, opts.maxReconnectInterval, ec.onEvent) return db, listener diff --git a/codebase/app/postgres_worker/option.go b/codebase/app/postgres_worker/option.go index c412d49..660bc7e 100644 --- a/codebase/app/postgres_worker/option.go +++ b/codebase/app/postgres_worker/option.go @@ -1,6 +1,7 @@ package postgresworker import ( + "database/sql" "time" "github.com/golangid/candi/candiutils" @@ -16,6 +17,7 @@ type ( minReconnectInterval time.Duration maxReconnectInterval time.Duration onErrorConnectionFunc func(error) + dbOption func(*sql.DB) sources map[string]*PostgresSource } @@ -97,3 +99,10 @@ func SetOnErrorConnectionCallback(callback func(error)) OptionFunc { o.onErrorConnectionFunc = callback } } + +// SetDBOption option func +func SetDBOption(dbOption func(*sql.DB)) OptionFunc { + return func(o *option) { + o.dbOption = dbOption + } +}