postgre worker: add db option

golangid · Apr 3, 2024 · bf71105 · bf71105
1 parent 3e5e24d
commit bf71105
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 0 deletions.
diff --git 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
@@ -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
+	}
+}