diff --git a/libsql/internal/http/driver.go b/libsql/internal/http/driver.go index 059fb3a..9752da5 100644 --- a/libsql/internal/http/driver.go +++ b/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) } diff --git a/libsql/internal/http/hranaV2/hranaV2.go b/libsql/internal/http/hranaV2/hranaV2.go index 5bcd052..89c430b 100644 --- a/libsql/internal/http/hranaV2/hranaV2.go +++ b/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 = "" } } diff --git a/libsql/sql.go b/libsql/sql.go index ec12f6a..67ae20f 100644 --- a/libsql/sql.go +++ b/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)