fix(driver,liner): after mr discussion

Cianidos · Cianidos · commit 4ab4614e833d · 2026-02-04T19:58:44.000+03:00
More generic way to pass dialfunc to driver.
updateDialler -&gt; Driver.Configure

Explanations about the way dialfunc passed that way.

pkg/driver package now have inverted dependency. Drivers should add
themselves to driver dispatcher.

Closes RNDSTROPPY-70
diff --git a/.golangci.yml b/.golangci.yml
@@ -36,8 +36,6 @@ linters:
     - forcetypeassert
     - funlen
     - gocheckcompilerdirectives
-    - gochecknoglobals
-    - gochecknoinits
     - gochecksumtype
     - gocognit
     - goconst
@@ -150,7 +148,6 @@ linters:
         - gocognit
         - lll
         - mnd
-        - gochecknoglobals
         - err113
         - ireturn
       allow-unused: true
@@ -226,7 +223,6 @@ linters:
           - wrapcheck
         path: _test\.go
       - linters:
-          - gochecknoglobals
           - unused
         path: bind.go
       - linters:
diff --git a/cmd/stroppy/commands/gen/new.go b/cmd/stroppy/commands/gen/new.go
@@ -22,7 +22,7 @@ const (
 	presetFlagName  = "preset"
 )
 
