diff --git a/core/src/main/java/org/apache/iceberg/rest/RESTCatalog.java b/core/src/main/java/org/apache/iceberg/rest/RESTCatalog.java index 0176128ed576..47f96dcb16b0 100644 --- a/core/src/main/java/org/apache/iceberg/rest/RESTCatalog.java +++ b/core/src/main/java/org/apache/iceberg/rest/RESTCatalog.java @@ -69,13 +69,28 @@ public RESTCatalog(Function, RESTClient> clientBuilder) { public RESTCatalog( SessionCatalog.SessionContext context, Function, RESTClient> clientBuilder) { - this.sessionCatalog = new RESTSessionCatalog(clientBuilder, null); + this.sessionCatalog = newSessionCatalog(clientBuilder); this.delegate = sessionCatalog.asCatalog(context); this.nsDelegate = (SupportsNamespaces) delegate; this.context = context; this.viewSessionCatalog = sessionCatalog.asViewCatalog(context); } + /** + * Create a new {@link RESTSessionCatalog} instance. + * + *

This method can be overridden in subclasses to provide custom session catalog + * implementations, which in turn can provide custom table and view operations by overriding the + * protected methods in {@link RESTSessionCatalog}. + * + * @param clientBuilder a function to build REST clients + * @return a new RESTSessionCatalog instance + */ + protected RESTSessionCatalog newSessionCatalog( + Function, RESTClient> clientBuilder) { + return new RESTSessionCatalog(clientBuilder, null); + } + @Override public void initialize(String name, Map props) { Preconditions.checkArgument(props != null, "Invalid configuration: null"); diff --git a/core/src/main/java/org/apache/iceberg/rest/RESTSessionCatalog.java b/core/src/main/java/org/apache/iceberg/rest/RESTSessionCatalog.java index b903f13adc09..85b04f3868cd 100644 --- a/core/src/main/java/org/apache/iceberg/rest/RESTSessionCatalog.java +++ b/core/src/main/java/org/apache/iceberg/rest/RESTSessionCatalog.java @@ -27,6 +27,7 @@ import java.util.Set; import java.util.function.BiFunction; import java.util.function.Function; +import java.util.function.Supplier; import java.util.stream.Collectors; import org.apache.iceberg.BaseTable; import org.apache.iceberg.CatalogProperties; @@ -450,7 +451,7 @@ public Table loadTable(SessionContext context, TableIdentifier identifier) { RESTClient tableClient = client.withAuthSession(tableSession); RESTTableOperations ops = - new RESTTableOperations( + newTableOps( tableClient, paths.table(finalIdentifier), Map::of, @@ -529,7 +530,7 @@ public Table registerTable( AuthSession tableSession = authManager.tableSession(ident, tableConf, contextualSession); RESTClient tableClient = client.withAuthSession(tableSession); RESTTableOperations ops = - new RESTTableOperations( + newTableOps( tableClient, paths.table(ident), Map::of, @@ -788,7 +789,7 @@ public Table create() { AuthSession tableSession = authManager.tableSession(ident, tableConf, contextualSession); RESTClient tableClient = client.withAuthSession(tableSession); RESTTableOperations ops = - new RESTTableOperations( + newTableOps( tableClient, paths.table(ident), Map::of, @@ -815,7 +816,7 @@ public Transaction createTransaction() { RESTClient tableClient = client.withAuthSession(tableSession); RESTTableOperations ops = - new RESTTableOperations( + newTableOps( tableClient, paths.table(ident), Map::of, @@ -878,7 +879,7 @@ public Transaction replaceTransaction() { RESTClient tableClient = client.withAuthSession(tableSession); RESTTableOperations ops = - new RESTTableOperations( + newTableOps( tableClient, paths.table(ident), Map::of, @@ -1010,6 +1011,82 @@ private FileIO tableFileIO( return newFileIO(context, fullConf, storageCredentials); } + /** + * Create a new {@link RESTTableOperations} instance for simple table operations. + * + *

This method can be overridden in subclasses to provide custom table operations + * implementations. + * + * @param restClient the REST client to use for communicating with the catalog server + * @param path the REST path for the table + * @param headers a supplier for additional HTTP headers to include in requests + * @param fileIO the FileIO implementation for reading and writing table metadata and data files + * @param current the current table metadata + * @param supportedEndpoints the set of supported REST endpoints + * @return a new RESTTableOperations instance + */ + protected RESTTableOperations newTableOps( + RESTClient restClient, + String path, + Supplier> headers, + FileIO fileIO, + TableMetadata current, + Set supportedEndpoints) { + return new RESTTableOperations(restClient, path, headers, fileIO, current, supportedEndpoints); + } + + /** + * Create a new {@link RESTTableOperations} instance for transaction-based operations (create or + * replace). + * + *

This method can be overridden in subclasses to provide custom table operations + * implementations for transaction-based operations. + * + * @param restClient the REST client to use for communicating with the catalog server + * @param path the REST path for the table + * @param headers a supplier for additional HTTP headers to include in requests + * @param fileIO the FileIO implementation for reading and writing table metadata and data files + * @param updateType the {@link RESTTableOperations.UpdateType} being performed + * @param createChanges the list of metadata updates to apply during table creation or replacement + * @param current the current table metadata (may be null for CREATE operations) + * @param supportedEndpoints the set of supported REST endpoints + * @return a new RESTTableOperations instance + */ + protected RESTTableOperations newTableOps( + RESTClient restClient, + String path, + Supplier> headers, + FileIO fileIO, + RESTTableOperations.UpdateType updateType, + List createChanges, + TableMetadata current, + Set supportedEndpoints) { + return new RESTTableOperations( + restClient, path, headers, fileIO, updateType, createChanges, current, supportedEndpoints); + } + + /** + * Create a new {@link RESTViewOperations} instance. + * + *

This method can be overridden in subclasses to provide custom view operations + * implementations. + * + * @param restClient the REST client to use for communicating with the catalog server + * @param path the REST path for the view + * @param headers a supplier for additional HTTP headers to include in requests + * @param current the current view metadata + * @param supportedEndpoints the set of supported REST endpoints + * @return a new RESTViewOperations instance + */ + protected RESTViewOperations newViewOps( + RESTClient restClient, + String path, + Supplier> headers, + ViewMetadata current, + Set supportedEndpoints) { + return new RESTViewOperations(restClient, path, headers, current, supportedEndpoints); + } + private static ConfigResponse fetchConfig( RESTClient client, AuthSession initialAuth, Map properties) { // send the client's warehouse location to the service to keep in sync @@ -1154,7 +1231,7 @@ public View loadView(SessionContext context, TableIdentifier identifier) { ViewMetadata metadata = response.metadata(); RESTViewOperations ops = - new RESTViewOperations( + newViewOps( client.withAuthSession(tableSession), paths.view(identifier), Map::of, @@ -1333,7 +1410,7 @@ public View create() { Map tableConf = response.config(); AuthSession tableSession = authManager.tableSession(identifier, tableConf, contextualSession); RESTViewOperations ops = - new RESTViewOperations( + newViewOps( client.withAuthSession(tableSession), paths.view(identifier), Map::of, @@ -1424,7 +1501,7 @@ private View replace(LoadViewResponse response) { AuthSession contextualSession = authManager.contextualSession(context, catalogAuth); AuthSession tableSession = authManager.tableSession(identifier, tableConf, contextualSession); RESTViewOperations ops = - new RESTViewOperations( + newViewOps( client.withAuthSession(tableSession), paths.view(identifier), Map::of, diff --git a/core/src/test/java/org/apache/iceberg/rest/TestRESTCatalog.java b/core/src/test/java/org/apache/iceberg/rest/TestRESTCatalog.java index 6f7af7ae758a..780b1c89f3ca 100644 --- a/core/src/test/java/org/apache/iceberg/rest/TestRESTCatalog.java +++ b/core/src/test/java/org/apache/iceberg/rest/TestRESTCatalog.java @@ -41,9 +41,14 @@ import java.util.Map; import java.util.Objects; import java.util.Optional; +import java.util.Set; import java.util.UUID; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.function.BiFunction; import java.util.function.Consumer; +import java.util.function.Function; +import java.util.function.Supplier; import org.apache.hadoop.conf.Configuration; import org.apache.http.HttpHeaders; import org.apache.iceberg.BaseTable; @@ -56,6 +61,7 @@ import org.apache.iceberg.Schema; import org.apache.iceberg.Snapshot; import org.apache.iceberg.Table; +import org.apache.iceberg.TableMetadata; import org.apache.iceberg.Transaction; import org.apache.iceberg.UpdatePartitionSpec; import org.apache.iceberg.UpdateSchema; @@ -71,6 +77,7 @@ import org.apache.iceberg.exceptions.ServiceFailureException; import org.apache.iceberg.expressions.Expressions; import org.apache.iceberg.inmemory.InMemoryCatalog; +import org.apache.iceberg.io.FileIO; import org.apache.iceberg.relocated.com.google.common.collect.ImmutableList; import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap; import org.apache.iceberg.relocated.com.google.common.collect.Lists; @@ -3066,6 +3073,101 @@ public void testCommitStateUnknownNotReconciled() { .satisfies(ex -> assertThat(((CommitStateUnknownException) ex).getSuppressed()).isEmpty()); } + @Test + public void testCustomTableOperationsInjection() throws IOException { + AtomicBoolean customTableOpsCalled = new AtomicBoolean(); + AtomicBoolean customTransactionTableOpsCalled = new AtomicBoolean(); + + // Custom RESTSessionCatalog that overrides table operations creation + class CustomRESTSessionCatalog extends RESTSessionCatalog { + CustomRESTSessionCatalog( + Function, RESTClient> clientBuilder, + BiFunction, FileIO> ioBuilder) { + super(clientBuilder, ioBuilder); + } + + @Override + protected RESTTableOperations newTableOps( + RESTClient restClient, + String path, + Supplier> headers, + FileIO fileIO, + TableMetadata current, + Set supportedEndpoints) { + customTableOpsCalled.set(true); + return super.newTableOps(restClient, path, headers, fileIO, current, supportedEndpoints); + } + + @Override + protected RESTTableOperations newTableOps( + RESTClient restClient, + String path, + Supplier> headers, + FileIO fileIO, + RESTTableOperations.UpdateType updateType, + List createChanges, + TableMetadata current, + Set supportedEndpoints) { + customTransactionTableOpsCalled.set(true); + return super.newTableOps( + restClient, + path, + headers, + fileIO, + updateType, + createChanges, + current, + supportedEndpoints); + } + } + + // Custom RESTCatalog that provides the custom session catalog + class CustomRESTCatalog extends RESTCatalog { + CustomRESTCatalog( + SessionCatalog.SessionContext context, + Function, RESTClient> clientBuilder) { + super(context, clientBuilder); + } + + @Override + protected RESTSessionCatalog newSessionCatalog( + Function, RESTClient> clientBuilder) { + return new CustomRESTSessionCatalog(clientBuilder, null); + } + } + + try (CustomRESTCatalog catalog = + new CustomRESTCatalog( + SessionCatalog.SessionContext.createEmpty(), + (config) -> new RESTCatalogAdapter(backendCatalog))) { + catalog.setConf(new Configuration()); + catalog.initialize( + "test", + ImmutableMap.of( + CatalogProperties.FILE_IO_IMPL, "org.apache.iceberg.inmemory.InMemoryFileIO")); + + catalog.createNamespace(NS); + + // Test table operations without UpdateType + assertThat(customTableOpsCalled).isFalse(); + assertThat(customTransactionTableOpsCalled).isFalse(); + + catalog.createTable(TABLE, SCHEMA); + assertThat(customTableOpsCalled).isTrue(); + assertThat(customTransactionTableOpsCalled).isFalse(); + + // Test table operations with UpdateType and createChanges + customTableOpsCalled.set(false); + assertThat(customTableOpsCalled).isFalse(); + assertThat(customTransactionTableOpsCalled).isFalse(); + + TableIdentifier table2 = TableIdentifier.of(NS, "table2"); + catalog.buildTable(table2, SCHEMA).createTransaction().commitTransaction(); + assertThat(customTableOpsCalled).isFalse(); + assertThat(customTransactionTableOpsCalled).isTrue(); + } + } + private RESTCatalog catalog(RESTCatalogAdapter adapter) { RESTCatalog catalog = new RESTCatalog(SessionCatalog.SessionContext.createEmpty(), (config) -> adapter); diff --git a/core/src/test/java/org/apache/iceberg/rest/TestRESTViewCatalog.java b/core/src/test/java/org/apache/iceberg/rest/TestRESTViewCatalog.java index f3ad68c0020a..e2dcd170fde9 100644 --- a/core/src/test/java/org/apache/iceberg/rest/TestRESTViewCatalog.java +++ b/core/src/test/java/org/apache/iceberg/rest/TestRESTViewCatalog.java @@ -32,14 +32,20 @@ import java.nio.file.Path; import java.util.List; import java.util.Map; +import java.util.Set; import java.util.UUID; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.function.BiFunction; import java.util.function.Consumer; +import java.util.function.Function; +import java.util.function.Supplier; import org.apache.iceberg.CatalogProperties; import org.apache.iceberg.catalog.Catalog; import org.apache.iceberg.catalog.Namespace; import org.apache.iceberg.catalog.SessionCatalog; import org.apache.iceberg.catalog.TableIdentifier; import org.apache.iceberg.inmemory.InMemoryCatalog; +import org.apache.iceberg.io.FileIO; import org.apache.iceberg.relocated.com.google.common.collect.ImmutableList; import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap; import org.apache.iceberg.rest.HTTPRequest.HTTPMethod; @@ -48,6 +54,7 @@ import org.apache.iceberg.rest.responses.ListTablesResponse; import org.apache.iceberg.rest.responses.LoadViewResponse; import org.apache.iceberg.view.ViewCatalogTests; +import org.apache.iceberg.view.ViewMetadata; import org.eclipse.jetty.server.Server; import org.eclipse.jetty.server.handler.gzip.GzipHandler; import org.eclipse.jetty.servlet.ServletContextHandler; @@ -308,6 +315,70 @@ public void viewExistsFallbackToGETRequestWithLegacyServer() { ImmutableMap.of(RESTCatalogProperties.VIEW_ENDPOINTS_SUPPORTED, "true")); } + @Test + public void testCustomViewOperationsInjection() throws Exception { + AtomicBoolean customViewOpsCalled = new AtomicBoolean(); + + // Custom RESTSessionCatalog that overrides view operations creation + class CustomRESTSessionCatalog extends RESTSessionCatalog { + CustomRESTSessionCatalog( + Function, RESTClient> clientBuilder, + BiFunction, FileIO> ioBuilder) { + super(clientBuilder, ioBuilder); + } + + @Override + protected RESTViewOperations newViewOps( + RESTClient restClient, + String path, + Supplier> headers, + ViewMetadata current, + Set supportedEndpoints) { + customViewOpsCalled.set(true); + return super.newViewOps(restClient, path, headers, current, supportedEndpoints); + } + } + + // Custom RESTCatalog that provides the custom session catalog + class CustomRESTCatalog extends RESTCatalog { + CustomRESTCatalog( + SessionCatalog.SessionContext context, + Function, RESTClient> clientBuilder) { + super(context, clientBuilder); + } + + @Override + protected RESTSessionCatalog newSessionCatalog( + Function, RESTClient> clientBuilder) { + return new CustomRESTSessionCatalog(clientBuilder, null); + } + } + + try (CustomRESTCatalog catalog = + new CustomRESTCatalog( + SessionCatalog.SessionContext.createEmpty(), + (config) -> new RESTCatalogAdapter(backendCatalog))) { + catalog.initialize( + "test", + ImmutableMap.of( + CatalogProperties.FILE_IO_IMPL, "org.apache.iceberg.inmemory.InMemoryFileIO")); + + Namespace namespace = Namespace.of("ns"); + catalog.createNamespace(namespace); + + // Test view operations + assertThat(customViewOpsCalled).isFalse(); + TableIdentifier viewIdentifier = TableIdentifier.of(namespace, "view1"); + catalog + .buildView(viewIdentifier) + .withSchema(SCHEMA) + .withDefaultNamespace(namespace) + .withQuery("spark", "select * from ns.table") + .create(); + assertThat(customViewOpsCalled).isTrue(); + } + } + @Override protected RESTCatalog catalog() { return restCatalog;