IcebergIO and AddFiles multi bucket catalog validation

.github/trigger_files/IO_Iceberg_Integration_Tests.json

@@ -1,4 +1,4 @@
 {
     "comment": "Modify this file in a trivial way to cause this test suite to run.",
-    "modification": 10
+    "modification": 11
 }

sdks/java/io/iceberg/build.gradle

@@ -179,10 +179,15 @@ task integrationTest(type: Test) {
         "--project=${gcpProject}",
         "--tempLocation=${gcpTempLocation}",
     ])
-    // Warehouse (= catalog) used by the BigLake REST catalog tests; overridable for runs
-    // against a non-default project's catalog.
-    systemProperty "beam.iceberg.biglake.warehouse",
-        project.findProperty('biglakeWarehouse') ?: 'gs://managed-iceberg-biglake-its'
+    // Multiple-bucket Lakehouse REST catalog used by RESTCatalogBLMSIT and AddFilesIT (see
+    // LakehouseTestCatalog). Locations: the catalog's default location first, then a restricted
+    // location in another bucket. Override for runs against another project's catalog:
+    //   -PlakehouseWarehouse=bl://projects/PROJECT/catalogs/CATALOG
+    //   -PlakehouseLocations=gs://default-bucket/path,gs://other-bucket/path
+    systemProperty "beam.iceberg.lakehouse.warehouse",
+        project.findProperty('lakehouseWarehouse') ?: 'bl://projects/apache-beam-testing/catalogs/beam-lakehouse-it'
+    systemProperty "beam.iceberg.lakehouse.locations",
+        project.findProperty('lakehouseLocations') ?: 'gs://beam-lakehouse-it,gs://beam-lakehouse-it-added-path'
     // Connection + storage root for BigQueryManagedTableCrossEngineIT.
     if (project.findProperty('bqImtConnection') != null) {
         systemProperty "beam.bq.imt.connection", project.findProperty('bqImtConnection')

sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/LakehouseTestCatalog.java

@@ -0,0 +1,198 @@
+/*
+ * 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 org.apache.beam.sdk.io.iceberg;
+
+import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.api.services.storage.model.StorageObject;
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+import org.apache.beam.sdk.extensions.gcp.options.GcpOptions;
+import org.apache.beam.sdk.extensions.gcp.options.GcsOptions;
+import org.apache.beam.sdk.extensions.gcp.util.GcsUtil;
+import org.apache.beam.sdk.testing.TestPipeline;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Splitter;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
+import org.apache.iceberg.catalog.Catalog;
+import org.apache.iceberg.catalog.Namespace;
+import org.apache.iceberg.catalog.SupportsNamespaces;
+import org.apache.iceberg.catalog.TableIdentifier;
+
+/**
+ * Test-side description of the Lakehouse Iceberg REST catalog the ITs run against.
+ *
+ * <p>The catalog is a multiple-bucket catalog addressed by a {@code
+ * bl://projects/PROJECT/catalogs/CATALOG} warehouse and allowed to store resources under a fixed
+ * set of Cloud Storage locations: the catalog's default location plus any restricted locations.
+ * Configured (defaults and overrides) by the integrationTest task in build.gradle through the
+ * system properties
+ *
+ * <ul>
+ *   <li>{@code beam.iceberg.lakehouse.warehouse}: the {@code bl://} warehouse URI
+ *   <li>{@code beam.iceberg.lakehouse.locations}: comma-separated {@code gs://} prefixes the
+ *       catalog may write to; the first is the catalog's default location, the rest are additional
+ *       restricted locations
+ * </ul>
+ */
+public final class LakehouseTestCatalog {
+  private static final Pattern WAREHOUSE_PATTERN =
+      Pattern.compile("bl://projects/([^/]+)/catalogs/([^/]+)");
+
+  public static final String WAREHOUSE = requiredProperty("beam.iceberg.lakehouse.warehouse");
+
+  public static final List<String> LOCATIONS =
+      ImmutableList.copyOf(
+          Splitter.on(',')
+              .trimResults()
+              .omitEmptyStrings()
+              .split(requiredProperty("beam.iceberg.lakehouse.locations")));
+
+  /** Catalog id, which is also the second segment of BigQuery's 4-part table reference. */
+  public static final String CATALOG_ID = parseCatalogId(WAREHOUSE);
+
+  private static final String PROJECT =
+      TestPipeline.testingPipelineOptions().as(GcpOptions.class).getProject();
+
+  private LakehouseTestCatalog() {}
+
+  private static String requiredProperty(String name) {
+    String value = System.getProperty(name);
+    checkArgument(
+        value != null && !value.isEmpty(),
+        "System property %s is not set; run through the integrationTest gradle task, which"
+            + " sets it, or pass -D%s=...",
+        name,
+        name);
+    return value;
+  }
+
+  private static String parseCatalogId(String warehouse) {
+    // Legacy single-bucket catalogs are addressed by their bucket and named after it.
+    if (warehouse.startsWith("gs://")) {
+      return bucketOf(warehouse);
+    }
+    Matcher matcher = WAREHOUSE_PATTERN.matcher(warehouse);
+    checkArgument(
+        matcher.matches(),
+        "Expected a bl://projects/PROJECT/catalogs/CATALOG (or legacy gs://BUCKET) warehouse, got '%s'",
+        warehouse);
+    return matcher.group(2);
+  }
+
+  /** The catalog's default location; tables land here unless created with an explicit location. */
+  public static String defaultLocation() {
+    return LOCATIONS.get(0);
+  }
+
+  /** A restricted location outside the default one, i.e. in a second bucket. */
+  public static String additionalLocation() {
+    checkArgument(
+        LOCATIONS.size() >= 2,
+        "beam.iceberg.lakehouse.locations must list at least two locations, got %s",
+        LOCATIONS);
+    return LOCATIONS.get(1);
+  }
+
+  public static String bucketOf(String gcsLocation) {
+    checkArgument(gcsLocation.startsWith("gs://"), "Not a gs:// location: %s", gcsLocation);
+    String withoutScheme = gcsLocation.substring("gs://".length());
+    int slash = withoutScheme.indexOf('/');
+    return slash < 0 ? withoutScheme : withoutScheme.substring(0, slash);
+  }
+
+  /** Object-name prefix (no bucket, no leading slash) of a {@code gs://bucket/path} location. */
+  public static String prefixOf(String gcsLocation) {
+    String withoutScheme = gcsLocation.substring("gs://".length());
+    int slash = withoutScheme.indexOf('/');
+    return slash < 0 ? "" : withoutScheme.substring(slash + 1);
+  }
+
+  public static Map<String, String> catalogProperties() {
+    return ImmutableMap.<String, String>builder()
+        .put("type", "rest")
+        .put("uri", "https://biglake.googleapis.com/iceberg/v1/restcatalog")
+        .put("warehouse", WAREHOUSE)
+        .put("header.x-goog-user-project", PROJECT)
+        // Required by catalogs in vended-credentials mode; ignored in end-user mode.
+        .put("header.X-Iceberg-Access-Delegation", "vended-credentials")
+        .put("io-impl", "org.apache.iceberg.gcp.gcs.GCSFileIO")
+        .put("rest.auth.type", "org.apache.iceberg.gcp.auth.GoogleAuthManager")
+        .build();
+  }
+
+  /**
+   * Drops every table in the namespaces and then the namespaces themselves, and deletes the tables'
+   * files from Cloud Storage.
+   */
+  public static void dropNamespacesAndFiles(Catalog catalog, List<String> namespaces)
+      throws IOException {
+    for (String name : namespaces) {
+      Namespace namespace = Namespace.of(name);
+      if (!((SupportsNamespaces) catalog).namespaceExists(namespace)) {
+        continue;
+      }
+      dropTablesAndFiles(catalog, name);
+      ((SupportsNamespaces) catalog).dropNamespace(namespace);
+    }
+  }
+
+  /**
+   * Drops every table in the namespace and deletes the tables' files from Cloud Storage: Lakehouse
+   * keeps a dropped table's data and metadata (even with purge), and table locations carry a random
+   * suffix, so they are captured before the drop.
+   */
+  public static void dropTablesAndFiles(Catalog catalog, String namespaceName) throws IOException {
+    Namespace namespace = Namespace.of(namespaceName);
+    if (!((SupportsNamespaces) catalog).namespaceExists(namespace)) {
+      return;
+    }
+    List<String> tableLocations = new ArrayList<>();
+    for (TableIdentifier identifier : catalog.listTables(namespace)) {
+      tableLocations.add(catalog.loadTable(identifier).location());
+      catalog.dropTable(identifier);
+    }
+    for (String location : tableLocations) {
+      deleteObjects(location);
+    }
+  }
+
+  /** Deletes every object under a {@code gs://bucket/prefix} location. */
+  public static void deleteObjects(String gcsLocation) throws IOException {
+    GcsUtil gcsUtil = TestPipeline.testingPipelineOptions().as(GcsOptions.class).getGcsUtil();
+    List<StorageObject> objects =
+        gcsUtil.listObjects(bucketOf(gcsLocation), prefixOf(gcsLocation), null).getItems();
+    if (objects == null || objects.isEmpty()) {
+      return;
+    }
+    List<String> paths = new ArrayList<>();
+    for (StorageObject object : objects) {
+      paths.add("gs://" + object.getBucket() + "/" + object.getName());
+    }
+    gcsUtil.remove(paths);
+  }
+
+  /** BigQuery's 4-part {@code project.catalog.namespace.table} reference. */
+  public static String bigQueryTableSpec(String namespace, String table) {
+    return String.format("%s.%s.%s.%s", PROJECT, CATALOG_ID, namespace, table);
+  }
+}

sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/catalog/IcebergCatalogBaseIT.java

@@ -175,10 +175,9 @@ public abstract class IcebergCatalogBaseIT implements Serializable {
   /**
    * Catalogs whose tables are also queryable with BigQuery return the BigQuery table reference for
    * the given Iceberg table id: either the 4-part {@code project.catalog.namespace.table} form for
-   * Lakehouse runtime catalog (BigLake metastore REST) tables, or the 3-part {@code
-   * project.dataset.table} form for the BigQuery metastore federation, where namespaces surface as
-   * datasets. Returning null (the default) disables the cross-engine read checks in {@link
-   * #testReadWithBigQueryIO()}.
+   * Lakehouse runtime catalog (Iceberg REST) tables, or the 3-part {@code project.dataset.table}
+   * form for the BigQuery metastore federation, where namespaces surface as datasets. Returning
+   * null (the default) disables the cross-engine read checks in {@link #testReadWithBigQueryIO()}.
    */
   public @Nullable String bigQueryTableSpec(String tableId) {
     return null;
@@ -238,14 +237,14 @@ public void cleanUp() throws Exception {
     try {
       GcsUtil gcsUtil = OPTIONS.as(GcsOptions.class).getGcsUtil();
       GcsPath path = GcsPath.fromUri(warehouse);
+      // The warehouse may be a bare bucket (no object path), where getFileName() throws.
+      String prefix =
+          path.getObject().isEmpty()
+              ? getClass().getSimpleName()
+              : getClass().getSimpleName() + "/" + path.getFileName();
 
       @Nullable List<StorageObject> objects =
-          gcsUtil
-              .listObjects(
-                  path.getBucket(),
-                  getClass().getSimpleName() + "/" + path.getFileName().toString(),
-                  null)
-              .getItems();
+          gcsUtil.listObjects(path.getBucket(), prefix, null).getItems();
 
       // sometimes a catalog's cleanup will take care of all the files.
       // If any files are left though, manually delete them with GCS utils
@@ -432,7 +431,7 @@ private List<Row> populateTable(Table table, @Nullable String charOverride) thro
     }
   }
 
-  private List<Record> readRecords(Table table) throws IOException {
+  protected List<Record> readRecords(Table table) throws IOException {
     org.apache.iceberg.Schema tableSchema = table.schema();
     TableScan tableScan = table.newScan().project(tableSchema);
     List<Record> writtenRecords = new ArrayList<>();

sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/catalog/IcebergCdcWriteBaseIT.java

@@ -236,13 +236,13 @@ public void cleanUp() throws Exception {
     try {
       GcsUtil gcsUtil = OPTIONS.as(GcsOptions.class).getGcsUtil();
       GcsPath path = GcsPath.fromUri(warehouse);
+      // The warehouse may be a bare bucket (no object path), where getFileName() throws.
+      String prefix =
+          path.getObject().isEmpty()
+              ? getClass().getSimpleName()
+              : getClass().getSimpleName() + "/" + path.getFileName();
       @Nullable List<StorageObject> objects =
-          gcsUtil
-              .listObjects(
-                  path.getBucket(),
-                  getClass().getSimpleName() + "/" + path.getFileName().toString(),
-                  null)
-              .getItems();
+          gcsUtil.listObjects(path.getBucket(), prefix, null).getItems();
       // A catalog's cleanup sometimes removes every file; delete whatever is left.
       if (objects != null) {
         gcsUtil.remove(

sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/catalog/LakehouseCatalogCdcWriteIT.java

@@ -17,7 +17,9 @@
  */
 package org.apache.beam.sdk.io.iceberg.catalog;
 
+import java.io.IOException;
 import java.util.Map;
+import org.apache.beam.sdk.io.iceberg.LakehouseTestCatalog;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
 import org.apache.iceberg.catalog.Catalog;
 import org.apache.iceberg.rest.RESTCatalog;
@@ -28,27 +30,19 @@
 public class LakehouseCatalogCdcWriteIT extends IcebergCdcWriteBaseIT {
   private static Map<String, String> catalogProps;
 
-  private static final String LAKEHOUSE_WAREHOUSE =
-      System.getProperty("beam.iceberg.biglake.warehouse", "gs://managed-iceberg-biglake-its");
-
   @BeforeClass
   public static void setup() {
-    warehouse = LAKEHOUSE_WAREHOUSE;
-    catalogProps =
-        ImmutableMap.<String, String>builder()
-            .put("type", "rest")
-            .put("uri", "https://biglake.googleapis.com/iceberg/v1/restcatalog")
-            .put("warehouse", LAKEHOUSE_WAREHOUSE)
-            .put("header.x-goog-user-project", OPTIONS.getProject())
-            .put("io-impl", "org.apache.iceberg.gcp.gcs.GCSFileIO")
-            .put("rest.auth.type", "org.apache.iceberg.gcp.auth.GoogleAuthManager")
-            .build();
+    warehouse = LakehouseTestCatalog.defaultLocation();
+    catalogProps = LakehouseTestCatalog.catalogProperties();
   }
 
   @After
-  public void after() {
+  public void after() throws IOException {
+    // Lakehouse keeps a dropped table's files, so remove them before the base class drops the
+    // namespace.
+    LakehouseTestCatalog.dropTablesAndFiles(catalog, namespace());
     // The base class points its cleanup at this warehouse.
-    warehouse = LAKEHOUSE_WAREHOUSE;
+    warehouse = LakehouseTestCatalog.defaultLocation();
   }
 
   @Override