JWT Auth

SpencerTorres · SpencerTorres · commit c24ae5761050 · 2025-04-11T21:14:19.000-04:00
diff --git a/clickhouse.go b/clickhouse.go
@@ -375,3 +375,8 @@ func (ch *clickhouse) Close() error {
 		}
 	}
 }
+
+func (ch *clickhouse) UpdateJWT(jwt string) error {
+	ch.opt.Auth.JWT = jwt
+	return nil
+}
diff --git a/clickhouse_options.go b/clickhouse_options.go
@@ -77,8 +77,12 @@ var compressionMap = map[string]CompressionMethod{
 
 type Auth struct { // has_control_character
 	Database string
+
 	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 {
diff --git a/clickhouse_std.go b/clickhouse_std.go
@@ -54,7 +54,10 @@ func (o *stdConnOpener) Driver() driver.Driver {
 			debugf = log.New(os.Stdout, "[clickhouse-std] ", 0).Printf
 		}
 	}
-	return &stdDriver{debugf: debugf}
+	return &stdDriver{
+		opt:    o.opt,
+		debugf: debugf,
+	}
 }
 
 func (o *stdConnOpener) Connect(ctx context.Context) (_ driver.Conn, err error) {
@@ -201,6 +204,7 @@ type stdConnect interface {
 }
 
 type stdDriver struct {
+	opt    *Options
 	conn   stdConnect
 	commit func() error
 	debugf func(format string, v ...any)
@@ -382,6 +386,11 @@ 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)
diff --git a/conn.go b/conn.go
@@ -110,7 +110,7 @@ func dial(ctx context.Context, addr string, num int, opt *Options) (*connect, er
 		}
 	)
 
-	if err := connect.handshake(opt.Auth.Database, opt.Auth.Username, opt.Auth.Password); err != nil {
+	if err := connect.handshake(opt.Auth); err != nil {
 		return nil, err
 	}
 
diff --git a/conn_handshake.go b/conn_handshake.go
@@ -25,7 +25,11 @@ import (
 	"github.com/ClickHouse/clickhouse-go/v2/lib/proto"
 )
 
-func (c *connect) handshake(database, username, password string) error {
+// 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{})
 	// set a read deadline - alternative to context.Read operation will fail if no data is received after deadline.
@@ -43,9 +47,15 @@ func (c *connect) handshake(database, username, password string) error {
 		}
 		handshake.Encode(c.buffer)
 		{
-			c.buffer.PutString(database)
-			c.buffer.PutString(username)
-			c.buffer.PutString(password)
+			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)
+			}
 		}
 		if err := c.flush(); err != nil {
 			return err
diff --git a/conn_http.go b/conn_http.go
@@ -119,6 +119,34 @@ 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)
+	} 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 {
+			req.Header.Set("X-ClickHouse-Key", opt.Auth.Password)
+			req.Header.Set("X-ClickHouse-SSL-Certificate-Auth", "off")
+		} else {
+			req.Header.Set("X-ClickHouse-SSL-Certificate-Auth", "on")
+		}
+	} else if opt.TLS == nil && len(opt.Auth.Username) > 0 {
+		if len(opt.Auth.Password) > 0 {
+			req.URL.User = url.UserPassword(opt.Auth.Username, opt.Auth.Password)
+
+		} else {
+			req.URL.User = url.User(opt.Auth.Username)
+		}
+	}
+
+	req.Header.Set("User-Agent", opt.ClientInfo.String())
+
+	for k, v := range opt.HttpHeaders {
+		req.Header.Set(k, v)
+	}
+}
+
 func dialHttp(ctx context.Context, addr string, num int, opt *Options) (*httpConnect, error) {
 	var debugf = func(format string, v ...any) {}
 	if opt.Debug {
@@ -151,29 +179,6 @@ func dialHttp(ctx context.Context, addr string, num int, opt *Options) (*httpCon
 		Path:   opt.HttpUrlPath,
 	}
 
-	headers := make(map[string]string)
-	for k, v := range opt.HttpHeaders {
-		headers[k] = v
-	}
-
-	if opt.TLS == nil && len(opt.Auth.Username) > 0 {
-		if len(opt.Auth.Password) > 0 {
-			u.User = url.UserPassword(opt.Auth.Username, opt.Auth.Password)
-		} else {
-			u.User = url.User(opt.Auth.Username)
-		}
-	} else if opt.TLS != nil && len(opt.Auth.Username) > 0 {
-		headers["X-ClickHouse-User"] = opt.Auth.Username
-		if len(opt.Auth.Password) > 0 {
-			headers["X-ClickHouse-Key"] = opt.Auth.Password
-			headers["X-ClickHouse-SSL-Certificate-Auth"] = "off"
-		} else {
-			headers["X-ClickHouse-SSL-Certificate-Auth"] = "on"
-		}
-	}
-
-	headers["User-Agent"] = opt.ClientInfo.String()
-
 	query := u.Query()
 	if len(opt.Auth.Database) > 0 {
 		query.Set("database", opt.Auth.Database)
@@ -225,6 +230,7 @@ func dialHttp(ctx context.Context, addr string, num int, opt *Options) (*httpCon
 	}
 
 	conn := &httpConnect{
+		opt: opt,
 		client: &http.Client{
 			Transport: t,
 		},
@@ -234,7 +240,6 @@ func dialHttp(ctx context.Context, addr string, num int, opt *Options) (*httpCon
 		blockCompressor: compress.NewWriter(compress.Level(opt.Compression.Level), compress.Method(opt.Compression.Method)),
 		compressionPool: compressionPool,
 		blockBufferSize: opt.BlockBufferSize,
-		headers:         headers,
 	}
 	location, err := conn.readTimeZone(ctx)
 	if err != nil {
@@ -251,6 +256,7 @@ func dialHttp(ctx context.Context, addr string, num int, opt *Options) (*httpCon
 	}
 
 	return &httpConnect{
+		opt: opt,
 		client: &http.Client{
 			Transport: t,
 		},
@@ -261,11 +267,11 @@ func dialHttp(ctx context.Context, addr string, num int, opt *Options) (*httpCon
 		compressionPool: compressionPool,
 		location:        location,
 		blockBufferSize: opt.BlockBufferSize,
-		headers:         headers,
 	}, nil
 }
 
 type httpConnect struct {
+	opt             *Options
 	url             *url.URL
 	client          *http.Client
 	location        *time.Location
@@ -274,7 +280,6 @@ type httpConnect struct {
 	blockCompressor *compress.Writer
 	compressionPool Pool[HTTPReaderWriter]
 	blockBufferSize uint8
-	headers         map[string]string
 }
 
 func (h *httpConnect) isBad() bool {
@@ -456,9 +461,12 @@ func (h *httpConnect) createRequest(ctx context.Context, requestUrl string, read
 	if err != nil {
 		return nil, err
 	}
+
+	applyOptionsToRequest(req, h.opt)
 	for k, v := range headers {
 		req.Header.Add(k, v)
 	}
+
 	var query url.Values
 	if options != nil {
 		query = req.URL.Query()
diff --git a/conn_http_async_insert.go b/conn_http_async_insert.go
@@ -38,7 +38,7 @@ func (h *httpConnect) asyncInsert(ctx context.Context, query string, wait bool,
 		}
 	}
 
-	res, err := h.sendQuery(ctx, query, &options, h.headers)
+	res, err := h.sendQuery(ctx, query, &options, nil)
 	if res != nil {
 		defer res.Body.Close()
 		// we don't care about result, so just discard it to reuse connection
diff --git a/conn_http_batch.go b/conn_http_batch.go
@@ -207,9 +207,7 @@ func (b *httpBatch) Send() (err error) {
 
 	options.settings["query"] = b.query
 	headers["Content-Type"] = "application/octet-stream"
-	for k, v := range b.conn.headers {
-		headers[k] = v
-	}
+
 	res, err := b.conn.sendStreamQuery(b.ctx, r, &options, headers)
 
 	if res != nil {
diff --git a/conn_http_exec.go b/conn_http_exec.go
@@ -29,7 +29,7 @@ func (h *httpConnect) exec(ctx context.Context, query string, args ...any) error
 		return err
 	}
 
-	res, err := h.sendQuery(ctx, query, &options, h.headers)
+	res, err := h.sendQuery(ctx, query, &options, nil)
 	if res != nil {
 		defer res.Body.Close()
 		// we don't care about result, so just discard it to reuse connection
diff --git a/conn_http_query.go b/conn_http_query.go
@@ -42,10 +42,6 @@ func (h *httpConnect) query(ctx context.Context, release func(*connect, error),
 		headers["Accept-Encoding"] = h.compression.String()
 	}
 
-	for k, v := range h.headers {
-		headers[k] = v
-	}
-
 	res, err := h.sendQuery(ctx, query, &options, headers)
 	if err != nil {
 		return nil, err
diff --git a/jwt.go b/jwt.go
@@ -0,0 +1,46 @@
+// Licensed to ClickHouse, Inc. under one or more contributor
+// license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright
+// ownership. ClickHouse, Inc. licenses this file to you under
+// the Apache License, Version 2.0 (the "License"); you may
+// not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package clickhouse
+
+import (
+	"database/sql"
+	"fmt"
+)
+
+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
+	}
+
+	chDriver, ok := db.Driver().(jwtUpdater)
+	if !ok {
+		return fmt.Errorf("failed to update JWT: db instance must be ClickHouse")
+	}
+
+	return chDriver.UpdateJWT(jwt)
+}
diff --git a/lib/driver/driver.go b/lib/driver/driver.go
@@ -61,6 +61,11 @@ 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
diff --git a/tests/conn_test.go b/tests/conn_test.go
@@ -455,3 +455,24 @@ func TestFreeBufOnConnRelease(t *testing.T) {
 	err = conn.Exec(context.Background(), "DROP TABLE TestFreeBufOnConnRelease")
 	require.NoError(t, err)
 }
+
+func TestNativeJWTAuth(t *testing.T) {
+	SkipNotCloud(t)
+
+	conn, err := GetJWTConnection(testSet, nil, &tls.Config{}, 1000*time.Millisecond)
+	require.NoError(t, err)
+
+	// Token works
+	require.NoError(t, conn.Ping(context.Background()))
+
+	// Wait for connection to timeout
+	time.Sleep(1500 * time.Millisecond)
+
+	// Break the token
+	require.NoError(t, conn.UpdateJWT("broken_jwt"))
+
+	// Next ping should fail
+	require.Error(t, conn.Ping(context.Background()))
+
+	require.NoError(t, conn.Close())
+}
diff --git a/tests/package.go b/tests/package.go
@@ -23,6 +23,7 @@ func init() {
 	}
 }
 
+// SkipOnCloud skips the test if it's run on ClickHouse Cloud
 func SkipOnCloud(t *testing.T, reasons ...string) {
 	if CloudClickHouse {
 		t.Skip(append(
@@ -31,3 +32,13 @@ func SkipOnCloud(t *testing.T, reasons ...string) {
 		))
 	}
 }
+
+// SkipNotCloud skips the test if it's not run on ClickHouse Cloud
+func SkipNotCloud(t *testing.T, reasons ...string) {
+	if !CloudClickHouse {
+		t.Skip(append(
+			[]string{"Skipping test for non-cloud ClickHouse"},
+			reasons...,
+		))
+	}
+}
diff --git a/tests/std/conn_test.go b/tests/std/conn_test.go
@@ -457,3 +457,33 @@ func TestCustomSettings(t *testing.T) {
 		})
 	}
 }
+
+func TestStdJWTAuth(t *testing.T) {
+	clickhouse_tests.SkipNotCloud(t)
+
+	protocols := map[string]clickhouse.Protocol{"Native": clickhouse.Native, "Http": clickhouse.HTTP}
+	for name, protocol := range protocols {
+		t.Run(fmt.Sprintf("%s Protocol", name), func(t *testing.T) {
+			conn, err := GetOpenDBConnectionJWT(testSet, protocol, nil, &tls.Config{})
+			require.NoError(t, err)
+			conn.SetMaxOpenConns(1)
+			conn.SetConnMaxLifetime(1000 * time.Millisecond)
+			conn.SetConnMaxIdleTime(1000 * time.Millisecond)
+			conn.SetMaxIdleConns(1)
+
+			// Token works
+			require.NoError(t, conn.PingContext(context.Background()))
+
+			// Wait for connection to timeout
+			time.Sleep(1500 * time.Millisecond)
+
+			// Break the token
+			require.NoError(t, clickhouse.UpdateSqlJWT(conn, "broken_jwt"))
+
+			// Next ping should fail
+			require.Error(t, conn.PingContext(context.Background()))
+
+			require.NoError(t, conn.Close())
+		})
+	}
+}
diff --git a/tests/std/utils.go b/tests/std/utils.go
diff --git a/tests/utils.go b/tests/utils.go

Original file line number	Diff line number	Diff line change
`@@ -375,3 +375,8 @@ func (ch *clickhouse) Close() error {`
`375`	`375`	`}`
`376`	`376`	`}`
`377`	`377`	`}`
	`378`	`+`
	`379`	`+func (ch *clickhouse) UpdateJWT(jwt string) error {`
	`380`	`+ ch.opt.Auth.JWT = jwt`
	`381`	`+ return nil`
	`382`	`+}`
Original file line number	Diff line number	Diff line change
`@@ -54,7 +54,10 @@ func (o *stdConnOpener) Driver() driver.Driver {`
`54`	`54`	`debugf = log.New(os.Stdout, "[clickhouse-std] ", 0).Printf`
`55`	`55`	`}`
`56`	`56`	`}`
`57`		`- return &stdDriver{debugf: debugf}`
	`57`	`+ return &stdDriver{`
	`58`	`+ opt: o.opt,`
	`59`	`+ debugf: debugf,`
	`60`	`+ }`
`58`	`61`	`}`
`59`	`62`
`60`	`63`	`func (o *stdConnOpener) Connect(ctx context.Context) (_ driver.Conn, err error) {`
`@@ -201,6 +204,7 @@ type stdConnect interface {`
`201`	`204`	`}`
`202`	`205`
`203`	`206`	`type stdDriver struct {`
	`207`	`+ opt *Options`
`204`	`208`	`conn stdConnect`
`205`	`209`	`commit func() error`
`206`	`210`	`debugf func(format string, v ...any)`
`@@ -382,6 +386,11 @@ func (std *stdDriver) Close() error {`
`382`	`386`	`return err`
`383`	`387`	`}`
`384`	`388`
	`389`	`+func (std *stdDriver) UpdateJWT(jwt string) error {`
	`390`	`+ std.opt.Auth.JWT = jwt`
	`391`	`+ return nil`
	`392`	`+}`
	`393`	`+`
`385`	`394`	`type stdBatch struct {`
`386`	`395`	`batch ldriver.Batch`
`387`	`396`	`debugf func(format string, v ...any)`
Original file line number	Diff line number	Diff line change
`@@ -110,7 +110,7 @@ func dial(ctx context.Context, addr string, num int, opt Options) (connect, er`
`110`	`110`	`}`
`111`	`111`	`)`
`112`	`112`
`113`		`- if err := connect.handshake(opt.Auth.Database, opt.Auth.Username, opt.Auth.Password); err != nil {`
	`113`	`+ if err := connect.handshake(opt.Auth); err != nil {`
`114`	`114`	`return nil, err`
`115`	`115`	`}`
`116`	`116`
Original file line number	Diff line number	Diff line change
`@@ -38,7 +38,7 @@ func (h *httpConnect) asyncInsert(ctx context.Context, query string, wait bool,`
`38`	`38`	`}`
`39`	`39`	`}`
`40`	`40`
`41`		`- res, err := h.sendQuery(ctx, query, &options, h.headers)`
	`41`	`+ res, err := h.sendQuery(ctx, query, &options, nil)`
`42`	`42`	`if res != nil {`
`43`	`43`	`defer res.Body.Close()`
`44`	`44`	`// we don't care about result, so just discard it to reuse connection`
Original file line number	Diff line number	Diff line change
`@@ -29,7 +29,7 @@ func (h *httpConnect) exec(ctx context.Context, query string, args ...any) error`
`29`	`29`	`return err`
`30`	`30`	`}`
`31`	`31`
`32`		`- res, err := h.sendQuery(ctx, query, &options, h.headers)`
	`32`	`+ res, err := h.sendQuery(ctx, query, &options, nil)`
`33`	`33`	`if res != nil {`
`34`	`34`	`defer res.Body.Close()`
`35`	`35`	`// we don't care about result, so just discard it to reuse connection`