From 5e888bc85e157b40175c91ca748beeb30dd200cd Mon Sep 17 00:00:00 2001
From: Christoph Hartmann <chris@lollyrock.com>
Date: Sat, 11 Nov 2023 17:17:45 +0100
Subject: [PATCH] =?UTF-8?q?=E2=AD=90=EF=B8=8F=20use=20endpoint=20url=20fro?=
 =?UTF-8?q?m=20service=20account?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 internal/signer/signer.go |  8 ++++----
 option/option.go          | 22 +++++++++++++---------
 2 files changed, 17 insertions(+), 13 deletions(-)

diff --git a/internal/signer/signer.go b/internal/signer/signer.go
index a7ce5a7..18480de 100644
--- a/internal/signer/signer.go
+++ b/internal/signer/signer.go
@@ -52,17 +52,17 @@ func privateKeyFromBytes(bytes []byte) (*ecdsa.PrivateKey, error) {
 	}
 }
 
-func NewServiceAccountTokenSource(data []byte) (*serviceAccountTokenSource, error) {
+func NewServiceAccountTokenSource(data []byte) (*serviceAccountTokenSource, *serviceAccountCredentials, error) {
 	var credentials *serviceAccountCredentials
 	err := json.Unmarshal(data, &credentials)
 	if credentials == nil || err != nil {
-		return nil, errors.New("valid service account needs to be provided")
+		return nil, nil, errors.New("valid service account needs to be provided")
 	}
 
 	// verify that we can read the private key
 	privateKey, err := privateKeyFromBytes([]byte(credentials.PrivateKey))
 	if err != nil {
-		return nil, errors.New("cannot load retrieved key: " + err.Error())
+		return nil, nil, errors.New("cannot load retrieved key: " + err.Error())
 	}
 
 	// configure authentication plugin, since the server only accepts authenticated calls
@@ -75,7 +75,7 @@ func NewServiceAccountTokenSource(data []byte) (*serviceAccountTokenSource, erro
 
 	return &serviceAccountTokenSource{
 		cfg: cfg,
-	}, nil
+	}, credentials, nil
 }
 
 type tokenSourceConfig struct {
diff --git a/option/option.go b/option/option.go
index eb9f4d0..74a2308 100644
--- a/option/option.go
+++ b/option/option.go
@@ -45,17 +45,21 @@ func (w withHTTPClient) Apply(o *internal.DialSettings) {
 
 // WithTokenSource returns a ClientOption that specifies the oauth2.TokenSource
 func WithTokenSource(s oauth2.TokenSource) ClientOption {
-	return withTokenSource{s, nil}
+	return withTokenSource{"", s, nil}
 }
 
 type withTokenSource struct {
-	ts  oauth2.TokenSource
-	err error
+	endpoint string
+	ts       oauth2.TokenSource
+	err      error
 }
 
 func (w withTokenSource) Apply(o *internal.DialSettings) {
 	o.TokenSource = w.ts
 	o.TokenError = w.err
+	if w.endpoint != "" {
+		o.Endpoint = w.endpoint + "/query"
+	}
 }
 
 // WithAPIToken returns a ClientOption that specifies the oauth2.TokenSource with the given token.
@@ -63,24 +67,24 @@ func WithAPIToken(token string) ClientOption {
 	src := oauth2.StaticTokenSource(
 		&oauth2.Token{AccessToken: token},
 	)
-	return withTokenSource{src, nil}
+	return withTokenSource{"", src, nil}
 }
 
 // WithServiceAccount returns a ClientOption that specifies the credentials file to use.
 func WithServiceAccount(data []byte) ClientOption {
-	ts, err := signer.NewServiceAccountTokenSource(data)
-	return withTokenSource{ts, err}
+	ts, sa, err := signer.NewServiceAccountTokenSource(data)
+	return withTokenSource{sa.ApiEndpoint, ts, err}
 }
 
 // WithServiceAccountFile returns a ClientOption that specifies the credentials file to use.
 func WithServiceAccountFile(filename string) ClientOption {
 	data, err := os.ReadFile(filename)
 	if err != nil {
-		return withTokenSource{nil, err}
+		return withTokenSource{"", nil, err}
 	}
 
-	ts, err := signer.NewServiceAccountTokenSource(data)
-	return withTokenSource{ts, err}
+	ts, sa, err := signer.NewServiceAccountTokenSource(data)
+	return withTokenSource{sa.ApiEndpoint, ts, err}
 }
 
 // WithoutAuthentication returns a ClientOption that disables authentication.