This is an automated email from the ASF dual-hosted git repository.
laskoviymishka 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 a15bea708 feat(rest): support isolated IO registries (#2366)
a15bea708 is described below
commit a15bea70891abc79a3eea7bdf7db5ec9f5acc97f
Author: Minh Vu <[email protected]>
AuthorDate: Wed Nov 4 18:39:10 2026 +0100
feat(rest): support isolated IO registries (#2366)
---
catalog/rest/io_registry_test.go | 218 +++++++++++++++++++++++++++++++++++++
catalog/rest/options.go | 8 ++
catalog/rest/rest.go | 9 +-
catalog/rest/scan_planning.go | 5 +-
catalog/rest/scan_planning_test.go | 12 +-
catalog/rest/vended_creds.go | 28 ++++-
6 files changed, 267 insertions(+), 13 deletions(-)
diff --git a/catalog/rest/io_registry_test.go b/catalog/rest/io_registry_test.go
new file mode 100644
index 000000000..7390fde95
--- /dev/null
+++ b/catalog/rest/io_registry_test.go
@@ -0,0 +1,218 @@
+// 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
+
+import (
+ "context"
+ "fmt"
+ "net/http"
+ "net/http/httptest"
+ "net/url"
+ "testing"
+
+ "github.com/apache/iceberg-go"
+ iceio "github.com/apache/iceberg-go/io"
+ "github.com/stretchr/testify/require"
+ "golang.org/x/sync/semaphore"
+)
+
+func TestIsolatedIORegistryForVendedCredentials(t *testing.T) {
+ registry := &iceio.Registry{}
+ fs := iceio.NewMemFS()
+ var gotProps map[string]string
+ require.NoError(t, registry.Register("isolated", func(_
context.Context, _ *url.URL, props map[string]string) (iceio.IO, error) {
+ gotProps = props
+
+ return fs, nil
+ }))
+ r := &vendedCredentialRefresher{
+ mu: semaphore.NewWeighted(1), ioRegistry: registry,
+ location: "isolated://bucket/metadata", props:
iceberg.Properties{"base": "value", "token": "vended"},
+ vended: iceberg.Properties{"token": "vended"},
+ }
+ loaded, err := r.loadFS(t.Context())
+ require.NoError(t, err)
+ require.Same(t, fs, loaded)
+ require.Equal(t, "value", gotProps["base"])
+ require.Equal(t, "vended", gotProps["token"])
+ require.NoError(t, r.close())
+}
+
+func TestIsolatedIORegistryForPlanCredentials(t *testing.T) {
+ registry := &iceio.Registry{}
+ fs := iceio.NewMemFS()
+ var gotProps map[string]string
+ require.NoError(t, registry.Register("isolated", func(_
context.Context, _ *url.URL, props map[string]string) (iceio.IO, error) {
+ gotProps = props
+
+ return fs, nil
+ }))
+ p := planIOFromCredentials([]StorageCredential{{Prefix:
"isolated://bucket/", Config: iceberg.Properties{"token": "vended"}}},
"isolated://bucket/metadata", iceberg.Properties{"base": "value"}, registry)
+ require.NotNil(t, p)
+ loaded, err := p.Load(t.Context())
+ require.NoError(t, err)
+ resolver, ok := loaded.(*prefixScopedIO)
+ require.True(t, ok)
+ actual, err := resolver.filesystemFor("isolated://bucket/data")
+ require.NoError(t, err)
+ require.Same(t, fs, actual)
+ require.Equal(t, "value", gotProps["base"])
+ require.Equal(t, "vended", gotProps["token"])
+ require.NoError(t, p.Close())
+}
+
+func TestCatalogPlanFilesUsesIsolatedIORegistry(t *testing.T) {
+ registry := &iceio.Registry{}
+ fs := iceio.NewMemFS()
+ var gotProps map[string]string
+ require.NoError(t, registry.Register("isolated", func(_
context.Context, _ *url.URL, props map[string]string) (iceio.IO, error) {
+ gotProps = props
+
+ return fs, nil
+ }))
+
+ cat := newScanPlanningTestCatalog(t, []endpoint{endpointPlanTableScan},
func(mux *http.ServeMux) {
+ mux.HandleFunc("/v1/namespaces/db/tables/tbl/plan", func(w
http.ResponseWriter, _ *http.Request) {
+ _, _ =
w.Write([]byte(`{"status":"completed","plan-id":"plan-1","storage-credentials":[{"prefix":"isolated://bucket/","config":{"token":"vended"}}]}`))
+ })
+ }, WithIORegistry(registry))
+
+ req := planFilesReq()
+ req.MetadataLocation = "isolated://bucket/metadata.json"
+ req.FileIOProperties = iceberg.Properties{"base": "value", "token":
"table"}
+ result, err := cat.PlanFiles(t.Context(), req)
+ require.NoError(t, err)
+ require.NotNil(t, result.IO)
+ t.Cleanup(func() { require.NoError(t, result.IO.Close()) })
+
+ loaded, err := result.IO.Load(t.Context())
+ require.NoError(t, err)
+ resolver, ok := loaded.(*prefixScopedIO)
+ require.True(t, ok)
+ actual, err := resolver.filesystemFor("isolated://bucket/data.parquet")
+ require.NoError(t, err)
+ require.Same(t, fs, actual)
+ require.Equal(t, "value", gotProps["base"])
+ require.Equal(t, "vended", gotProps["token"])
+}
+
+// A catalog with an explicit empty registry must not silently use the global
+// file backend when a remote plan supplies credentials for that scheme.
+func TestCatalogPlanFilesDoesNotFallBackToGlobalIO(t *testing.T) {
+ cat := newScanPlanningTestCatalog(t, []endpoint{endpointPlanTableScan},
func(mux *http.ServeMux) {
+ mux.HandleFunc("/v1/namespaces/db/tables/tbl/plan", func(w
http.ResponseWriter, _ *http.Request) {
+ _, _ =
w.Write([]byte(`{"status":"completed","plan-id":"plan-1","storage-credentials":[{"prefix":"file:///warehouse/","config":{"token":"vended"}}]}`))
+ })
+ }, WithIORegistry(&iceio.Registry{}))
+
+ req := planFilesReq()
+ req.MetadataLocation = "file:///warehouse/metadata.json"
+ result, err := cat.PlanFiles(t.Context(), req)
+ require.NoError(t, err)
+ require.NotNil(t, result.IO)
+ t.Cleanup(func() { require.NoError(t, result.IO.Close()) })
+
+ loaded, err := result.IO.Load(t.Context())
+ require.NoError(t, err)
+ resolver, ok := loaded.(*prefixScopedIO)
+ require.True(t, ok)
+
+ _, err = resolver.filesystemFor("file:///warehouse/data.parquet")
+ require.ErrorIs(t, err, iceio.ErrIOSchemeNotFound)
+}
+
+func TestCatalogUsesIsolatedIORegistry(t *testing.T) {
+ mux := http.NewServeMux()
+ mux.HandleFunc("/v1/config", serveEmptyConfig)
+ server := httptest.NewServer(mux)
+ defer server.Close()
+
+ registry := &iceio.Registry{}
+ fs := iceio.NewMemFS()
+ require.NoError(t, registry.Register("isolated", func(context.Context,
*url.URL, map[string]string) (iceio.IO, error) { return fs, nil }))
+ cat, err := NewCatalog(t.Context(), "test", server.URL,
WithIORegistry(registry))
+ require.NoError(t, err)
+ defer cat.Close()
+ for _, vended := range []bool{false, true} {
+ t.Run(fmt.Sprintf("vended=%t", vended), func(t *testing.T) {
+ tbl, err := cat.tableFromResponse(t.Context(),
[]string{"ns", "tbl"}, nil, "isolated://bucket/metadata", nil, nil, vended,
nil, nil)
+ require.NoError(t, err)
+ loaded, err := tbl.FS(t.Context())
+ require.NoError(t, err)
+ require.Same(t, fs, loaded)
+ })
+ }
+}
+
+func TestIsolatedIORegistryDoesNotUseGlobalFactories(t *testing.T) {
+ registry := &iceio.Registry{}
+ _, err := loadRegistryFS(t.Context(), registry, nil,
"file:///metadata.json")
+ require.ErrorIs(t, err, iceio.ErrIOSchemeNotFound)
+ _, err = loadRegistryFSFunc(registry, nil,
"file:///metadata.json")(t.Context())
+ require.ErrorIs(t, err, iceio.ErrIOSchemeNotFound)
+
+ fs, err := loadRegistryFS(t.Context(), nil, nil,
"file:///metadata.json")
+ require.NoError(t, err)
+ require.IsType(t, iceio.LocalFS{}, fs)
+ fs, err = loadRegistryFSFunc(nil, nil,
"file:///metadata.json")(t.Context())
+ require.NoError(t, err)
+ require.IsType(t, iceio.LocalFS{}, fs)
+}
+
+func TestCatalogsCanUseDifferentFactoriesForSameScheme(t *testing.T) {
+ mux := http.NewServeMux()
+ mux.HandleFunc("/v1/config", serveEmptyConfig)
+ server := httptest.NewServer(mux)
+ defer server.Close()
+
+ cases := []struct {
+ name string
+ fs *iceio.MemFS
+ }{
+ {name: "first", fs: iceio.NewMemFS()},
+ {name: "second", fs: iceio.NewMemFS()},
+ }
+ catalogs := make([]*Catalog, len(cases))
+ for i, tc := range cases {
+ registry := &iceio.Registry{}
+ fs := tc.fs
+ require.NoError(t, registry.Register("same-scheme",
func(context.Context, *url.URL, map[string]string) (iceio.IO, error) {
+ return fs, nil
+ }))
+
+ cat, err := NewCatalog(t.Context(), tc.name, server.URL,
WithIORegistry(registry))
+ require.NoError(t, err)
+ defer cat.Close()
+ catalogs[i] = cat
+ }
+
+ // Both catalogs must remain alive while the same URI scheme resolves to
+ // different factories, including through vended-credential refreshers.
+ for i, tc := range cases {
+ t.Run(tc.name, func(t *testing.T) {
+ for _, vended := range []bool{false, true} {
+ tbl, err := catalogs[i].tableFromResponse(
+ t.Context(), []string{"ns", "tbl"},
nil, "same-scheme://bucket/metadata", nil, nil, vended, nil, nil)
+ require.NoError(t, err)
+ loaded, err := tbl.FS(t.Context())
+ require.NoError(t, err)
+ require.Same(t, tc.fs, loaded)
+ }
+ })
+ }
+}
diff --git a/catalog/rest/options.go b/catalog/rest/options.go
index 7c13f51c7..a5e321ba5 100644
--- a/catalog/rest/options.go
+++ b/catalog/rest/options.go
@@ -23,10 +23,17 @@ import (
"net/url"
"github.com/apache/iceberg-go"
+ iceio "github.com/apache/iceberg-go/io"
)
type Option func(*options)
+// WithIORegistry selects the registry used to load table and scan FileIOs.
+// A nil registry uses the package-level default.
+func WithIORegistry(registry *iceio.Registry) Option {
+ return func(o *options) { o.ioRegistry = registry }
+}
+
// TransportFactory creates a transport owned by the REST catalog and returns
// the cleanup function to run when the catalog is closed. The TLS config is
// the one configured with WithTLSConfig.
@@ -213,6 +220,7 @@ func WithTransportFactory(factory TransportFactory) Option {
}
type options struct {
+ ioRegistry *iceio.Registry
signer RequestSigner
signerFactory SignerFactory
tlsConfig *tls.Config
diff --git a/catalog/rest/rest.go b/catalog/rest/rest.go
index 304fa03df..d99368246 100644
--- a/catalog/rest/rest.go
+++ b/catalog/rest/rest.go
@@ -907,8 +907,9 @@ func toProps(o *options) iceberg.Properties {
}
type Catalog struct {
- baseURI *url.URL
- cl *http.Client
+ ioRegistry *iceio.Registry
+ baseURI *url.URL
+ cl *http.Client
// closeSession releases transports owned by the catalog.
Caller-provided
// transports are excluded when the session is constructed.
closeSession func()
@@ -1073,6 +1074,7 @@ func setupOAuthManager(r *Catalog, cl *http.Client, opts
*options) (AuthManager,
}
func (r *Catalog) init(ctx context.Context, ops *options, uri string) error {
+ r.ioRegistry = ops.ioRegistry
baseuri, err := url.Parse(uri)
if err != nil {
return err
@@ -1340,6 +1342,7 @@ func (r *Catalog) tableFromResponse(
refresher := &vendedCredentialRefresher{
mu: semaphore.NewWeighted(1),
identifier: identifier,
+ ioRegistry: r.ioRegistry,
location: loc,
props: config,
vended: vended,
@@ -1349,7 +1352,7 @@ func (r *Catalog) tableFromResponse(
}
fsF = refresher.loadFS
} else {
- fsF = iceio.LoadFSFunc(config, loc)
+ fsF = loadRegistryFSFunc(r.ioRegistry, config, loc)
}
// Resolve the reporter from the client-supplied properties only (see
the
diff --git a/catalog/rest/scan_planning.go b/catalog/rest/scan_planning.go
index 920d9ce01..150aad825 100644
--- a/catalog/rest/scan_planning.go
+++ b/catalog/rest/scan_planning.go
@@ -266,7 +266,7 @@ func (r *Catalog) PlanFiles(ctx context.Context, req
table.ScanPlanningRequest)
result := table.ScanPlanningResult{
Tasks: tasks,
- IO: planIOFromCredentials(completed.StorageCredentials,
req.MetadataLocation, r.planIOBaseProps(req)),
+ IO: planIOFromCredentials(completed.StorageCredentials,
req.MetadataLocation, r.planIOBaseProps(req), r.ioRegistry),
}
if result.IO == nil {
// Without plan-scoped credentials, all task materialization is
complete
@@ -471,12 +471,13 @@ func remoteScanTasks(envelopes []ScanTasks, req
table.ScanPlanningRequest) ([]ta
// for each data or delete file it opens, overlaid on the table's base IO props
// so vended keys win on collision. With none vended it returns a nil PlanIO,
// and the scan falls back to the table's own FileIO.
-func planIOFromCredentials(creds []StorageCredential, location string,
baseProps iceberg.Properties) table.PlanIO {
+func planIOFromCredentials(creds []StorageCredential, location string,
baseProps iceberg.Properties, registry *iceio.Registry) table.PlanIO {
if len(creds) == 0 {
return nil
}
return &planScopedIO{refresher: &vendedCredentialRefresher{
+ ioRegistry: registry,
mu: semaphore.NewWeighted(1),
location: location,
props: maps.Clone(baseProps),
diff --git a/catalog/rest/scan_planning_test.go
b/catalog/rest/scan_planning_test.go
index 88f733a10..2254097f8 100644
--- a/catalog/rest/scan_planning_test.go
+++ b/catalog/rest/scan_planning_test.go
@@ -967,7 +967,7 @@ func TestScanPlanningEndpointGating(t *testing.T) {
}
}
-func newScanPlanningTestCatalog(t *testing.T, endpoints []endpoint, register
func(*http.ServeMux)) *Catalog {
+func newScanPlanningTestCatalog(t *testing.T, endpoints []endpoint, register
func(*http.ServeMux), opts ...Option) *Catalog {
t.Helper()
mux := http.NewServeMux()
@@ -991,7 +991,7 @@ func newScanPlanningTestCatalog(t *testing.T, endpoints
[]endpoint, register fun
srv := httptest.NewServer(mux)
t.Cleanup(srv.Close)
- cat, err := NewCatalog(context.Background(), "rest", srv.URL)
+ cat, err := NewCatalog(context.Background(), "rest", srv.URL, opts...)
require.NoError(t, err)
return cat
@@ -1932,7 +1932,7 @@ func TestPlanScopedIOExpiredCredentials(t *testing.T) {
},
}}
- planIO := planIOFromCredentials(creds,
"file:///bucket/db/tbl/metadata/v1.json", nil)
+ planIO := planIOFromCredentials(creds,
"file:///bucket/db/tbl/metadata/v1.json", nil, nil)
require.NotNil(t, planIO)
p, ok := planIO.(*planScopedIO)
@@ -1971,7 +1971,7 @@ func TestPlanScopedIOAlreadyExpiredCredentials(t
*testing.T) {
},
}}
- p, ok := planIOFromCredentials(creds,
"file:///bucket/db/tbl/metadata/v1.json", nil).(*planScopedIO)
+ p, ok := planIOFromCredentials(creds,
"file:///bucket/db/tbl/metadata/v1.json", nil, nil).(*planScopedIO)
require.True(t, ok)
p.refresher.nowFunc = func() time.Time { return now }
@@ -2000,7 +2000,7 @@ func
TestPlanScopedIOIgnoresExpiredCredentialForUnrelatedPrefix(t *testing.T) {
},
}
- p, ok := planIOFromCredentials(creds, "file:///metadata/v1.json",
nil).(*planScopedIO)
+ p, ok := planIOFromCredentials(creds, "file:///metadata/v1.json", nil,
nil).(*planScopedIO)
require.True(t, ok)
p.refresher.nowFunc = func() time.Time { return now }
@@ -2022,7 +2022,7 @@ func TestPlanScopedIOCredentialsWithoutExpiry(t
*testing.T) {
creds := []StorageCredential{{Prefix: "file:///bucket/", Config:
iceberg.Properties{"s3.access-key-id": "vended"}}}
- p, ok := planIOFromCredentials(creds,
"file:///bucket/db/tbl/metadata/v1.json", nil).(*planScopedIO)
+ p, ok := planIOFromCredentials(creds,
"file:///bucket/db/tbl/metadata/v1.json", nil, nil).(*planScopedIO)
require.True(t, ok)
_, err := p.Load(context.Background())
diff --git a/catalog/rest/vended_creds.go b/catalog/rest/vended_creds.go
index d3b4a8229..aecb595e5 100644
--- a/catalog/rest/vended_creds.go
+++ b/catalog/rest/vended_creds.go
@@ -110,6 +110,7 @@ func parseCredentialExpiry(config iceberg.Properties)
(time.Time, bool) {
}
type vendedCredentialRefresher struct {
+ ioRegistry *iceio.Registry
// Use a weighted semaphore with a single unit to use as an exclusive
lock
// but cancellation (via context) is supported. This is important as we
do IO
// while holding this lock and we want to allow others to cancel during
acquisition.
@@ -211,11 +212,12 @@ func (v *vendedCredentialRefresher) loadFS(ctx
context.Context) (iceio.IO, error
if len(v.credentials) > 0 {
prefixIO := newPrefixScopedIO(ioCtx, v.props, v.credentials)
prefixIO.nowFunc = v.now
+ prefixIO.ioRegistry = v.ioRegistry
newIO = prefixIO
// Expiry is enforced by prefixScopedIO against only the
credential
// selected for each object location.
} else {
- loaded, err := iceio.LoadFS(ioCtx, config, v.location)
+ loaded, err := loadRegistryFS(ioCtx, v.ioRegistry, config,
v.location)
if err != nil {
ioCancel()
@@ -340,6 +342,7 @@ func (v *vendedCredentialRefresher) close() error {
// files in different storage prefixes, so resolving credentials once at plan
// creation would be incorrect.
type prefixScopedIO struct {
+ ioRegistry *iceio.Registry
ctx context.Context
baseProps iceberg.Properties
credentials []StorageCredential
@@ -408,7 +411,7 @@ func (p *prefixScopedIO) filesystemFor(name string)
(iceio.IO, error) {
props := p.propertiesForLocation(name)
- fs, err := iceio.LoadFS(p.ctx, props, name)
+ fs, err := loadRegistryFS(p.ctx, p.ioRegistry, props, name)
if err != nil {
return nil, err
}
@@ -534,3 +537,24 @@ func (p *prefixScopedIO) Close() error {
return closeErr
}
+
+// loadRegistryFS preserves the global registry when none is configured.
+func loadRegistryFS(ctx context.Context, registry *iceio.Registry, props
iceberg.Properties, location string) (iceio.IO, error) {
+ if registry == nil {
+ return iceio.LoadFS(ctx, props, location)
+ }
+
+ return registry.LoadFS(ctx, props, location)
+}
+
+// loadRegistryFSFunc preserves the global registry for lazy loads when none
is configured.
+func loadRegistryFSFunc(registry *iceio.Registry, props iceberg.Properties,
location string) func(context.Context) (iceio.IO, error) {
+ return func(ctx context.Context) (iceio.IO, error) {
+ fs, err := loadRegistryFS(ctx, registry, props, location)
+ if err != nil {
+ return nil, fmt.Errorf("failed to load metadata file at
%s: %w", location, err)
+ }
+
+ return fs, nil
+ }
+}
