Configurable interface to check query type to detect RETURNING

db.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	"database/sql"
 	"database/sql/driver"
-	"strings"
 	"time"
 
 	"go.uber.org/multierr"
@@ -57,6 +56,7 @@ type sqlDB struct {
 	replicas         []*sql.DB
 	loadBalancer     DBLoadBalancer
 	stmtLoadBalancer StmtLoadBalancer
+	queryTypeChecker QueryTypeChecker
 }
 
 // PrimaryDBs return all the active primary DB
@@ -185,8 +185,7 @@ func (db *sqlDB) Query(query string, args ...interface{}) (*sql.Rows, error) {
 // The args are for any placeholder parameters in the query.
 func (db *sqlDB) QueryContext(ctx context.Context, query string, args ...interface{}) (rows *sql.Rows, err error) {
 	var curDB *sql.DB
-	_query := strings.ToUpper(query)
-	writeFlag := strings.Contains(_query, "RETURNING")
+	writeFlag := db.queryTypeChecker.Check(query) == QueryTypeWrite
 
 	if writeFlag {
 		curDB = db.ReadWrite()
@@ -213,8 +212,7 @@ func (db *sqlDB) QueryRow(query string, args ...interface{}) *sql.Row {
 // Errors are deferred until Row's Scan method is called.
 func (db *sqlDB) QueryRowContext(ctx context.Context, query string, args ...interface{}) *sql.Row {
 	var curDB *sql.DB
-	_query := strings.ToUpper(query)
-	writeFlag := strings.Contains(_query, "RETURNING")
+	writeFlag := db.queryTypeChecker.Check(query) == QueryTypeWrite
 
 	if writeFlag {
 		curDB = db.ReadWrite()

options.go

@@ -16,10 +16,11 @@ const (
 
 // Option define the option property
 type Option struct {
-	PrimaryDBs []*sql.DB
-	ReplicaDBs []*sql.DB
-	StmtLB     StmtLoadBalancer
-	DBLB       DBLoadBalancer
+	PrimaryDBs       []*sql.DB
+	ReplicaDBs       []*sql.DB
+	StmtLB           StmtLoadBalancer
+	DBLB             DBLoadBalancer
+	QueryTypeChecker QueryTypeChecker
 }
 
 // OptionFunc used for option chaining
@@ -39,6 +40,14 @@ func WithReplicaDBs(replicaDBs ...*sql.DB) OptionFunc {
 	}
 }
 
+// WithQueryTypeChecker sets the query type checker instance.
+// The default one just checks for the presence of the string "RETURNING" in the uppercase query.
+func WithQueryTypeChecker(checker QueryTypeChecker) OptionFunc {
+	return func(opt *Option) {
+		opt.QueryTypeChecker = checker
+	}
+}
+
 // WithLoadBalancer configure the loadbalancer for the resolver
 func WithLoadBalancer(lb LoadBalancerPolicy) OptionFunc {
 	return func(opt *Option) {
@@ -61,7 +70,8 @@ func WithLoadBalancer(lb LoadBalancerPolicy) OptionFunc {
 
 func defaultOption() *Option {
 	return &Option{
-		DBLB:   &RoundRobinLoadBalancer[*sql.DB]{},
-		StmtLB: &RoundRobinLoadBalancer[*sql.Stmt]{},
+		DBLB:             &RoundRobinLoadBalancer[*sql.DB]{},
+		StmtLB:           &RoundRobinLoadBalancer[*sql.Stmt]{},
+		QueryTypeChecker: &DefaultQueryTypeChecker{},
 	}
 }

query.go

@@ -0,0 +1,28 @@
+package dbresolver
+
+import "strings"
+
+type QueryType int
+
+const (
+	QueryTypeUnknown QueryType = iota
+	QueryTypeRead
+	QueryTypeWrite
+)
+
+// QueryTypeChecker is used to try to detect the query type, like for detecting RETURNING clauses in
+// INSERT/UPDATE clauses.
+type QueryTypeChecker interface {
+	Check(query string) QueryType
+}
+
+// DefaultQueryTypeChecker searches for a "RETURNING" string inside the query to detect a write query.
+type DefaultQueryTypeChecker struct {
+}
+
+func (c DefaultQueryTypeChecker) Check(query string) QueryType {
+	if strings.Contains(strings.ToUpper(query), "RETURNING") {
+		return QueryTypeWrite
+	}
+	return QueryTypeUnknown
+}

resolver.go

@@ -16,5 +16,6 @@ func New(opts ...OptionFunc) DB {
 		replicas:         opt.ReplicaDBs,
 		loadBalancer:     opt.DBLB,
 		stmtLoadBalancer: opt.StmtLB,
+		queryTypeChecker: opt.QueryTypeChecker,
 	}
 }