update JWT interface

Author: SpencerTorres

clickhouse.go

@@ -375,8 +375,3 @@ func (ch *clickhouse) Close() error {
 		}
 	}
 }
-
-func (ch *clickhouse) UpdateJWT(jwt string) error {
-	ch.opt.Auth.JWT = jwt
-	return nil
-}

clickhouse_options.go

@@ -80,9 +80,6 @@ type Auth struct { // has_control_character
 
 	Username string
 	Password string
-
-	// JWT for ClickHouse Cloud. Use this instead of Username and Password if you're using JWT auth.
-	JWT string
 }
 
 type Compression struct {
@@ -160,6 +157,11 @@ type Options struct {
 	// HTTPProxy specifies an HTTP proxy URL to use for requests made by the client.
 	HTTPProxyURL *url.URL
 
+	// GetJWT should return a JWT for authentication with ClickHouse Cloud.
+	// This is called per connection/request, so you may cache the token in your app if needed.
+	// Use this instead of Auth.Username and Auth.Password if you're using JWT auth.
+	GetJWT GetJWTFunc
+
 	scheme      string
 	ReadTimeout time.Duration
 }

clickhouse_std.go

@@ -386,11 +386,6 @@ func (std *stdDriver) Close() error {
 	return err
 }
 
-func (std *stdDriver) UpdateJWT(jwt string) error {
-	std.opt.Auth.JWT = jwt
-	return nil
-}
-
 type stdBatch struct {
 	batch  ldriver.Batch
 	debugf func(format string, v ...any)

conn.go

@@ -110,7 +110,18 @@ func dial(ctx context.Context, addr string, num int, opt *Options) (*connect, er
 		}
 	)
 
-	if err := connect.handshake(opt.Auth); err != nil {
+	auth := opt.Auth
+	if useJWTAuth(opt) {
+		jwt, err := opt.GetJWT(ctx)
+		if err != nil {
+			return nil, fmt.Errorf("failed to get JWT: %w", err)
+		}
+
+		auth.Username = jwtAuthMarker
+		auth.Password = jwt
+	}
+
+	if err := connect.handshake(auth); err != nil {
 		return nil, err
 	}
 

conn_handshake.go

@@ -25,10 +25,6 @@ import (
 	"github.com/ClickHouse/clickhouse-go/v2/lib/proto"
 )
 
-// jwtAuthMarker is the marker for JSON Web Token authentication in ClickHouse Cloud.
-// At the protocol level this is used in place of a username.
-const jwtAuthMarker = " JWT AUTHENTICATION "
-
 func (c *connect) handshake(auth Auth) error {
 	defer c.buffer.Reset()
 	c.debugf("[handshake] -> %s", proto.ClientHandshake{})
@@ -48,14 +44,8 @@ func (c *connect) handshake(auth Auth) error {
 		handshake.Encode(c.buffer)
 		{
 			c.buffer.PutString(auth.Database)
-
-			if auth.JWT != "" {
-				c.buffer.PutString(jwtAuthMarker)
-				c.buffer.PutString(auth.JWT)
-			} else {
-				c.buffer.PutString(auth.Username)
-				c.buffer.PutString(auth.Password)
-			}
+			c.buffer.PutString(auth.Username)
+			c.buffer.PutString(auth.Password)
 		}
 		if err := c.flush(); err != nil {
 			return err

conn_http.go

@@ -120,9 +120,20 @@ func (rw *HTTPReaderWriter) reset(pw *io.PipeWriter) io.WriteCloser {
 }
 
 // applyOptionsToRequest applies the client Options (such as auth, headers, client info) to the given http.Request
-func applyOptionsToRequest(req *http.Request, opt *Options) {
-	if opt.TLS != nil && len(opt.Auth.JWT) > 0 {
-		req.Header.Set("Authorization", "Bearer "+opt.Auth.JWT)
+func applyOptionsToRequest(ctx context.Context, req *http.Request, opt *Options) error {
+	jwt := queryOptionsJWT(ctx)
+	useJWT := jwt != "" || useJWTAuth(opt)
+
+	if opt.TLS != nil && useJWT {
+		if jwt == "" {
+			var err error
+			jwt, err = opt.GetJWT(ctx)
+			if err != nil {
+				return fmt.Errorf("failed to get JWT: %w", err)
+			}
+		}
+
+		req.Header.Set("Authorization", "Bearer "+jwt)
 	} else if opt.TLS != nil && len(opt.Auth.Username) > 0 {
 		req.Header.Set("X-ClickHouse-User", opt.Auth.Username)
 		if len(opt.Auth.Password) > 0 {
@@ -145,6 +156,8 @@ func applyOptionsToRequest(req *http.Request, opt *Options) {
 	for k, v := range opt.HttpHeaders {
 		req.Header.Set(k, v)
 	}
+
+	return nil
 }
 
 func dialHttp(ctx context.Context, addr string, num int, opt *Options) (*httpConnect, error) {
@@ -462,7 +475,11 @@ func (h *httpConnect) createRequest(ctx context.Context, requestUrl string, read
 		return nil, err
 	}
 
-	applyOptionsToRequest(req, h.opt)
+	err = applyOptionsToRequest(ctx, req, h.opt)
+	if err != nil {
+		return nil, err
+	}
+
 	for k, v := range headers {
 		req.Header.Add(k, v)
 	}

context.go

@@ -53,6 +53,7 @@ type (
 		async    AsyncOptions
 		queryID  string
 		quotaKey string
+		jwt      string
 		events   struct {
 			logs          func(*Log)
 			progress      func(*Progress)
@@ -95,6 +96,15 @@ func WithQuotaKey(quotaKey string) QueryOption {
 	}
 }
 
+// WithJWT overrides the existing authentication with the given JWT.
+// This only applies for clients connected with HTTPS to ClickHouse Cloud.
+func WithJWT(jwt string) QueryOption {
+	return func(o *QueryOptions) error {
+		o.jwt = jwt
+		return nil
+	}
+}
+
 func WithSettings(settings Settings) QueryOption {
 	return func(o *QueryOptions) error {
 		o.settings = settings
@@ -211,6 +221,16 @@ func queryOptions(ctx context.Context) QueryOptions {
 	return opt
 }
 
+// queryOptionsJWT returns the JWT within the given context's QueryOptions.
+// Empty string if not present.
+func queryOptionsJWT(ctx context.Context) string {
+	if opt, ok := ctx.Value(_contextOptionKey).(QueryOptions); ok {
+		return opt.jwt
+	}
+
+	return ""
+}
+
 // queryOptionsAsync returns the AsyncOptions struct within the given context's QueryOptions.
 func queryOptionsAsync(ctx context.Context) AsyncOptions {
 	if opt, ok := ctx.Value(_contextOptionKey).(QueryOptions); ok {

jwt.go

@@ -18,29 +18,16 @@
 package clickhouse
 
 import (
-	"database/sql"
-	"fmt"
+	"context"
 )
 
-type jwtUpdater interface {
-	UpdateJWT(jwt string) error
-}
-
-// UpdateSqlJWT is a helper function that updates the JWT within the given sql.DB instance, useful for
-// updating expired tokens.
-// For the Native interface, the JWT is only updated for new connections.
-// For the HTTP interface, the JWT is updated immediately for subsequent requests.
-// Existing Native connections are unaffected, but may be forcibly closed by the server upon token expiry.
-// For a completely fresh set of connections you should open a new instance.
-func UpdateSqlJWT(db *sql.DB, jwt string) error {
-	if db == nil {
-		return nil
-	}
+// jwtAuthMarker is the marker for JSON Web Token authentication in ClickHouse Cloud.
+// At the protocol level this is used in place of a username.
+const jwtAuthMarker = " JWT AUTHENTICATION "
 
-	chDriver, ok := db.Driver().(jwtUpdater)
-	if !ok {
-		return fmt.Errorf("failed to update JWT: db instance must be ClickHouse")
-	}
+type GetJWTFunc = func(ctx context.Context) (string, error)
 
-	return chDriver.UpdateJWT(jwt)
+// useJWTAuth returns true if the client should use JWT auth
+func useJWTAuth(opt *Options) bool {
+	return opt.GetJWT != nil
 }

lib/driver/driver.go

@@ -61,11 +61,6 @@ type (
 		Ping(context.Context) error
 		Stats() Stats
 		Close() error
-
-		// UpdateJWT updates the JWT used for new connections, useful for swapping expired tokens.
-		// Existing connections are unaffected, but may be forcibly closed by the server upon token expiry.
-		// For a completely fresh set of connections you should create a new Conn instance.
-		UpdateJWT(jwt string) error
 	}
 	Row interface {
 		Err() error

tests/conn_test.go

@@ -456,10 +456,25 @@ func TestFreeBufOnConnRelease(t *testing.T) {
 	require.NoError(t, err)
 }
 
+func TestJWTError(t *testing.T) {
+	getJWT := func(ctx context.Context) (string, error) {
+		return "", fmt.Errorf("test error")
+	}
+
+	conn, err := GetJWTConnection(testSet, nil, nil, 1000*time.Millisecond, getJWT)
+	require.NoError(t, err)
+	require.ErrorContains(t, conn.Ping(context.Background()), "test error")
+}
+
 func TestNativeJWTAuth(t *testing.T) {
 	SkipNotCloud(t)
 
-	conn, err := GetJWTConnection(testSet, nil, &tls.Config{}, 1000*time.Millisecond)
+	jwt := GetEnv("CLICKHOUSE_JWT", "")
+	getJWT := func(ctx context.Context) (string, error) {
+		return jwt, nil
+	}
+
+	conn, err := GetJWTConnection(testSet, nil, &tls.Config{}, 1000*time.Millisecond, getJWT)
 	require.NoError(t, err)
 
 	// Token works
@@ -469,7 +484,7 @@ func TestNativeJWTAuth(t *testing.T) {
 	time.Sleep(1500 * time.Millisecond)
 
 	// Break the token
-	require.NoError(t, conn.UpdateJWT("broken_jwt"))
+	jwt = "broken_jwt"
 
 	// Next ping should fail
 	require.Error(t, conn.Ping(context.Background()))