-var Cmd = &cobra.Command{ //nolint: gochecknoglobals
+var Cmd = &cobra.Command{
 	Use:   "gen",
 	Short: "Generate stroppy development environment",
 	Long: fmt.Sprintf(`
@@ -133,7 +133,7 @@ Examples:
 	},
 }
 
-func init() { //nolint: gochecknoinits // allow in cmd
+func init() {
 	Cmd.PersistentFlags().String(
 		workdirFlagName,
 		".",
diff --git a/cmd/stroppy/commands/root.go b/cmd/stroppy/commands/root.go
@@ -12,7 +12,7 @@ import (
 	"github.com/stroppy-io/stroppy/pkg/common/logger"
 )
 
-var rootCmd = &cobra.Command{ //nolint: gochecknoglobals
+var rootCmd = &cobra.Command{
 	Use:   "stroppy",
 	Short: "Tool to generate and run stress tests (e.g benchmarking) for databases",
 	Long: `
@@ -26,7 +26,7 @@ For more information see https://github.com/stroppy-io/stroppy`,
 	SilenceUsage: true,
 }
 
-var versionCmd = &cobra.Command{ //nolint: gochecknoglobals
+var versionCmd = &cobra.Command{
 	Use:   "version",
 	Short: "Print versions of stroppy components",
 	Long:  ``,
@@ -42,7 +42,7 @@ func Execute() {
 	}
 }
 
-func init() { //nolint: gochecknoinits // allow in cmd
+func init() {
 	cobra.EnableCommandSorting = false
 	rootCmd.CompletionOptions.HiddenDefaultCmd = true
 	rootCmd.SetVersionTemplate(`{{with .Name}}{{printf "%s " .}}{{end}}{{printf "%s" .Version}}`)
diff --git a/cmd/stroppy/commands/run/run.go b/cmd/stroppy/commands/run/run.go
@@ -14,7 +14,7 @@ const (
 	maxArgs = 2 // sql_file.sql is optional
 )
 
-var Cmd = &cobra.Command{ //nolint: gochecknoglobals
+var Cmd = &cobra.Command{
 	Use:   "run <script.ts> [sql_file.sql]",
 	Short: "Run benchmark script with k6",
 	Long: `Run a TypeScript benchmark script with k6.
diff --git a/cmd/xk6air/driver_wrapper.go b/cmd/xk6air/driver_wrapper.go
@@ -12,17 +12,35 @@ import (
 )
 
 type DriverWrapper struct {
-	vu                modules.VU
-	lg                *zap.Logger
-	drv               driver.Driver
-	onceUpdateDialler sync.Once
+	vu  modules.VU
+	lg  *zap.Logger
+	drv driver.Driver
+
+	configureOnce sync.Once
 }
 
-func (d *DriverWrapper) RunQuery(sql string, args map[string]any) any {
-	d.onceUpdateDialler.Do(
-		func() { d.drv.UpdateDialler(d.vu.Context(), d.vu.State().Dialer.DialContext) },
+// This is a custom "VU setup" hook.
+//
+// NOTE: k6 have no option to make per VU setup code execution by itself.
+// Check https://github.com/grafana/k6/issues/785
+// https://github.com/grafana/k6/issues/1638
+//
+// Unfortunatly it's impossible to pass DialFunc at [Instance.NewDriverByConfigBin]
+// because there is nil [modules.VU.State]. It may be fixed in the feature:
+// https://github.com/grafana/k6/issues?q=is%3Aopen+is%3Aissue+label%3Anew-http
+// https://github.com/grafana/k6/issues/2293
+func (d *DriverWrapper) configure() {
+	d.configureOnce.Do(
+		func() {
+			d.drv.Configure(d.vu.Context(), driver.Options{
+				DialFunc: d.vu.State().Dialer.DialContext,
+			})
+		},
 	)
+}
 
+func (d *DriverWrapper) RunQuery(sql string, args map[string]any) any {
+	d.configure()
 	stats, err := d.drv.RunQuery(d.vu.Context(), sql, args)
 	if err != nil {
 		return fmt.Errorf("error while executing sql query: %w", err)
@@ -31,13 +49,8 @@ func (d *DriverWrapper) RunQuery(sql string, args map[string]any) any {
 }
 
 // InsertValuesBin starts bulk insert blocking operation on driver.
-func (d *DriverWrapper) InsertValuesBin(
-	insertMsg []byte,
-	count int64,
-) any {
-	d.onceUpdateDialler.Do(
-		func() { d.drv.UpdateDialler(d.vu.Context(), d.vu.State().Dialer.DialContext) },
-	)
+func (d *DriverWrapper) InsertValuesBin(insertMsg []byte, count int64) any {
+	d.configure()
 	var descriptor stroppy.InsertDescriptor
 	err := proto.Unmarshal(insertMsg, &descriptor)
 	if err != nil {
diff --git a/cmd/xk6air/instance.go b/cmd/xk6air/instance.go
@@ -7,6 +7,7 @@ import (
 	"github.com/stroppy-io/stroppy/pkg/common/logger"
 	"github.com/stroppy-io/stroppy/pkg/common/proto/stroppy"
 	"github.com/stroppy-io/stroppy/pkg/driver"
+	_ "github.com/stroppy-io/stroppy/pkg/driver/postgres"
 	"go.k6.io/k6/js/modules"
 	"go.uber.org/zap"
 	"google.golang.org/protobuf/proto"
diff --git a/cmd/xk6air/module.go b/cmd/xk6air/module.go
@@ -20,7 +20,7 @@ var rootModule *RootModule
 var _ modules.Module = new(RootModule)
 
 // rootModule initialization.
-func init() { //nolint:gochecknoinits // allow for xk6
+func init() {
 	lg := logger.
 		NewFromEnv().
 		Named("k6-module").
diff --git a/internal/static/embed.go b/internal/static/embed.go
@@ -30,7 +30,7 @@ const (
 )
 
 // StaticFiles are copied to temp dir for k6 execution.
-var StaticFiles = []FileName{ //nolint: gochecknoglobals
+var StaticFiles = []FileName{
 	ProtoJSFileName,
 	K6PluginFileName,
 	HelpersFileName,
@@ -40,14 +40,14 @@ var StaticFiles = []FileName{ //nolint: gochecknoglobals
 }
 
 // DevStaticFiles are additional files for development environment.
-var DevStaticFiles = []FileName{ //nolint: gochecknoglobals
+var DevStaticFiles = []FileName{
 	PackageJSONFileName,
 	ParseSQLTSFileName,
 	ProtoTSFileName,
 	TSConfigFileName,
 }
 
-var Binaries = []FileName{ //nolint: gochecknoglobals
+var Binaries = []FileName{
 	K6PluginFileName,
 }
 
diff --git a/internal/version/version.go b/internal/version/version.go
@@ -1,3 +1,3 @@
 package version
 
-var Version = "0.0.0" //nolint: gochecknoglobals
+var Version = "0.0.0"
diff --git a/pkg/common/generate/randstr/string.go b/pkg/common/generate/randstr/string.go
@@ -12,7 +12,7 @@ func (sg *StringGenerator) Next() string {
 	return sg.cutter.Cut()
 }
 
-var DefaultEnglishAlphabet = [][2]int32{{65, 90}, {97, 122}} //nolint: gochecknoglobals
+var DefaultEnglishAlphabet = [][2]int32{{65, 90}, {97, 122}}
 
 func NewStringGenerator(
 	seed uint64,
diff --git a/pkg/common/logger/logger.go b/pkg/common/logger/logger.go
@@ -22,9 +22,9 @@ type Config struct {
 	LogLevel string `default:"info"       mapstructure:"level" validate:"oneof=debug info warn error"`
 }
 
-var globalLogger = atomic.Pointer[zap.Logger]{} //nolint:gochecknoglobals // global logger needed for all app.
+var globalLogger = atomic.Pointer[zap.Logger]{}
 
-func init() { //nolint: gochecknoinits // allow
+func init() {
 	setGlobalLogger(newDefault())
 }
 
diff --git a/pkg/common/shutdown/shortcut.go b/pkg/common/shutdown/shortcut.go
@@ -2,7 +2,7 @@ package shutdown
 
 import "os"
 
-var globalShutdown = New() //nolint:gochecknoglobals // used for main app
+var globalShutdown = New()
 
 // Global returns the global shutdown manager.
 //
diff --git a/pkg/driver/dispatcher.go b/pkg/driver/dispatcher.go
@@ -2,44 +2,56 @@ package driver
 
 import (
 	"context"
+	"errors"
+	"fmt"
 	"net"
 
 	"go.uber.org/zap"
 
 	stroppy "github.com/stroppy-io/stroppy/pkg/common/proto/stroppy"
-	"github.com/stroppy-io/stroppy/pkg/driver/postgres"
 	"github.com/stroppy-io/stroppy/pkg/driver/stats"
 )
 
-type Driver interface {
-	InsertValues(ctx context.Context, unit *stroppy.InsertDescriptor) (*stats.Query, error)
-	RunQuery(ctx context.Context, sql string, args map[string]any) (*stats.Query, error)
-	Teardown(ctx context.Context) error
+type (
+	Options struct {
+		// Allows to pass k6 DialFunc to driver for proper network metrics.
+		DialFunc func(ctx context.Context, network, addr string) (net.Conn, error)
+	}
+
+	Driver interface {
+		InsertValues(ctx context.Context, unit *stroppy.InsertDescriptor) (*stats.Query, error)
+		RunQuery(ctx context.Context, sql string, args map[string]any) (*stats.Query, error)
+		Teardown(ctx context.Context) error
+		Configure(ctx context.Context, opts Options) error
+	}
 
-	UpdateDialler(
+	driverConstructor = func(
 		ctx context.Context,
-		dialFunc func(ctx context.Context, network, addr string) (net.Conn, error),
-	) (err error)
+		lg *zap.Logger,
+		config *stroppy.DriverConfig,
+	) (Driver, error)
+)
+
+var ErrNoRegisteredDriver = errors.New("no registered driver")
+
+var registry = map[stroppy.DriverConfig_DriverType]driverConstructor{}
+
+func RegisterDriver(
+	driverType stroppy.DriverConfig_DriverType,
+	constructor driverConstructor,
+) {
+	registry[driverType] = constructor
 }
 
 func Dispatch( //nolint: ireturn // better than return any
 	ctx context.Context,
 	lg *zap.Logger,
 	config *stroppy.DriverConfig,
 ) (Driver, error) {
-	switch drvType := config.GetDriverType(); drvType {
-	case stroppy.DriverConfig_DRIVER_TYPE_UNSPECIFIED:
-		lg.Sugar().
-			Warnf("driver type UNSPECIFIED, fall back to %s", stroppy.DriverConfig_DRIVER_TYPE_POSTGRES)
-
-		fallthrough // as good suggestion
-	case stroppy.DriverConfig_DRIVER_TYPE_POSTGRES:
-		drv, err := postgres.NewDriver(ctx, lg, config)
-
-		return drv, err
-	default:
-		lg.Sugar().Panicf("driver type '%s' not dispatchable", drvType.String())
-
-		return nil, nil //nolint:nilnil // unreachable after panic
+	drvType := config.GetDriverType()
+	if constructor, ok := registry[drvType]; ok {
+		return constructor(ctx, lg, config)
 	}
+
+	return nil, fmt.Errorf("driver type '%s': %w", drvType.String(), ErrNoRegisteredDriver)
 }
diff --git a/pkg/driver/postgres/driver.go b/pkg/driver/postgres/driver.go
@@ -2,7 +2,7 @@ package postgres
 
 import (
 	"context"
-	"net"
+	"fmt"
 
 	"github.com/jackc/pgx/v5"
 	"github.com/jackc/pgx/v5/pgconn"
@@ -11,9 +11,19 @@ import (
 
 	"github.com/stroppy-io/stroppy/pkg/common/logger"
 	stroppy "github.com/stroppy-io/stroppy/pkg/common/proto/stroppy"
+	"github.com/stroppy-io/stroppy/pkg/driver"
 	"github.com/stroppy-io/stroppy/pkg/driver/postgres/pool"
 )
 
+func init() {
+	driver.RegisterDriver(
+		stroppy.DriverConfig_DRIVER_TYPE_POSTGRES,
+		func(ctx context.Context, lg *zap.Logger, config *stroppy.DriverConfig) (driver.Driver, error) {
+			return NewDriver(ctx, lg, config)
+		},
+	)
+}
+
 type Executor interface {
 	Exec(ctx context.Context, sql string, arguments ...any) (pgconn.CommandTag, error)
 	CopyFrom(
@@ -32,6 +42,8 @@ type Driver struct {
 	pgxPool Executor
 }
 
+var _ driver.Driver = new(Driver)
+
 func NewDriver(
 	ctx context.Context,
 	lg *zap.Logger,
@@ -72,17 +84,20 @@ func NewDriver(
 	return d, nil
 }
 
-func (d *Driver) UpdateDialler(
-	ctx context.Context,
-	dialFunc func(ctx context.Context, network, addr string) (net.Conn, error),
-) (err error) {
-	cfg := d.pgxPool.Config()
-	cfg.ConnConfig.DialFunc = dialFunc
+func (d *Driver) Configure(ctx context.Context, opts driver.Options) (err error) {
+	if opts.DialFunc != nil {
+		cfg := d.pgxPool.Config()
+		cfg.ConnConfig.DialFunc = opts.DialFunc
 
-	d.pgxPool.Close()
-	d.pgxPool, err = pgxpool.NewWithConfig(ctx, cfg)
+		d.pgxPool.Close()
 
-	return err
+		d.pgxPool, err = pgxpool.NewWithConfig(ctx, cfg)
+		if err != nil {
+			return fmt.Errorf("can't start reconfigured pgxpool: %w", err)
+		}
+	}
+
+	return nil
 }
 
 func (d *Driver) Teardown(_ context.Context) error {
diff --git a/pkg/driver/postgres/pool/pool.go b/pkg/driver/postgres/pool/pool.go
diff --git a/pkg/utils/protoyaml/protoyaml.go b/pkg/utils/protoyaml/protoyaml.go

Original file line number	Diff line number	Diff line change
`@@ -12,7 +12,7 @@ import (`
`12`	`12`	`"github.com/stroppy-io/stroppy/pkg/common/logger"`
`13`	`13`	`)`
`14`	`14`
`15`		`-var rootCmd = &cobra.Command{ //nolint: gochecknoglobals`
	`15`	`+var rootCmd = &cobra.Command{`
`16`	`16`	`Use: "stroppy",`
`17`	`17`	`Short: "Tool to generate and run stress tests (e.g benchmarking) for databases",`
`18`	`18`	Long: `
@@ -26,7 +26,7 @@ For more information see https://github.com/stroppy-io/stroppy`,
`26`	`26`	`SilenceUsage: true,`
`27`	`27`	`}`
`28`	`28`
`29`		`-var versionCmd = &cobra.Command{ //nolint: gochecknoglobals`
	`29`	`+var versionCmd = &cobra.Command{`
`30`	`30`	`Use: "version",`
`31`	`31`	`Short: "Print versions of stroppy components",`
`32`	`32`	Long: ``,
`@@ -42,7 +42,7 @@ func Execute() {`
`42`	`42`	`}`
`43`	`43`	`}`
`44`	`44`
`45`		`-func init() { //nolint: gochecknoinits // allow in cmd`
	`45`	`+func init() {`
`46`	`46`	`cobra.EnableCommandSorting = false`
`47`	`47`	`rootCmd.CompletionOptions.HiddenDefaultCmd = true`
`48`	`48`	rootCmd.SetVersionTemplate(`{{with .Name}}{{printf "%s " .}}{{end}}{{printf "%s" .Version}}`)
Original file line number	Diff line number	Diff line change
`@@ -14,7 +14,7 @@ const (`
`14`	`14`	`maxArgs = 2 // sql_file.sql is optional`
`15`	`15`	`)`
`16`	`16`
`17`		`-var Cmd = &cobra.Command{ //nolint: gochecknoglobals`
	`17`	`+var Cmd = &cobra.Command{`
`18`	`18`	`Use: "run <script.ts> [sql_file.sql]",`
`19`	`19`	`Short: "Run benchmark script with k6",`
`20`	`20`	Long: `Run a TypeScript benchmark script with k6.
Original file line number	Diff line number	Diff line change
`@@ -30,7 +30,7 @@ const (`
`30`	`30`	`)`
`31`	`31`
`32`	`32`	`// StaticFiles are copied to temp dir for k6 execution.`
`33`		`-var StaticFiles = []FileName{ //nolint: gochecknoglobals`
	`33`	`+var StaticFiles = []FileName{`
`34`	`34`	`ProtoJSFileName,`
`35`	`35`	`K6PluginFileName,`
`36`	`36`	`HelpersFileName,`
`@@ -40,14 +40,14 @@ var StaticFiles = []FileName{ //nolint: gochecknoglobals`
`40`	`40`	`}`
`41`	`41`
`42`	`42`	`// DevStaticFiles are additional files for development environment.`
`43`		`-var DevStaticFiles = []FileName{ //nolint: gochecknoglobals`
	`43`	`+var DevStaticFiles = []FileName{`
`44`	`44`	`PackageJSONFileName,`
`45`	`45`	`ParseSQLTSFileName,`
`46`	`46`	`ProtoTSFileName,`
`47`	`47`	`TSConfigFileName,`
`48`	`48`	`}`
`49`	`49`
`50`		`-var Binaries = []FileName{ //nolint: gochecknoglobals`
	`50`	`+var Binaries = []FileName{`
`51`	`51`	`K6PluginFileName,`
`52`	`52`	`}`
`53`	`53`
Original file line number	Diff line number	Diff line change
`@@ -1,3 +1,3 @@`
`1`	`1`	`package version`
`2`	`2`
`3`		`-var Version = "0.0.0" //nolint: gochecknoglobals`
	`3`	`+var Version = "0.0.0"`
Original file line number	Diff line number	Diff line change
`@@ -12,7 +12,7 @@ func (sg *StringGenerator) Next() string {`
`12`	`12`	`return sg.cutter.Cut()`
`13`	`13`	`}`
`14`	`14`
`15`		`-var DefaultEnglishAlphabet = [][2]int32{{65, 90}, {97, 122}} //nolint: gochecknoglobals`
	`15`	`+var DefaultEnglishAlphabet = [][2]int32{{65, 90}, {97, 122}}`
`16`	`16`
`17`	`17`	`func NewStringGenerator(`
`18`	`18`	`seed uint64,`