diff --git a/CHANGES.md b/CHANGES.md index f0d5d06b9d9f..bcdcaca8cdb5 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -92,6 +92,7 @@ ## Bugfixes +* (Java) Fixed BigQuery Storage Write API failed-row timestamp conversion, preserving microsecond precision and UTC formatting ([#40110](https://github.com/apache/beam/issues/40110)). * (Java) Fixed the Spark runner firing processing-time timers in reverse timestamp order ([#39824](https://github.com/apache/beam/issues/39824)). * (Python) Fixed incorrect profiler options handling on portable runners ([#39613](https://github.com/apache/beam/issues/39613)). * (Java) KafkaIO dynamic reads no longer require the obsolete `beam_fn_api` experiment ([#29998](https://github.com/apache/beam/issues/29998)). diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryUtils.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryUtils.java index c8530f5ce678..289e7d54de43 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryUtils.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryUtils.java @@ -230,6 +230,16 @@ public abstract static class Builder { .appendZoneRegionId() .toFormatter(); + // TIMESTAMP_FORMATTER accepts alternate separators and must not be used for printing. + static final java.time.format.DateTimeFormatter BIGQUERY_TIMESTAMP_MICROS_FORMATTER = + new java.time.format.DateTimeFormatterBuilder() + .appendPattern("uuuu-MM-dd HH:mm:ss") + .appendFraction(java.time.temporal.ChronoField.NANO_OF_SECOND, 0, 6, true) + .appendLiteral(" UTC") + .toFormatter() + .withResolverStyle(java.time.format.ResolverStyle.STRICT) + .withZone(ZoneOffset.UTC); + private static final DateTimeFormatter BIGQUERY_TIMESTAMP_PRINTER; /** @@ -936,7 +946,11 @@ public static Row toBeamRow(Schema rowSchema, TableSchema bqSchema, TableRow jso long nanos = (micros % 1_000_000) * 1_000; return java.time.Instant.ofEpochSecond(seconds, nanos); } catch (NumberFormatException e) { - return java.time.Instant.parse(jsonBQString); + try { + return java.time.Instant.parse(jsonBQString); + } catch (DateTimeParseException e2) { + return BIGQUERY_TIMESTAMP_MICROS_FORMATTER.parse(jsonBQString, java.time.Instant::from); + } } } else if (fieldType.isLogicalType(Timestamp.IDENTIFIER)) { if (!jsonBQString.contains("UTC")) { diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/TableRowToStorageApiProto.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/TableRowToStorageApiProto.java index ec1daed48b96..f376d7c30a57 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/TableRowToStorageApiProto.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/TableRowToStorageApiProto.java @@ -19,6 +19,7 @@ import static java.util.stream.Collectors.toList; import static java.util.stream.Collectors.toSet; +import static org.apache.beam.sdk.io.gcp.bigquery.BigQueryUtils.BIGQUERY_TIMESTAMP_MICROS_FORMATTER; import static org.apache.beam.sdk.io.gcp.bigquery.BigQueryUtils.DATETIME_SPACE_FORMATTER; import static org.apache.beam.sdk.io.gcp.bigquery.BigQueryUtils.TIMESTAMP_FORMATTER; @@ -56,6 +57,7 @@ import java.time.LocalTime; import java.time.ZoneOffset; import java.time.format.DateTimeParseException; +import java.time.temporal.ChronoUnit; import java.util.AbstractMap; import java.util.ArrayList; import java.util.Collections; @@ -1990,7 +1992,7 @@ public static Object jsonValueFromMessageValue( long epochSeconds = epochMicros / 1_000_000L; long nanoAdjustment = (epochMicros % 1_000_000L) * 1_000L; Instant instant = Instant.ofEpochSecond(epochSeconds, nanoAdjustment); - return LocalDateTime.ofInstant(instant, ZoneOffset.UTC).format(TIMESTAMP_FORMATTER); + return BIGQUERY_TIMESTAMP_MICROS_FORMATTER.format(instant); } else if (fieldDescriptor.getType().equals(FieldDescriptor.Type.MESSAGE)) { Message message = (Message) fieldValue; String messageName = fieldDescriptor.getMessageType().getName(); @@ -2000,7 +2002,8 @@ public static Object jsonValueFromMessageValue( long seconds = (long) message.getField(descriptor.findFieldByName("seconds")); int nanos = (int) message.getField(descriptor.findFieldByName("nanos")); Instant instant = Instant.ofEpochSecond(seconds, nanos); - return LocalDateTime.ofInstant(instant, ZoneOffset.UTC).format(TIMESTAMP_FORMATTER); + return BIGQUERY_TIMESTAMP_MICROS_FORMATTER.format( + instant.truncatedTo(ChronoUnit.MICROS)); } else if (messageName.equals("TimestampPicos")) { Descriptor descriptor = message.getDescriptorForType(); long seconds = (long) message.getField(descriptor.findFieldByName("seconds")); diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOWriteTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOWriteTest.java index 7310e2f850e5..a88a533d6005 100644 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOWriteTest.java +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOWriteTest.java @@ -1335,20 +1335,12 @@ public void runTestWriteAvro(boolean schemaFromView) throws Exception { .set("strval", "test") .set("longval", "1") .set("doubleval", "1.0") - .set( - "instantval", - useStorageApi || useStorageApiApproximate - ? "2019-01-01 T00:00:00" - : "2019-01-01 00:00:00 UTC"), + .set("instantval", "2019-01-01 00:00:00 UTC"), new TableRow() .set("strval", "test2") .set("longval", "2") .set("doubleval", "2.0") - .set( - "instantval", - useStorageApi || useStorageApiApproximate - ? "2019-02-01 T00:00:00" - : "2019-02-01 00:00:00 UTC"))); + .set("instantval", "2019-02-01 00:00:00 UTC"))); } @Test @@ -3876,7 +3868,7 @@ public void testStorageApiErrorsWriteTableRows() throws Exception { TableRow goodNested = new TableRow() .set("number", "42") - .set("timestamp", "1970-01-01 T00:00:00.000043") + .set("timestamp", "1970-01-01 00:00:00.000043 UTC") .set("time", "00:52:07.123456") .set("datetime", "2019-08-16T00:52:07.123456") .set("date", "2019-08-16") @@ -4427,7 +4419,7 @@ public void testWriteProtosEncodedValues(boolean directWrite) throws Exception { .setMode("REPEATED") .setFields(tableSchema.getFields()))); - final String timestamp = "1970-01-01 T00:00:00.000043"; + final String timestamp = "1970-01-01 00:00:00.000043 UTC"; final String date = "2019-08-16"; final String numeric = "23"; final String bignumeric = "123456789012345678"; @@ -4585,7 +4577,7 @@ public void testWriteProtosUnEncodedValues(boolean directWrite) throws Exception .setMode("REPEATED") .setFields(tableSchema.getFields()))); - final String timestamp = "1970-01-01 T00:00:00.000043"; + final String timestamp = "1970-01-01 00:00:00.000043 UTC"; final String date = "2019-08-16"; final String numeric = "23"; final String bignumeric = "123456789012345678"; @@ -4715,7 +4707,7 @@ public void testWriteProtosWrappedValues(boolean directWrite) throws Exception { .setMode("REPEATED") .setFields(tableSchema.getFields()))); - final String timestamp = "1970-01-01 T00:00:00.000043"; + final String timestamp = "1970-01-01 00:00:00.000043 UTC"; long timestampMicros = (long) TYPE_MAP_PROTO_CONVERTERS diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryUtilsTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryUtilsTest.java index 52dbef55286f..f4266022828c 100644 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryUtilsTest.java +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryUtilsTest.java @@ -1182,6 +1182,67 @@ public void testToBeamRow_timestamp_micros() { assertEquals(expectedRow, beamRowMicros); } + @Test + public void testTimestampInputCompatibility() { + Schema schema = Schema.builder().addLogicalTypeField("ts", SqlTypes.TIMESTAMP).build(); + java.time.Instant expected = java.time.Instant.parse("2026-09-03T18:51:43.417123Z"); + for (Object input : + Arrays.asList( + "2026-09-03 18:51:43.417123 UTC", + "2026-09-03T18:51:43.417123Z", + "2026-09-03t18:51:43.417123z", + "2026-09-03T20:51:43.417123+02:00", + "1788461503417123", + 1788461503417123L)) { + assertEquals( + input.toString(), + expected, + BigQueryUtils.toBeamRow(schema, new TableRow().set("ts", input)).getValue("ts")); + } + assertEquals( + java.time.Instant.parse("1969-12-31T23:59:59.999999Z"), + BigQueryUtils.toBeamRow(schema, new TableRow().set("ts", -1L)).getValue("ts")); + assertEquals( + java.time.Instant.EPOCH, + BigQueryUtils.toBeamRow(schema, new TableRow().set("ts", "1970-01-01 00:00:00 UTC")) + .getValue("ts")); + assertThrows( + java.time.format.DateTimeParseException.class, + () -> BigQueryUtils.toBeamRow(schema, new TableRow().set("ts", "not a timestamp"))); + assertThrows( + java.time.format.DateTimeParseException.class, + () -> BigQueryUtils.toBeamRow(schema, new TableRow().set("ts", "2026-02-30 00:00:00 UTC"))); + } + + @Test + public void testDateTimeTimestampInputCompatibility() { + Schema schema = Schema.builder().addDateTimeField("ts").build(); + for (Object input : + Arrays.asList("2026-09-03 18:51:43.417123 UTC", "1788461503.417123", 1788461503.417123)) { + assertEquals( + input.toString(), + 1788461503417L, + BigQueryUtils.toBeamRow(schema, new TableRow().set("ts", input)) + .getDateTime("ts") + .getMillis()); + } + assertEquals( + 1000L, + BigQueryUtils.toBeamRow(schema, new TableRow().set("ts", 1L)) + .getDateTime("ts") + .getMillis()); + assertEquals( + 0L, + BigQueryUtils.toBeamRow(schema, new TableRow().set("ts", "-0.000001")) + .getDateTime("ts") + .getMillis()); + assertThrows( + IllegalArgumentException.class, + () -> + BigQueryUtils.toBeamRow( + schema, new TableRow().set("ts", "2026-09-03 18:51:43.417123456 UTC"))); + } + @Test public void testToTableSpec() { TableReference withProject = diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiTimestampWriteTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiTimestampWriteTest.java new file mode 100644 index 000000000000..25694a88ef5b --- /dev/null +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiTimestampWriteTest.java @@ -0,0 +1,277 @@ +/* + * 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.hamcrest.Matchers.containsInAnyOrder; +import static org.hamcrest.Matchers.hasItems; + +import com.google.api.core.ApiFuture; +import com.google.api.core.ApiFutures; +import com.google.api.services.bigquery.model.Table; +import com.google.api.services.bigquery.model.TableFieldSchema; +import com.google.api.services.bigquery.model.TableReference; +import com.google.api.services.bigquery.model.TableRow; +import com.google.api.services.bigquery.model.TableSchema; +import com.google.cloud.bigquery.storage.v1.AppendRowsRequest; +import com.google.cloud.bigquery.storage.v1.AppendRowsResponse; +import com.google.cloud.bigquery.storage.v1.Exceptions; +import com.google.cloud.bigquery.storage.v1.ProtoRows; +import com.google.protobuf.DescriptorProtos; +import com.google.protobuf.Descriptors.Descriptor; +import com.google.protobuf.DynamicMessage; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import org.apache.beam.runners.direct.DirectOptions; +import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO.Write.CreateDisposition; +import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO.Write.Method; +import org.apache.beam.sdk.io.gcp.testing.FakeBigQueryServices; +import org.apache.beam.sdk.io.gcp.testing.FakeDatasetService; +import org.apache.beam.sdk.io.gcp.testing.FakeJobService; +import org.apache.beam.sdk.schemas.Schema; +import org.apache.beam.sdk.schemas.logicaltypes.SqlTypes; +import org.apache.beam.sdk.testing.PAssert; +import org.apache.beam.sdk.testing.TestPipeline; +import org.apache.beam.sdk.testing.TestStream; +import org.apache.beam.sdk.transforms.Create; +import org.apache.beam.sdk.transforms.FlatMapElements; +import org.apache.beam.sdk.transforms.MapElements; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.Row; +import org.apache.beam.sdk.values.TypeDescriptor; +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.checkerframework.checker.nullness.qual.Nullable; +import org.joda.time.Duration; +import org.junit.After; +import org.junit.Before; +import org.junit.Rule; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.Parameterized; +import org.junit.runners.Parameterized.Parameter; +import org.junit.runners.Parameterized.Parameters; + +/** Tests timestamp decoding after a real BigQueryIO append rejection, without a cloud service. */ +@RunWith(Parameterized.class) +public class StorageApiTimestampWriteTest { + private static final Schema MICROS_SCHEMA = + Schema.builder().addStringField("id").addLogicalTypeField("ts", SqlTypes.TIMESTAMP).build(); + private static final Schema MILLIS_SCHEMA = + Schema.builder().addStringField("id").addDateTimeField("ts").build(); + private static final Schema ERROR_SCHEMA = + Schema.builder() + .addStringField("error_message") + .addRowField("failed_row", MICROS_SCHEMA) + .build(); + private static final List APPEND_SIZES = Collections.synchronizedList(new ArrayList<>()); + + @Parameters(name = "method={0}, streaming={1}") + public static Iterable parameters() { + return ImmutableList.of( + new Object[] {Method.STORAGE_WRITE_API, false}, + new Object[] {Method.STORAGE_WRITE_API, true}, + new Object[] {Method.STORAGE_API_AT_LEAST_ONCE, false}); + } + + @Parameter public Method method; + + @Parameter(1) + public boolean streaming; + + @Rule public final transient TestPipeline pipeline = TestPipeline.create(); + + @Before + public void setUp() throws Exception { + FakeDatasetService.setUp(); + BigQueryIO.clearStaticCaches(); + APPEND_SIZES.clear(); + pipeline.getOptions().as(BigQueryOptions.class).setProject("project-id"); + pipeline.getOptions().as(DirectOptions.class).setTargetParallelism(1); + } + + @After + public void tearDown() throws Exception { + BigQueryIO.clearStaticCaches(); + } + + @Test + public void testFailedRowsRetainTimestampsAndErrors() throws Exception { + TableReference table = + new TableReference() + .setProjectId("project-id") + .setDatasetId("dataset-id") + .setTableId("table-id"); + TableSchema schema = + new TableSchema() + .setFields( + ImmutableList.of( + new TableFieldSchema().setName("id").setType("STRING"), + new TableFieldSchema().setName("ts").setType("TIMESTAMP"))); + FakeDatasetService dataset = new RejectOneRowPerAppend(); + dataset.createDataset("project-id", "dataset-id", "", "", null); + dataset.createTable(new Table().setTableReference(table).setSchema(schema)); + PCollection rows; + if (streaming) { + rows = + pipeline.apply( + TestStream.create(TableRowJsonCoder.of()) + .addElements(row("A"), row("B"), row("C")) + .advanceProcessingTime(Duration.standardSeconds(2)) + .advanceWatermarkToInfinity()); + } else { + rows = + pipeline + .apply(Create.of(0)) + .apply( + FlatMapElements.into(TypeDescriptor.of(TableRow.class)) + .via(ignored -> ImmutableList.of(row("A"), row("B"), row("C")))) + .setCoder(TableRowJsonCoder.of()); + } + BigQueryIO.Write write = + BigQueryIO.writeTableRows() + .to(table) + .withMethod(method) + .withCreateDisposition(CreateDisposition.CREATE_NEVER) + .withPropagateSuccessfulStorageApiWrites(true) + .withTestServices( + new FakeBigQueryServices() + .withDatasetService(dataset) + .withJobService(new FakeJobService())) + .withoutValidation(); + if (streaming) { + write = + write + .withNumStorageWriteApiStreams(1) + .withTriggeringFrequency(Duration.standardSeconds(1)); + } else { + write = write.withNumStorageWriteApiStreams(0); + } + WriteResult result = rows.apply(write); + // No failure formatter: exercise the proto-to-TableRow fallback used by failed-row consumers. + PCollection errors = + result + .getFailedStorageApiInserts() + .apply( + "Decode failed rows", + MapElements.into(TypeDescriptor.of(Row.class)) + .via( + error -> + Row.withSchema(ERROR_SCHEMA) + .addValues( + error.getErrorMessage(), + BigQueryUtils.toBeamRow(MICROS_SCHEMA, error.getRow())) + .build())) + .setRowSchema(ERROR_SCHEMA); + PAssert.that(errors).containsInAnyOrder(error("A"), error("B")); + PCollection millis = + result + .getFailedStorageApiInserts() + .apply( + "Decode millisecond rows", + MapElements.into(TypeDescriptor.of(Row.class)) + .via(error -> BigQueryUtils.toBeamRow(MILLIS_SCHEMA, error.getRow()))) + .setRowSchema(MILLIS_SCHEMA); + PAssert.that(millis) + .containsInAnyOrder( + Row.withSchema(MILLIS_SCHEMA) + .addValues("A", org.joda.time.Instant.ofEpochMilli(1788461503417L)) + .build(), + Row.withSchema(MILLIS_SCHEMA) + .addValues("B", org.joda.time.Instant.ofEpochMilli(-1L)) + .build()); + PAssert.that(result.getSuccessfulStorageApiInserts()).containsInAnyOrder(row("C")); + pipeline.run().waitUntilFinish(); + assertThat( + dataset.getAllRows("project-id", "dataset-id", "table-id"), containsInAnyOrder(row("C"))); + assertThat(APPEND_SIZES, hasItems(3, 2, 1)); + } + + private static TableRow row(String id) { + return new TableRow() + .set("id", id) + .set( + "ts", + id.equals("B") ? "1969-12-31 23:59:59.999999 UTC" : "2026-09-03 18:51:43.417123 UTC"); + } + + private static Row error(String id) { + java.time.Instant timestamp = + java.time.Instant.parse( + id.equals("B") ? "1969-12-31T23:59:59.999999Z" : "2026-09-03T18:51:43.417123Z"); + return Row.withSchema(ERROR_SCHEMA) + .addValues( + "Invalid id: " + id, Row.withSchema(MICROS_SCHEMA).addValues(id, timestamp).build()) + .build(); + } + + private static class RejectOneRowPerAppend extends FakeDatasetService { + @Override + public BigQueryServices.StreamAppendClient getStreamAppendClient( + String streamName, + DescriptorProtos.DescriptorProto descriptor, + boolean useConnectionPool, + AppendRowsRequest.MissingValueInterpretation missingValueInterpretation) + throws Exception { + BigQueryServices.StreamAppendClient delegate = + super.getStreamAppendClient( + streamName, descriptor, useConnectionPool, missingValueInterpretation); + Descriptor protoDescriptor = TableRowToStorageApiProto.wrapDescriptorProto(descriptor); + return new BigQueryServices.StreamAppendClient() { + @Override + public ApiFuture appendRows(long offset, ProtoRows rows) + throws Exception { + APPEND_SIZES.add(rows.getSerializedRowsCount()); + for (int i = 0; i < rows.getSerializedRowsCount(); i++) { + String id = + (String) + DynamicMessage.parseFrom(protoDescriptor, rows.getSerializedRows(i)) + .getField(protoDescriptor.findFieldByName("id")); + if (!id.equals("C")) { + return ApiFutures.immediateFailedFuture( + new Exceptions.AppendSerializationError( + 400, "Invalid row", streamName, ImmutableMap.of(i, "Invalid id: " + id))); + } + } + return delegate.appendRows(offset, rows); + } + + @Override + public com.google.cloud.bigquery.storage.v1.@Nullable TableSchema getUpdatedSchema() { + return delegate.getUpdatedSchema(); + } + + @Override + public void pin() { + delegate.pin(); + } + + @Override + public void unpin() { + delegate.unpin(); + } + + @Override + public void close() throws Exception { + delegate.close(); + } + }; + } + } +} diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/TableRowToStorageApiProtoTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/TableRowToStorageApiProtoTest.java index 90735bae1261..02065d6dac3c 100644 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/TableRowToStorageApiProtoTest.java +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/TableRowToStorageApiProtoTest.java @@ -17,7 +17,6 @@ */ package org.apache.beam.sdk.io.gcp.bigquery; -import static org.apache.beam.sdk.io.gcp.bigquery.BigQueryUtils.TIMESTAMP_FORMATTER; import static org.apache.beam.sdk.io.gcp.bigquery.TableRowToStorageApiProto.TYPE_MAP_PROTO_CONVERTERS; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; @@ -42,13 +41,12 @@ import com.google.protobuf.Descriptors.FieldDescriptor; import com.google.protobuf.DynamicMessage; import com.google.protobuf.Int64Value; +import com.google.protobuf.Message; import java.math.BigDecimal; import java.math.BigInteger; import java.nio.charset.StandardCharsets; import java.time.Instant; import java.time.LocalDate; -import java.time.LocalDateTime; -import java.time.ZoneOffset; import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; @@ -59,6 +57,9 @@ import javax.annotation.Nullable; import org.apache.beam.sdk.io.gcp.bigquery.TableRowToStorageApiProto.SchemaConversionException; import org.apache.beam.sdk.io.gcp.bigquery.TableRowToStorageApiProto.SchemaInformation; +import org.apache.beam.sdk.schemas.Schema; +import org.apache.beam.sdk.schemas.logicaltypes.SqlTypes; +import org.apache.beam.sdk.values.Row; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Functions; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Predicates; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; @@ -1507,11 +1508,20 @@ private TableRow normalizeTableRowF( TYPE_MAP_PROTO_CONVERTERS.get(schemaInformation.getType()).apply("", value); return BaseEncoding.base64().encode(byteString.toByteArray()); case TIMESTAMP: - long timestampLongValue = (long) convertedValue; - long epochSeconds = timestampLongValue / 1_000_000L; - long nanoAdjustment = (timestampLongValue % 1_000_000L) * 1_000L; - Instant instant = Instant.ofEpochSecond(epochSeconds, nanoAdjustment); - return LocalDateTime.ofInstant(instant, ZoneOffset.UTC).format(TIMESTAMP_FORMATTER); + Map expectedTimestamps = + ImmutableMap.builder() + .put(43L, "1970-01-01 00:00:00.000043 UTC") + .put(-3600000000L, "1969-12-31 23:00:00 UTC") + .put(1234567000L, "1970-01-01 00:20:34.567 UTC") + .put(343L, "1970-01-01 00:00:00.000343 UTC") + .put(18000123456L, "1970-01-01 05:00:00.123456 UTC") + .put(123000L, "1970-01-01 00:00:00.123 UTC") + .put(253402300799999999L, "9999-12-31 23:59:59.999999 UTC") + .build(); + assertTrue( + "Missing literal timestamp expectation: " + value, + expectedTimestamps.containsKey(convertedValue)); + return expectedTimestamps.get(convertedValue); case DATE: int daysInt = (int) convertedValue; return LocalDate.ofEpochDay(daysInt).toString(); @@ -1599,6 +1609,235 @@ private void assertBaseRecord(DynamicMessage msg, boolean withF) { overriddenNames); } + @Test + public void testTimestampRoundTrip() throws Exception { + TableSchema schema = + new TableSchema() + .setFields(ImmutableList.of(new TableFieldSchema().setName("ts").setType("TIMESTAMP"))); + Descriptor descriptor = + TableRowToStorageApiProto.getDescriptorFromTableSchema(schema, true, false); + SchemaInformation information = SchemaInformation.fromTableSchema(schema); + DynamicMessage message = + DynamicMessage.newBuilder(descriptor) + .setField(descriptor.findFieldByName("ts"), 1788461503417123L) + .build(); + TableRow recovered = + TableRowToStorageApiProto.tableRowFromMessage( + information, message, true, Predicates.alwaysTrue()); + assertEquals("2026-09-03 18:51:43.417123 UTC", recovered.get("ts")); + Schema beamSchema = Schema.builder().addLogicalTypeField("ts", SqlTypes.TIMESTAMP).build(); + assertEquals( + Instant.parse("2026-09-03T18:51:43.417123Z"), + BigQueryUtils.toBeamRow(beamSchema, recovered).getValue("ts")); + } + + @Test + public void testIntegerTimestampBoundaries() throws Exception { + TableSchema schema = + new TableSchema() + .setFields(ImmutableList.of(new TableFieldSchema().setName("ts").setType("TIMESTAMP"))); + Descriptor descriptor = + TableRowToStorageApiProto.getDescriptorFromTableSchema(schema, true, false); + long[] micros = { + 0L, -1L, -1000001L, 43L, 1788461503417000L, -62135596800000000L, 253402300799999999L + }; + String[] expected = { + "1970-01-01 00:00:00 UTC", + "1969-12-31 23:59:59.999999 UTC", + "1969-12-31 23:59:58.999999 UTC", + "1970-01-01 00:00:00.000043 UTC", + "2026-09-03 18:51:43.417 UTC", + "0001-01-01 00:00:00 UTC", + "9999-12-31 23:59:59.999999 UTC" + }; + for (int i = 0; i < micros.length; i++) { + Message message = + DynamicMessage.newBuilder(descriptor) + .setField(descriptor.findFieldByName("ts"), micros[i]) + .build(); + assertTimestampConsumers( + schema, + message, + expected[i], + Instant.ofEpochSecond( + Math.floorDiv(micros[i], 1000000L), Math.floorMod(micros[i], 1000000L) * 1000L)); + } + } + + @Test + public void testProtoTimestampTruncatesToMicros() throws Exception { + DescriptorProto proto = + DescriptorProto.newBuilder() + .setName("Row") + .addField( + FieldDescriptorProto.newBuilder() + .setName("ts") + .setNumber(1) + .setType(FieldDescriptorProto.Type.TYPE_MESSAGE) + .setTypeName(".google.protobuf.Timestamp")) + .build(); + Descriptor descriptor = + com.google.protobuf.Descriptors.FileDescriptor.buildFrom( + DescriptorProtos.FileDescriptorProto.newBuilder() + .setName("timestamp_test.proto") + .addDependency("google/protobuf/timestamp.proto") + .addMessageType(proto) + .build(), + new com.google.protobuf.Descriptors.FileDescriptor[] { + com.google.protobuf.Timestamp.getDescriptor().getFile() + }) + .findMessageTypeByName("Row"); + TableSchema schema = + new TableSchema() + .setFields(ImmutableList.of(new TableFieldSchema().setName("ts").setType("TIMESTAMP"))); + long[] seconds = {1788461503L, -1L, 0L}; + int[] nanos = {417123456, 999999999, 999}; + String[] expected = { + "2026-09-03 18:51:43.417123 UTC", "1969-12-31 23:59:59.999999 UTC", "1970-01-01 00:00:00 UTC" + }; + for (int i = 0; i < seconds.length; i++) { + Message message = + DynamicMessage.newBuilder(descriptor) + .setField( + descriptor.findFieldByName("ts"), + com.google.protobuf.Timestamp.newBuilder() + .setSeconds(seconds[i]) + .setNanos(nanos[i]) + .build()) + .build(); + assertTimestampConsumers( + schema, message, expected[i], Instant.ofEpochSecond(seconds[i], nanos[i] / 1000 * 1000)); + } + } + + private static void assertTimestampConsumers( + TableSchema schema, Message message, String expected, Instant instant) throws Exception { + SchemaInformation information = SchemaInformation.fromTableSchema(schema); + TableRow recovered = + TableRowToStorageApiProto.tableRowFromMessage( + information, message, true, Predicates.alwaysTrue()); + assertEquals(expected, recovered.get("ts")); + Schema microsSchema = Schema.builder().addLogicalTypeField("ts", SqlTypes.TIMESTAMP).build(); + Schema millisSchema = Schema.builder().addDateTimeField("ts").build(); + assertEquals(instant, BigQueryUtils.toBeamRow(microsSchema, recovered).getValue("ts")); + assertEquals( + instant.toEpochMilli(), + BigQueryUtils.toBeamRow(millisSchema, recovered).getDateTime("ts").getMillis()); + Descriptor descriptor = + TableRowToStorageApiProto.getDescriptorFromTableSchema(schema, true, false); + DynamicMessage roundTrip = + TableRowToStorageApiProto.messageFromTableRow( + information, + descriptor, + recovered, + false, + false, + null, + null, + -1, + TableRowToStorageApiProto.ErrorCollector.DONT_COLLECT); + assertEquals(toEpochMicros(instant), roundTrip.getField(descriptor.findFieldByName("ts"))); + } + + @Test + public void testTimestampPicosPreservesPrecision() throws Exception { + TableSchema schema = + new TableSchema() + .setFields( + ImmutableList.of( + new TableFieldSchema() + .setName("ts") + .setType("TIMESTAMP") + .setTimestampPrecision(12L))); + Descriptor descriptor = + TableRowToStorageApiProto.getDescriptorFromTableSchema(schema, true, false); + SchemaInformation information = SchemaInformation.fromTableSchema(schema); + TableRow input = new TableRow().set("ts", "2024-01-15T10:30:45.123456789012Z"); + DynamicMessage message = + TableRowToStorageApiProto.messageFromTableRow( + information, + descriptor, + input, + false, + false, + null, + null, + -1, + TableRowToStorageApiProto.ErrorCollector.DONT_COLLECT); + assertEquals( + input, + TableRowToStorageApiProto.tableRowFromMessage( + information, message, true, Predicates.alwaysTrue())); + } + + @Test + public void testNestedRepeatedAndNullTimestamps() throws Exception { + TableSchema nested = + new TableSchema() + .setFields( + ImmutableList.of( + new TableFieldSchema().setName("ts").setType("TIMESTAMP"), + new TableFieldSchema().setName("missing").setType("TIMESTAMP"), + new TableFieldSchema() + .setName("times") + .setType("TIMESTAMP") + .setMode("REPEATED"))); + TableSchema schema = + new TableSchema() + .setFields( + ImmutableList.of( + new TableFieldSchema() + .setName("nested") + .setType("STRUCT") + .setFields(nested.getFields()))); + Schema nestedBeam = + Schema.builder() + .addLogicalTypeField("ts", SqlTypes.TIMESTAMP) + .addNullableField("missing", Schema.FieldType.logicalType(SqlTypes.TIMESTAMP)) + .addArrayField("times", Schema.FieldType.logicalType(SqlTypes.TIMESTAMP)) + .build(); + Schema beam = Schema.builder().addRowField("nested", nestedBeam).build(); + TableRow input = + new TableRow() + .set( + "nested", + new TableRow() + .set("ts", "2026-09-03T18:51:43.417123Z") + .set("missing", null) + .set( + "times", + ImmutableList.of("1969-12-31T23:59:59.999999Z", "1970-01-01T00:00:00Z"))); + SchemaInformation information = SchemaInformation.fromTableSchema(schema); + Descriptor descriptor = + TableRowToStorageApiProto.getDescriptorFromTableSchema(schema, true, false); + DynamicMessage message = + TableRowToStorageApiProto.messageFromTableRow( + information, + descriptor, + input, + false, + false, + null, + null, + -1, + TableRowToStorageApiProto.ErrorCollector.DONT_COLLECT); + TableRow recovered = + TableRowToStorageApiProto.tableRowFromMessage( + information, message, true, Predicates.alwaysTrue()); + assertEquals( + Row.withSchema(beam) + .addValue( + Row.withSchema(nestedBeam) + .addValues( + Instant.parse("2026-09-03T18:51:43.417123Z"), + null, + ImmutableList.of( + Instant.parse("1969-12-31T23:59:59.999999Z"), Instant.EPOCH)) + .build()) + .build(), + BigQueryUtils.toBeamRow(beam, recovered)); + } + @Test public void testMessageFromTableRow() throws Exception { TableRow tableRow =