MSSQL: Simplify proxy logic (#86278)

* Simplify proxy dialer creation

- Set new dialer on connector
- Create MSSQL connector in a similar fashion to postgres

* Update test

* Fix lint

* More lint

* Use correct driver name

pkg/tsdb/mssql/mssql.go

@@ -19,7 +19,7 @@ import (
 	"github.com/grafana/grafana-plugin-sdk-go/data"
 	"github.com/grafana/grafana-plugin-sdk-go/data/sqlutil"
 	mssql "github.com/microsoft/go-mssqldb"
-	_ "github.com/microsoft/go-mssqldb/azuread"
+	"github.com/microsoft/go-mssqldb/azuread"
 	_ "github.com/microsoft/go-mssqldb/integratedauth/krb5"
 
 	"github.com/grafana/grafana-plugin-sdk-go/backend/log"
@@ -70,6 +70,72 @@ func (s *Service) QueryData(ctx context.Context, req *backend.QueryDataRequest)
 	return dsHandler.QueryData(ctx, req)
 }
 
+func newMSSQL(ctx context.Context, driverName string, userFacingDefaultError string, rowLimit int64, dsInfo sqleng.DataSourceInfo, cnnstr string, logger log.Logger, settings backend.DataSourceInstanceSettings) (*sql.DB, *sqleng.DataSourceHandler, error) {
+	var connector *mssql.Connector
+	var err error
+	if driverName == "azuresql" {
+		connector, err = azuread.NewConnector(cnnstr)
+	} else {
+		connector, err = mssql.NewConnector(cnnstr)
+	}
+
+	if err != nil {
+		logger.Error("mssql connector creation failed", "error", err)
+		return nil, nil, fmt.Errorf("mssql connector creation failed")
+	}
+
+	proxyClient, err := settings.ProxyClient(ctx)
+	if err != nil {
+		logger.Error("mssql proxy creation failed", "error", err)
+		return nil, nil, fmt.Errorf("mssql proxy creation failed")
+	}
+
+	if proxyClient.SecureSocksProxyEnabled() {
+		dialer, err := proxyClient.NewSecureSocksProxyContextDialer()
+		if err != nil {
+			logger.Error("mssql proxy creation failed", "error", err)
+			return nil, nil, fmt.Errorf("mssql proxy creation failed")
+		}
+		URL, err := ParseURL(dsInfo.URL, logger)
+		if err != nil {
+			return nil, nil, err
+		}
+
+		mssqlDialer, err := newMSSQLProxyDialer(URL.Hostname(), dialer)
+		if err != nil {
+			return nil, nil, err
+		}
+		// update the mssql dialer with the proxy dialer
+		connector.Dialer = (mssqlDialer)
+	}
+
+	config := sqleng.DataPluginConfiguration{
+		DSInfo:            dsInfo,
+		MetricColumnTypes: []string{"VARCHAR", "CHAR", "NVARCHAR", "NCHAR"},
+		RowLimit:          rowLimit,
+	}
+
+	queryResultTransformer := mssqlQueryResultTransformer{
+		userError: userFacingDefaultError,
+	}
+
+	db := sql.OpenDB(connector)
+
+	db.SetMaxOpenConns(config.DSInfo.JsonData.MaxOpenConns)
+	db.SetMaxIdleConns(config.DSInfo.JsonData.MaxIdleConns)
+	db.SetConnMaxLifetime(time.Duration(config.DSInfo.JsonData.ConnMaxLifetime) * time.Second)
+
+	handler, err := sqleng.NewQueryDataHandler(userFacingDefaultError, db, config, &queryResultTransformer, newMssqlMacroEngine(),
+		logger)
+	if err != nil {
+		logger.Error("Failed connecting to Postgres", "err", err)
+		return nil, nil, err
+	}
+
+	logger.Debug("Successfully connected to Postgres")
+	return db, handler, nil
+}
+
 func newInstanceSettings(cfg *setting.Cfg, logger log.Logger) datasource.InstanceFactoryFunc {
 	return func(ctx context.Context, settings backend.DataSourceInstanceSettings) (instancemgmt.Instance, error) {
 		jsonData := sqleng.JsonData{
@@ -123,47 +189,15 @@ func newInstanceSettings(cfg *setting.Cfg, logger log.Logger) datasource.Instanc
 			driverName = "azuresql"
 		}
 
-		proxyClient, err := settings.ProxyClient(ctx)
+		_, handler, err := newMSSQL(ctx, driverName, cfg.UserFacingDefaultError, cfg.DataProxyRowLimit, dsInfo, cnnstr, logger, settings)
+
 		if err != nil {
+			logger.Error("Failed connecting to MSSQL", "err", err)
 			return nil, err
 		}
 
-		// register a new proxy driver if the secure socks proxy is enabled
+		logger.Debug("Successfully connected to MSSQL")
-		if proxyClient.SecureSocksProxyEnabled() {
+		return handler, nil
-			dialer, err := proxyClient.NewSecureSocksProxyContextDialer()
-			if err != nil {
-				return nil, err
-			}
-			URL, err := ParseURL(dsInfo.URL, logger)
-			if err != nil {
-				return nil, err
-			}
-			driverName, err = createMSSQLProxyDriver(cnnstr, URL.Hostname(), dialer)
-			if err != nil {
-				return nil, err
-			}
-		}
-
-		config := sqleng.DataPluginConfiguration{
-			DSInfo:            dsInfo,
-			MetricColumnTypes: []string{"VARCHAR", "CHAR", "NVARCHAR", "NCHAR"},
-			RowLimit:          cfg.DataProxyRowLimit,
-		}
-
-		queryResultTransformer := mssqlQueryResultTransformer{
-			userError: cfg.UserFacingDefaultError,
-		}
-
-		db, err := sql.Open(driverName, cnnstr)
-		if err != nil {
-			return nil, err
-		}
-
-		db.SetMaxOpenConns(config.DSInfo.JsonData.MaxOpenConns)
-		db.SetMaxIdleConns(config.DSInfo.JsonData.MaxIdleConns)
-		db.SetConnMaxLifetime(time.Duration(config.DSInfo.JsonData.ConnMaxLifetime) * time.Second)
-
-		return sqleng.NewQueryDataHandler(cfg.UserFacingDefaultError, db, config, &queryResultTransformer, newMssqlMacroEngine(), logger)
 	}
 }
 

pkg/tsdb/mssql/proxy.go

@@ -2,42 +2,13 @@ package mssql
 
 import (
 	"context"
-	"crypto/md5"
-	"database/sql"
-	"database/sql/driver"
 	"errors"
-	"fmt"
 	"net"
-	"slices"
 
 	mssql "github.com/microsoft/go-mssqldb"
 	"golang.org/x/net/proxy"
 )
 
-// createMSSQLProxyDriver creates and registers a new sql driver that uses a mssql connector and updates the dialer to
-// route connections through the secure socks proxy
-func createMSSQLProxyDriver(cnnstr string, hostName string, dialer proxy.Dialer) (string, error) {
-	// create a unique driver per connection string
-	hash := fmt.Sprintf("%x", md5.Sum([]byte(cnnstr)))
-	driverName := "mssql-proxy-" + hash
-
-	// only register the driver once
-	if !slices.Contains(sql.Drivers(), driverName) {
-		connector, err := mssql.NewConnector(cnnstr)
-		if err != nil {
-			return "", err
-		}
-
-		driver, err := newMSSQLProxyDriver(connector, hostName, dialer)
-		if err != nil {
-			return "", err
-		}
-		sql.Register(driverName, driver)
-	}
-
-	return driverName, nil
-}
-
 type HostTransportDialer struct {
 	Dialer proxy.ContextDialer
 	Host   string
@@ -47,36 +18,12 @@ func (m HostTransportDialer) DialContext(ctx context.Context, network string, ad
 	return m.Dialer.DialContext(ctx, network, addr)
 }
 
-func (m HostTransportDialer) HostName() string {
+// // we wrap the proxy.Dialer to become dialer that the mssql module accepts
-	return m.Host
+func newMSSQLProxyDialer(hostName string, dialer proxy.Dialer) (mssql.Dialer, error) {
-}
-
-// mssqlProxyDriver is a regular mssql driver with an updated dialer.
-// This is needed because there is no way to save a dialer to the mssql driver in xorm
-type mssqlProxyDriver struct {
-	c *mssql.Connector
-}
-
-var _ driver.DriverContext = (*mssqlProxyDriver)(nil)
-
-// newMSSQLProxyDriver updates the dialer for a mssql connector with a dialer that proxys connections through the secure socks proxy
-// and returns a new mssql driver to register
-func newMSSQLProxyDriver(connector *mssql.Connector, hostName string, dialer proxy.Dialer) (*mssqlProxyDriver, error) {
 	contextDialer, ok := dialer.(proxy.ContextDialer)
 	if !ok {
 		return nil, errors.New("unable to cast socks proxy dialer to context proxy dialer")
 	}
 
-	connector.Dialer = HostTransportDialer{contextDialer, hostName}
+	return &HostTransportDialer{contextDialer, hostName}, nil
-	return &mssqlProxyDriver{c: connector}, nil
-}
-
-// OpenConnector returns the normal mssql connector that has the updated dialer context
-func (d *mssqlProxyDriver) OpenConnector(name string) (driver.Connector, error) {
-	return d.c, nil
-}
-
-// Open uses the connector with the updated dialer context to open a new connection
-func (d *mssqlProxyDriver) Open(dsn string) (driver.Conn, error) {
-	return d.c.Connect(context.Background())
 }

pkg/tsdb/mssql/proxy_test.go

@@ -2,6 +2,7 @@ package mssql
 
 import (
 	"context"
+	"database/sql"
 	"fmt"
 	"net"
 	"testing"
@@ -23,50 +24,21 @@ func (d *testDialer) DialContext(ctx context.Context, network, address string) (
 }
 
 var _ proxy.Dialer = (&testDialer{})
-var _ proxy.ContextDialer = (&testDialer{})
-
-func newTestDialer() proxy.Dialer {
-	d := testDialer{}
-	return &d
-}
 
 func TestMSSQLProxyDriver(t *testing.T) {
 	cnnstr := "server=127.0.0.1;port=1433;user id=sa;password=yourStrong(!)Password;database=db"
-	driverName, err := createMSSQLProxyDriver(cnnstr, "127.0.0.1", newTestDialer())
-	require.NoError(t, err)
-
-	t.Run("Driver should not be registered more than once", func(t *testing.T) {
-		testDriver, err := createMSSQLProxyDriver(cnnstr, "127.0.0.1", newTestDialer())
-		require.NoError(t, err)
-		require.Equal(t, driverName, testDriver)
-	})
-
-	t.Run("A new driver should be created for a new connection string", func(t *testing.T) {
-		testDriver, err := createMSSQLProxyDriver("server=localhost;user id=sa;password=yourStrong(!)Password;database=db2", "localhost", newTestDialer())
-		require.NoError(t, err)
-		require.NotEqual(t, driverName, testDriver)
-	})
 
 	t.Run("Connector should use dialer context that routes through the socks proxy to db", func(t *testing.T) {
 		connector, err := mssql.NewConnector(cnnstr)
 		require.NoError(t, err)
-		driver, err := newMSSQLProxyDriver(connector, "127.0.0.1", newTestDialer())
+		dialer, err := newMSSQLProxyDialer("127.0.0.1", &testDialer{})
 		require.NoError(t, err)
 
-		conn, err := driver.OpenConnector(cnnstr)
+		connector.Dialer = (dialer)
-		require.NoError(t, err)
 
-		_, err = conn.Connect(context.Background())
+		db := sql.OpenDB(connector)
-		require.Contains(t, err.Error(), "test-dialer: DialContext is not functional")
+		err = db.Ping()
-	})
 
-	t.Run("Open should use the connector that routes through the socks proxy to db", func(t *testing.T) {
-		connector, err := mssql.NewConnector(cnnstr)
-		require.NoError(t, err)
-		driver, err := newMSSQLProxyDriver(connector, "127.0.0.1", newTestDialer())
-		require.NoError(t, err)
-
-		_, err = driver.Open(cnnstr)
 		require.Contains(t, err.Error(), "test-dialer: DialContext is not functional")
 	})
 }