From 53aa54b9c2a18d9f73b5ae584085603a05a3ad17 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 27 Aug 2026 08:36:09 -0400 Subject: [PATCH 1/4] Validate multi region biglake catalog --- sdks/java/io/iceberg/build.gradle | 15 +- .../beam/sdk/io/iceberg/AddFilesIT.java | 220 ++++++++++++++---- .../sdk/io/iceberg/BigLakeTestCatalog.java | 146 ++++++++++++ .../iceberg/catalog/IcebergCatalogBaseIT.java | 2 +- .../io/iceberg/catalog/RESTCatalogBLMSIT.java | 108 ++++++--- 5 files changed, 422 insertions(+), 69 deletions(-) create mode 100644 sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/BigLakeTestCatalog.java diff --git a/sdks/java/io/iceberg/build.gradle b/sdks/java/io/iceberg/build.gradle index e2e8a12d01eb..a186ef8000cf 100644 --- a/sdks/java/io/iceberg/build.gradle +++ b/sdks/java/io/iceberg/build.gradle @@ -169,10 +169,19 @@ 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. + // Multiple-bucket BigLake REST catalog used by RESTCatalogBLMSIT and AddFilesIT (see + // BigLakeTestCatalog). Locations: the catalog's default location first, then a restricted + // location in another bucket. The cross-region location must be a bucket in a different + // region than the catalog. Override for runs against another project's catalog: + // -PbiglakeWarehouse=bl://projects/PROJECT/catalogs/CATALOG + // -PbiglakeLocations=gs://default-bucket/path,gs://other-bucket/path + // -PbiglakeCrossRegionLocation=gs://bucket-in-another-region/path systemProperty "beam.iceberg.biglake.warehouse", - project.findProperty('biglakeWarehouse') ?: 'gs://managed-iceberg-biglake-its' + project.findProperty('biglakeWarehouse') ?: 'bl://projects/apache-beam-testing/catalogs/beam-lakehouse-it' + systemProperty "beam.iceberg.biglake.locations", + project.findProperty('biglakeLocations') ?: 'gs://beam-lakehouse-it,gs://beam-lakehouse-it-added-path' + systemProperty "beam.iceberg.biglake.crossRegionLocation", + project.findProperty('biglakeCrossRegionLocation') ?: 'gs://managed-iceberg-biglake-its/biglake_cross_region' // Connection + storage root for BigQueryManagedTableCrossEngineIT. if (project.findProperty('bqImtConnection') != null) { systemProperty "beam.bq.imt.connection", project.findProperty('bqImtConnection') diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/AddFilesIT.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/AddFilesIT.java index 39a1ab4d427e..eea844222bbd 100644 --- a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/AddFilesIT.java +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/AddFilesIT.java @@ -22,8 +22,12 @@ import static org.apache.beam.sdk.io.iceberg.IcebergUtils.beamSchemaToIcebergSchema; import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; import static org.apache.beam.sdk.values.TypeDescriptors.strings; +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.startsWith; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotEquals; +import static org.junit.Assert.assertThrows; import static org.junit.Assert.assertTrue; import com.google.api.services.storage.model.StorageObject; @@ -76,6 +80,7 @@ import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables; import org.apache.hadoop.util.Lists; +import org.apache.iceberg.BaseTable; import org.apache.iceberg.DataFile; import org.apache.iceberg.PartitionSpec; import org.apache.iceberg.Snapshot; @@ -102,12 +107,12 @@ public class AddFilesIT { private static final Logger LOG = LoggerFactory.getLogger(AddFilesIT.class); - // Bucket-backed BigLake catalogs are named after their bucket. Overridable for local runs - // against a different project's catalog: -Dbeam.iceberg.biglake.warehouse=gs://my-bucket - private static final String CATALOG_NAME = - System.getProperty("beam.iceberg.biglake.warehouse", "gs://managed-iceberg-biglake-its") - .replace("gs://", ""); - private static final String WAREHOUSE = "gs://" + CATALOG_NAME; + // Multiple-bucket BigLake catalog (see BigLakeTestCatalog). Source parquet files, and the + // GCS notifications announcing them, live under the catalog's default location. + private static final String DATA_LOCATION = BigLakeTestCatalog.defaultLocation(); + private static final String DATA_BUCKET = BigLakeTestCatalog.bucketOf(DATA_LOCATION); + private static final String DATA_PREFIX = BigLakeTestCatalog.prefixOf(DATA_LOCATION); + private static final String CROSS_REGION_LOCATION = BigLakeTestCatalog.CROSS_REGION_LOCATION; private static final String PROJECT = TestPipeline.testingPipelineOptions().as(GcpOptions.class).getProject(); @Rule public TestName testName = new TestName(); @@ -124,18 +129,13 @@ public class AddFilesIT { .addStringField("name") .addStringField("kind") .build(); - private static final Map BIGLAKE_PROPS = - Map.of( - "type", "rest", - "uri", "https://biglake.googleapis.com/iceberg/v1/restcatalog", - "warehouse", WAREHOUSE, - "header.x-goog-user-project", PROJECT, - "rest.auth.type", "google", - "io-impl", "org.apache.iceberg.gcp.gcs.GCSFileIO"); + private static final Map BIGLAKE_PROPS = BigLakeTestCatalog.catalogProperties(); private Storage storage; private PubsubClient pubsub; private Notification notification; private final String namespace = getClass().getSimpleName() + "_" + System.currentTimeMillis(); + // Namespace placed in the catalog's additional location (second bucket). + private final String altNamespace = namespace + "_alt"; private String srcTableName; private String destTableName; private TableIdentifier srcTableId; @@ -171,29 +171,30 @@ public void setup() throws IOException { .setPayloadFormat(NotificationInfo.PayloadFormat.JSON_API_V1) .build(); try { - notification = storage.createNotification(WAREHOUSE.replace("gs://", ""), notificationInfo); + notification = storage.createNotification(DATA_BUCKET, notificationInfo); } catch (StorageException e) { if (e.getMessage().contains("Too many overlapping notifications")) { - List existing = storage.listNotifications(WAREHOUSE.replace("gs://", "")); + List existing = storage.listNotifications(DATA_BUCKET); LOG.warn( "Too many notifications on bucket {}: {}. Deleting existing notifications to make room: {}", - WAREHOUSE, + DATA_BUCKET, e, existing.stream() .map(NotificationInfo::getNotificationId) .collect(Collectors.toList())); - existing.forEach( - n -> storage.deleteNotification(WAREHOUSE.replace("gs://", ""), n.getNotificationId())); + existing.forEach(n -> storage.deleteNotification(DATA_BUCKET, n.getNotificationId())); // try creating it again - notification = storage.createNotification(WAREHOUSE.replace("gs://", ""), notificationInfo); + notification = storage.createNotification(DATA_BUCKET, notificationInfo); } else { throw e; } } salt = System.currentTimeMillis(); - dirName = format("%s-%s/%s", getClass().getSimpleName(), salt, testName.getMethodName()); + dirName = + format( + "%s/%s-%s/%s", DATA_PREFIX, getClass().getSimpleName(), salt, testName.getMethodName()); srcTableName = "src_" + testName.getMethodName() + "_" + salt; destTableName = "dest_" + testName.getMethodName() + "_" + salt; srcTableId = TableIdentifier.of(namespace, srcTableName); @@ -205,10 +206,12 @@ public void setup() throws IOException { } private void cleanupCatalog() { - Namespace ns = Namespace.of(namespace); - if (catalog.namespaceExists(ns)) { - catalog.listTables(ns).forEach(catalog::dropTable); - catalog.dropNamespace(ns); + for (String name : Arrays.asList(namespace, altNamespace)) { + Namespace ns = Namespace.of(name); + if (catalog.namespaceExists(ns)) { + catalog.listTables(ns).forEach(catalog::dropTable); + catalog.dropNamespace(ns); + } } } @@ -222,7 +225,7 @@ public void cleanup() { } try { - storage.deleteNotification(WAREHOUSE.replace("gs://", ""), notification.getNotificationId()); + storage.deleteNotification(DATA_BUCKET, notification.getNotificationId()); storage.close(); } catch (Exception e) { LOG.warn("Failed to clean up GCS notifications", e); @@ -234,14 +237,19 @@ public void cleanup() { LOG.warn("Failed to clean up Iceberg catalog", e); } + deleteBlobs(DATA_BUCKET, dirName); + deleteBlobs( + BigLakeTestCatalog.bucketOf(CROSS_REGION_LOCATION), + BigLakeTestCatalog.prefixOf(CROSS_REGION_LOCATION) + "/" + dirName); + } + + private void deleteBlobs(String bucket, String prefix) { try { Iterable blobs = - storage - .list(WAREHOUSE.replace("gs://", ""), Storage.BlobListOption.prefix(dirName)) - .getValues(); + storage.list(bucket, Storage.BlobListOption.prefix(prefix)).getValues(); blobs.forEach(b -> storage.delete(b.getBlobId())); } catch (Exception e) { - LOG.warn("Failed to clean up GCS bucket", e); + LOG.warn("Failed to clean up gs://{}/{}", bucket, prefix, e); } } @@ -347,8 +355,8 @@ public void testStreamingParquetImport() throws InterruptedException, TimeoutException, IOException { // start with a table that does not exist - String parquetDir = format("%s/%s/", WAREHOUSE, dirName); - String tempDir = format("%s/%s-tmp/", WAREHOUSE, dirName); + String parquetDir = format("gs://%s/%s/", DATA_BUCKET, dirName); + String tempDir = format("gs://%s/%s-tmp/", DATA_BUCKET, dirName); // let the add files pipeline run in the background PipelineResult addFilesPipeline = startAddFilesListener(dirName); @@ -380,8 +388,7 @@ record -> GcsUtil gcsUtil = TestPipeline.testingPipelineOptions().as(GcsOptions.class).getGcsUtil(); - Iterable objects = - gcsUtil.listObjects(WAREHOUSE.replace("gs://", ""), dirName, null).getItems(); + Iterable objects = gcsUtil.listObjects(DATA_BUCKET, dirName, null).getItems(); List writtenFilePaths = Lists.newArrayList(objects).stream() .map(o -> format("gs://%s/%s", o.getBucket(), o.getName())) @@ -421,11 +428,147 @@ public void testBatchParquetImportToUIT() throws IOException { testBatchParquetImport(true); } + /** + * The catalog does not police where added files live: data in a bucket it does not manage, even + * in another region, registers and reads fine through Iceberg. BigQuery, however, cannot read + * data files outside the catalog's region, so the cross-engine read of such a table fails. (Files + * in an unmanaged bucket in the catalog's region read fine from BigQuery.) Documents the + * trade-off of adding files in place from a cross-region bucket. + */ + @Test + public void testBatchParquetImportFromCrossRegionBucket() throws IOException { + String crossRegionBucket = BigLakeTestCatalog.bucketOf(CROSS_REGION_LOCATION); + for (String location : BigLakeTestCatalog.LOCATIONS) { + assertNotEquals( + "Test needs a bucket outside the catalog: " + CROSS_REGION_LOCATION, + BigLakeTestCatalog.bucketOf(location), + crossRegionBucket); + } + List writtenFilePaths = + writeParquetFiles( + crossRegionBucket, BigLakeTestCatalog.prefixOf(CROSS_REGION_LOCATION) + "/" + dirName); + + Pipeline p = Pipeline.create(); + PCollectionRowTuple tuple = + p.apply(Create.of(writtenFilePaths)) + .apply( + new AddFiles( + IcebergCatalogConfig.builder().setCatalogProperties(BIGLAKE_PROPS).build(), + destTableId.toString(), + null, + PARTITION_FIELDS, + null, + TABLE_PROPS, + null, + null)); + PAssert.that(tuple.get("errors")).empty(); + p.run().waitUntilFinish(); + + assertTrue(checkTableHasRegisteredParquetFiles(writtenFilePaths)); + checkRecordsInDestinationTable(/* alsoCheckWithBigQueryIO= */ false); + + Pipeline bq = Pipeline.create(); + bq.apply( + Managed.read(Managed.BIGQUERY) + .withConfig( + ImmutableMap.of( + "table", + BigLakeTestCatalog.bigQueryTableSpec( + destTableId.namespace().toString(), destTableId.name())))) + .getSinglePCollection(); + assertThrows(Pipeline.PipelineExecutionException.class, () -> bq.run().waitUntilFinish()); + } + + /** + * The destination table lives in the catalog's additional location (a second bucket) while the + * source parquet files stay in the default one. BigLake pins tables under their namespace's + * location, so the table is created in a namespace placed in the second bucket; AddFiles must + * commit metadata there and reference the files in place across buckets. + */ + @Test + public void testBatchParquetImportToTableInAdditionalLocation() throws IOException { + String namespaceLocation = BigLakeTestCatalog.additionalLocation() + "/" + altNamespace; + assertNotEquals( + "Test needs two distinct buckets", + DATA_BUCKET, + BigLakeTestCatalog.bucketOf(namespaceLocation)); + catalog.createNamespace( + Namespace.of(altNamespace), ImmutableMap.of("location", namespaceLocation)); + destTableId = TableIdentifier.of(altNamespace, destTableName); + catalog.createTable(destTableId, beamSchemaToIcebergSchema(ROW_SCHEMA), SPEC); + assertThat(catalog.loadTable(destTableId).location(), startsWith(namespaceLocation)); + + List writtenFilePaths = writeParquetFiles(); + Pipeline p = Pipeline.create(); + PCollectionRowTuple tuple = + p.apply(Create.of(writtenFilePaths)) + .apply( + new AddFiles( + IcebergCatalogConfig.builder().setCatalogProperties(BIGLAKE_PROPS).build(), + destTableId.toString(), + null, + PARTITION_FIELDS, + null, + TABLE_PROPS, + null, + null)); + PAssert.that(tuple.get("errors")).empty(); + p.run().waitUntilFinish(); + + assertTrue(checkTableHasRegisteredParquetFiles(writtenFilePaths)); + Table destTable = catalog.loadTable(destTableId); + String metadataLocation = ((BaseTable) destTable).operations().current().metadataFileLocation(); + assertThat(metadataLocation, startsWith(BigLakeTestCatalog.additionalLocation())); + for (String path : writtenFilePaths) { + assertThat(path, startsWith("gs://" + DATA_BUCKET + "/")); + } + checkRecordsInDestinationTable(/* alsoCheckWithBigQueryIO= */ true); + } + + private List writeParquetFiles() throws IOException { + return writeParquetFiles(DATA_BUCKET, dirName); + } + + /** Writes TEST_ROWS as parquet under gs://{bucket}/{dir}/ and returns the written file paths. */ + private List writeParquetFiles(String bucket, String dir) throws IOException { + String parquetDir = format("gs://%s/%s/", bucket, dir); + String tempDir = format("gs://%s/%s-tmp/", bucket, dir); + LOG.info("Writing records to the parquet dir"); + Pipeline q = Pipeline.create(); + org.apache.avro.Schema avroSchema = AvroUtils.toAvroSchema(ROW_SCHEMA); + q.apply(Create.of(TEST_ROWS)) + .setRowSchema(ROW_SCHEMA) + .apply( + MapElements.into(TypeDescriptor.of(GenericRecord.class)) + .via(AvroUtils.getRowToGenericRecordFunction(avroSchema))) + .setCoder(AvroCoder.of(avroSchema)) + .apply( + FileIO.writeDynamic() + .by( + record -> + format("%s-%s-%s", record.get("id"), record.get("name"), record.get("age"))) + .via(ParquetIO.sink(avroSchema)) + .withNaming(name -> defaultNaming(name, ".parquet")) + .withTempDirectory(tempDir) + .to(parquetDir) + .withDestinationCoder(StringUtf8Coder.of())); + q.run().waitUntilFinish(); + + GcsUtil gcsUtil = TestPipeline.testingPipelineOptions().as(GcsOptions.class).getGcsUtil(); + Iterable objects = gcsUtil.listObjects(bucket, dir, null).getItems(); + List writtenFilePaths = + Lists.newArrayList(objects).stream() + .map(o -> format("gs://%s/%s", o.getBucket(), o.getName())) + .collect(Collectors.toList()); + LOG.info("Written file paths: {}", writtenFilePaths); + return writtenFilePaths; + } + private void testBatchParquetImport(boolean isUIT) throws IOException { // start with a table that does not exist - String parquetDir = format("%s/%s/", WAREHOUSE, dirName); - String tempDir = format("%s/%s-tmp/", WAREHOUSE, dirName); + String parquetDir = format("gs://%s/%s/", DATA_BUCKET, dirName); + String tempDir = format("gs://%s/%s-tmp/", DATA_BUCKET, dirName); // write some parquet files LOG.info("Writing records to the parquet dir"); @@ -451,8 +594,7 @@ record -> GcsUtil gcsUtil = TestPipeline.testingPipelineOptions().as(GcsOptions.class).getGcsUtil(); - Iterable objects = - gcsUtil.listObjects(WAREHOUSE.replace("gs://", ""), dirName, null).getItems(); + Iterable objects = gcsUtil.listObjects(DATA_BUCKET, dirName, null).getItems(); List writtenFilePaths = Lists.newArrayList(objects).stream() .map(o -> format("gs://%s/%s", o.getBucket(), o.getName())) @@ -522,7 +664,7 @@ private void checkRecordsInDestinationTable(boolean alsoCheckWithBigQueryIO) { format( "%s.%s.%s.%s", PROJECT, - CATALOG_NAME, + BigLakeTestCatalog.CATALOG_ID, destTableId.namespace(), destTableId.name())))) .getSinglePCollection() diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/BigLakeTestCatalog.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/BigLakeTestCatalog.java new file mode 100644 index 000000000000..061babb7d0d5 --- /dev/null +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/BigLakeTestCatalog.java @@ -0,0 +1,146 @@ +/* + * 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 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.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; + +/** + * Test-side description of the BigLake (Lakehouse) Iceberg REST catalog the ITs run against. + * + *

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 + * + *

    + *
  • {@code beam.iceberg.biglake.warehouse}: the {@code bl://} warehouse URI + *
  • {@code beam.iceberg.biglake.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 + *
  • {@code beam.iceberg.biglake.crossRegionLocation}: a {@code gs://} prefix in a bucket + * located in a different region than the catalog + *
+ */ +public final class BigLakeTestCatalog { + private static final Pattern WAREHOUSE_PATTERN = + Pattern.compile("bl://projects/([^/]+)/catalogs/([^/]+)"); + + public static final String WAREHOUSE = requiredProperty("beam.iceberg.biglake.warehouse"); + + public static final List LOCATIONS = + ImmutableList.copyOf( + Splitter.on(',') + .trimResults() + .omitEmptyStrings() + .split(requiredProperty("beam.iceberg.biglake.locations"))); + + /** + * A location in a bucket the catalog does not manage and that is in a different region than the + * catalog (BigQuery cannot read cross-region data files; see + * AddFilesIT#testBatchParquetImportFromCrossRegionBucket). + */ + public static final String CROSS_REGION_LOCATION = + requiredProperty("beam.iceberg.biglake.crossRegionLocation"); + + /** 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 BigLakeTestCatalog() {} + + 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.biglake.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 catalogProperties() { + return ImmutableMap.builder() + .put("type", "rest") + .put("uri", "https://biglake.googleapis.com/iceberg/v1/restcatalog") + .put("warehouse", WAREHOUSE) + .put("header.x-goog-user-project", PROJECT) + .put("io-impl", "org.apache.iceberg.gcp.gcs.GCSFileIO") + .put("rest.auth.type", "org.apache.iceberg.gcp.auth.GoogleAuthManager") + .build(); + } + + /** 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); + } +} diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/catalog/IcebergCatalogBaseIT.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/catalog/IcebergCatalogBaseIT.java index 9eca20189e5e..12da329e115f 100644 --- a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/catalog/IcebergCatalogBaseIT.java +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/catalog/IcebergCatalogBaseIT.java @@ -433,7 +433,7 @@ private List populateTable(Table table, @Nullable String charOverride) thro } } - private List readRecords(Table table) throws IOException { + protected List readRecords(Table table) throws IOException { org.apache.iceberg.Schema tableSchema = table.schema(); TableScan tableScan = table.newScan().project(tableSchema); List writtenRecords = new ArrayList<>(); diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/catalog/RESTCatalogBLMSIT.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/catalog/RESTCatalogBLMSIT.java index aa83d2b7db2d..95442455132f 100644 --- a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/catalog/RESTCatalogBLMSIT.java +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/catalog/RESTCatalogBLMSIT.java @@ -17,44 +17,53 @@ */ package org.apache.beam.sdk.io.iceberg.catalog; +import static org.apache.beam.sdk.managed.Managed.ICEBERG; +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.containsInAnyOrder; +import static org.hamcrest.Matchers.startsWith; +import static org.junit.Assert.assertFalse; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.List; import java.util.Map; +import org.apache.beam.sdk.io.iceberg.BigLakeTestCatalog; +import org.apache.beam.sdk.managed.Managed; +import org.apache.beam.sdk.transforms.Create; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; +import org.apache.iceberg.BaseTable; +import org.apache.iceberg.DataFile; +import org.apache.iceberg.Snapshot; +import org.apache.iceberg.Table; import org.apache.iceberg.catalog.Catalog; +import org.apache.iceberg.catalog.Namespace; +import org.apache.iceberg.catalog.SupportsNamespaces; import org.apache.iceberg.catalog.TableIdentifier; +import org.apache.iceberg.data.Record; import org.apache.iceberg.rest.RESTCatalog; import org.junit.After; import org.junit.BeforeClass; +import org.junit.Test; -/** Tests for {@link org.apache.iceberg.rest.RESTCatalog} using BigLake Metastore. */ +/** + * Tests for {@link org.apache.iceberg.rest.RESTCatalog} using a multiple-bucket BigLake Metastore + * catalog (see {@link BigLakeTestCatalog}). + */ public class RESTCatalogBLMSIT extends IcebergCatalogBaseIT { private static Map catalogProps; - // Using a special bucket for this test class because - // BigLake does not support using subfolders as a warehouse (yet). - // Overridable for local runs against a different project's catalog, e.g. - // -Dbeam.iceberg.biglake.warehouse=gs://my-bucket (bucket-backed catalogs are named after - // their bucket). - private static final String BIGLAKE_WAREHOUSE = - System.getProperty("beam.iceberg.biglake.warehouse", "gs://managed-iceberg-biglake-its"); - @BeforeClass public static void setup() { - warehouse = BIGLAKE_WAREHOUSE; - catalogProps = - ImmutableMap.builder() - .put("type", "rest") - .put("uri", "https://biglake.googleapis.com/iceberg/v1/restcatalog") - .put("warehouse", BIGLAKE_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(); + // The catalog decides where tables go (its default location); the base class only uses + // `warehouse` to sweep leftover files, so point it at that location. + warehouse = BigLakeTestCatalog.defaultLocation(); + catalogProps = BigLakeTestCatalog.catalogProperties(); } @After public void after() { // making sure the cleanup path is directed at the correct warehouse - warehouse = BIGLAKE_WAREHOUSE; + warehouse = BigLakeTestCatalog.defaultLocation(); } @Override @@ -65,13 +74,11 @@ public String type() { @Override public String bigQueryTableSpec(String tableId) { // BigQuery surfaces Lakehouse runtime catalog (BigLake metastore REST) tables via 4-part - // project.catalog.namespace.table identifiers; the catalog id of a bucket-backed catalog is - // the bucket name. Requires the caller to hold biglake.* read permissions (e.g. - // roles/biglake.viewer) in addition to the usual BigQuery roles. + // project.catalog.namespace.table identifiers. Requires the caller to hold biglake.* read + // permissions (e.g. roles/biglake.viewer) in addition to the usual BigQuery roles. TableIdentifier identifier = TableIdentifier.parse(tableId); - String catalogId = BIGLAKE_WAREHOUSE.replace("gs://", ""); - return String.format( - "%s.%s.%s.%s", OPTIONS.getProject(), catalogId, identifier.namespace(), identifier.name()); + return BigLakeTestCatalog.bigQueryTableSpec( + identifier.namespace().toString(), identifier.name()); } @Override @@ -88,4 +95,53 @@ public Map managedIcebergConfig(String tableId) { .put("catalog_properties", catalogProps) .build(); } + + /** + * A multiple-bucket catalog may place resources under any of its restricted locations, not just + * its default one. BigLake pins a table under its namespace's location, so the namespace is + * created in a second bucket; Beam only receives the table through the catalog and must write + * wherever it was placed. + */ + @Test + public void testWriteReadTableInAdditionalLocation() throws IOException { + String altNamespace = namespace() + "_alt"; + String altTableId = altNamespace + ".test_table"; + String namespaceLocation = BigLakeTestCatalog.additionalLocation() + "/" + altNamespace; + assertFalse( + "Test needs two distinct buckets", + BigLakeTestCatalog.bucketOf(namespaceLocation) + .equals(BigLakeTestCatalog.bucketOf(BigLakeTestCatalog.defaultLocation()))); + namespacesToCleanup.add(altNamespace); + ((SupportsNamespaces) catalog) + .createNamespace( + Namespace.of(altNamespace), ImmutableMap.of("location", namespaceLocation)); + Table table = catalog.createTable(TableIdentifier.parse(altTableId), ICEBERG_SCHEMA); + assertThat(table.location(), startsWith(namespaceLocation)); + + pipeline + .apply(Create.of(inputRows)) + .setRowSchema(BEAM_SCHEMA) + .apply(Managed.write(ICEBERG).withConfig(managedIcebergConfig(altTableId))); + pipeline.run().waitUntilFinish(); + + table.refresh(); + List returnedRecords = readRecords(table); + assertThat( + returnedRecords, containsInAnyOrder(inputRows.stream().map(RECORD_FUNC::apply).toArray())); + + // Both the data files Beam wrote and the metadata the catalog committed live in the + // additional location, not in the catalog's default bucket. + List dataFileLocations = new ArrayList<>(); + for (Snapshot snapshot : table.snapshots()) { + for (DataFile dataFile : snapshot.addedDataFiles(table.io())) { + dataFileLocations.add(dataFile.location()); + } + } + assertFalse("No data files were written", dataFileLocations.isEmpty()); + for (String location : dataFileLocations) { + assertThat(location, startsWith(BigLakeTestCatalog.additionalLocation())); + } + String metadataLocation = ((BaseTable) table).operations().current().metadataFileLocation(); + assertThat(metadataLocation, startsWith(BigLakeTestCatalog.additionalLocation())); + } } From 57550c68dcba494013fc6b440ebf3bb6271bc4f3 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 27 Aug 2026 08:45:15 -0400 Subject: [PATCH 2/4] trigger tests --- .github/trigger_files/IO_Iceberg_Integration_Tests.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/trigger_files/IO_Iceberg_Integration_Tests.json b/.github/trigger_files/IO_Iceberg_Integration_Tests.json index 34a6e02150e7..5d04b2c0a8c7 100644 --- a/.github/trigger_files/IO_Iceberg_Integration_Tests.json +++ b/.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": 4 + "modification": 5 } From c810f0170b17f0e0e68bd8914f11c8d47ab851db Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 27 Aug 2026 09:22:56 -0400 Subject: [PATCH 3/4] remove x-region test --- sdks/java/io/iceberg/build.gradle | 6 +- .../beam/sdk/io/iceberg/AddFilesIT.java | 76 ++----------------- .../sdk/io/iceberg/BigLakeTestCatalog.java | 10 --- 3 files changed, 7 insertions(+), 85 deletions(-) diff --git a/sdks/java/io/iceberg/build.gradle b/sdks/java/io/iceberg/build.gradle index a186ef8000cf..11e71f340ec1 100644 --- a/sdks/java/io/iceberg/build.gradle +++ b/sdks/java/io/iceberg/build.gradle @@ -171,17 +171,13 @@ task integrationTest(type: Test) { ]) // Multiple-bucket BigLake REST catalog used by RESTCatalogBLMSIT and AddFilesIT (see // BigLakeTestCatalog). Locations: the catalog's default location first, then a restricted - // location in another bucket. The cross-region location must be a bucket in a different - // region than the catalog. Override for runs against another project's catalog: + // location in another bucket. Override for runs against another project's catalog: // -PbiglakeWarehouse=bl://projects/PROJECT/catalogs/CATALOG // -PbiglakeLocations=gs://default-bucket/path,gs://other-bucket/path - // -PbiglakeCrossRegionLocation=gs://bucket-in-another-region/path systemProperty "beam.iceberg.biglake.warehouse", project.findProperty('biglakeWarehouse') ?: 'bl://projects/apache-beam-testing/catalogs/beam-lakehouse-it' systemProperty "beam.iceberg.biglake.locations", project.findProperty('biglakeLocations') ?: 'gs://beam-lakehouse-it,gs://beam-lakehouse-it-added-path' - systemProperty "beam.iceberg.biglake.crossRegionLocation", - project.findProperty('biglakeCrossRegionLocation') ?: 'gs://managed-iceberg-biglake-its/biglake_cross_region' // Connection + storage root for BigQueryManagedTableCrossEngineIT. if (project.findProperty('bqImtConnection') != null) { systemProperty "beam.bq.imt.connection", project.findProperty('bqImtConnection') diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/AddFilesIT.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/AddFilesIT.java index eea844222bbd..d2874d6df273 100644 --- a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/AddFilesIT.java +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/AddFilesIT.java @@ -27,7 +27,6 @@ import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertNotEquals; -import static org.junit.Assert.assertThrows; import static org.junit.Assert.assertTrue; import com.google.api.services.storage.model.StorageObject; @@ -112,7 +111,6 @@ public class AddFilesIT { private static final String DATA_LOCATION = BigLakeTestCatalog.defaultLocation(); private static final String DATA_BUCKET = BigLakeTestCatalog.bucketOf(DATA_LOCATION); private static final String DATA_PREFIX = BigLakeTestCatalog.prefixOf(DATA_LOCATION); - private static final String CROSS_REGION_LOCATION = BigLakeTestCatalog.CROSS_REGION_LOCATION; private static final String PROJECT = TestPipeline.testingPipelineOptions().as(GcpOptions.class).getProject(); @Rule public TestName testName = new TestName(); @@ -237,19 +235,12 @@ public void cleanup() { LOG.warn("Failed to clean up Iceberg catalog", e); } - deleteBlobs(DATA_BUCKET, dirName); - deleteBlobs( - BigLakeTestCatalog.bucketOf(CROSS_REGION_LOCATION), - BigLakeTestCatalog.prefixOf(CROSS_REGION_LOCATION) + "/" + dirName); - } - - private void deleteBlobs(String bucket, String prefix) { try { Iterable blobs = - storage.list(bucket, Storage.BlobListOption.prefix(prefix)).getValues(); + storage.list(DATA_BUCKET, Storage.BlobListOption.prefix(dirName)).getValues(); blobs.forEach(b -> storage.delete(b.getBlobId())); } catch (Exception e) { - LOG.warn("Failed to clean up gs://{}/{}", bucket, prefix, e); + LOG.warn("Failed to clean up GCS bucket", e); } } @@ -428,57 +419,6 @@ public void testBatchParquetImportToUIT() throws IOException { testBatchParquetImport(true); } - /** - * The catalog does not police where added files live: data in a bucket it does not manage, even - * in another region, registers and reads fine through Iceberg. BigQuery, however, cannot read - * data files outside the catalog's region, so the cross-engine read of such a table fails. (Files - * in an unmanaged bucket in the catalog's region read fine from BigQuery.) Documents the - * trade-off of adding files in place from a cross-region bucket. - */ - @Test - public void testBatchParquetImportFromCrossRegionBucket() throws IOException { - String crossRegionBucket = BigLakeTestCatalog.bucketOf(CROSS_REGION_LOCATION); - for (String location : BigLakeTestCatalog.LOCATIONS) { - assertNotEquals( - "Test needs a bucket outside the catalog: " + CROSS_REGION_LOCATION, - BigLakeTestCatalog.bucketOf(location), - crossRegionBucket); - } - List writtenFilePaths = - writeParquetFiles( - crossRegionBucket, BigLakeTestCatalog.prefixOf(CROSS_REGION_LOCATION) + "/" + dirName); - - Pipeline p = Pipeline.create(); - PCollectionRowTuple tuple = - p.apply(Create.of(writtenFilePaths)) - .apply( - new AddFiles( - IcebergCatalogConfig.builder().setCatalogProperties(BIGLAKE_PROPS).build(), - destTableId.toString(), - null, - PARTITION_FIELDS, - null, - TABLE_PROPS, - null, - null)); - PAssert.that(tuple.get("errors")).empty(); - p.run().waitUntilFinish(); - - assertTrue(checkTableHasRegisteredParquetFiles(writtenFilePaths)); - checkRecordsInDestinationTable(/* alsoCheckWithBigQueryIO= */ false); - - Pipeline bq = Pipeline.create(); - bq.apply( - Managed.read(Managed.BIGQUERY) - .withConfig( - ImmutableMap.of( - "table", - BigLakeTestCatalog.bigQueryTableSpec( - destTableId.namespace().toString(), destTableId.name())))) - .getSinglePCollection(); - assertThrows(Pipeline.PipelineExecutionException.class, () -> bq.run().waitUntilFinish()); - } - /** * The destination table lives in the catalog's additional location (a second bucket) while the * source parquet files stay in the default one. BigLake pins tables under their namespace's @@ -525,14 +465,10 @@ public void testBatchParquetImportToTableInAdditionalLocation() throws IOExcepti checkRecordsInDestinationTable(/* alsoCheckWithBigQueryIO= */ true); } + /** Writes TEST_ROWS as parquet under the test's data dir and returns the written file paths. */ private List writeParquetFiles() throws IOException { - return writeParquetFiles(DATA_BUCKET, dirName); - } - - /** Writes TEST_ROWS as parquet under gs://{bucket}/{dir}/ and returns the written file paths. */ - private List writeParquetFiles(String bucket, String dir) throws IOException { - String parquetDir = format("gs://%s/%s/", bucket, dir); - String tempDir = format("gs://%s/%s-tmp/", bucket, dir); + String parquetDir = format("gs://%s/%s/", DATA_BUCKET, dirName); + String tempDir = format("gs://%s/%s-tmp/", DATA_BUCKET, dirName); LOG.info("Writing records to the parquet dir"); Pipeline q = Pipeline.create(); org.apache.avro.Schema avroSchema = AvroUtils.toAvroSchema(ROW_SCHEMA); @@ -555,7 +491,7 @@ record -> q.run().waitUntilFinish(); GcsUtil gcsUtil = TestPipeline.testingPipelineOptions().as(GcsOptions.class).getGcsUtil(); - Iterable objects = gcsUtil.listObjects(bucket, dir, null).getItems(); + Iterable objects = gcsUtil.listObjects(DATA_BUCKET, dirName, null).getItems(); List writtenFilePaths = Lists.newArrayList(objects).stream() .map(o -> format("gs://%s/%s", o.getBucket(), o.getName())) diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/BigLakeTestCatalog.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/BigLakeTestCatalog.java index 061babb7d0d5..b1f43bd996c0 100644 --- a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/BigLakeTestCatalog.java +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/BigLakeTestCatalog.java @@ -43,8 +43,6 @@ *
  • {@code beam.iceberg.biglake.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 - *
  • {@code beam.iceberg.biglake.crossRegionLocation}: a {@code gs://} prefix in a bucket - * located in a different region than the catalog * */ public final class BigLakeTestCatalog { @@ -60,14 +58,6 @@ public final class BigLakeTestCatalog { .omitEmptyStrings() .split(requiredProperty("beam.iceberg.biglake.locations"))); - /** - * A location in a bucket the catalog does not manage and that is in a different region than the - * catalog (BigQuery cannot read cross-region data files; see - * AddFilesIT#testBatchParquetImportFromCrossRegionBucket). - */ - public static final String CROSS_REGION_LOCATION = - requiredProperty("beam.iceberg.biglake.crossRegionLocation"); - /** Catalog id, which is also the second segment of BigQuery's 4-part table reference. */ public static final String CATALOG_ID = parseCatalogId(WAREHOUSE); From cdd7a9c4c9d83358128b8c48f12fb8a8d41c5a83 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 27 Aug 2026 11:08:07 -0400 Subject: [PATCH 4/4] fix tests, cleanup test artifacts --- .../beam/sdk/io/iceberg/AddFilesIT.java | 17 +++---- .../sdk/io/iceberg/BigLakeTestCatalog.java | 50 +++++++++++++++++++ .../iceberg/catalog/IcebergCatalogBaseIT.java | 13 +++-- .../io/iceberg/catalog/RESTCatalogBLMSIT.java | 9 ++++ 4 files changed, 73 insertions(+), 16 deletions(-) diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/AddFilesIT.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/AddFilesIT.java index d2874d6df273..b12075e99de5 100644 --- a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/AddFilesIT.java +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/AddFilesIT.java @@ -190,9 +190,14 @@ public void setup() throws IOException { } salt = System.currentTimeMillis(); + // Object-name prefix of this test's parquet files; DATA_PREFIX is empty for a bare bucket. dirName = format( - "%s/%s-%s/%s", DATA_PREFIX, getClass().getSimpleName(), salt, testName.getMethodName()); + "%s%s-%s/%s", + DATA_PREFIX.isEmpty() ? "" : DATA_PREFIX + "/", + getClass().getSimpleName(), + salt, + testName.getMethodName()); srcTableName = "src_" + testName.getMethodName() + "_" + salt; destTableName = "dest_" + testName.getMethodName() + "_" + salt; srcTableId = TableIdentifier.of(namespace, srcTableName); @@ -203,14 +208,8 @@ public void setup() throws IOException { catalog.createNamespace(Namespace.of(namespace)); } - private void cleanupCatalog() { - for (String name : Arrays.asList(namespace, altNamespace)) { - Namespace ns = Namespace.of(name); - if (catalog.namespaceExists(ns)) { - catalog.listTables(ns).forEach(catalog::dropTable); - catalog.dropNamespace(ns); - } - } + private void cleanupCatalog() throws IOException { + BigLakeTestCatalog.dropNamespacesAndFiles(catalog, Arrays.asList(namespace, altNamespace)); } @After diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/BigLakeTestCatalog.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/BigLakeTestCatalog.java index b1f43bd996c0..a9cdd6745e19 100644 --- a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/BigLakeTestCatalog.java +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/BigLakeTestCatalog.java @@ -19,15 +19,24 @@ 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 BigLake (Lakehouse) Iceberg REST catalog the ITs run against. @@ -124,11 +133,52 @@ public static Map catalogProperties() { .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: BigLake 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 dropNamespacesAndFiles(Catalog catalog, List namespaces) + throws IOException { + List tableLocations = new ArrayList<>(); + for (String name : namespaces) { + Namespace namespace = Namespace.of(name); + if (!((SupportsNamespaces) catalog).namespaceExists(namespace)) { + continue; + } + for (TableIdentifier identifier : catalog.listTables(namespace)) { + tableLocations.add(catalog.loadTable(identifier).location()); + catalog.dropTable(identifier); + } + ((SupportsNamespaces) catalog).dropNamespace(namespace); + } + 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 objects = + gcsUtil.listObjects(bucketOf(gcsLocation), prefixOf(gcsLocation), null).getItems(); + if (objects == null || objects.isEmpty()) { + return; + } + List 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); diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/catalog/IcebergCatalogBaseIT.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/catalog/IcebergCatalogBaseIT.java index 12da329e115f..11a6cd3306c5 100644 --- a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/catalog/IcebergCatalogBaseIT.java +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/catalog/IcebergCatalogBaseIT.java @@ -238,15 +238,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 objects = - gcsUtil - .listObjects( - path.getBucket(), - getClass().getSimpleName() + "/" + path.getFileName().toString(), - null) - .getItems(); + List objects = 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 diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/catalog/RESTCatalogBLMSIT.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/catalog/RESTCatalogBLMSIT.java index 95442455132f..1df2b145fde4 100644 --- a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/catalog/RESTCatalogBLMSIT.java +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/catalog/RESTCatalogBLMSIT.java @@ -81,6 +81,15 @@ public String bigQueryTableSpec(String tableId) { identifier.namespace().toString(), identifier.name()); } + @Override + public void catalogCleanup(List namespaces) throws IOException { + List names = new ArrayList<>(); + for (Namespace namespace : namespaces) { + names.add(namespace.toString()); + } + BigLakeTestCatalog.dropNamespacesAndFiles(catalog, names); + } + @Override public Catalog createCatalog() { RESTCatalog restCatalog = new RESTCatalog();