Merge pull request #144 from tursodatabase/request-headers

allow client to pass arbitrary http headers in the sql-over-http requests

Author: sivukhin

libsql/internal/http/driver.go

@@ -6,6 +6,6 @@ import (
 	"github.com/tursodatabase/libsql-client-go/libsql/internal/http/hranaV2"
 )
 
-func Connect(url, jwt, host string, schemaDb bool, remoteEncryptionKey string) driver.Conn {
-	return hranaV2.Connect(url, jwt, host, schemaDb, remoteEncryptionKey)
+func Connect(url, jwt, host string, schemaDb bool, remoteEncryptionKey string, requestHeaders map[string]string) driver.Conn {
+	return hranaV2.Connect(url, jwt, host, schemaDb, remoteEncryptionKey, requestHeaders)
 }

libsql/internal/http/hranaV2/hranaV2.go

@@ -37,8 +37,8 @@ func init() {
 	commitHash = "unknown"
 }
 
-func Connect(url, jwt, host string, schemaDb bool, encryptionKey string) driver.Conn {
-	return &hranaV2Conn{url, jwt, host, schemaDb, encryptionKey, "", false, 0}
+func Connect(url, jwt, host string, schemaDb bool, encryptionKey string, requestHeaders map[string]string) driver.Conn {
+	return &hranaV2Conn{url, jwt, host, schemaDb, encryptionKey, requestHeaders, "", false, 0}
 }
 
 type hranaV2Stmt struct {
@@ -88,6 +88,7 @@ type hranaV2Conn struct {
 	host                string
 	schemaDb            bool
 	remoteEncryptionKey string
+	requestHeaders      map[string]string
 	baton               string
 	streamClosed        bool
 	replicationIndex    uint64
@@ -123,11 +124,11 @@ func (h *hranaV2Conn) PrepareContext(ctx context.Context, query string) (driver.
 
 func (h *hranaV2Conn) Close() error {
 	if h.baton != "" {
-		go func(baton, url, jwt, host, encryptionKey string) {
+		go func(baton, url, jwt, host, encryptionKey string, requestHeaders map[string]string) {
 			msg := hrana.PipelineRequest{Baton: baton}
 			msg.Add(hrana.CloseStream())
-			_, _, _ = sendPipelineRequest(context.Background(), &msg, url, jwt, host, encryptionKey)
-		}(h.baton, h.url, h.jwt, h.host, h.remoteEncryptionKey)
+			_, _, _ = sendPipelineRequest(context.Background(), &msg, url, jwt, host, encryptionKey, requestHeaders)
+		}(h.baton, h.url, h.jwt, h.host, h.remoteEncryptionKey, h.requestHeaders)
 	}
 	return nil
 }
@@ -175,7 +176,7 @@ func (h *hranaV2Conn) sendPipelineRequest(ctx context.Context, msg *hrana.Pipeli
 	if h.replicationIndex > 0 {
 		addReplicationIndex(msg, h.replicationIndex)
 	}
-	result, streamClosed, err := sendPipelineRequest(ctx, msg, h.url, h.jwt, h.host, h.remoteEncryptionKey)
+	result, streamClosed, err := sendPipelineRequest(ctx, msg, h.url, h.jwt, h.host, h.remoteEncryptionKey, h.requestHeaders)
 	if streamClosed {
 		h.streamClosed = true
 	}
@@ -232,7 +233,7 @@ func getReplicationIndex(response *hrana.PipelineResponse) uint64 {
 	return replicationIndex
 }
 
-func sendPipelineRequest(ctx context.Context, msg *hrana.PipelineRequest, url string, jwt string, host string, remoteEncryptionKey string) (result hrana.PipelineResponse, streamClosed bool, err error) {
+func sendPipelineRequest(ctx context.Context, msg *hrana.PipelineRequest, url string, jwt string, host string, remoteEncryptionKey string, requestHeaders map[string]string) (result hrana.PipelineResponse, streamClosed bool, err error) {
 	reqBody, err := json.Marshal(msg)
 	if err != nil {
 		return hrana.PipelineResponse{}, false, err
@@ -254,6 +255,9 @@ func sendPipelineRequest(ctx context.Context, msg *hrana.PipelineRequest, url st
 	}
 
 	req.Host = host
+	for name, value := range requestHeaders {
+		req.Header.Set(name, value)
+	}
 	resp, err := http.DefaultClient.Do(req)
 	if err != nil {
 		return hrana.PipelineResponse{}, false, err
@@ -597,11 +601,11 @@ func (h *hranaV2Conn) QueryContext(ctx context.Context, query string, args []dri
 
 func (h *hranaV2Conn) closeStream() {
 	if h.baton != "" {
-		go func(baton, url, jwt, host, encryptionKey string) {
+		go func(baton, url, jwt, host, encryptionKey string, requestHeaders map[string]string) {
 			msg := hrana.PipelineRequest{Baton: baton}
 			msg.Add(hrana.CloseStream())
-			_, _, _ = sendPipelineRequest(context.Background(), &msg, url, jwt, host, encryptionKey)
-		}(h.baton, h.url, h.jwt, h.host, h.remoteEncryptionKey)
+			_, _, _ = sendPipelineRequest(context.Background(), &msg, url, jwt, host, encryptionKey, requestHeaders)
+		}(h.baton, h.url, h.jwt, h.host, h.remoteEncryptionKey, h.requestHeaders)
 		h.baton = ""
 	}
 }

libsql/sql.go

@@ -19,6 +19,7 @@ type config struct {
 	proxy               *string
 	schemaDb            *bool
 	remoteEncryptionKey *string
+	requestHeaders      map[string]string
 }
 
 type Option interface {
@@ -90,6 +91,26 @@ func WithRemoteEncryptionKey(key string) Option {
 	})
 }
 
+// WithRequestHeaders attaches arbitrary HTTP headers to every request the
+// driver sends to the remote server. Passing the `Host` key (case-insensitive)
+// has no effect.
+func WithRequestHeaders(headers map[string]string) Option {
+	return option(func(o *config) error {
+		if o.requestHeaders != nil {
+			return fmt.Errorf("requestHeaders already set")
+		}
+		if len(headers) == 0 {
+			return fmt.Errorf("requestHeaders must not be empty")
+		}
+		copied := make(map[string]string, len(headers))
+		for k, v := range headers {
+			copied[k] = v
+		}
+		o.requestHeaders = copied
+		return nil
+	})
+}
+
 func (c config) connector(dbPath string) (driver.Connector, error) {
 	u, err := url.Parse(dbPath)
 	if err != nil {
@@ -182,7 +203,7 @@ func (c config) connector(dbPath string) (driver.Connector, error) {
 		return wsConnector{url: u.String(), authToken: authToken}, nil
 	}
 	if u.Scheme == "https" || u.Scheme == "http" {
-		return httpConnector{url: u.String(), authToken: authToken, host: host, schemaDb: schemaDb, remoteEncryptionKey: encryptionKey}, nil
+		return httpConnector{url: u.String(), authToken: authToken, host: host, schemaDb: schemaDb, remoteEncryptionKey: encryptionKey, requestHeaders: c.requestHeaders}, nil
 	}
 
 	return nil, fmt.Errorf("unsupported URL scheme: %s\nThis driver supports only URLs that start with libsql://, file://, https://, http://, wss:// and ws://", u.Scheme)
@@ -208,10 +229,11 @@ type httpConnector struct {
 	host                string
 	schemaDb            bool
 	remoteEncryptionKey string
+	requestHeaders      map[string]string
 }
 
 func (c httpConnector) Connect(_ctx context.Context) (driver.Conn, error) {
-	return http.Connect(c.url, c.authToken, c.host, c.schemaDb, c.remoteEncryptionKey), nil
+	return http.Connect(c.url, c.authToken, c.host, c.schemaDb, c.remoteEncryptionKey, c.requestHeaders), nil
 }
 
 func (c httpConnector) Driver() driver.Driver {
@@ -360,7 +382,7 @@ func (d Driver) Open(dbUrl string) (driver.Conn, error) {
 		return ws.Connect(u.String(), jwt)
 	}
 	if u.Scheme == "https" || u.Scheme == "http" {
-		return http.Connect(u.String(), jwt, u.Host, false, ""), nil
+		return http.Connect(u.String(), jwt, u.Host, false, "", nil), nil
 	}
 
 	return nil, fmt.Errorf("unsupported URL scheme: %s\nThis driver supports only URLs that start with libsql://, file://, https://, http://, wss:// and ws://", u.Scheme)