From 0329cceeebe556202fe0fbb86acb2c9902042eb5 Mon Sep 17 00:00:00 2001
From: Brad Lovering <425979+bradlo@users.noreply.github.com>
Date: Fri, 24 Feb 2023 10:19:52 -0700
Subject: [PATCH] add PreRequestHook to client options so clients can modify
 the request just prior to dispatch, for example to set/modify headers. (#63)

---
 rai/client.go  | 20 ++++++++++++++------
 rai/version.go |  2 +-
 2 files changed, 15 insertions(+), 7 deletions(-)

diff --git a/rai/client.go b/rai/client.go
index 6ac9caa..7c63be1 100644
--- a/rai/client.go
+++ b/rai/client.go
@@ -37,10 +37,13 @@ import (
 
 const userAgent = "raictl/" + Version
 
+type PreRequestHook func(*http.Request) *http.Request
+
 type ClientOptions struct {
 	Config
 	HTTPClient         *http.Client
 	AccessTokenHandler AccessTokenHandler
+	PreRequestHook     PreRequestHook
 }
 
 func NewClientOptions(cfg *Config) *ClientOptions {
@@ -55,6 +58,7 @@ type Client struct {
 	Port               string
 	HttpClient         *http.Client
 	accessTokenHandler AccessTokenHandler
+	preRequestHook     PreRequestHook
 }
 
 const DefaultHost = "azure.relationalai.com"
@@ -86,12 +90,13 @@ func NewClient(ctx context.Context, opts *ClientOptions) *Client {
 		opts.HTTPClient = &http.Client{}
 	}
 	client := &Client{
-		ctx:        ctx,
-		Region:     region,
-		Scheme:     scheme,
-		Host:       host,
-		Port:       port,
-		HttpClient: opts.HTTPClient}
+		ctx:            ctx,
+		Region:         region,
+		Scheme:         scheme,
+		Host:           host,
+		Port:           port,
+		preRequestHook: opts.PreRequestHook,
+		HttpClient:     opts.HTTPClient}
 	if opts.AccessTokenHandler != nil {
 		client.accessTokenHandler = opts.AccessTokenHandler
 	} else if opts.Credentials == nil {
@@ -336,6 +341,9 @@ func isErrorStatus(rsp *http.Response) bool {
 // Execute the given request and return the response or error.
 func (c *Client) Do(req *http.Request) (*http.Response, error) {
 	req = req.WithContext(c.ctx)
+	if c.preRequestHook != nil {
+		req = c.preRequestHook(req)
+	}
 	rsp, err := c.HttpClient.Do(req)
 	if err != nil {
 		return nil, err
diff --git a/rai/version.go b/rai/version.go
index 330ef5e..89777a1 100644
--- a/rai/version.go
+++ b/rai/version.go
@@ -14,4 +14,4 @@
 
 package rai
 
-const Version = "0.4.0-alpha"
+const Version = "0.4.1-alpha"