This is an automated email from the ASF dual-hosted git repository.
zeroshade pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/iceberg-go.git
The following commit(s) were added to refs/heads/main by this push:
new 5c233d439 feat(rest): export NewOauth2AuthManager for custom token
sources (#2433)
5c233d439 is described below
commit 5c233d43981d997696756bd906412d72d2b8b594
Author: Revanth <[email protected]>
AuthorDate: Thu Nov 5 11:18:15 2026 -0600
feat(rest): export NewOauth2AuthManager for custom token sources (#2433)
---
catalog/rest/auth.go | 26 ++++++++
catalog/rest/auth_test.go | 78 ++++++++++++++++++++++
catalog/rest/auth_token_source_test.go | 117 +++++++++++++++++++++++++++++++++
website/src/configuration.md | 27 +++++++-
4 files changed, 247 insertions(+), 1 deletion(-)
diff --git a/catalog/rest/auth.go b/catalog/rest/auth.go
index 4789eb78d..0b50b6c88 100644
--- a/catalog/rest/auth.go
+++ b/catalog/rest/auth.go
@@ -57,6 +57,32 @@ type Oauth2AuthManager struct {
var _ ContextAuthManager = (*Oauth2AuthManager)(nil)
+// NewOauth2AuthManager returns an AuthManager that sends tokens from ts, for
+// identity providers the built-in client-credentials flow can't reach, such as
+// Google Application Default Credentials for BigLake. Pass the result to
+// WithAuthManager.
+//
+// Like oauth2.NewClient, it wraps ts in oauth2.ReuseTokenSource, so ts is
+// called only when the cached token is missing or expired. A token without an
+// expiry is reused indefinitely. The header uses the token's type, and an
+// *oauth2.RetrieveError from ts is reported as ErrOAuthError.
+//
+// The built-in flows bound each refresh with their own HTTP client timeout.
+// A refresh through ts is bounded only by the HTTP client ts uses; for
+// sources built from a context, such as google.DefaultTokenSource, that is
+// the client in the context's oauth2.HTTPClient value, or
+// http.DefaultClient, which has no timeout. AuthHeaderWithContext still
+// returns when its context is done.
+//
+// NewOauth2AuthManager panics if ts is nil.
+func NewOauth2AuthManager(ts oauth2.TokenSource) *Oauth2AuthManager {
+ if ts == nil {
+ panic("rest: NewOauth2AuthManager called with a nil
oauth2.TokenSource")
+ }
+
+ return &Oauth2AuthManager{tokenSource: oauth2.ReuseTokenSource(nil, ts)}
+}
+
// AuthHeader returns the authorization header with the bearer token.
func (o *Oauth2AuthManager) AuthHeader() (string, string, error) {
tok, err, _ := o.refresh.Do("token", o.fetchToken)
diff --git a/catalog/rest/auth_test.go b/catalog/rest/auth_test.go
index 165be30bc..543db6170 100644
--- a/catalog/rest/auth_test.go
+++ b/catalog/rest/auth_test.go
@@ -190,3 +190,81 @@ func
TestOauth2AuthManager_AuthHeaderWithContext_SharesRefresh(t *testing.T) {
assert.Equal(t, "Authorization", key)
assert.Equal(t, "Bearer shared_token", value)
}
+
+// sequenceTokenSource returns its tokens in order, repeating the last one, and
+// counts Token calls.
+type sequenceTokenSource struct {
+ tokens []*oauth2.Token
+ err error
+ calls atomic.Int32
+}
+
+func (s *sequenceTokenSource) Token() (*oauth2.Token, error) {
+ n := int(s.calls.Add(1))
+ if s.err != nil {
+ return nil, s.err
+ }
+
+ return s.tokens[min(n, len(s.tokens))-1], nil
+}
+
+func TestNewOauth2AuthManager_ReusesValidToken(t *testing.T) {
+ // An empty TokenType defaults to Bearer, as oauth2.Token.Type does.
+ source := &sequenceTokenSource{tokens: []*oauth2.Token{
+ {AccessToken: "a", Expiry: time.Now().Add(time.Hour)},
+ }}
+ manager := NewOauth2AuthManager(source)
+
+ for range 3 {
+ key, value, err := manager.AuthHeader()
+ require.NoError(t, err)
+ assert.Equal(t, "Authorization", key)
+ assert.Equal(t, "Bearer a", value)
+ }
+ assert.EqualValues(t, 1, source.calls.Load(), "a valid token must be
reused, not refetched per request")
+}
+
+func TestNewOauth2AuthManager_RefreshesExpiredToken(t *testing.T) {
+ source := &sequenceTokenSource{tokens: []*oauth2.Token{
+ {AccessToken: "a", TokenType: "Bearer", Expiry:
time.Now().Add(-time.Minute)},
+ {AccessToken: "b", TokenType: "Bearer", Expiry:
time.Now().Add(time.Hour)},
+ }}
+ manager := NewOauth2AuthManager(source)
+
+ for _, want := range []string{"Bearer a", "Bearer b", "Bearer b"} {
+ _, value, err :=
manager.AuthHeaderWithContext(context.Background())
+ require.NoError(t, err)
+ assert.Equal(t, want, value)
+ }
+ assert.EqualValues(t, 2, source.calls.Load())
+}
+
+func TestNewOauth2AuthManager_RetrieveErrorIsOAuthError(t *testing.T) {
+ manager := NewOauth2AuthManager(&sequenceTokenSource{err:
&oauth2.RetrieveError{
+ ErrorCode: "invalid_grant",
+ ErrorDescription: "Token has been expired or revoked.",
+ }})
+
+ _, _, err := manager.AuthHeader()
+ require.ErrorIs(t, err, ErrOAuthError)
+ assert.Contains(t, err.Error(), "invalid_grant")
+ assert.Contains(t, err.Error(), "Token has been expired or revoked.")
+}
+
+func TestNewOauth2AuthManager_CanceledContextSkipsSource(t *testing.T) {
+ source := &sequenceTokenSource{tokens: []*oauth2.Token{{AccessToken:
"a"}}}
+ manager := NewOauth2AuthManager(source)
+
+ ctx, cancel := context.WithCancel(context.Background())
+ cancel()
+ _, _, err := manager.AuthHeaderWithContext(ctx)
+
+ require.ErrorIs(t, err, context.Canceled)
+ assert.Zero(t, source.calls.Load(), "a canceled caller must not trigger
a refresh")
+}
+
+func TestNewOauth2AuthManager_NilSourcePanics(t *testing.T) {
+ assert.PanicsWithValue(t, "rest: NewOauth2AuthManager called with a nil
oauth2.TokenSource", func() {
+ NewOauth2AuthManager(nil)
+ })
+}
diff --git a/catalog/rest/auth_token_source_test.go
b/catalog/rest/auth_token_source_test.go
new file mode 100644
index 000000000..a0ce87527
--- /dev/null
+++ b/catalog/rest/auth_token_source_test.go
@@ -0,0 +1,117 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF 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 rest_test
+
+import (
+ "context"
+ "encoding/json"
+ "log"
+ "net/http"
+ "net/http/httptest"
+ "sync"
+ "testing"
+ "time"
+
+ "github.com/apache/iceberg-go/catalog/rest"
+ "github.com/stretchr/testify/require"
+ "golang.org/x/oauth2"
+ "golang.org/x/oauth2/google"
+)
+
+// rotatingTokenSource returns an already expired token first and a valid one
+// after that, standing in for a source whose first token has run out.
+type rotatingTokenSource struct {
+ mu sync.Mutex
+ calls int
+}
+
+func (s *rotatingTokenSource) Token() (*oauth2.Token, error) {
+ s.mu.Lock()
+ defer s.mu.Unlock()
+ s.calls++
+ if s.calls == 1 {
+ return &oauth2.Token{AccessToken: "first", Expiry:
time.Now().Add(-time.Minute)}, nil
+ }
+
+ return &oauth2.Token{AccessToken: "second", Expiry:
time.Now().Add(time.Hour)}, nil
+}
+
+func TestNewOauth2AuthManagerAuthenticatesCatalogRequests(t *testing.T) {
+ var (
+ mu sync.Mutex
+ headers []string
+ )
+ record := func(r *http.Request) {
+ mu.Lock()
+ defer mu.Unlock()
+ headers = append(headers, r.URL.Path+"
"+r.Header.Get("Authorization"))
+ }
+
+ mux := http.NewServeMux()
+ mux.HandleFunc("/v1/config", func(w http.ResponseWriter, r
*http.Request) {
+ record(r)
+ json.NewEncoder(w).Encode(map[string]any{
+ "defaults": map[string]any{},
+ "overrides": map[string]any{},
+ "endpoints": rest.AllEndpointStrings,
+ })
+ })
+ mux.HandleFunc("/v1/namespaces", func(w http.ResponseWriter, r
*http.Request) {
+ record(r)
+ json.NewEncoder(w).Encode(map[string]any{"namespaces": []any{}})
+ })
+ srv := httptest.NewServer(mux)
+ defer srv.Close()
+
+ ctx := context.Background()
+ cat, err := rest.NewCatalog(ctx, "rest", srv.URL,
+
rest.WithAuthManager(rest.NewOauth2AuthManager(&rotatingTokenSource{})))
+ require.NoError(t, err)
+ _, err = cat.ListNamespaces(ctx, nil)
+ require.NoError(t, err)
+
+ // The expired first token is sent once, then replaced on the next
request.
+ require.Equal(t, []string{
+ "/v1/config Bearer first",
+ "/v1/namespaces Bearer second",
+ }, headers)
+}
+
+// ExampleNewOauth2AuthManager connects to Google's BigLake Iceberg REST
+// catalog with Application Default Credentials.
+func ExampleNewOauth2AuthManager() {
+ ctx := context.Background()
+
+ // Application Default Credentials: a service-account key file, gcloud
user
+ // credentials, or workload identity from the metadata server.
+ ts, err := google.DefaultTokenSource(ctx,
"https://www.googleapis.com/auth/cloud-platform";)
+ if err != nil {
+ log.Fatal(err)
+ }
+
+ cat, err := rest.NewCatalog(ctx, "biglake",
"https://biglake.googleapis.com/iceberg/v1/restcatalog";,
+ rest.WithAuthManager(rest.NewOauth2AuthManager(ts)),
+ rest.WithWarehouseLocation("gs://lakehouse-public-data"),
+ // The project BigLake bills the requests to.
+ rest.WithHeaders(map[string]string{"x-goog-user-project":
"my-project"}),
+ )
+ if err != nil {
+ log.Fatal(err)
+ }
+ _ = cat
+}
diff --git a/website/src/configuration.md b/website/src/configuration.md
index bda15dcb0..43b23aeec 100644
--- a/website/src/configuration.md
+++ b/website/src/configuration.md
@@ -70,7 +70,7 @@ The most option-rich surface. Source:
[`catalog/rest/options.go`](https://github
| Group | Options |
|---|---|
-| Authentication | `WithCredential`, `WithOAuthToken`, `WithAuthManager`,
`WithAuthURI`, `WithScope`, `WithAudience`, `WithResource` |
+| Authentication | `WithCredential`, `WithOAuthToken`, `WithAuthManager`,
`NewOauth2AuthManager`, `WithAuthURI`, `WithScope`, `WithAudience`,
`WithResource` |
| AWS SigV4 | `WithSigV4`, `WithSigV4RegionSvc`, `WithSigner` |
| HTTP | `WithHeaders`, `WithTLSConfig`, `WithOAuthTLSConfig`,
`WithCustomTransport` |
| Catalog routing | `WithPrefix`, `WithWarehouseLocation`,
`WithMetadataLocation` |
@@ -99,6 +99,31 @@ The most option-rich surface. Source:
[`catalog/rest/options.go`](https://github
> until the `catalog/rest/sigv4` backend is blank-imported; the error names the
> import to add.
+#### Custom token sources (Google BigLake)
+
+`WithCredential` uses the OAuth2 client-credentials flow against the catalog's
+token endpoint. For other identity providers, wrap any `oauth2.TokenSource`
with
+`NewOauth2AuthManager` and pass it to `WithAuthManager`. Google's BigLake
Iceberg
+REST catalog, for example, takes Application Default Credentials:
+
+```go
+ts, err := google.DefaultTokenSource(ctx,
"https://www.googleapis.com/auth/cloud-platform";)
+if err != nil {
+ return err
+}
+cat, err := rest.NewCatalog(ctx, "biglake",
"https://biglake.googleapis.com/iceberg/v1/restcatalog";,
+ rest.WithAuthManager(rest.NewOauth2AuthManager(ts)),
+ rest.WithWarehouseLocation("gs://lakehouse-public-data"),
+ rest.WithHeaders(map[string]string{"x-goog-user-project":
"my-project"}))
+```
+
+The manager reuses a token until it expires, then asks the source for a new
one.
+That refresh runs on the HTTP client the source was built with. For
+`google.DefaultTokenSource` that is the client in the context's
+`oauth2.HTTPClient` value, or `http.DefaultClient`, which has no timeout.
+Unlike Java's `rest.auth.type=google`, this is configured in code; iceberg-go
+has no `rest.auth.type` property.
+
#### Metrics reporting
The REST catalog can POST scan and commit metrics to the catalog's
