Add test for QuerySample collector

go.mod

@@ -8,6 +8,7 @@ require (
 	github.com/Azure/azure-sdk-for-go/sdk/azcore v1.13.0
 	github.com/Azure/azure-sdk-for-go/sdk/azidentity v1.7.0
 	github.com/Azure/go-autorest/autorest v0.11.29
+	github.com/DATA-DOG/go-sqlmock v1.5.2
 	github.com/IBM/sarama v1.43.3
 	github.com/KimMachineGun/automemlimit v0.6.0
 	github.com/Lusitaniae/apache_exporter v0.11.1-0.20220518131644-f9522724dab4
@@ -768,7 +769,7 @@ require (
 	github.com/xeipuuv/gojsonschema v1.2.0 // indirect
 	github.com/xhit/go-str2duration/v2 v2.1.0 // indirect
 	github.com/xo/dburl v0.20.0 // indirect
-	github.com/xwb1989/sqlparser v0.0.0-20180606152119-120387863bf2 // indirect
+	github.com/xwb1989/sqlparser v0.0.0-20180606152119-120387863bf2
 	github.com/yl2chen/cidranger v1.0.2 // indirect
 	github.com/youmark/pkcs8 v0.0.0-20240424034433-3c2c7870ae76 // indirect
 	github.com/yusufpapurcu/wmi v1.2.4 // indirect

go.sum

@@ -1646,6 +1646,7 @@ github.com/kisielk/errcheck v1.1.0/go.mod h1:EZBBE59ingxPouuu3KfxchcWSUPOHkagtvW
 github.com/kisielk/errcheck v1.2.0/go.mod h1:/BMXB+zMLi60iA8Vv6Ksmxu/1UDYcXs4uQLJ+jE2L00=
 github.com/kisielk/errcheck v1.5.0/go.mod h1:pFxgyoBC7bSaBwPgfKdkLd5X25qrDl4LWUI2bnpBCr8=
 github.com/kisielk/gotool v1.0.0/go.mod h1:XhKaO+MFFWcvkIS/tQcRk01m1F5IRFswLeQ+oQHNcck=
+github.com/kisielk/sqlstruct v0.0.0-20201105191214-5f3e10d3ab46/go.mod h1:yyMNCyc/Ib3bDTKd379tNMpB/7/H5TjM2Y9QJ5THLbE=
 github.com/klauspost/asmfmt v1.3.2 h1:4Ri7ox3EwapiOjCki+hw14RyKk201CN4rzyCJRFLpK4=
 github.com/klauspost/asmfmt v1.3.2/go.mod h1:AG8TuvYojzulgDAMCnYn50l/5QV3Bs/tp6j0HLHbNSE=
 github.com/klauspost/compress v1.11.0/go.mod h1:aoV0uJVorq1K+umq18yTdKaF57EivdYsUV+/s2qKfXs=

internal/component/database_observability/mysql/collector/query_sample.go

@@ -3,7 +3,6 @@ package collector
 import (
 	"context"
 	"database/sql"
-	"errors"
 	"fmt"
 	"strings"
 	"time"
@@ -32,7 +31,7 @@ const selectQuerySamples = `
 	WHERE last_seen > DATE_SUB(NOW(), INTERVAL 1 DAY)`
 
 type QuerySampleArguments struct {
-	DSN            string
+	DB             *sql.DB
 	ScrapeInterval time.Duration
 	EntryHandler   loki.EntryHandler
 
@@ -51,21 +50,8 @@ type QuerySample struct {
 }
 
 func NewQuerySample(args QuerySampleArguments) (*QuerySample, error) {
-	dbConnection, err := sql.Open("mysql", args.DSN)
-	if err != nil {
-		return nil, err
-	}
-
-	if dbConnection == nil {
-		return nil, errors.New("nil DB connection")
-	}
-
-	if err = dbConnection.Ping(); err != nil {
-		return nil, err
-	}
-
 	return &QuerySample{
-		dbConnection:   dbConnection,
+		dbConnection:   args.DB,
 		scrapeInterval: args.ScrapeInterval,
 		entryHandler:   args.EntryHandler,
 		logger:         args.Logger,

internal/component/database_observability/mysql/collector/query_sample_test.go

@@ -0,0 +1,65 @@
+package collector
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	loki_fake "github.com/grafana/alloy/internal/component/common/loki/client/fake"
+	"github.com/prometheus/common/model"
+
+	"github.com/DATA-DOG/go-sqlmock"
+	"github.com/go-kit/log"
+	"github.com/stretchr/testify/require"
+)
+
+func TestQuerySampleRun(t *testing.T) {
+	db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(sqlmock.QueryMatcherEqual))
+	require.NoError(t, err)
+	defer db.Close()
+
+	lokiClient := loki_fake.NewClient(func() {})
+
+	collector, err := NewQuerySample(QuerySampleArguments{
+		DB:             db,
+		ScrapeInterval: time.Second,
+		EntryHandler:   lokiClient,
+		Logger:         log.NewNopLogger(),
+	})
+	require.NoError(t, err)
+	require.NotNil(t, collector)
+
+	mock.ExpectQuery(selectQuerySamples).WithoutArgs().WillReturnRows(
+		sqlmock.NewRows([]string{
+			"digest",
+			"query_sample_text",
+			"query_sample_seen",
+			"query_sample_timer_wait",
+		}).AddRow(
+			"abc123",
+			"select * from table_name where id = 1",
+			"2024-01-01T00:00:00.000Z",
+			"1000",
+		),
+	)
+
+	err = collector.Run(context.Background())
+	require.NoError(t, err)
+
+	require.Eventually(t, func() bool {
+		return len(lokiClient.Received()) == 2
+	}, 5*time.Second, 100*time.Millisecond)
+
+	collector.Stop()
+	lokiClient.Stop()
+
+	lokiEntries := lokiClient.Received()
+	for _, entry := range lokiEntries {
+		require.Equal(t, model.LabelSet{"job": "integrations/db-o11y"}, entry.Labels)
+	}
+	require.Equal(t, `level=info msg="query samples fetched" op="query_sample" digest="abc123" query_sample_text="select * from table_name where id = 1" query_sample_seen="2024-01-01T00:00:00.000Z" query_sample_timer_wait="1000" query_redacted="select * from table_name where id = :redacted1"`, lokiEntries[0].Line)
+	require.Equal(t, `level=info msg="table name parsed" op="query_parsed_table_name" digest="abc123" table="table_name"`, lokiEntries[1].Line)
+
+	err = mock.ExpectationsWereMet()
+	require.NoError(t, err)
+}

internal/component/database_observability/mysql/component.go

@@ -2,6 +2,8 @@ package mysql
 
 import (
 	"context"
+	"database/sql"
+	"errors"
 	"fmt"
 	"net/http"
 	"path"
@@ -169,10 +171,26 @@ func (c *Component) Update(args component.Arguments) error {
 	newArgs := args.(Arguments)
 	c.args = newArgs
 
+	dbConnection, err := sql.Open("mysql", string(newArgs.DataSourceName))
+	if err != nil {
+		return err
+	}
+
+	if dbConnection == nil {
+		return errors.New("nil DB connection")
+	}
+
+	if err = dbConnection.Ping(); err != nil {
+		return err
+	}
+
 	entryHandler := loki.NewEntryHandler(c.handler.Chan(), func() {})
 
+	// dbConnection.Close()
+	// entryHandler.Stop()
+
 	qsCollector, err := collector.NewQuerySample(collector.QuerySampleArguments{
-		DSN:            string(newArgs.DataSourceName),
+		DB:             dbConnection,
 		ScrapeInterval: newArgs.ScrapeInterval,
 		EntryHandler:   entryHandler,
 		Logger:         c.opts.Logger,