From f6956de0b6854ae52c603939a2afb9aa8bf18881 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 21 Jul 2026 09:42:56 -0400 Subject: [PATCH 1/2] initial --- CHANGES.md | 1 + .../sdk/io/gcp/bigquery/BigQueryHelpers.java | 80 +++- .../beam/sdk/io/gcp/bigquery/BigQueryIO.java | 13 + .../bigquery/BigQueryStorageSourceBase.java | 7 +- .../bigquery/BigQueryStorageTableSource.java | 3 +- .../io/gcp/bigquery/BigQueryTableSource.java | 5 + .../io/gcp/bigquery/BigQueryHelpersTest.java | 359 +++++++++++++++++- sdks/java/io/iceberg/build.gradle | 4 + .../beam/sdk/io/iceberg/AddFilesIT.java | 71 +++- .../catalog/BigQueryMetastoreCatalogIT.java | 10 + .../iceberg/catalog/IcebergCatalogBaseIT.java | 261 +++++++++++++ .../io/iceberg/catalog/RESTCatalogBLMSIT.java | 21 +- 12 files changed, 817 insertions(+), 18 deletions(-) diff --git a/CHANGES.md b/CHANGES.md index d853314a0ad3..dacc90a60e6f 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -67,6 +67,7 @@ * Upgraded Iceberg dependency to 1.11.0 (Java) ([#38925](https://github.com/apache/beam/issues/38925)). * Support for X source added (Java/Python) ([#X](https://github.com/apache/beam/issues/X)). * Add ArrowFlight IO (Java) ([#20116](https://github.com/apache/beam/issues/20116)). +* BigQueryIO now supports reading BigQuery Lakehouse runtime catalog (BigLake metastore) Iceberg tables with the Storage Read API, using 4-part `project.catalog.namespace.table` identifiers (or a `TableReference` with a composite `catalog.namespace` dataset id). Previously such references were silently mis-parsed, and tables without storage statistics failed with a `NullPointerException` (Java) ([#39597](https://github.com/apache/beam/issues/39597)) . ## New Features / Improvements diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryHelpers.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryHelpers.java index 55c703438f02..20a6724dad49 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryHelpers.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryHelpers.java @@ -45,6 +45,7 @@ import java.util.Map; import java.util.UUID; import java.util.regex.Matcher; +import java.util.regex.Pattern; import java.util.stream.Collectors; import org.apache.beam.sdk.extensions.gcp.util.BackOffAdapter; import org.apache.beam.sdk.io.FileSystems; @@ -79,6 +80,9 @@ public class BigQueryHelpers { private static final Logger LOG = LoggerFactory.getLogger(BigQueryHelpers.class); + /** Matches valid Project ID patterns. */ + private static final Pattern PROJECT_NAME_SEGMENT_PATTERN = Pattern.compile("[-a-z0-9]*[a-z0-9]"); + // Given a potential failure and a current job-id, return the next job-id to be used on retry. // Algorithm is as follows (given input of job_id_prefix-N) // If BigQuery has no status for job_id_prefix-n, we should retry with the same id. @@ -458,6 +462,13 @@ static List getOrCreateMapListValue(Map> map, K key) { * Parse a table specification in the form {@code "[project_id]:[dataset_id].[table_id]"} or * {@code "[project_id].[dataset_id].[table_id]"} or {@code "[dataset_id].[table_id]"}. * + *

Lakehouse runtime catalog (BigLake metastore) tables are referenced with four parts, {@code + * "[project_id].[catalog_id].[namespace_id].[table_id]"} (or {@code + * "[project_id]:[catalog_id].[namespace_id].[table_id]"}); these parse to a composite {@code + * "[catalog_id].[namespace_id]"} dataset id, which is the form the BigQuery APIs accept for such + * tables. More generally, when a specification contains more than three segments, everything + * between the project id and the final (table) segment becomes the dataset id. + * *

If the project id is omitted, the default project id is used. */ @SuppressWarnings({ @@ -471,14 +482,75 @@ public static TableReference parseTableSpec(String tableSpec) { "Table specification [%s] is not in one of the expected formats (" + " [project_id]:[dataset_id].[table_id]," + " [project_id].[dataset_id].[table_id]," - + " [dataset_id].[table_id])", + + " [dataset_id].[table_id]," + + " [project_id]:[catalog_id].[namespace_id].[table_id]," + + " [project_id].[catalog_id].[namespace_id].[table_id])", tableSpec)); } - TableReference ref = new TableReference(); - ref.setProjectId(match.group("PROJECT")); + // Table ids cannot contain '.', so the table is always the segment after + // the last dot. + int lastDot = tableSpec.lastIndexOf('.'); + String table = tableSpec.substring(lastDot + 1); + String prefix = tableSpec.substring(0, lastDot); + + String project = null; + String dataset; + long colonCount = prefix.chars().filter(c -> c == ':').count(); + if (colonCount == 0) { + // No colon means the purely dotted form ("p.d.t", "d.t", "p.catalog.ns.t"): the + // leading segment is the project id when it is a plausible project id. + // (Dataset ids may contain characters such as '_' that project ids may + // not, in which case the whole prefix is the dataset id.) + // The firstDot < length-1 guard keeps degenerate trailing-dot specs + // ("pp..t", accepted by the character-set gate with dataset "pp.") + // instead of producing an empty dataset id. + int firstDot = prefix.indexOf('.'); + if (firstDot >= 0 + && firstDot < prefix.length() - 1 + && BigQueryIO.PROJECT_ID_PATTERN.matcher(prefix.substring(0, firstDot)).matches()) { + project = prefix.substring(0, firstDot); + dataset = prefix.substring(firstDot + 1); + } else { + dataset = prefix; + } + } else if (colonCount == 1) { + // One colon ("p:d.t", "p:catalog.ns.t", "example.com:proj.ds.t"). If the + // project part is dotted, it is a legacy domain-scoped id written with + // a '.' separator after the project: the first dataset segment completes + // the project id, and any remaining middle segments bind as a (possibly + // composite) dataset. (Domain-scoped project names cannot contain dots) + int colon = prefix.indexOf(':'); + project = prefix.substring(0, colon); + dataset = prefix.substring(colon + 1); + int firstDot = dataset.indexOf('.'); + // Absorb the first dataset segment into a dotted (domain-scoped) project + // only when the split leaves a non-empty dataset. + if (firstDot >= 0 + && firstDot < dataset.length() - 1 + && project.indexOf('.') >= 0 + && PROJECT_NAME_SEGMENT_PATTERN.matcher(dataset.substring(0, firstDot)).matches()) { + project = project + ":" + dataset.substring(0, firstDot); + dataset = dataset.substring(firstDot + 1); + } + } else { + // Two colons - the last colon is an explicit project terminator. This is + // the canonical spelling for a domain-scoped project, whose id itself + // contains a colon ("example.com:proj:ds.t"), including with a composite + // Lakehouse catalog dataset ("example.com:proj:catalog.ns.t"). Both + // domain-scoped spellings keep toTableSpec/parseTableSpec a round trip + // for composite dataset ids. (More than two colons cannot form a valid + // reference - project ids contain at most one colon, but such specs pass + // the character-set gate, so they bind here too and the impossible + // project id is rejected by the service.) + int lastColon = prefix.lastIndexOf(':'); + project = prefix.substring(0, lastColon); + dataset = prefix.substring(lastColon + 1); + } - return ref.setDatasetId(match.group("DATASET")).setTableId(match.group("TABLE")); + TableReference ref = new TableReference(); + ref.setProjectId(project); + return ref.setDatasetId(dataset).setTableId(table); } @SuppressWarnings({ diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java index c2ac4efb5b5d..56b0200e4ba8 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java @@ -589,6 +589,13 @@ public class BigQueryIO { */ private static final String PROJECT_ID_REGEXP = "[a-z][-a-z0-9:.]{0,61}[a-z0-9]"; + /** + * Matches a whole string against {@link #PROJECT_ID_REGEXP}. Used by {@link + * BigQueryHelpers#parseTableSpec} to decide whether the leading segment of a dotted table + * specification is a project id. + */ + static final Pattern PROJECT_ID_PATTERN = Pattern.compile(PROJECT_ID_REGEXP); + /** Regular expression that matches Dataset IDs. */ private static final String DATASET_REGEXP = "[-\\w.]{1,1024}"; @@ -604,6 +611,12 @@ public class BigQueryIO { /** * Matches table specifications in the form {@code "[project_id]:[dataset_id].[table_id]"}, {@code * "[project_id].[dataset_id].[table_id]"}, or {@code "[dataset_id].[table_id]"}. + * + *

This pattern is used for syntactic validation only; the assignment of the matched string's + * segments to the project/dataset/table fields is done by {@link BigQueryHelpers#parseTableSpec}, + * which additionally understands 4-part Lakehouse runtime catalog (BigLake metastore) references + * {@code "[project_id].[catalog_id].[namespace_id].[table_id]"}, mapping them to a composite + * {@code "[catalog_id].[namespace_id]"} dataset id. */ private static final String DATASET_TABLE_REGEXP = String.format( diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryStorageSourceBase.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryStorageSourceBase.java index 402de619464a..7f14985fc90f 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryStorageSourceBase.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryStorageSourceBase.java @@ -152,7 +152,12 @@ public List> split( int streamCount = 0; if (!bqOptions.getEnableStorageReadApiV2()) { if (desiredBundleSizeBytes > 0) { - long tableSizeBytes = (targetTable != null) ? targetTable.getNumBytes() : 0; + // numBytes is null for tables that don't report storage statistics, e.g. Lakehouse + // runtime catalog (BigLake metastore) tables. + long tableSizeBytes = + (targetTable != null && targetTable.getNumBytes() != null) + ? targetTable.getNumBytes() + : 0; streamCount = (int) Math.min(tableSizeBytes / desiredBundleSizeBytes, MAX_SPLIT_COUNT); } diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryStorageTableSource.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryStorageTableSource.java index 8b7240158dc1..6887e663b39f 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryStorageTableSource.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryStorageTableSource.java @@ -163,12 +163,13 @@ public void populateDisplayData(DisplayData.Builder builder) { @Override public long getEstimatedSizeBytes(PipelineOptions options) throws Exception { Table table = getTargetTable(options.as(BigQueryOptions.class)); - if (table != null) { + if (table != null && table.getNumBytes() != null) { return table.getNumBytes(); } // If the table does not exist, then it will be null. // Avoid the NullPointerException here, allow a more meaningful table "not_found" // error to be shown to the user, upon table read. + // Lakehouse runtime catalog (BigLake metastore) tables exist but report no numBytes. return 0; } diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryTableSource.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryTableSource.java index 1b6aedf8cb17..ba9a8a15610d 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryTableSource.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryTableSource.java @@ -82,6 +82,11 @@ public synchronized long getEstimatedSizeBytes(PipelineOptions options) throws E } Long numBytes = table.getNumBytes(); + if (numBytes == null) { + // Tables that don't report storage statistics, e.g. Lakehouse runtime catalog + // (BigLake metastore) tables. + numBytes = 0L; + } if (table.getStreamingBuffer() != null && table.getStreamingBuffer().getEstimatedBytes() != null) { numBytes += table.getStreamingBuffer().getEstimatedBytes().longValue(); diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryHelpersTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryHelpersTest.java index 0d36d7bb46d0..e6039781bc3d 100644 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryHelpersTest.java +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryHelpersTest.java @@ -18,6 +18,7 @@ package org.apache.beam.sdk.io.gcp.bigquery; import static org.junit.Assert.assertEquals; +import static org.junit.Assert.fail; import com.google.api.client.util.Data; import com.google.api.services.bigquery.model.Clustering; @@ -28,6 +29,7 @@ import com.google.api.services.bigquery.model.TableReference; import com.google.api.services.bigquery.model.TableRow; import java.util.Arrays; +import java.util.List; import java.util.Optional; import java.util.Random; import java.util.Set; @@ -101,13 +103,368 @@ public void testTableParsing_noProjectId() { assertEquals("table_name", ref.getTableId()); } + @Test + public void testTableParsing_lakehouseCatalogDotted() { + // 4-part Lakehouse runtime catalog reference: project.catalog.namespace.table. The + // catalog+namespace form a composite dataset id, including when the catalog name uses the + // GCS-bucket charset (lowercase, digits, dashes). + TableReference ref = BigQueryHelpers.parseTableSpec("my-project.my-bucket-catalog.my_ns.tbl"); + assertEquals("my-project", ref.getProjectId()); + assertEquals("my-bucket-catalog.my_ns", ref.getDatasetId()); + assertEquals("tbl", ref.getTableId()); + } + + @Test + public void testTableParsing_lakehouseCatalogColon() { + TableReference ref = BigQueryHelpers.parseTableSpec("my-project:my-catalog.my_ns.tbl"); + assertEquals("my-project", ref.getProjectId()); + assertEquals("my-catalog.my_ns", ref.getDatasetId()); + assertEquals("tbl", ref.getTableId()); + } + + @Test + public void testTableParsing_lakehouseCatalogNoProject() { + // Dataset ids may contain characters that project ids may not (e.g. '_'), in which case the + // whole prefix is the (composite) dataset id. + TableReference ref = BigQueryHelpers.parseTableSpec("my_catalog.my_ns.tbl"); + assertEquals(null, ref.getProjectId()); + assertEquals("my_catalog.my_ns", ref.getDatasetId()); + assertEquals("tbl", ref.getTableId()); + } + + @Test + public void testTableParsing_multiLevelNamespace() { + // More than four segments: everything between the project and the table becomes the dataset. + TableReference ref = BigQueryHelpers.parseTableSpec("my-project.cat.ns1.ns2.tbl"); + assertEquals("my-project", ref.getProjectId()); + assertEquals("cat.ns1.ns2", ref.getDatasetId()); + assertEquals("tbl", ref.getTableId()); + } + + @Test + public void testTableParsing_lakehouseWithPartitionDecorator() { + TableReference ref = BigQueryHelpers.parseTableSpec("my-project.my-catalog.ns.tbl$20260101"); + assertEquals("my-project", ref.getProjectId()); + assertEquals("my-catalog.ns", ref.getDatasetId()); + assertEquals("tbl$20260101", ref.getTableId()); + } + + @Test + public void testTableParsing_domainScopedProjectPreserved() { + // Legacy domain-scoped projects keep their historical binding. + TableReference ref = BigQueryHelpers.parseTableSpec("example.com:project:data_set.tbl"); + assertEquals("example.com:project", ref.getProjectId()); + assertEquals("data_set", ref.getDatasetId()); + assertEquals("tbl", ref.getTableId()); + + ref = BigQueryHelpers.parseTableSpec("example.com:project.data_set.tbl"); + assertEquals("example.com:project", ref.getProjectId()); + assertEquals("data_set", ref.getDatasetId()); + assertEquals("tbl", ref.getTableId()); + } + + @Test + public void testTableParsing_domainScopedProjectWithCompositeDataset() { + // With two colons, the last colon is an explicit project terminator, so a domain-scoped + // project can address a Lakehouse catalog table: the remainder binds as a composite + // dataset. + TableReference ref = BigQueryHelpers.parseTableSpec("example.com:project:cat.ns.tbl"); + assertEquals("example.com:project", ref.getProjectId()); + assertEquals("cat.ns", ref.getDatasetId()); + assertEquals("tbl", ref.getTableId()); + + // The dotted spelling binds consistently: the first segment after the colon completes the + // domain-scoped project id; further middle segments form the composite dataset. (Project + // names cannot contain dots, so the pre-fix greedy binding of this string was invalid.) + ref = BigQueryHelpers.parseTableSpec("example.com:project.cat.ns.tbl"); + assertEquals("example.com:project", ref.getProjectId()); + assertEquals("cat.ns", ref.getDatasetId()); + assertEquals("tbl", ref.getTableId()); + } + + @Test + public void testTableParsing_partitionDecoratorColonForm() { + TableReference ref = BigQueryHelpers.parseTableSpec("my-project:my-catalog.ns.tbl$20260101"); + assertEquals("my-project", ref.getProjectId()); + assertEquals("my-catalog.ns", ref.getDatasetId()); + assertEquals("tbl$20260101", ref.getTableId()); + } + + @Test + public void testTableParsing_tableIdSpecialCharacters() { + // Table ids may contain spaces, '@', '$', dashes, and unicode letters, none of which + // affect segment binding (only '.' and ':' are structural). + TableReference ref = BigQueryHelpers.parseTableSpec("my-project.data_set.my table@x-1"); + assertEquals("my-project", ref.getProjectId()); + assertEquals("data_set", ref.getDatasetId()); + assertEquals("my table@x-1", ref.getTableId()); + + ref = BigQueryHelpers.parseTableSpec("my-project.my-catalog.ns.ग्राहक"); + assertEquals("my-project", ref.getProjectId()); + assertEquals("my-catalog.ns", ref.getDatasetId()); + assertEquals("ग्राहक", ref.getTableId()); + } + + @Test + public void testTableParsing_colonFormMultiLevelNamespace() { + TableReference ref = BigQueryHelpers.parseTableSpec("my-project:cat.ns1.ns2.tbl"); + assertEquals("my-project", ref.getProjectId()); + assertEquals("cat.ns1.ns2", ref.getDatasetId()); + assertEquals("tbl", ref.getTableId()); + } + + @Test + public void testTableParsing_domainScopedMultiLevelNamespace() { + // Single-colon domain-scoped spelling with a multi-level composite dataset: the first + // segment after the colon completes the project id; everything else up to the table binds + // as the dataset. + TableReference ref = BigQueryHelpers.parseTableSpec("example.com:proj.cat.ns1.ns2.tbl"); + assertEquals("example.com:proj", ref.getProjectId()); + assertEquals("cat.ns1.ns2", ref.getDatasetId()); + assertEquals("tbl", ref.getTableId()); + } + + @Test + public void testTableParsing_moreThanTwoColonsBindsAtLastColon() { + // More than two colons cannot form a valid reference (project ids contain at most one + // colon), but such specs pass the character-set gate; they bind at the last colon so the + // impossible project id is rejected by the service rather than producing a malformed + // dataset id. + TableReference ref = BigQueryHelpers.parseTableSpec("d1:d2:d3:data_set.tbl"); + assertEquals("d1:d2:d3", ref.getProjectId()); + assertEquals("data_set", ref.getDatasetId()); + assertEquals("tbl", ref.getTableId()); + } + + @Test + public void testTableParsing_consecutiveDotsPreservedInDataset() { + // Dataset content is preserved verbatim; degenerate consecutive dots are not collapsed + // (the service rejects the invalid dataset id; the parser must not mangle it). + TableReference ref = BigQueryHelpers.parseTableSpec("my-project:a..b.tbl"); + assertEquals("my-project", ref.getProjectId()); + assertEquals("a..b", ref.getDatasetId()); + assertEquals("tbl", ref.getTableId()); + } + + @Test + public void testTableParsing_trailingDotStaysOnLegacyBinding() { + // Degenerate trailing-dot specs pass the character-set gate with the dot inside the dataset + // id. They must keep their historical binding; the parser must never produce an empty + // dataset id. + TableReference ref = BigQueryHelpers.parseTableSpec("pp..t"); + assertEquals(null, ref.getProjectId()); + assertEquals("pp.", ref.getDatasetId()); + assertEquals("t", ref.getTableId()); + + ref = BigQueryHelpers.parseTableSpec("google.com:proj..t"); + assertEquals("google.com", ref.getProjectId()); + assertEquals("proj.", ref.getDatasetId()); + assertEquals("t", ref.getTableId()); + } + + @Test + public void testTableParsing_absorptionRequiresProjectCharset() { + // The segment absorbed into a domain-scoped project must be legal in the project charset, + // exactly what the historical greedy regex could absorb. + // '_' and uppercase were never absorbable: + TableReference ref = BigQueryHelpers.parseTableSpec("google.com:My_Cat.n.x.t"); + assertEquals("google.com", ref.getProjectId()); + assertEquals("My_Cat.n.x", ref.getDatasetId()); + assertEquals("t", ref.getTableId()); + + // ...but a single-character segment was: + ref = BigQueryHelpers.parseTableSpec("example.com:a.b.t"); + assertEquals("example.com:a", ref.getProjectId()); + assertEquals("b", ref.getDatasetId()); + assertEquals("t", ref.getTableId()); + } + + @Test + public void testTableParsing_nonProjectLeadingSegmentsFoldIntoDataset() { + // Leading segments that cannot be project ids (uppercase, '_', too short, leading '_') + // fold into the (composite) dataset with a null project, regardless of depth. + String[][] cases = { + {"My-Ds.ns.t", "My-Ds.ns"}, + {"my_cat.a.b.t", "my_cat.a.b"}, + {"a.c.n.t", "a.c.n"}, + {"_x.a.b.c.t", "_x.a.b.c"}, + }; + for (String[] c : cases) { + TableReference ref = BigQueryHelpers.parseTableSpec(c[0]); + assertEquals(c[0], null, ref.getProjectId()); + assertEquals(c[0], c[1], ref.getDatasetId()); + assertEquals(c[0], "t", ref.getTableId()); + } + } + + @Test + public void testTableParsing_domainScopedNoAbsorptionWithoutDottedDataset() { + // Dot-less post-colon remainder: nothing to absorb; the pre-colon text alone is the project. + TableReference ref = BigQueryHelpers.parseTableSpec("google.com:data_set.tbl"); + assertEquals("google.com", ref.getProjectId()); + assertEquals("data_set", ref.getDatasetId()); + assertEquals("tbl", ref.getTableId()); + } + + @Test + public void testTableParsing_twoColonsDotlessDomain() { + // The explicit-terminator rule does not require a dotted (domain-like) project. + TableReference ref = BigQueryHelpers.parseTableSpec("p1:x2:data_set.tbl"); + assertEquals("p1:x2", ref.getProjectId()); + assertEquals("data_set", ref.getDatasetId()); + assertEquals("tbl", ref.getTableId()); + + ref = BigQueryHelpers.parseTableSpec("p1:x2:cat.ns.tbl"); + assertEquals("p1:x2", ref.getProjectId()); + assertEquals("cat.ns", ref.getDatasetId()); + assertEquals("tbl", ref.getTableId()); + } + + @Test + public void testTableParsing_uppercaseCatalogUnderNormalProject() { + TableReference ref = BigQueryHelpers.parseTableSpec("my-project:My_Cat.ns.tbl"); + assertEquals("my-project", ref.getProjectId()); + assertEquals("My_Cat.ns", ref.getDatasetId()); + assertEquals("tbl", ref.getTableId()); + } + + @Test + public void testTableParsingError_rejectedForms() { + String[] rejected = { + "p1:t2", // colon form requires dataset.table + "ds.", // empty table + ".t", // empty dataset + "p1:d2.t3.", // trailing separator + "p1::ds.t", // colon cannot follow the separator colon + "MyProj:ds.t", // uppercase cannot precede a colon + }; + for (String spec : rejected) { + try { + BigQueryHelpers.parseTableSpec(spec); + fail("Expected IllegalArgumentException for spec: " + spec); + } catch (IllegalArgumentException expected) { + // expected + } + } + } + + @Test + public void testParseTableSpecIdempotentThroughToTableSpec() { + // parse(toTableSpec(parse(s))) == parse(s) for every accepted spec, the string-level + // consequence of the round-trip fixed point, covering all rebinding families. + String[] specs = { + "my-project:data_set.tbl", + "my-project.cat.ns.tbl", + "my-project:cat.ns.tbl", + "my-project:cat.ns1.ns2.tbl", + "example.com:proj:cat.ns.tbl", + "example.com:proj.cat.ns.tbl", + "my_cat.ns.tbl", + "my-project:My_Cat.ns.tbl", + "my-project:a..b.tbl", + "pp..t", + "google.com:proj..t", + "my-project:data_set.tbl$20260101", + }; + for (String spec : specs) { + TableReference once = BigQueryHelpers.parseTableSpec(spec); + TableReference twice = BigQueryHelpers.parseTableSpec(BigQueryHelpers.toTableSpec(once)); + assertEquals(spec, once.getProjectId(), twice.getProjectId()); + assertEquals(spec, once.getDatasetId(), twice.getDatasetId()); + assertEquals(spec, once.getTableId(), twice.getTableId()); + } + } + + @Test + public void testTableParsingError_colonIsNotATableSeparator() { + // The dataset/table separator must be a dot; a spec with no dot at all is rejected even + // if it contains colons. + thrown.expect(IllegalArgumentException.class); + BigQueryHelpers.parseTableSpec("my-project:data_set:tbl"); + } + + @Test + public void testTableParsing_projectlessCatalogSpecIsAmbiguous() { + // A project-less catalog reference is indistinguishable from project.dataset.table when the + // catalog name fits the project-id charset: the project interpretation wins. Users must + // write the full 4-part name (or use a TableReference) for such catalogs; only catalog + // names that are illegal as project ids (e.g. containing '_') parse as a composite dataset + // with the project left to be defaulted. + TableReference ref = BigQueryHelpers.parseTableSpec("my-bucket-catalog.my_ns.tbl"); + assertEquals("my-bucket-catalog", ref.getProjectId()); + assertEquals("my_ns", ref.getDatasetId()); + assertEquals("tbl", ref.getTableId()); + } + + @Test + public void testTableParsing_informationSchema() { + // INFORMATION_SCHEMA views ride along with the multi-segment rule: the pseudo-schema folds + // into the dataset id. (Neither the Storage Read API nor extract jobs support reading + // INFORMATION_SCHEMA views, so this reference form never reaches a table API; pinning the + // binding here documents that it at least keeps a valid, undotted project id.) + TableReference ref = + BigQueryHelpers.parseTableSpec("my-project.data_set.INFORMATION_SCHEMA.TABLES"); + assertEquals("my-project", ref.getProjectId()); + assertEquals("data_set.INFORMATION_SCHEMA", ref.getDatasetId()); + assertEquals("TABLES", ref.getTableId()); + } + + @Test + public void testTableParsing_shortFirstSegmentIsNotAProject() { + // A single character cannot be a project id, so the prefix folds into the dataset id + // (historical behavior). + TableReference ref = BigQueryHelpers.parseTableSpec("a.b.c"); + assertEquals(null, ref.getProjectId()); + assertEquals("a.b", ref.getDatasetId()); + assertEquals("c", ref.getTableId()); + } + + @Test + public void testToTableSpecParseTableSpecRoundTrip() { + List refs = + Arrays.asList( + new TableReference() + .setProjectId("my-project") + .setDatasetId("data_set") + .setTableId("tbl"), + new TableReference() + .setProjectId("my-project") + .setDatasetId("my-catalog.my_ns") + .setTableId("tbl"), + new TableReference() + .setProjectId("example.com:project") + .setDatasetId("data_set") + .setTableId("tbl"), + new TableReference() + .setProjectId("example.com:project") + .setDatasetId("my-catalog.my_ns") + .setTableId("tbl"), + new TableReference() + .setProjectId("my-project") + .setDatasetId("cat.ns1.ns2") + .setTableId("tbl"), + new TableReference() + .setProjectId("my-project") + .setDatasetId("my-catalog.ns") + .setTableId("tbl$20260101"), + new TableReference().setDatasetId("data_set").setTableId("tbl")); + for (TableReference ref : refs) { + TableReference reparsed = BigQueryHelpers.parseTableSpec(BigQueryHelpers.toTableSpec(ref)); + assertEquals(BigQueryHelpers.toTableSpec(ref), ref.getProjectId(), reparsed.getProjectId()); + assertEquals(BigQueryHelpers.toTableSpec(ref), ref.getDatasetId(), reparsed.getDatasetId()); + assertEquals(BigQueryHelpers.toTableSpec(ref), ref.getTableId(), reparsed.getTableId()); + } + } + @Test public void testTableParsingError0() { String expectedMessage = "Table specification [foo_bar_baz] is not in one of the expected formats (" + " [project_id]:[dataset_id].[table_id]," + " [project_id].[dataset_id].[table_id]," - + " [dataset_id].[table_id])"; + + " [dataset_id].[table_id]," + + " [project_id]:[catalog_id].[namespace_id].[table_id]," + + " [project_id].[catalog_id].[namespace_id].[table_id])"; thrown.expect(IllegalArgumentException.class); thrown.expectMessage(expectedMessage); diff --git a/sdks/java/io/iceberg/build.gradle b/sdks/java/io/iceberg/build.gradle index 983ebd07fefc..fe91533227fa 100644 --- a/sdks/java/io/iceberg/build.gradle +++ b/sdks/java/io/iceberg/build.gradle @@ -169,6 +169,10 @@ 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' // Disable Gradle cache: these ITs interact with live service that should always be considered "out of date" outputs.upToDateWhen { false } 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 54c9cb8dc93f..4964f90cc2d0 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 @@ -102,7 +102,12 @@ public class AddFilesIT { private static final Logger LOG = LoggerFactory.getLogger(AddFilesIT.class); - private static final String WAREHOUSE = "gs://managed-iceberg-biglake-its"; + // 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; private static final String PROJECT = TestPipeline.testingPipelineOptions().as(GcpOptions.class).getProject(); @Rule public TestName testName = new TestName(); @@ -131,7 +136,7 @@ public class AddFilesIT { private Storage storage; private PubsubClient pubsub; private Notification notification; - private final String namespace = getClass().getSimpleName(); + private final String namespace = getClass().getSimpleName() + "_" + System.currentTimeMillis(); private String srcTableName; private String destTableName; private TableIdentifier srcTableId; @@ -303,7 +308,7 @@ public void testStreamingImportFromExistingIcebergTable() addFilesPipeline.cancel(); // check all records are there - checkRecordsInDestinationTable(); + checkRecordsInDestinationTable(/* alsoCheckWithBigQueryIO= */ false); } /** @@ -404,11 +409,20 @@ record -> addFilesPipeline.cancel(); // check all records are there - checkRecordsInDestinationTable(); + checkRecordsInDestinationTable(/* alsoCheckWithBigQueryIO= */ false); } @Test public void testBatchParquetImport() throws IOException { + testBatchParquetImport(false); + } + + @Test + public void testBatchParquetImportToUIT() throws IOException { + testBatchParquetImport(true); + } + + private void testBatchParquetImport(boolean isUIT) throws IOException { // start with a table that does not exist String parquetDir = format("%s/%s/", WAREHOUSE, dirName); @@ -449,6 +463,11 @@ record -> // before adding, confirm the destination table still does not exist assertFalse(catalog.tableExists(destTableId)); + Map tableProps = new HashMap<>(TABLE_PROPS); + if (isUIT) { + tableProps.put("gcp.biglake.bigquery-dml.enabled", "true"); + } + // run batch AddFiles Pipeline p = Pipeline.create(); PCollectionRowTuple tuple = @@ -458,9 +477,9 @@ record -> IcebergCatalogConfig.builder().setCatalogProperties(BIGLAKE_PROPS).build(), namespace + "." + destTableName, null, - PARTITION_FIELDS, + isUIT ? null : PARTITION_FIELDS, null, - TABLE_PROPS, + tableProps, null, null)); PAssert.that(tuple.get("errors")).empty(); @@ -475,11 +494,11 @@ record -> LOG.info( "Destination table has registered all source files ({} files).", writtenFilePaths.size()); - // check all records are there - checkRecordsInDestinationTable(); + // check all records are there. + checkRecordsInDestinationTable(/* alsoCheckWithBigQueryIO= */ true); } - private void checkRecordsInDestinationTable() { + private void checkRecordsInDestinationTable(boolean alsoCheckWithBigQueryIO) { Pipeline s = Pipeline.create(); PCollection destRows = s.apply( @@ -489,9 +508,43 @@ private void checkRecordsInDestinationTable() { "table", destTableId.toString(), "catalog_properties", BIGLAKE_PROPS))) .getSinglePCollection(); PAssert.that(destRows).containsInAnyOrder(TEST_ROWS); + + if (alsoCheckWithBigQueryIO) { + // Cross-engine check: the same rows must be readable with BigQueryIO via the 4-part + // project.catalog.namespace.table reference. Rows are compared on a canonical string + // because BigQuery widens int32 (age) to INT64, so whole-row equality does not hold. + PCollection bqRows = + s.apply( + "read with BigQueryIO", + Managed.read(Managed.BIGQUERY) + .withConfig( + ImmutableMap.of( + "table", + format( + "%s.%s.%s.%s", + PROJECT, + CATALOG_NAME, + destTableId.namespace(), + destTableId.name())))) + .getSinglePCollection() + .apply( + "canonicalize bq rows", + MapElements.into(strings()).via(AddFilesIT::canonicalRecord)); + PAssert.that(bqRows) + .containsInAnyOrder( + TEST_ROWS.stream().map(AddFilesIT::canonicalRecord).collect(Collectors.toList())); + } s.run().waitUntilFinish(); } + private static String canonicalRecord(Row row) { + return String.valueOf((Object) row.getValue("id")) + + "|" + + row.getValue("name") + + "|" + + String.valueOf((Object) row.getValue("age")); + } + private boolean checkTableHasRegisteredParquetFiles(List parquetFiles) { Table destTable = catalog.loadTable(destTableId); diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/catalog/BigQueryMetastoreCatalogIT.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/catalog/BigQueryMetastoreCatalogIT.java index a34e580d29b7..a0c412596014 100644 --- a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/catalog/BigQueryMetastoreCatalogIT.java +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/catalog/BigQueryMetastoreCatalogIT.java @@ -49,6 +49,16 @@ public String type() { return "bqms"; } + @Override + public String bigQueryTableSpec(String tableId) { + // The BigQuery metastore federation surfaces Iceberg namespaces as ordinary BigQuery + // datasets, so tables use plain 3-part project.dataset.table references, with the namespace + // as the dataset. + TableIdentifier identifier = TableIdentifier.parse(tableId); + return String.format( + "%s.%s.%s", OPTIONS.getProject(), identifier.namespace(), identifier.name()); + } + @Override public Catalog createCatalog() { return CatalogUtil.loadCatalog( 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 5c28f0192a61..5c4e5c76c19b 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 @@ -55,6 +55,8 @@ import org.apache.beam.sdk.extensions.gcp.options.GcsOptions; import org.apache.beam.sdk.extensions.gcp.util.GcsUtil; import org.apache.beam.sdk.extensions.gcp.util.gcsfs.GcsPath; +import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO; +import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO.TypedRead.Method; import org.apache.beam.sdk.io.iceberg.IcebergUtils; import org.apache.beam.sdk.managed.Managed; import org.apache.beam.sdk.schemas.Schema; @@ -91,10 +93,12 @@ import org.apache.iceberg.catalog.Namespace; import org.apache.iceberg.catalog.SupportsNamespaces; import org.apache.iceberg.catalog.TableIdentifier; +import org.apache.iceberg.data.GenericRecord; import org.apache.iceberg.data.IdentityPartitionConverters; import org.apache.iceberg.data.Record; import org.apache.iceberg.data.parquet.GenericParquetReaders; import org.apache.iceberg.data.parquet.GenericParquetWriter; +import org.apache.iceberg.deletes.EqualityDeleteWriter; import org.apache.iceberg.encryption.InputFilesDecryptor; import org.apache.iceberg.io.CloseableIterable; import org.apache.iceberg.io.DataWriter; @@ -111,6 +115,7 @@ import org.joda.time.Instant; import org.joda.time.LocalDate; import org.joda.time.LocalTime; +import org.joda.time.ReadableInstant; import org.junit.After; import org.junit.Before; import org.junit.Rule; @@ -153,6 +158,18 @@ public abstract class IcebergCatalogBaseIT implements Serializable { public abstract String type(); + /** + * 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()}. + */ + public @Nullable String bigQueryTableSpec(String tableId) { + return null; + } + public void catalogSetup() { ((SupportsNamespaces) catalog).createNamespace(Namespace.of(namespace())); } @@ -592,6 +609,250 @@ public void testWriteRead() throws IOException { containsInAnyOrder(expectedRows.stream().map(RECORD_FUNC::apply).toArray())); } + /** + * Cross-engine consistency: rows written through the Iceberg catalog must be readable with + * BigQueryIO's Storage Read API using the catalog's BigQuery table reference (see {@link + * #bigQueryTableSpec(String)}). Exercises the full read, server-side projection + filtering + * push-down, and a query read with the reference embedded in SQL. + * + *

Rows are compared on a projection of fields whose types survive the Iceberg-to-BigQuery + * mapping losslessly; BigQuery widens e.g. {@code int32} to {@code INT64}, so whole-row equality + * against the Iceberg schema does not hold by design. + */ + @Test + public void testReadWithBigQueryIO() throws Exception { + String tableSpec = bigQueryTableSpec(tableId()); + assumeTrue("Catalog does not surface its tables in BigQuery", tableSpec != null); + Table table = catalog.createTable(TableIdentifier.parse(tableId()), ICEBERG_SCHEMA); + List expectedRows = populateTable(table); + int intCutoff = numRecords() / 2; + + List expectedFull = + expectedRows.stream() + .map( + r -> + canonical( + r.getString("str"), + r.getInt64("modulo_5"), + r.getBoolean("bool_field"), + r.getInt32("int_field"))) + .collect(Collectors.toList()); + List expectedFiltered = + expectedRows.stream() + .filter(r -> checkStateNotNull(r.getInt32("int_field")) < intCutoff) + .map(r -> canonical(r.getString("str"), r.getInt32("int_field"))) + .collect(Collectors.toList()); + + PCollection fullRead = + pipeline + .apply( + "BQ direct read", + BigQueryIO.readTableRows().from(tableSpec).withMethod(Method.DIRECT_READ)) + .apply( + "canonicalize full", + MapElements.into(TypeDescriptors.strings()) + .via( + tr -> + canonical( + tr.get("str"), + tr.get("modulo_5"), + tr.get("bool_field"), + tr.get("int_field")))); + PAssert.that(fullRead).containsInAnyOrder(expectedFull); + + PCollection pushdownRead = + pipeline + .apply( + "BQ pushdown read", + BigQueryIO.readTableRows() + .from(tableSpec) + .withMethod(Method.DIRECT_READ) + .withSelectedFields(Arrays.asList("str", "int_field")) + .withRowRestriction("int_field < " + intCutoff)) + .apply( + "canonicalize pushdown", + MapElements.into(TypeDescriptors.strings()) + .via(tr -> canonical(tr.get("str"), tr.get("int_field")))); + PAssert.that(pushdownRead).containsInAnyOrder(expectedFiltered); + + PCollection queryRead = + pipeline + .apply( + "BQ query read", + BigQueryIO.readTableRows() + .fromQuery( + String.format( + "SELECT str, int_field FROM `%s` WHERE int_field < %d", + tableSpec, intCutoff)) + .usingStandardSql() + .withMethod(Method.DIRECT_READ)) + .apply( + "canonicalize query", + MapElements.into(TypeDescriptors.strings()) + .via(tr -> canonical(tr.get("str"), tr.get("int_field")))); + PAssert.that(queryRead).containsInAnyOrder(expectedFiltered); + + pipeline.run().waitUntilFinish(); + } + + /** + * String canonicalization for cross-engine row comparison: {@code String.valueOf} normalizes + * representation differences (e.g. {@code Long} vs {@code Integer} vs numeric strings in {@code + * TableRow}). + */ + private static String canonical(Object... values) { + StringBuilder sb = new StringBuilder(); + for (Object value : values) { + sb.append(String.valueOf(normalize(value))).append('|'); + } + return sb.toString(); + } + + /** + * Normalizes representation differences that are inherent to the Iceberg-to-BigQuery type mapping + * so that values compare equal across engines: BigQuery widens {@code float32} to {@code + * FLOAT64}, and the schema'd Beam Row conversion used by the Managed BigQuery read maps TIMESTAMP + * to Beam's DATETIME field type, a millisecond Joda instant, while the Iceberg-side rows carry + * microsecond {@code java.time.Instant}s. (Plain {@code readTableRows()} is not affected: its + * {@code TableRow} values are strings with full microsecond precision.) + */ + private static Object normalize(Object value) { + if (value instanceof Float) { + return ((Float) value).doubleValue(); + } + if (value instanceof java.time.Instant) { + return ((java.time.Instant) value).toEpochMilli(); + } + if (value instanceof ReadableInstant) { + return ((ReadableInstant) value).getMillis(); + } + return value; + } + + /** + * Full-schema canonicalization for the type-fidelity check. Works on both the Iceberg-side + * expected rows and the BigQuery-lens rows returned by the Managed BigQuery read; {@link + * #normalize(Object)} absorbs the documented type widenings. + */ + private static String fidelityCanonical(Row row) { + Row nested = checkStateNotNull(row.getRow("row")); + @Nullable Row nullableNested = row.getRow("nullable_row"); + return canonical( + row.getValue("str"), + row.getValue("char"), + row.getValue("modulo_5"), + row.getValue("bool_field"), + row.getValue("int_field"), + nestedCanonical(nested), + (Object) row.getValue("arr_long"), + nullableNested == null ? "null" : nestedCanonical(nullableNested), + row.getValue("nullable_long"), + row.getValue("datetime_tz"), + row.getValue("datetime"), + row.getValue("date"), + row.getValue("time")); + } + + private static String nestedCanonical(Row nested) { + Row doubly = checkStateNotNull(nested.getRow("nested_row")); + return canonical( + nested.getValue("nested_str"), + doubly.getValue("doubly_nested_str"), + doubly.getValue("doubly_nested_float"), + nested.getValue("nested_int"), + nested.getValue("nested_float")); + } + + /** + * Type fidelity across the Iceberg-to-BigQuery read: every field of {@link #BEAM_SCHEMA} + * (including timestamps, datetime/date/time, arrays, and nullable nested rows) must survive the + * round trip with its value intact, modulo the documented widenings handled by {@link + * #normalize(Object)}. + */ + @Test + public void testReadWithBigQueryIOTypeFidelity() throws Exception { + String tableSpec = bigQueryTableSpec(tableId()); + assumeTrue("Catalog does not surface its tables in BigQuery", tableSpec != null); + Table table = catalog.createTable(TableIdentifier.parse(tableId()), ICEBERG_SCHEMA); + List expectedRows = populateTable(table); + List expected = + expectedRows.stream() + .map(IcebergCatalogBaseIT::fidelityCanonical) + .collect(Collectors.toList()); + + PCollection actual = + pipeline + .apply( + "BQ managed read", + Managed.read(Managed.BIGQUERY).withConfig(ImmutableMap.of("table", tableSpec))) + .getSinglePCollection() + .apply( + "fidelity canonicalize", + MapElements.into(TypeDescriptors.strings()) + .via(IcebergCatalogBaseIT::fidelityCanonical)); + PAssert.that(actual).containsInAnyOrder(expected); + pipeline.run().waitUntilFinish(); + } + + /** + * Merge-on-read: BigQuery reads must apply Iceberg v2 row-level delete files. Writes an + * equality-delete file with the Iceberg API (as an external engine would), commits it as a row + * delta, and asserts the deleted rows are absent from the BigQueryIO read. + */ + @Test + public void testReadWithBigQueryIOAfterRowLevelDeletes() throws Exception { + String tableSpec = bigQueryTableSpec(tableId()); + assumeTrue("Catalog does not surface its tables in BigQuery", tableSpec != null); + Table table = + catalog.createTable( + TableIdentifier.parse(tableId()), + ICEBERG_SCHEMA, + PartitionSpec.unpartitioned(), + ImmutableMap.of("format-version", "2")); + List expectedRows = populateTable(table); + int deleteBelow = numRecords() / 2; + + org.apache.iceberg.Schema deleteRowSchema = ICEBERG_SCHEMA.select("int_field"); + OutputFile deleteOutputFile; + try (FileIO io = table.io()) { + deleteOutputFile = io.newOutputFile(table.location() + "/deletes-" + UUID.randomUUID()); + } + EqualityDeleteWriter deleteWriter = + Parquet.writeDeletes(deleteOutputFile) + .createWriterFunc(GenericParquetWriter::create) + .overwrite() + .rowSchema(deleteRowSchema) + .withSpec(PartitionSpec.unpartitioned()) + .equalityFieldIds(ICEBERG_SCHEMA.findField("int_field").fieldId()) + .buildEqualityWriter(); + try (EqualityDeleteWriter closingWriter = deleteWriter) { + for (int i = 0; i < deleteBelow; i++) { + GenericRecord deleteRecord = GenericRecord.create(deleteRowSchema); + deleteRecord.setField("int_field", i); + closingWriter.write(deleteRecord); + } + } + table.newRowDelta().addDeletes(deleteWriter.toDeleteFile()).commit(); + + List expectedRemaining = + expectedRows.stream() + .filter(r -> checkStateNotNull(r.getInt32("int_field")) >= deleteBelow) + .map(r -> canonical(r.getString("str"), r.getInt32("int_field"))) + .collect(Collectors.toList()); + + PCollection afterDeletes = + pipeline + .apply( + "BQ read after deletes", + BigQueryIO.readTableRows().from(tableSpec).withMethod(Method.DIRECT_READ)) + .apply( + "canonicalize after deletes", + MapElements.into(TypeDescriptors.strings()) + .via(tr -> canonical(tr.get("str"), tr.get("int_field")))); + PAssert.that(afterDeletes).containsInAnyOrder(expectedRemaining); + pipeline.run().waitUntilFinish(); + } + @Test public void testWriteReadWithFilter() throws IOException { Table table = catalog.createTable(TableIdentifier.parse(tableId()), ICEBERG_SCHEMA); 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 c16df763333f..bb56e89f75ca 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 @@ -20,6 +20,7 @@ import java.util.Map; 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.TableIdentifier; import org.apache.iceberg.rest.RESTCatalog; import org.junit.After; import org.junit.BeforeClass; @@ -29,8 +30,12 @@ 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) - private static final String BIGLAKE_WAREHOUSE = "gs://managed-iceberg-biglake-its"; + // 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() { @@ -58,6 +63,18 @@ public String type() { return "biglake"; } + @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. + 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()); + } + @Override public Catalog createCatalog() { RESTCatalog restCatalog = new RESTCatalog(); From 02350006b1e417d74696c5dabc6a15d478535d04 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 4 Aug 2026 17:10:09 -0400 Subject: [PATCH 2/2] imt tests --- .../io/google-cloud-platform/build.gradle | 32 ++ .../BigQueryIOIcebergManagedTableIT.java | 408 ++++++++++++++++++ sdks/java/io/iceberg/build.gradle | 7 + .../BigQueryManagedTableCrossEngineIT.java | 159 +++++++ 4 files changed, 606 insertions(+) create mode 100644 sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOIcebergManagedTableIT.java create mode 100644 sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/BigQueryManagedTableCrossEngineIT.java diff --git a/sdks/java/io/google-cloud-platform/build.gradle b/sdks/java/io/google-cloud-platform/build.gradle index 43e16288348e..94fd8c27229a 100644 --- a/sdks/java/io/google-cloud-platform/build.gradle +++ b/sdks/java/io/google-cloud-platform/build.gradle @@ -202,6 +202,7 @@ task integrationTest(type: Test, dependsOn: processTestResources) { outputs.upToDateWhen { false } include '**/*IT.class' + exclude '**/BigQueryIOIcebergManagedTableIT.class' exclude '**/BigQueryIOReadIT.class' exclude '**/BigQueryIOStorageQueryIT.class' exclude '**/BigQueryIOStorageReadIT.class' @@ -251,6 +252,37 @@ task integrationTestKms(type: Test) { } } +/* + Integration tests for BigQueryIO against managed Apache Iceberg tables. Kept out of the + generic integrationTest task because they need a provisioned CLOUD_RESOURCE connection; + override the apache-beam-testing defaults with -PgcpProject, -PbqImtConnection and + -PbqImtStorageUri. + */ +task bigQueryManagedIcebergIntegrationTest(type: Test, dependsOn: processTestResources) { + group = "Verification" + def gcpProject = project.findProperty('gcpProject') ?: 'apache-beam-testing' + def gcpTempRoot = project.findProperty('gcpTempRoot') ?: 'gs://temp-storage-for-end-to-end-tests' + def bqImtConnection = project.findProperty('bqImtConnection') ?: + 'projects/apache-beam-testing/locations/us/connections/apache-beam-testing-storageapi-biglake-nodelete' + def bqImtStorageUri = project.findProperty('bqImtStorageUri') ?: 'gs://apache-beam-testing-bq-biglake' + systemProperty "beamTestPipelineOptions", JsonOutput.toJson([ + "--runner=DirectRunner", + "--project=${gcpProject}", + "--tempRoot=${gcpTempRoot}", + "--tempLocation=${gcpTempRoot}", + ]) + systemProperty "beam.bq.imt.connection", bqImtConnection + systemProperty "beam.bq.imt.storageUri", bqImtStorageUri + + outputs.upToDateWhen { false } + + include '**/BigQueryIOIcebergManagedTableIT.class' + + maxParallelForks 4 + classpath = sourceSets.test.runtimeClasspath + testClassesDirs = sourceSets.test.output.classesDirs +} + /* Integration tests for BigQueryIO that run on BigQuery's early rollout region (us-east7) with the intended purpose of catching breaking changes from new BigQuery releases. diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOIcebergManagedTableIT.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOIcebergManagedTableIT.java new file mode 100644 index 000000000000..83d47ab57e0b --- /dev/null +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOIcebergManagedTableIT.java @@ -0,0 +1,408 @@ +/* + * 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.gcp.bigquery; + +import static org.hamcrest.MatcherAssert.assertThat; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotNull; + +import com.google.api.services.bigquery.Bigquery; +import com.google.api.services.bigquery.model.QueryRequest; +import com.google.api.services.bigquery.model.QueryResponse; +import com.google.api.services.bigquery.model.Table; +import com.google.api.services.bigquery.model.TableCell; +import com.google.api.services.bigquery.model.TableFieldSchema; +import com.google.api.services.bigquery.model.TableRow; +import com.google.api.services.bigquery.model.TableSchema; +import java.io.IOException; +import java.util.List; +import java.util.stream.Collectors; +import java.util.stream.LongStream; +import org.apache.beam.sdk.Pipeline; +import org.apache.beam.sdk.extensions.gcp.options.GcpOptions; +import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO.TypedRead.Method; +import org.apache.beam.sdk.io.gcp.testing.BigqueryClient; +import org.apache.beam.sdk.testing.PAssert; +import org.apache.beam.sdk.testing.TestPipeline; +import org.apache.beam.sdk.transforms.Create; +import org.apache.beam.sdk.transforms.MapElements; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.TypeDescriptors; +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.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.hamcrest.Matchers; +import org.joda.time.Duration; +import org.junit.AfterClass; +import org.junit.BeforeClass; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** + * Integration tests for BigQueryIO against managed Apache Iceberg tables in BigQuery. + * + *

Requires a CLOUD_RESOURCE connection whose service account can administer the storage bucket. + * Defaults target standing apache-beam-testing resources (shared with {@link + * StorageApiSinkCreateIfNeededIT}); override with the {@code beam.bq.imt.connection} and {@code + * beam.bq.imt.storageUri} system properties. + */ +@RunWith(JUnit4.class) +public class BigQueryIOIcebergManagedTableIT { + + private static final BigqueryClient BQ_CLIENT = + new BigqueryClient("BigQueryIOIcebergManagedTableIT"); + // BigqueryClient's query helpers stage a destination table (DDL rejects that) and run in the + // default location, so SQL goes through a raw client instead. + private static final Bigquery RAW_BQ = + BigqueryClient.getNewBigqueryClient("BigQueryIOIcebergManagedTableIT"); + private static final String PROJECT = + TestPipeline.testingPipelineOptions().as(GcpOptions.class).getProject(); + + // Connection in "projects/{project}/locations/{location}/connections/{connection}" form. + private static final String CONNECTION = + System.getProperty( + "beam.bq.imt.connection", + "projects/apache-beam-testing/locations/us/connections/apache-beam-testing-storageapi-biglake-nodelete"); + private static final String STORAGE_URI_ROOT = + System.getProperty("beam.bq.imt.storageUri", "gs://apache-beam-testing-bq-biglake") + + "/BigQueryIOIcebergManagedTableIT"; + + private static final String DATASET_ID = "bq_imt_it_" + System.nanoTime(); + + private static final TableSchema BASE_SCHEMA = + new TableSchema() + .setFields( + ImmutableList.of( + new TableFieldSchema().setName("id").setType("INT64").setMode("REQUIRED"), + new TableFieldSchema().setName("name").setType("STRING"))); + + @BeforeClass + public static void setup() throws IOException, InterruptedException { + // The dataset must be colocated with the connection. + BQ_CLIENT.createNewDataset( + PROJECT, DATASET_ID, /* defaultTableExpirationMs= */ null, connectionLocation()); + } + + @AfterClass + public static void cleanup() { + BQ_CLIENT.deleteDataset(PROJECT, DATASET_ID); + } + + private static String connectionLocation() { + return Splitter.on('/').splitToList(CONNECTION).get(3); + } + + /** Connection reference in the dotted form BigQuery DDL accepts. */ + private static String connectionDotted() { + List parts = Splitter.on('/').splitToList(CONNECTION); + return String.format("%s.%s.%s", parts.get(1), parts.get(3), parts.get(5)); + } + + private static String tableSpec(String table) { + return String.format("%s.%s.%s", PROJECT, DATASET_ID, table); + } + + private static List rows(long from, long to) { + return LongStream.range(from, to) + .mapToObj(i -> new TableRow().set("id", i).set("name", "row_" + i)) + .collect(Collectors.toList()); + } + + private static String canonical(TableRow row) { + return row.get("id") + "|" + row.get("name"); + } + + private static List canonical(List rows) { + return rows.stream() + .map(BigQueryIOIcebergManagedTableIT::canonical) + .collect(Collectors.toList()); + } + + private ImmutableMap bigLakeConfig() { + return ImmutableMap.of( + BigQueryIO.CONNECTION_ID, CONNECTION, + BigQueryIO.STORAGE_URI, STORAGE_URI_ROOT); + } + + /** Runs SQL in the connection's location and returns the result rows. */ + private static List runSql(String sql) throws IOException { + QueryResponse response = + RAW_BQ + .jobs() + .query( + PROJECT, + new QueryRequest() + .setQuery(sql) + .setUseLegacySql(false) + .setLocation(connectionLocation()) + .setTimeoutMs(180_000L)) + .execute(); + if (!Boolean.TRUE.equals(response.getJobComplete())) { + throw new IOException("Query did not complete in time: " + sql); + } + return response.getRows(); + } + + private static String firstCell(List rows) { + TableCell cell = (TableCell) Iterables.getOnlyElement(rows).getF().get(0); + return (String) cell.getV(); + } + + private void createManagedTableViaDdl(String table, String columns) throws IOException { + runSql( + String.format( + "CREATE TABLE `%s` (%s) WITH CONNECTION `%s` " + + "OPTIONS (file_format='PARQUET', table_format='ICEBERG', storage_uri='%s/%s/%s')", + tableSpec(table), columns, connectionDotted(), STORAGE_URI_ROOT, DATASET_ID, table)); + } + + private void runWrite( + BigQueryIO.Write.Method method, + String table, + List input, + BigQueryIO.Write.CreateDisposition createDisposition, + @Nullable TableSchema schema, + boolean withBigLakeConfiguration) { + Pipeline p = Pipeline.create(TestPipeline.testingPipelineOptions()); + BigQueryIO.Write write = + BigQueryIO.writeTableRows() + .to(tableSpec(table)) + .withMethod(method) + .withCreateDisposition(createDisposition) + .withWriteDisposition(BigQueryIO.Write.WriteDisposition.WRITE_APPEND); + if (schema != null) { + write = write.withSchema(schema); + } + if (withBigLakeConfiguration) { + write = write.withBigLakeConfiguration(bigLakeConfig()); + } + p.apply(Create.of(input).withCoder(TableRowJsonCoder.of())).apply(write); + p.run().waitUntilFinish(); + } + + private long countRows(String table) throws IOException { + return Long.parseLong( + firstCell(runSql(String.format("SELECT COUNT(*) FROM `%s`", tableSpec(table))))); + } + + @Test + public void testAtLeastOnceWriteThenDirectRead() throws IOException, InterruptedException { + String table = "alo_create_" + System.nanoTime(); + List input = rows(0, 20); + runWrite( + BigQueryIO.Write.Method.STORAGE_API_AT_LEAST_ONCE, + table, + input, + BigQueryIO.Write.CreateDisposition.CREATE_IF_NEEDED, + BASE_SCHEMA, + /* withBigLakeConfiguration= */ true); + + Table created = BQ_CLIENT.getTableResource(PROJECT, DATASET_ID, table); + assertNotNull(created.getBiglakeConfiguration()); + assertEquals("ICEBERG", created.getBiglakeConfiguration().getTableFormat()); + // Stats can lag behind recent writes and report zero, but are never null for these tables. + assertNotNull(created.getNumBytes()); + + Pipeline p = Pipeline.create(TestPipeline.testingPipelineOptions()); + PCollection full = + p.apply( + "ReadFull", + BigQueryIO.readTableRows().from(tableSpec(table)).withMethod(Method.DIRECT_READ)) + .apply( + "CanonicalFull", + MapElements.into(TypeDescriptors.strings()) + .via(BigQueryIOIcebergManagedTableIT::canonical)); + PAssert.that(full).containsInAnyOrder(canonical(input)); + + PCollection pushdown = + p.apply( + "ReadPushdown", + BigQueryIO.readTableRows() + .from(tableSpec(table)) + .withMethod(Method.DIRECT_READ) + .withSelectedFields(ImmutableList.of("id")) + .withRowRestriction("id < 5")) + .apply( + "CanonicalPushdown", + MapElements.into(TypeDescriptors.strings()).via(r -> String.valueOf(r.get("id")))); + PAssert.that(pushdown).containsInAnyOrder("0", "1", "2", "3", "4"); + p.run().waitUntilFinish(); + } + + /** Exactly-once batch write: bounded STORAGE_WRITE_API (PENDING streams + batch commit). */ + @Test + public void testExactlyOnceBatchWrite() throws IOException, InterruptedException { + String table = "eo_batch_" + System.nanoTime(); + runWrite( + BigQueryIO.Write.Method.STORAGE_WRITE_API, + table, + rows(0, 50), + BigQueryIO.Write.CreateDisposition.CREATE_IF_NEEDED, + BASE_SCHEMA, + /* withBigLakeConfiguration= */ true); + assertEquals(50L, countRows(table)); + } + + /** Exactly-once streaming write: unbounded STORAGE_WRITE_API with triggering frequency. */ + @Test + public void testExactlyOnceStreamingWrite() throws IOException, InterruptedException { + String table = "eo_streaming_" + System.nanoTime(); + Pipeline p = Pipeline.create(TestPipeline.testingPipelineOptions()); + // The streaming exactly-once path is chosen by input boundedness; flip the bounded Create. + p.apply(Create.of(rows(0, 30)).withCoder(TableRowJsonCoder.of())) + .setIsBoundedInternal(PCollection.IsBounded.UNBOUNDED) + .apply( + BigQueryIO.writeTableRows() + .to(tableSpec(table)) + .withMethod(BigQueryIO.Write.Method.STORAGE_WRITE_API) + .withCreateDisposition(BigQueryIO.Write.CreateDisposition.CREATE_IF_NEEDED) + .withWriteDisposition(BigQueryIO.Write.WriteDisposition.WRITE_APPEND) + .withSchema(BASE_SCHEMA) + .withBigLakeConfiguration(bigLakeConfig()) + .withNumStorageWriteApiStreams(1) + .withTriggeringFrequency(Duration.standardSeconds(1))); + p.run().waitUntilFinish(); + assertEquals(30L, countRows(table)); + } + + /** The common production case: CREATE_NEVER write into a pre-existing managed Iceberg table. */ + @Test + public void testCreateNeverToPreExistingTable() throws IOException, InterruptedException { + String table = "create_never_" + System.nanoTime(); + createManagedTableViaDdl(table, "id INT64 NOT NULL, name STRING"); + runWrite( + BigQueryIO.Write.Method.STORAGE_API_AT_LEAST_ONCE, + table, + rows(0, 10), + BigQueryIO.Write.CreateDisposition.CREATE_NEVER, + null, + /* withBigLakeConfiguration= */ false); + assertThat(countRows(table), Matchers.greaterThanOrEqualTo(10L)); + } + + /** + * Schema widening on a managed Iceberg table: tables.patch accepts a new nullable column (unlike + * Iceberg tables in catalogs, which refuse patch), and subsequent writes can use it. + */ + @Test + public void testSchemaFieldAdditionThenWrite() throws IOException, InterruptedException { + String table = "schema_update_" + System.nanoTime(); + createManagedTableViaDdl(table, "id INT64 NOT NULL, name STRING"); + + TableSchema widened = + new TableSchema() + .setFields( + ImmutableList.of( + new TableFieldSchema().setName("id").setType("INT64").setMode("REQUIRED"), + new TableFieldSchema().setName("name").setType("STRING"), + new TableFieldSchema().setName("extra").setType("STRING"))); + BQ_CLIENT.updateTableSchema(PROJECT, DATASET_ID, table, widened); + + List input = + ImmutableList.of( + new TableRow().set("id", 1L).set("name", "a").set("extra", "x"), + new TableRow().set("id", 2L).set("name", "b").set("extra", "y")); + runWrite( + BigQueryIO.Write.Method.STORAGE_API_AT_LEAST_ONCE, + table, + input, + BigQueryIO.Write.CreateDisposition.CREATE_NEVER, + null, + /* withBigLakeConfiguration= */ false); + + long withExtra = + Long.parseLong( + firstCell( + runSql( + String.format( + "SELECT COUNT(*) FROM `%s` WHERE extra IS NOT NULL", tableSpec(table))))); + assertThat(withExtra, Matchers.greaterThanOrEqualTo(2L)); + } + + /** + * FILE_LOADS append into a pre-existing table. withBigLakeConfiguration is rejected outside the + * Storage Write API, so FILE_LOADS cannot create these tables itself. + */ + @Test + public void testFileLoadsSinglePartition() throws IOException, InterruptedException { + String table = "file_loads_" + System.nanoTime(); + createManagedTableViaDdl(table, "id INT64 NOT NULL, name STRING"); + runWrite( + BigQueryIO.Write.Method.FILE_LOADS, + table, + rows(0, 25), + BigQueryIO.Write.CreateDisposition.CREATE_NEVER, + BASE_SCHEMA, + /* withBigLakeConfiguration= */ false); + assertEquals(25L, countRows(table)); + } + + /** Legacy streaming (insertAll) is accepted; rows arrive through the streaming buffer. */ + @Test + public void testStreamingInsertsWrite() throws IOException, InterruptedException { + String table = "streaming_inserts_" + System.nanoTime(); + createManagedTableViaDdl(table, "id INT64 NOT NULL, name STRING"); + runWrite( + BigQueryIO.Write.Method.STREAMING_INSERTS, + table, + rows(0, 5), + BigQueryIO.Write.CreateDisposition.CREATE_NEVER, + BASE_SCHEMA, + /* withBigLakeConfiguration= */ false); + assertThat(countRows(table), Matchers.greaterThanOrEqualTo(5L)); + } + + /** Query read with time travel: managed Iceberg tables support FOR SYSTEM_TIME AS OF. */ + @Test + public void testQueryReadWithTimeTravel() throws IOException, InterruptedException { + String table = "time_travel_" + System.nanoTime(); + runWrite( + BigQueryIO.Write.Method.STORAGE_WRITE_API, + table, + rows(0, 10), + BigQueryIO.Write.CreateDisposition.CREATE_IF_NEEDED, + BASE_SCHEMA, + /* withBigLakeConfiguration= */ true); + // Server-side timestamp between the two writes. + String asOf = firstCell(runSql("SELECT STRING(CURRENT_TIMESTAMP())")); + runWrite( + BigQueryIO.Write.Method.STORAGE_WRITE_API, + table, + rows(10, 20), + BigQueryIO.Write.CreateDisposition.CREATE_NEVER, + null, + /* withBigLakeConfiguration= */ false); + + Pipeline p = Pipeline.create(TestPipeline.testingPipelineOptions()); + PCollection counted = + p.apply( + BigQueryIO.readTableRows() + .fromQuery( + String.format( + "SELECT COUNT(*) AS c FROM `%s` FOR SYSTEM_TIME AS OF TIMESTAMP '%s'", + tableSpec(table), asOf)) + .usingStandardSql()) + .apply( + MapElements.into(TypeDescriptors.strings()).via(r -> String.valueOf(r.get("c")))); + PAssert.that(counted).containsInAnyOrder("10"); + p.run().waitUntilFinish(); + } +} diff --git a/sdks/java/io/iceberg/build.gradle b/sdks/java/io/iceberg/build.gradle index fe91533227fa..e2e8a12d01eb 100644 --- a/sdks/java/io/iceberg/build.gradle +++ b/sdks/java/io/iceberg/build.gradle @@ -173,6 +173,13 @@ task integrationTest(type: Test) { // against a non-default project's catalog. systemProperty "beam.iceberg.biglake.warehouse", project.findProperty('biglakeWarehouse') ?: 'gs://managed-iceberg-biglake-its' + // Connection + storage root for BigQueryManagedTableCrossEngineIT. + if (project.findProperty('bqImtConnection') != null) { + systemProperty "beam.bq.imt.connection", project.findProperty('bqImtConnection') + } + if (project.findProperty('bqImtStorageUri') != null) { + systemProperty "beam.bq.imt.storageUri", project.findProperty('bqImtStorageUri') + } // Disable Gradle cache: these ITs interact with live service that should always be considered "out of date" outputs.upToDateWhen { false } diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/BigQueryManagedTableCrossEngineIT.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/BigQueryManagedTableCrossEngineIT.java new file mode 100644 index 000000000000..bcea98c85c1e --- /dev/null +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/BigQueryManagedTableCrossEngineIT.java @@ -0,0 +1,159 @@ +/* + * 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 com.google.api.services.bigquery.Bigquery; +import com.google.api.services.bigquery.model.QueryRequest; +import com.google.api.services.bigquery.model.QueryResponse; +import java.io.IOException; +import java.util.List; +import java.util.Map; +import java.util.stream.Collectors; +import java.util.stream.LongStream; +import org.apache.beam.sdk.Pipeline; +import org.apache.beam.sdk.extensions.gcp.options.GcpOptions; +import org.apache.beam.sdk.io.gcp.testing.BigqueryClient; +import org.apache.beam.sdk.managed.Managed; +import org.apache.beam.sdk.testing.PAssert; +import org.apache.beam.sdk.testing.TestPipeline; +import org.apache.beam.sdk.transforms.MapElements; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.TypeDescriptors; +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.ImmutableMap; +import org.junit.AfterClass; +import org.junit.BeforeClass; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** + * Reads a managed Apache Iceberg table in BigQuery through {@code Managed.ICEBERG} with the + * BigQueryMetastore catalog, addressing it by its ordinary 3-part name. + * + *

Shares the {@code beam.bq.imt.*} system properties with BigQueryIOIcebergManagedTableIT in the + * google-cloud-platform module. + */ +@RunWith(JUnit4.class) +public class BigQueryManagedTableCrossEngineIT { + + private static final BigqueryClient BQ_CLIENT = + new BigqueryClient("BigQueryManagedTableCrossEngineIT"); + // BigqueryClient's query helpers stage a destination table (DDL rejects that) and run in the + // default location, so SQL goes through a raw client instead. + private static final Bigquery RAW_BQ = + BigqueryClient.getNewBigqueryClient("BigQueryManagedTableCrossEngineIT"); + private static final String PROJECT = + TestPipeline.testingPipelineOptions().as(GcpOptions.class).getProject(); + + // Connection in "projects/{project}/locations/{location}/connections/{connection}" form. + private static final String CONNECTION = + System.getProperty( + "beam.bq.imt.connection", + "projects/apache-beam-testing/locations/us/connections/apache-beam-testing-storageapi-biglake-nodelete"); + private static final String STORAGE_URI_ROOT = + System.getProperty("beam.bq.imt.storageUri", "gs://apache-beam-testing-bq-biglake") + + "/BigQueryManagedTableCrossEngineIT"; + + private static final String DATASET_ID = "bq_imt_xengine_" + System.nanoTime(); + + @BeforeClass + public static void setup() throws IOException, InterruptedException { + // The dataset must be colocated with the connection. + BQ_CLIENT.createNewDataset( + PROJECT, DATASET_ID, /* defaultTableExpirationMs= */ null, connectionLocation()); + } + + @AfterClass + public static void cleanup() { + BQ_CLIENT.deleteDataset(PROJECT, DATASET_ID); + } + + private static String connectionLocation() { + return Splitter.on('/').splitToList(CONNECTION).get(3); + } + + /** Connection reference in the dotted form BigQuery DDL accepts. */ + private static String connectionDotted() { + List parts = Splitter.on('/').splitToList(CONNECTION); + return String.format("%s.%s.%s", parts.get(1), parts.get(3), parts.get(5)); + } + + /** Runs SQL in the connection's location. */ + private static void runSql(String sql) throws IOException { + QueryResponse response = + RAW_BQ + .jobs() + .query( + PROJECT, + new QueryRequest() + .setQuery(sql) + .setUseLegacySql(false) + .setLocation(connectionLocation()) + .setTimeoutMs(180_000L)) + .execute(); + if (!Boolean.TRUE.equals(response.getJobComplete())) { + throw new IOException("Query did not complete in time: " + sql); + } + } + + @Test + public void testManagedIcebergReadByThreePartName() throws IOException { + String table = "managed_read_" + System.nanoTime(); + runSql( + String.format( + "CREATE TABLE `%s.%s.%s` (id INT64, name STRING) WITH CONNECTION `%s` " + + "OPTIONS (file_format='PARQUET', table_format='ICEBERG', storage_uri='%s/%s/%s')", + PROJECT, DATASET_ID, table, connectionDotted(), STORAGE_URI_ROOT, DATASET_ID, table)); + runSql( + String.format( + "INSERT INTO `%s.%s.%s` " + + "SELECT id, CONCAT('row_', CAST(id AS STRING)) " + + "FROM UNNEST(GENERATE_ARRAY(0, 9)) id", + PROJECT, DATASET_ID, table)); + // The catalog resolves the table through its exported Iceberg metadata, and automatic exports + // can lag far behind recent writes; exporting makes the test deterministic. + runSql(String.format("EXPORT TABLE METADATA FROM `%s.%s.%s`", PROJECT, DATASET_ID, table)); + + Map config = + ImmutableMap.builder() + .put("table", DATASET_ID + "." + table) + .put( + "catalog_properties", + ImmutableMap.builder() + .put("gcp_project", PROJECT) + .put("gcp_location", connectionLocation()) + .put("warehouse", STORAGE_URI_ROOT) + .put("catalog-impl", "org.apache.iceberg.gcp.bigquery.BigQueryMetastoreCatalog") + .put("io-impl", "org.apache.iceberg.gcp.gcs.GCSFileIO") + .build()) + .build(); + + Pipeline p = Pipeline.create(TestPipeline.testingPipelineOptions()); + PCollection rows = + p.apply(Managed.read(Managed.ICEBERG).withConfig(config)) + .getSinglePCollection() + .apply( + MapElements.into(TypeDescriptors.strings()) + .via(row -> row.getInt64("id") + "|" + row.getString("name"))); + PAssert.that(rows) + .containsInAnyOrder( + LongStream.range(0, 10).mapToObj(i -> i + "|row_" + i).collect(Collectors.toList())); + p.run().waitUntilFinish(); + } +}