|
| 1 | +/* |
| 2 | + * Licensed to the Apache Software Foundation (ASF) under one |
| 3 | + * or more contributor license agreements. See the NOTICE file |
| 4 | + * distributed with this work for additional information |
| 5 | + * regarding copyright ownership. The ASF licenses this file |
| 6 | + * to you under the Apache License, Version 2.0 (the |
| 7 | + * "License"); you may not use this file except in compliance |
| 8 | + * with the License. You may obtain a copy of the License at |
| 9 | + * |
| 10 | + * http://www.apache.org/licenses/LICENSE-2.0 |
| 11 | + * |
| 12 | + * Unless required by applicable law or agreed to in writing, software |
| 13 | + * distributed under the License is distributed on an "AS IS" BASIS, |
| 14 | + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 15 | + * See the License for the specific language governing permissions and |
| 16 | + * limitations under the License. |
| 17 | + */ |
| 18 | +package org.apache.beam.sdk.io.iceberg; |
| 19 | + |
| 20 | +import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument; |
| 21 | + |
| 22 | +import com.google.api.services.storage.model.StorageObject; |
| 23 | +import java.io.IOException; |
| 24 | +import java.util.ArrayList; |
| 25 | +import java.util.List; |
| 26 | +import java.util.Map; |
| 27 | +import java.util.regex.Matcher; |
| 28 | +import java.util.regex.Pattern; |
| 29 | +import org.apache.beam.sdk.extensions.gcp.options.GcpOptions; |
| 30 | +import org.apache.beam.sdk.extensions.gcp.options.GcsOptions; |
| 31 | +import org.apache.beam.sdk.extensions.gcp.util.GcsUtil; |
| 32 | +import org.apache.beam.sdk.testing.TestPipeline; |
| 33 | +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Splitter; |
| 34 | +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; |
| 35 | +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; |
| 36 | +import org.apache.iceberg.catalog.Catalog; |
| 37 | +import org.apache.iceberg.catalog.Namespace; |
| 38 | +import org.apache.iceberg.catalog.SupportsNamespaces; |
| 39 | +import org.apache.iceberg.catalog.TableIdentifier; |
| 40 | + |
| 41 | +/** |
| 42 | + * Test-side description of the Lakehouse Iceberg REST catalog the ITs run against. |
| 43 | + * |
| 44 | + * <p>The catalog is a multiple-bucket catalog addressed by a {@code |
| 45 | + * bl://projects/PROJECT/catalogs/CATALOG} warehouse and allowed to store resources under a fixed |
| 46 | + * set of Cloud Storage locations: the catalog's default location plus any restricted locations. |
| 47 | + * Configured (defaults and overrides) by the integrationTest task in build.gradle through the |
| 48 | + * system properties |
| 49 | + * |
| 50 | + * <ul> |
| 51 | + * <li>{@code beam.iceberg.lakehouse.warehouse}: the {@code bl://} warehouse URI |
| 52 | + * <li>{@code beam.iceberg.lakehouse.locations}: comma-separated {@code gs://} prefixes the |
| 53 | + * catalog may write to; the first is the catalog's default location, the rest are additional |
| 54 | + * restricted locations |
| 55 | + * </ul> |
| 56 | + */ |
| 57 | +public final class LakehouseTestCatalog { |
| 58 | + private static final Pattern WAREHOUSE_PATTERN = |
| 59 | + Pattern.compile("bl://projects/([^/]+)/catalogs/([^/]+)"); |
| 60 | + |
| 61 | + public static final String WAREHOUSE = requiredProperty("beam.iceberg.lakehouse.warehouse"); |
| 62 | + |
| 63 | + public static final List<String> LOCATIONS = |
| 64 | + ImmutableList.copyOf( |
| 65 | + Splitter.on(',') |
| 66 | + .trimResults() |
| 67 | + .omitEmptyStrings() |
| 68 | + .split(requiredProperty("beam.iceberg.lakehouse.locations"))); |
| 69 | + |
| 70 | + /** Catalog id, which is also the second segment of BigQuery's 4-part table reference. */ |
| 71 | + public static final String CATALOG_ID = parseCatalogId(WAREHOUSE); |
| 72 | + |
| 73 | + private static final String PROJECT = |
| 74 | + TestPipeline.testingPipelineOptions().as(GcpOptions.class).getProject(); |
| 75 | + |
| 76 | + private LakehouseTestCatalog() {} |
| 77 | + |
| 78 | + private static String requiredProperty(String name) { |
| 79 | + String value = System.getProperty(name); |
| 80 | + checkArgument( |
| 81 | + value != null && !value.isEmpty(), |
| 82 | + "System property %s is not set; run through the integrationTest gradle task, which" |
| 83 | + + " sets it, or pass -D%s=...", |
| 84 | + name, |
| 85 | + name); |
| 86 | + return value; |
| 87 | + } |
| 88 | + |
| 89 | + private static String parseCatalogId(String warehouse) { |
| 90 | + // Legacy single-bucket catalogs are addressed by their bucket and named after it. |
| 91 | + if (warehouse.startsWith("gs://")) { |
| 92 | + return bucketOf(warehouse); |
| 93 | + } |
| 94 | + Matcher matcher = WAREHOUSE_PATTERN.matcher(warehouse); |
| 95 | + checkArgument( |
| 96 | + matcher.matches(), |
| 97 | + "Expected a bl://projects/PROJECT/catalogs/CATALOG (or legacy gs://BUCKET) warehouse, got '%s'", |
| 98 | + warehouse); |
| 99 | + return matcher.group(2); |
| 100 | + } |
| 101 | + |
| 102 | + /** The catalog's default location; tables land here unless created with an explicit location. */ |
| 103 | + public static String defaultLocation() { |
| 104 | + return LOCATIONS.get(0); |
| 105 | + } |
| 106 | + |
| 107 | + /** A restricted location outside the default one, i.e. in a second bucket. */ |
| 108 | + public static String additionalLocation() { |
| 109 | + checkArgument( |
| 110 | + LOCATIONS.size() >= 2, |
| 111 | + "beam.iceberg.lakehouse.locations must list at least two locations, got %s", |
| 112 | + LOCATIONS); |
| 113 | + return LOCATIONS.get(1); |
| 114 | + } |
| 115 | + |
| 116 | + public static String bucketOf(String gcsLocation) { |
| 117 | + checkArgument(gcsLocation.startsWith("gs://"), "Not a gs:// location: %s", gcsLocation); |
| 118 | + String withoutScheme = gcsLocation.substring("gs://".length()); |
| 119 | + int slash = withoutScheme.indexOf('/'); |
| 120 | + return slash < 0 ? withoutScheme : withoutScheme.substring(0, slash); |
| 121 | + } |
| 122 | + |
| 123 | + /** Object-name prefix (no bucket, no leading slash) of a {@code gs://bucket/path} location. */ |
| 124 | + public static String prefixOf(String gcsLocation) { |
| 125 | + String withoutScheme = gcsLocation.substring("gs://".length()); |
| 126 | + int slash = withoutScheme.indexOf('/'); |
| 127 | + return slash < 0 ? "" : withoutScheme.substring(slash + 1); |
| 128 | + } |
| 129 | + |
| 130 | + public static Map<String, String> catalogProperties() { |
| 131 | + return ImmutableMap.<String, String>builder() |
| 132 | + .put("type", "rest") |
| 133 | + .put("uri", "https://biglake.googleapis.com/iceberg/v1/restcatalog") |
| 134 | + .put("warehouse", WAREHOUSE) |
| 135 | + .put("header.x-goog-user-project", PROJECT) |
| 136 | + // Required by catalogs in vended-credentials mode; ignored in end-user mode. |
| 137 | + .put("header.X-Iceberg-Access-Delegation", "vended-credentials") |
| 138 | + .put("io-impl", "org.apache.iceberg.gcp.gcs.GCSFileIO") |
| 139 | + .put("rest.auth.type", "org.apache.iceberg.gcp.auth.GoogleAuthManager") |
| 140 | + .build(); |
| 141 | + } |
| 142 | + |
| 143 | + /** |
| 144 | + * Drops every table in the namespaces and then the namespaces themselves, and deletes the tables' |
| 145 | + * files from Cloud Storage. |
| 146 | + */ |
| 147 | + public static void dropNamespacesAndFiles(Catalog catalog, List<String> namespaces) |
| 148 | + throws IOException { |
| 149 | + for (String name : namespaces) { |
| 150 | + Namespace namespace = Namespace.of(name); |
| 151 | + if (!((SupportsNamespaces) catalog).namespaceExists(namespace)) { |
| 152 | + continue; |
| 153 | + } |
| 154 | + dropTablesAndFiles(catalog, name); |
| 155 | + ((SupportsNamespaces) catalog).dropNamespace(namespace); |
| 156 | + } |
| 157 | + } |
| 158 | + |
| 159 | + /** |
| 160 | + * Drops every table in the namespace and deletes the tables' files from Cloud Storage: Lakehouse |
| 161 | + * keeps a dropped table's data and metadata (even with purge), and table locations carry a random |
| 162 | + * suffix, so they are captured before the drop. |
| 163 | + */ |
| 164 | + public static void dropTablesAndFiles(Catalog catalog, String namespaceName) throws IOException { |
| 165 | + Namespace namespace = Namespace.of(namespaceName); |
| 166 | + if (!((SupportsNamespaces) catalog).namespaceExists(namespace)) { |
| 167 | + return; |
| 168 | + } |
| 169 | + List<String> tableLocations = new ArrayList<>(); |
| 170 | + for (TableIdentifier identifier : catalog.listTables(namespace)) { |
| 171 | + tableLocations.add(catalog.loadTable(identifier).location()); |
| 172 | + catalog.dropTable(identifier); |
| 173 | + } |
| 174 | + for (String location : tableLocations) { |
| 175 | + deleteObjects(location); |
| 176 | + } |
| 177 | + } |
| 178 | + |
| 179 | + /** Deletes every object under a {@code gs://bucket/prefix} location. */ |
| 180 | + public static void deleteObjects(String gcsLocation) throws IOException { |
| 181 | + GcsUtil gcsUtil = TestPipeline.testingPipelineOptions().as(GcsOptions.class).getGcsUtil(); |
| 182 | + List<StorageObject> objects = |
| 183 | + gcsUtil.listObjects(bucketOf(gcsLocation), prefixOf(gcsLocation), null).getItems(); |
| 184 | + if (objects == null || objects.isEmpty()) { |
| 185 | + return; |
| 186 | + } |
| 187 | + List<String> paths = new ArrayList<>(); |
| 188 | + for (StorageObject object : objects) { |
| 189 | + paths.add("gs://" + object.getBucket() + "/" + object.getName()); |
| 190 | + } |
| 191 | + gcsUtil.remove(paths); |
| 192 | + } |
| 193 | + |
| 194 | + /** BigQuery's 4-part {@code project.catalog.namespace.table} reference. */ |
| 195 | + public static String bigQueryTableSpec(String namespace, String table) { |
| 196 | + return String.format("%s.%s.%s.%s", PROJECT, CATALOG_ID, namespace, table); |
| 197 | + } |
| 198 | +} |
0 commit comments