Influx: Support PDC for Influx SQL (#103032)

* Support PDC for Influx SQL

* Correctly set the gRPC transport creds dial option

* Logging and refactor

* Allow URLs with no scheme

* Use passthrough resolver for socks proxy

* Update comment

* Extract url parsing and add test

* Check err

pkg/tsdb/influxdb/fsql/client.go

@@ -4,12 +4,14 @@ import (
 	"context"
 	"crypto/x509"
 	"fmt"
+	"net"
 	"sync"
 
 	"github.com/apache/arrow-go/v18/arrow/flight"
 	"github.com/apache/arrow-go/v18/arrow/flight/flightsql"
 	"github.com/apache/arrow-go/v18/arrow/ipc"
 	"github.com/apache/arrow-go/v18/arrow/memory"
+	"github.com/grafana/grafana-plugin-sdk-go/backend/proxy"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials"
 	"google.golang.org/grpc/credentials/insecure"
@@ -26,11 +28,17 @@ func (c *client) FlightClient() flight.Client {
 	return c.Client.Client
 }
 
-func newFlightSQLClient(addr string, metadata metadata.MD, secure bool) (*client, error) {
+func newFlightSQLClient(addr string, metadata metadata.MD, secure bool, proxyClient proxy.Client) (*client, error) {
-	dialOptions, err := grpcDialOptions(secure)
+	dialOptions, err := grpcDialOptions(secure, proxyClient)
 	if err != nil {
 		return nil, fmt.Errorf("grpc dial options: %s", err)
 	}
+
+	// If the secure socks proxy is enabled, we add the passthrough scheme. Otherwise, we use the raw address.
+	// This ensures the address is passed directly to the transport rather than trying to resolve via DNS.
+	if proxyClient.SecureSocksProxyEnabled() {
+		addr = fmt.Sprintf("passthrough:///%s", addr)
+	}
 	fsqlClient, err := flightsql.NewClient(addr, nil, nil, dialOptions...)
 	if err != nil {
 		return nil, err
@@ -38,21 +46,42 @@ func newFlightSQLClient(addr string, metadata metadata.MD, secure bool) (*client
 	return &client{Client: fsqlClient, md: metadata}, nil
 }
 
-func grpcDialOptions(secure bool) ([]grpc.DialOption, error) {
+func grpcDialOptions(secure bool, proxyClient proxy.Client) ([]grpc.DialOption, error) {
-	transport := grpc.WithTransportCredentials(insecure.NewCredentials())
+	dialOptions := []grpc.DialOption{}
+	secureDialOpt := grpc.WithTransportCredentials(insecure.NewCredentials())
+
 	if secure {
 		pool, err := x509.SystemCertPool()
 		if err != nil {
 			return nil, fmt.Errorf("x509: %s", err)
 		}
-		transport = grpc.WithTransportCredentials(credentials.NewClientTLSFromCert(pool, ""))
+		secureDialOpt = grpc.WithTransportCredentials(credentials.NewClientTLSFromCert(pool, ""))
+	}
+	dialOptions = append(dialOptions, secureDialOpt)
+
+	if proxyClient.SecureSocksProxyEnabled() {
+		dialer, err := proxyClient.NewSecureSocksProxyContextDialer()
+		if err != nil {
+			return nil, fmt.Errorf("failed to create influx proxy dialer: %s", err)
+		}
+
+		dialOptions = append(dialOptions, grpc.WithContextDialer(func(ctx context.Context, host string) (net.Conn, error) {
+			logger := glog.FromContext(ctx)
+			logger.Debug("Dialing secure socks proxy", "host", host)
+			conn, err := dialer.Dial("tcp", host)
+			if err != nil {
+				return nil, fmt.Errorf("not possible to dial secure socks proxy: %w", err)
+			}
+			select {
+			case <-ctx.Done():
+				return conn, fmt.Errorf("context canceled: %w", err)
+			default:
+				return conn, nil
+			}
+		}))
 	}
 
-	opts := []grpc.DialOption{
+	return dialOptions, nil
-		transport,
-	}
-
-	return opts, nil
 }
 
 // DoGetWithHeaderExtraction performs a normal DoGet, but wraps the stream in a

pkg/tsdb/influxdb/fsql/fsql.go

@@ -97,20 +97,39 @@ type runner struct {
 	client *client
 }
 
+func ParseURL(endpoint string) (string, error) {
+	if endpoint == "" {
+		return "", fmt.Errorf("missing URL from datasource configuration")
+	}
+
+	u, err := url.Parse(endpoint)
+	if err != nil {
+		return "", fmt.Errorf("bad URL : %s", err)
+	}
+
+	addr := u.Host
+	if u.Port() == "" {
+		addr += ":443"
+	}
+
+	// If the user has specified an address with no scheme it can still be valid
+	// So we use the raw URL value
+	if u.Host == "" {
+		addr = endpoint
+	}
+
+	return addr, nil
+}
+
 // runnerFromDataSource creates a runner from the datasource model (the datasource instance's configuration).
 func runnerFromDataSource(dsInfo *models.DatasourceInfo) (*runner, error) {
 	if dsInfo.URL == "" {
 		return nil, fmt.Errorf("missing URL from datasource configuration")
 	}
 
-	u, err := url.Parse(dsInfo.URL)
+	u, err := ParseURL(dsInfo.URL)
 	if err != nil {
-		return nil, fmt.Errorf("bad URL : %s", err)
+		return nil, err
-	}
-
-	addr := u.Host
-	if u.Port() == "" {
-		addr += ":443"
 	}
 
 	md := metadata.MD{}
@@ -121,7 +140,7 @@ func runnerFromDataSource(dsInfo *models.DatasourceInfo) (*runner, error) {
 		md.Set("Authorization", fmt.Sprintf("Bearer %s", dsInfo.Token))
 	}
 
-	fsqlClient, err := newFlightSQLClient(addr, md, !dsInfo.InsecureGrpc)
+	fsqlClient, err := newFlightSQLClient(u, md, !dsInfo.InsecureGrpc, dsInfo.ProxyClient)
 	if err != nil {
 		return nil, err
 	}

pkg/tsdb/influxdb/fsql/fsql_test.go

@@ -12,6 +12,7 @@ import (
 	"github.com/apache/arrow-go/v18/arrow/flight/flightsql/example"
 	"github.com/apache/arrow-go/v18/arrow/memory"
 	"github.com/grafana/grafana-plugin-sdk-go/backend"
+	"github.com/grafana/grafana-plugin-sdk-go/backend/proxy"
 	"github.com/stretchr/testify/require"
 	"github.com/stretchr/testify/suite"
 
@@ -70,6 +71,7 @@ func (suite *FSQLTestSuite) TestIntegration_QueryData() {
 				Version:      "test",
 				HTTPMode:     "proxy",
 				InsecureGrpc: true,
+				ProxyClient:  proxy.New(nil),
 			},
 			backend.QueryDataRequest{
 				Queries: []backend.DataQuery{
@@ -139,6 +141,7 @@ func TestInvalidSchema(t *testing.T) {
 			Version:      "test",
 			HTTPMode:     "proxy",
 			InsecureGrpc: true,
+			ProxyClient:  proxy.New(nil),
 		},
 		backend.QueryDataRequest{
 			Queries: []backend.DataQuery{
@@ -151,3 +154,61 @@ func TestInvalidSchema(t *testing.T) {
 	)
 	require.Equal(t, backend.ErrorSourceDownstream, resp.Responses["A"].ErrorSource)
 }
+
+func TestParseURL(t *testing.T) {
+	tests := []struct {
+		name     string
+		input    string
+		expected string
+		hasError bool
+	}{
+		{
+			name:     "empty URL",
+			input:    "",
+			expected: "",
+			hasError: true,
+		},
+		{
+			name:     "URL without scheme",
+			input:    "example.com",
+			expected: "example.com",
+			hasError: false,
+		},
+		{
+			name:     "URL without scheme and with port",
+			input:    "example.com:8181",
+			expected: "example.com:8181",
+			hasError: false,
+		},
+		{
+			name:     "URL without port",
+			input:    "http://example.com",
+			expected: "example.com:443",
+			hasError: false,
+		},
+		{
+			name:     "URL with http scheme",
+			input:    "http://example.com:8080",
+			expected: "example.com:8080",
+			hasError: false,
+		},
+		{
+			name:     "URL with https scheme",
+			input:    "https://example.com:8443",
+			expected: "example.com:8443",
+			hasError: false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			result, err := ParseURL(tt.input)
+			if tt.hasError {
+				require.Error(t, err)
+			} else {
+				require.NoError(t, err)
+				require.Equal(t, tt.expected, result)
+			}
+		})
+	}
+}

pkg/tsdb/influxdb/influxdb.go

@@ -72,6 +72,12 @@ func newInstanceSettings(httpClientProvider httpclient.Provider) datasource.Inst
 			database = settings.Database
 		}
 
+		proxyClient, err := settings.ProxyClient(ctx)
+		if err != nil {
+			logger.Error("influx proxy creation failed", "error", err)
+			return nil, fmt.Errorf("influx proxy creation failed")
+		}
+
 		model := &models.DatasourceInfo{
 			HTTPClient:    client,
 			URL:           settings.URL,
@@ -85,6 +91,7 @@ func newInstanceSettings(httpClientProvider httpclient.Provider) datasource.Inst
 			InsecureGrpc:  jsonData.InsecureGrpc,
 			Token:         settings.DecryptedSecureJSONData["token"],
 			Timeout:       opts.Timeouts.Timeout,
+			ProxyClient:   proxyClient,
 		}
 		return model, nil
 	}

pkg/tsdb/influxdb/models/datasource_info.go

@@ -3,6 +3,8 @@ package models
 import (
 	"net/http"
 	"time"
+
+	"github.com/grafana/grafana-plugin-sdk-go/backend/proxy"
 )
 
 type DatasourceInfo struct {
@@ -22,4 +24,6 @@ type DatasourceInfo struct {
 
 	// FlightSQL grpc connection
 	InsecureGrpc bool `json:"insecureGrpc"`
+
+	ProxyClient proxy.Client
 }