Unexport DB init helpers, add control for observability

database/config.go

@@ -1,4 +1,4 @@
 package database
 
 import "time"
 
@@ -11,4 +11,8 @@
 	MaxOpen         int           `split_words:"true" default:"5"`
 	InitConn        bool          `split_words:"true"`
 	ApplyMigrations bool          `split_words:"true"`
+
+	// MethodSkipPackages provides helper package paths to skip when identifying method name for observability.
+	// Item example: "github.com/jmoiron/sqlx".
+	MethodSkipPackages []string
 }

database/observability.go

@@ -14,13 +14,13 @@
 	"go.opencensus.io/trace"
 )
 
-// WithTracing instruments database connector with OpenCensus tracing.
-func WithTracing(dbConnector driver.Connector) driver.Connector {
+// withTracing instruments database connector with OpenCensus tracing.
+func withTracing(dbConnector driver.Connector) driver.Connector {
 	return ocsql.WrapConnector(dbConnector, tracingOptions()...)
 }
 
-// DriverNameWithTracing registers database driver name with OpenCensus tracing.
-func DriverNameWithTracing(driverName string) (string, error) {
+// driverNameWithTracing registers database driver name with OpenCensus tracing.
+func driverNameWithTracing(driverName string) (string, error) {
 	return ocsql.Register(driverName, tracingOptions()...)
 }
 
@@ -34,17 +34,17 @@
 	}
 }
 
-// WithQueriesLogging instruments database connector with query logging.
-func WithQueriesLogging(dbConnector driver.Connector, logger ctxd.Logger, statsTracker stats.Tracker) driver.Connector {
-	return dbwrap.WrapConnector(dbConnector, wrapOptions(logger, statsTracker)...)
+// withQueriesLogging instruments database connector with query logging.
+func withQueriesLogging(cfg Config, dbConnector driver.Connector, logger ctxd.Logger, statsTracker stats.Tracker) driver.Connector {
+	return dbwrap.WrapConnector(dbConnector, wrapOptions(logger, statsTracker, cfg.MethodSkipPackages)...)
 }
 
-// DriverNameWithQueriesLogging registers database driver name with query logging.
-func DriverNameWithQueriesLogging(driverName string, logger ctxd.Logger, statsTracker stats.Tracker) (string, error) {
-	return dbwrap.Register(driverName, wrapOptions(logger, statsTracker)...)
+// driverNameWithQueriesLogging registers database driver name with query logging.
+func driverNameWithQueriesLogging(cfg Config, driverName string, logger ctxd.Logger, statsTracker stats.Tracker) (string, error) {
+	return dbwrap.Register(driverName, wrapOptions(logger, statsTracker, cfg.MethodSkipPackages)...)
 }
 
-func wrapOptions(logger ctxd.Logger, statsTracker stats.Tracker) []dbwrap.Option {
+func wrapOptions(logger ctxd.Logger, statsTracker stats.Tracker, skipPackages []string) []dbwrap.Option {
 	if logger == nil {
 		logger = ctxd.NoOpLogger{}
 	}
@@ -53,11 +53,11 @@
 		statsTracker = stats.NoOp{}
 	}
 
-	skipPackages := []string{
+	skipPackages = append([]string{
 		"github.com/Masterminds/squirrel",
 		"github.com/bool64/sqluct",
 		"github.com/jmoiron/sqlx",
-	}
+	}, skipPackages...)
 
 	return []dbwrap.Option{
 		// This interceptor enables reverse debugging from DB side.
@@ -70,7 +70,7 @@
 		}),
 
 		// This option limits middleware applicability.
-		dbwrap.WithOperations(dbwrap.Query, dbwrap.StmtQuery, dbwrap.Exec, dbwrap.StmtExec),
+		dbwrap.WithOperations(dbwrap.Query, dbwrap.StmtQuery, dbwrap.Exec, dbwrap.StmtExec, dbwrap.RowsClose),
 
 		// This middleware logs statements with arguments at DEBUG level and counts stats.
 		dbwrap.WithMiddleware(observe(logger, statsTracker, skipPackages)),
@@ -87,6 +87,11 @@
 		// Closest caller in the stack with package not equal to listed and to "database/sql".
 		caller := dbwrap.Caller(skipPackages...)
 
+		if operation == dbwrap.RowsClose {
+			statsTracker.Add(ctx, "sql_storage_rows_close", 1, "method", caller)
+			return
+		}
+
 		ctx, span := trace.StartSpan(ctx, caller+":"+string(operation))
 		span.AddAttributes(
 			trace.StringAttribute("stmt", statement),

database/storage.go

@@ -20,8 +20,8 @@
 
 // SetupStorage initializes database pool and prepares storage.
 func SetupStorage(cfg Config, logger ctxd.Logger, statsTracker stats.Tracker, conn driver.Connector, migrations fs.FS) (*sqluct.Storage, error) {
-	conn = WithTracing(conn)
-	conn = WithQueriesLogging(conn, logger, statsTracker)
+	conn = withTracing(conn)
+	conn = withQueriesLogging(cfg, conn, logger, statsTracker)
 
 	db := sql.OpenDB(conn)
 
@@ -30,12 +30,12 @@
 
 // SetupStorageDSN initializes database pool and prepares storage.
 func SetupStorageDSN(cfg Config, logger ctxd.Logger, statsTracker stats.Tracker, migrations fs.FS) (*sqluct.Storage, error) {
-	wrapName, err := DriverNameWithTracing(cfg.DriverName)
+	wrapName, err := driverNameWithTracing(cfg.DriverName)
 	if err != nil {
 		return nil, err
 	}
 
-	wrapName, err = DriverNameWithQueriesLogging(wrapName, logger, statsTracker)
+	wrapName, err = driverNameWithQueriesLogging(cfg, wrapName, logger, statsTracker)
 	if err != nil {
 		return nil, err
 	}