Merge branch 'main' into fix-clickhouse-proxy-issue

kokizzu · web-flow · commit b3b56a53d6fe · 2025-04-19T02:39:46.000+08:00
diff --git a/clickhouse_options.go b/clickhouse_options.go
@@ -77,6 +77,7 @@ var compressionMap = map[string]CompressionMethod{
 
 type Auth struct { // has_control_character
 	Database string
+
 	Username string
 	Password string
 }
@@ -156,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
 }
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)
diff --git a/conn.go b/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.Database, opt.Auth.Username, opt.Auth.Password); 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
 	}
 
diff --git a/conn_handshake.go b/conn_handshake.go
@@ -25,7 +25,7 @@ import (
 	"github.com/ClickHouse/clickhouse-go/v2/lib/proto"
 )
 
-func (c *connect) handshake(database, username, password string) error {
+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 +43,9 @@ 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)
+			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,47 @@ 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(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 {
+			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)
+	}
+
+	return nil
+}
+
 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 +192,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 +243,7 @@ func dialHttp(ctx context.Context, addr string, num int, opt *Options) (*httpCon
 	}
 
 	conn := &httpConnect{
+		opt: opt,
 		client: &http.Client{
 			Transport: t,
 		},
@@ -234,7 +253,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 +269,7 @@ func dialHttp(ctx context.Context, addr string, num int, opt *Options) (*httpCon
 	}
 
 	return &httpConnect{
+		opt: opt,
 		client: &http.Client{
 			Transport: t,
 		},
@@ -261,11 +280,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 +293,6 @@ type httpConnect struct {
 	blockCompressor *compress.Writer
 	compressionPool Pool[HTTPReaderWriter]
 	blockBufferSize uint8
-	headers         map[string]string
 }
 
 func (h *httpConnect) isBad() bool {
@@ -456,9 +474,16 @@ func (h *httpConnect) createRequest(ctx context.Context, requestUrl string, read
 	if err != nil {
 		return nil, err
 	}
+
+	err = applyOptionsToRequest(ctx, req, h.opt)
+	if err != nil {
+		return nil, err
+	}
+
 	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/context.go b/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 {
diff --git a/go.mod b/go.mod
@@ -1,6 +1,7 @@
 module github.com/ClickHouse/clickhouse-go/v2
 
-go 1.22.0
+go 1.23.0
+
 toolchain go1.24.1
 
 require (
diff --git a/jwt.go b/jwt.go
@@ -0,0 +1,33 @@
+// 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 (
+	"context"
+)
+
+// 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 "
+
+type GetJWTFunc = func(ctx context.Context) (string, error)
+
+// useJWTAuth returns true if the client should use JWT auth
+func useJWTAuth(opt *Options) bool {
+	return opt.GetJWT != nil
+}
diff --git a/tests/conn_test.go b/tests/conn_test.go
@@ -455,3 +455,39 @@ func TestFreeBufOnConnRelease(t *testing.T) {
 	err = conn.Exec(context.Background(), "DROP TABLE TestFreeBufOnConnRelease")
 	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)
+
+	jwt := GetEnv("CLICKHOUSE_JWT", "")
+	getJWT := func(ctx context.Context) (string, error) {
+		return jwt, nil
+	}
+
+	conn, err := GetJWTConnection(testSet, nil, nil, 1000*time.Millisecond, getJWT)
+	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
+	jwt = "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
diff --git a/tests/std/conn_test.go b/tests/std/conn_test.go
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
`@@ -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)`
Original file line number	Diff line number	Diff line change
`@@ -110,7 +110,18 @@ 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`	`+ auth := opt.Auth`
	`114`	`+ if useJWTAuth(opt) {`
	`115`	`+ jwt, err := opt.GetJWT(ctx)`
	`116`	`+ if err != nil {`
	`117`	`+ return nil, fmt.Errorf("failed to get JWT: %w", err)`
	`118`	`+ }`
	`119`	`+`
	`120`	`+ auth.Username = jwtAuthMarker`
	`121`	`+ auth.Password = jwt`
	`122`	`+ }`
	`123`	`+`
	`124`	`+ if err := connect.handshake(auth); err != nil {`
`114`	`125`	`return nil, err`
`115`	`126`	`}`
`116`	`127`
Original file line number	Diff line number	Diff line change
`@@ -25,7 +25,7 @@ import (`
`25`	`25`	`"github.com/ClickHouse/clickhouse-go/v2/lib/proto"`
`26`	`26`	`)`
`27`	`27`
`28`		`-func (c *connect) handshake(database, username, password string) error {`
	`28`	`+func (c *connect) handshake(auth Auth) error {`
`29`	`29`	`defer c.buffer.Reset()`
`30`	`30`	`c.debugf("[handshake] -> %s", proto.ClientHandshake{})`
`31`	`31`	`// set a read deadline - alternative to context.Read operation will fail if no data is received after deadline.`
`@@ -43,9 +43,9 @@ func (c *connect) handshake(database, username, password string) error {`
`43`	`43`	`}`
`44`	`44`	`handshake.Encode(c.buffer)`
`45`	`45`	`{`
`46`		`- c.buffer.PutString(database)`
`47`		`- c.buffer.PutString(username)`
`48`		`- c.buffer.PutString(password)`
	`46`	`+ c.buffer.PutString(auth.Database)`
	`47`	`+ c.buffer.PutString(auth.Username)`
	`48`	`+ c.buffer.PutString(auth.Password)`
`49`	`49`	`}`
`50`	`50`	`if err := c.flush(); err != nil {`
`51`	`51`	`return err`
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`