fixed ydb_go_sdk_ydb_database_sql_conns metric + added integration test for check metric

Author: asmyasnikov

internal/xslices/keys.go

@@ -0,0 +1,18 @@
+package xslices
+
+import (
+	"cmp"
+	"slices"
+)
+
+func Keys[Key cmp.Ordered, T any](m map[Key]T) []Key {
+	keys := make([]Key, 0, len(m))
+
+	for key := range m {
+		keys = append(keys, key)
+	}
+
+	slices.Sort(keys)
+
+	return keys
+}

internal/xslices/keys_test.go

@@ -0,0 +1,18 @@
+package xslices
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestKeys(t *testing.T) {
+	require.Equal(t, []string{"1", "2"}, Keys(map[string]any{
+		"1": nil,
+		"2": nil,
+	}))
+	require.Equal(t, []int{1, 2}, Keys(map[int]any{
+		1: nil,
+		2: nil,
+	}))
+}

internal/xsql/connector.go

@@ -14,6 +14,7 @@ import (
 
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/bind"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/query"
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/stack"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xcontext"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xsql/legacy"
@@ -123,10 +124,17 @@ func (c *Connector) Open(name string) (driver.Conn, error) {
 	return nil, xerrors.WithStackTrace(driver.ErrSkip)
 }
 
-func (c *Connector) Connect(ctx context.Context) (driver.Conn, error) {
+func (c *Connector) Connect(ctx context.Context) (_ driver.Conn, finalErr error) { //nolint:funlen
+	onDone := trace.DatabaseSQLOnConnectorConnect(c.Trace(), &ctx,
+		stack.FunctionID("", stack.Package("database/sql")),
+	)
+
 	switch c.processor {
 	case QUERY_SERVICE:
 		s, err := query.CreateSession(ctx, c.Query())
+		defer func() {
+			onDone(s, finalErr)
+		}()
 		if err != nil {
 			return nil, xerrors.WithStackTrace(err)
 		}
@@ -152,6 +160,9 @@ func (c *Connector) Connect(ctx context.Context) (driver.Conn, error) {
 
 	case LEGACY:
 		s, err := c.Table().CreateSession(ctx) //nolint:staticcheck
+		defer func() {
+			onDone(s, finalErr)
+		}()
 		if err != nil {
 			return nil, xerrors.WithStackTrace(err)
 		}

metrics/sql.go

@@ -1,16 +1,16 @@
 package metrics
 
 import (
-	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xsync"
 	"time"
 
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xsync"
 	"github.com/ydb-platform/ydb-go-sdk/v3/trace"
 )
 
-// databaseSQL makes trace.DatabaseSQL with measuring `database/sql` events
+// DatabaseSQL makes trace.DatabaseSQL with measuring `database/sql` events
 //
 //nolint:funlen
-func databaseSQL(config Config) trace.DatabaseSQL {
+func DatabaseSQL(config Config) trace.DatabaseSQL {
 	config = config.WithSystem("database").WithSystem("sql")
 
 	conns := config.GaugeVec("conns")

metrics/traces.go

@@ -19,7 +19,7 @@ func WithTraces(config Config) ydb.Option {
 		ydb.WithTraceCoordination(coordination(config)),
 		ydb.WithTraceRatelimiter(ratelimiter(config)),
 		ydb.WithTraceDiscovery(discovery(config)),
-		ydb.WithTraceDatabaseSQL(databaseSQL(config)),
+		ydb.WithTraceDatabaseSQL(DatabaseSQL(config)),
 		ydb.WithTraceRetry(retry(config)),
 	)
 }

tests/integration/database_sql_metrics_test.go

@@ -0,0 +1,78 @@
+//go:build integration
+// +build integration
+
+package integration
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/ydb-platform/ydb-go-sdk/v3"
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xslices"
+	"github.com/ydb-platform/ydb-go-sdk/v3/metrics"
+	"github.com/ydb-platform/ydb-go-sdk/v3/trace"
+)
+
+func TestDatabaseSqlMetrics(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	var (
+		scope    = newScope(t)
+		registry = &registryConfig{
+			details:    trace.DatabaseSQLEvents,
+			gauges:     newVec[gaugeVec](),
+			counters:   newVec[counterVec](),
+			timers:     newVec[timerVec](),
+			histograms: newVec[histogramVec](),
+		}
+		db = scope.SQLDriver(ydb.WithDatabaseSQLTrace(metrics.DatabaseSQL(registry)))
+	)
+
+	require.Equal(t,
+		[]string{"database.sql.conns{}", "database.sql.tx{}"},
+		xslices.Keys(registry.gauges.data),
+	)
+	require.EqualValues(t, 1, registry.gauges.data["database.sql.conns{}"].gauges["database.sql.conns{}"].value)
+
+	cc1, err := db.Conn(ctx)
+	require.NoError(t, err)
+	require.NotNil(t, cc1)
+	require.EqualValues(t, 1, registry.gauges.data["database.sql.conns{}"].gauges["database.sql.conns{}"].value)
+	require.Empty(t, registry.gauges.data["database.sql.tx{}"].gauges)
+
+	cc2, err := db.Conn(ctx)
+	require.NoError(t, err)
+	require.NotNil(t, cc2)
+	require.EqualValues(t, 2, registry.gauges.data["database.sql.conns{}"].gauges["database.sql.conns{}"].value)
+	require.Empty(t, registry.gauges.data["database.sql.tx{}"].gauges)
+
+	tx1, err := cc1.BeginTx(ctx, nil)
+	require.NoError(t, err)
+	require.NotNil(t, tx1)
+	require.NotEmpty(t, registry.gauges.data["database.sql.tx{}"].gauges)
+	require.EqualValues(t, 1, registry.gauges.data["database.sql.tx{}"].gauges["database.sql.tx{}"].value)
+
+	require.NoError(t, tx1.Commit())
+	require.EqualValues(t, 0, registry.gauges.data["database.sql.tx{}"].gauges["database.sql.tx{}"].value)
+
+	require.NoError(t, cc1.Close())
+	require.EqualValues(t, 2, registry.gauges.data["database.sql.conns{}"].gauges["database.sql.conns{}"].value)
+
+	tx2, err := cc2.BeginTx(ctx, nil)
+	require.NoError(t, err)
+	require.NotNil(t, tx2)
+	require.NotEmpty(t, registry.gauges.data["database.sql.tx{}"].gauges)
+	require.EqualValues(t, 1, registry.gauges.data["database.sql.tx{}"].gauges["database.sql.tx{}"].value)
+
+	require.NoError(t, tx2.Rollback())
+	require.EqualValues(t, 0, registry.gauges.data["database.sql.tx{}"].gauges["database.sql.tx{}"].value)
+
+	require.NoError(t, cc2.Close())
+	require.EqualValues(t, 2, registry.gauges.data["database.sql.conns{}"].gauges["database.sql.conns{}"].value)
+
+	require.NoError(t, db.Close())
+	require.EqualValues(t, 0, registry.gauges.data["database.sql.conns{}"].gauges["database.sql.conns{}"].value)
+}

trace/sql.go

@@ -77,8 +77,8 @@ type (
 	}
 	// Internals: https://github.com/ydb-platform/ydb-go-sdk/blob/master/VERSIONING.md#internals
 	DatabaseSQLConnectorConnectDoneInfo struct {
-		Error   error
 		Session sessionInfo
+		Error   error
 	}
 	// Internals: https://github.com/ydb-platform/ydb-go-sdk/blob/master/VERSIONING.md#internals
 	DatabaseSQLConnPingStartInfo struct {