From ff1ebc66e458519deca37275ba91650133188683 Mon Sep 17 00:00:00 2001 From: Phong Chuong <147636638+PhongChuong@users.noreply.github.com> Date: Fri, 19 Apr 2024 10:13:45 -0400 Subject: [PATCH] feat: Add support for inserting Range values (#3246) * feat: Add support for inserting Range values * fix: IT assert to use size instead of const --- .../java/com/google/cloud/bigquery/Range.java | 13 + .../com/google/cloud/bigquery/RangeTest.java | 18 + .../cloud/bigquery/it/ITBigQueryTest.java | 323 ++++++++++++------ 3 files changed, 246 insertions(+), 108 deletions(-) diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/Range.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/Range.java index 198bd28fa..8d244fbeb 100644 --- a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/Range.java +++ b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/Range.java @@ -19,6 +19,7 @@ import com.google.auto.value.AutoValue; import com.google.cloud.bigquery.FieldValue.Attribute; +import com.google.common.collect.ImmutableMap; import java.io.Serializable; import javax.annotation.Nullable; @@ -44,6 +45,18 @@ public FieldValue getEnd() { @Nullable abstract String getEndInner(); + /** Returns the start and end values of this range. */ + public ImmutableMap getValues() { + ImmutableMap.Builder result = ImmutableMap.builder(); + if (!getStart().isNull()) { + result.put("start", getStart().getStringValue()); + } + if (!getEnd().isNull()) { + result.put("end", getEnd().getStringValue()); + } + return result.build(); + } + /** Returns the type of the range. */ @Nullable public abstract FieldElementType getType(); diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/RangeTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/RangeTest.java index 45bb83259..2d98376b3 100644 --- a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/RangeTest.java +++ b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/RangeTest.java @@ -18,6 +18,7 @@ import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertTrue; +import com.google.common.collect.ImmutableMap; import org.junit.Test; public class RangeTest { @@ -77,6 +78,17 @@ public void testToBuilder() { compareRange(RANGE_TIMESTAMP, RANGE_TIMESTAMP.toBuilder().build()); } + @Test + public void testGetValues() { + compareRange(null, null, Range.of("[null, NULL)").getValues()); + compareRange(null, null, Range.of("[unbounded, UNBOUNDED)").getValues()); + compareRange(null, null, Range.of("[nUlL, uNbOuNdEd)").getValues()); + + compareRange(null, "2020-12-31", Range.of("[null, 2020-12-31)").getValues()); + compareRange("2020-01-01", null, Range.of("[2020-01-01, null)").getValues()); + compareRange("2020-01-01", "2020-12-31", Range.of("[2020-01-01, 2020-12-31)").getValues()); + } + private static void compareRange(Range expected, Range value) { assertEquals(expected.getStart(), value.getStart()); assertEquals(expected.getEnd(), value.getEnd()); @@ -97,4 +109,10 @@ private static void compareRange(String expectedStart, String expectedEnd, Range assertEquals(expectedEnd, range.getEnd().getStringValue()); } } + + private static void compareRange( + String expectedStart, String expectedEnd, ImmutableMap values) { + assertEquals(expectedStart, values.get("start")); + assertEquals(expectedEnd, values.get("end")); + } } diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java index 430235fba..d78549f6e 100644 --- a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java +++ b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java @@ -838,6 +838,157 @@ public class ITBigQueryTest { + " \"universe_domain\": \"fake.domain\"\n" + "}"; + private static final Schema RANGE_SCHEMA = + Schema.of( + Field.newBuilder("name", StandardSQLTypeName.STRING) + .setMode(Field.Mode.NULLABLE) + .setDescription("Name of the row") + .build(), + Field.newBuilder("date", StandardSQLTypeName.RANGE) + .setMode(Field.Mode.NULLABLE) + .setDescription("Range field with DATE") + .setRangeElementType(FieldElementType.newBuilder().setType("DATE").build()) + .build(), + Field.newBuilder("datetime", StandardSQLTypeName.RANGE) + .setMode(Field.Mode.NULLABLE) + .setDescription("Range field with DATETIME") + .setRangeElementType(FieldElementType.newBuilder().setType("DATETIME").build()) + .build(), + Field.newBuilder("timestamp", StandardSQLTypeName.RANGE) + .setMode(Field.Mode.NULLABLE) + .setDescription("Range field with TIMESTAMP") + .setRangeElementType(FieldElementType.newBuilder().setType("TIMESTAMP").build()) + .build()); + + private static final ImmutableMap RANGE_TEST_VALUES_DATES = + new ImmutableMap.Builder() + .put( + "bounded", + Range.newBuilder() + .setStart("2020-01-01") + .setEnd("2020-12-31") + .setType(FieldElementType.newBuilder().setType("DATE").build()) + .build()) + .put( + "unboundedStart", + Range.newBuilder() + .setStart(null) + .setEnd("2020-12-31") + .setType(FieldElementType.newBuilder().setType("DATE").build()) + .build()) + .put( + "unboundedEnd", + Range.newBuilder() + .setStart("2020-01-01") + .setEnd(null) + .setType(FieldElementType.newBuilder().setType("DATE").build()) + .build()) + .put( + "unbounded", + Range.newBuilder() + .setStart(null) + .setEnd(null) + .setType(FieldElementType.newBuilder().setType("DATE").build()) + .build()) + .build(); + + private static final ImmutableMap RANGE_TEST_VALUES_DATETIME = + new ImmutableMap.Builder() + .put( + "bounded", + Range.newBuilder() + .setStart("2014-08-19T05:41:35.220000") + .setEnd("2015-09-20T06:41:35.220000") + .setType(FieldElementType.newBuilder().setType("DATETIME").build()) + .build()) + .put( + "unboundedStart", + Range.newBuilder() + .setStart(null) + .setEnd("2015-09-20T06:41:35.220000") + .setType(FieldElementType.newBuilder().setType("DATETIME").build()) + .build()) + .put( + "unboundedEnd", + Range.newBuilder() + .setStart("2014-08-19T05:41:35.220000") + .setEnd(null) + .setType(FieldElementType.newBuilder().setType("DATETIME").build()) + .build()) + .put( + "unbounded", + Range.newBuilder() + .setStart(null) + .setEnd(null) + .setType(FieldElementType.newBuilder().setType("DATETIME").build()) + .build()) + .build(); + + private static final ImmutableMap RANGE_TEST_VALUES_TIMESTAMP = + new ImmutableMap.Builder() + .put( + "bounded", + Range.newBuilder() + .setStart("2014-08-19 12:41:35.220000+00:00") + .setEnd("2015-09-20 13:41:35.220000+01:00") + .setType(FieldElementType.newBuilder().setType("TIMESTAMP").build()) + .build()) + .put( + "unboundedStart", + Range.newBuilder() + .setStart(null) + .setEnd("2015-09-20 13:41:35.220000+01:00") + .setType(FieldElementType.newBuilder().setType("TIMESTAMP").build()) + .build()) + .put( + "unboundedEnd", + Range.newBuilder() + .setStart("2014-08-19 12:41:35.220000+00:00") + .setEnd(null) + .setType(FieldElementType.newBuilder().setType("TIMESTAMP").build()) + .build()) + .put( + "unbounded", + Range.newBuilder() + .setStart(null) + .setEnd(null) + .setType(FieldElementType.newBuilder().setType("TIMESTAMP").build()) + .build()) + .build(); + + // timestamps are returned as seconds since epoch + private static final ImmutableMap RANGE_TEST_VALUES_EXPECTED_TIMESTAMP = + new ImmutableMap.Builder() + .put( + "bounded", + Range.newBuilder() + .setStart("1408452095.220000") + .setEnd("1442752895.220000") + .setType(FieldElementType.newBuilder().setType("TIMESTAMP").build()) + .build()) + .put( + "unboundedStart", + Range.newBuilder() + .setStart(null) + .setEnd("1442752895.220000") + .setType(FieldElementType.newBuilder().setType("TIMESTAMP").build()) + .build()) + .put( + "unboundedEnd", + Range.newBuilder() + .setStart("1408452095.220000") + .setEnd(null) + .setType(FieldElementType.newBuilder().setType("TIMESTAMP").build()) + .build()) + .put( + "unbounded", + Range.newBuilder() + .setStart(null) + .setEnd(null) + .setType(FieldElementType.newBuilder().setType("TIMESTAMP").build()) + .build()) + .build(); + private static BigQuery bigquery; private static Storage storage; @@ -1313,123 +1464,79 @@ public void testIntervalType() throws InterruptedException { } @Test - public void testRangeCreateTable() throws InterruptedException { - String tableName = "test_range_create_table"; + public void testRangeType() throws InterruptedException { + String tableName = "test_range_type_table"; TableId tableId = TableId.of(DATASET, tableName); - Field rangeFieldWithDate = - Field.newBuilder("rangeFieldDate", StandardSQLTypeName.RANGE) - .setMode(Field.Mode.NULLABLE) - .setDescription("Range field with DATE") - .setRangeElementType(FieldElementType.newBuilder().setType("DATE").build()) - .build(); - - Field rangeFieldWithDatetime = - Field.newBuilder("rangeFieldDatetime", StandardSQLTypeName.RANGE) - .setMode(Field.Mode.NULLABLE) - .setDescription("Range field with DATETIME") - .setRangeElementType(FieldElementType.newBuilder().setType("DATETIME").build()) - .build(); - - Field rangeFieldWithTimestamp = - Field.newBuilder("rangeFieldTimestamp", StandardSQLTypeName.RANGE) - .setMode(Field.Mode.NULLABLE) - .setDescription("Range field with TIMESTAMP") - .setRangeElementType(FieldElementType.newBuilder().setType("TIMESTAMP").build()) - .build(); - Schema schema = Schema.of(rangeFieldWithDate, rangeFieldWithDatetime, rangeFieldWithTimestamp); - StandardTableDefinition standardTableDefinition = StandardTableDefinition.of(schema); + StandardTableDefinition standardTableDefinition = StandardTableDefinition.of(RANGE_SCHEMA); try { - // Create a table with a RANGE column. + // Create a table with a RANGE columns and verify the result. Table createdTable = bigquery.create(TableInfo.of(tableId, standardTableDefinition)); assertNotNull(createdTable); - } finally { - assertTrue(bigquery.delete(tableId)); - } - } - - @Test - public void testRangeType() throws InterruptedException { - // TODO: Combine testRangeType test with testRangeCreateTable test. - String tableName = "test_range_type_table"; - QueryJobConfiguration createTable = - QueryJobConfiguration.newBuilder( - String.format( - "CREATE TABLE %s AS SELECT RANGE(DATE '2020-01-01', DATE '2020-12-31') as date, \n" - + "RANGE(DATETIME '2020-01-01T12:00:00', DATETIME '2020-12-31T12:00:00') as datetime, \n" - + "RANGE(TIMESTAMP '2014-01-01 07:00:00.000000+00:00', TIMESTAMP '2015-01-01 07:00:00.000000+00:00') as timestamp", - tableName)) - .setDefaultDataset(DatasetId.of(DATASET)) - .setUseLegacySql(false) - .build(); - bigquery.query(createTable); - String query = - String.format( - "SELECT date, datetime, timestamp\n" - + "FROM %s.%s\n" - + "WHERE date = @dateParam\n" - + "AND datetime = @datetimeParam\n" - + "AND timestamp = @timestampParam", - DATASET, tableName); + Table remoteTable = bigquery.getTable(DATASET, tableName); + Schema remoteSchema = remoteTable.getDefinition().getSchema(); + assertEquals(RANGE_SCHEMA, remoteSchema); + + // Insert range values to the table. + InsertAllRequest.Builder request = InsertAllRequest.newBuilder(tableId); + for (String name : RANGE_TEST_VALUES_DATES.keySet()) { + ImmutableMap.Builder builder = ImmutableMap.builder(); + builder.put("name", name); + builder.put("date", RANGE_TEST_VALUES_DATES.get(name).getValues()); + builder.put("datetime", RANGE_TEST_VALUES_DATETIME.get(name).getValues()); + builder.put("timestamp", RANGE_TEST_VALUES_TIMESTAMP.get(name).getValues()); + request.addRow(builder.build()); + } + bigquery.insertAll(request.build()); - Range dateRange = - Range.newBuilder() - .setType(FieldElementType.newBuilder().setType("DATE").build()) - .setStart("2020-01-01") - .setEnd("2020-12-31") - .build(); - Range datetimeRange = - Range.newBuilder() - .setType(FieldElementType.newBuilder().setType("DATETIME").build()) - .setStart("2020-01-01T12:00:00") - .setEnd("2020-12-31T12:00:00") - .build(); - Range timestampRange = - Range.newBuilder() - .setType(FieldElementType.newBuilder().setType("TIMESTAMP").build()) - .setStart("2014-01-01 07:00:00.000000+00:00") - .setEnd("2015-01-01 07:00:00.000000+00:00") - .build(); + // Test listTableData + TableResult result = bigquery.listTableData(DATASET, tableName, RANGE_SCHEMA); + assertEquals(RANGE_TEST_VALUES_DATES.size(), Iterables.size(result.getValues())); + for (FieldValueList values : result.iterateAll()) { + String name = values.get("name").getStringValue(); + assertEquals(RANGE_TEST_VALUES_DATES.get(name), values.get("date").getRangeValue()); + assertEquals(RANGE_TEST_VALUES_DATETIME.get(name), values.get("datetime").getRangeValue()); + assertEquals( + RANGE_TEST_VALUES_EXPECTED_TIMESTAMP.get(name), + values.get("timestamp").getRangeValue()); + } - // Test Query Parameter. - QueryJobConfiguration config = - QueryJobConfiguration.newBuilder(query) - .setDefaultDataset(DatasetId.of(DATASET)) - .addNamedParameter("dateParam", QueryParameterValue.range(dateRange)) - .addNamedParameter("datetimeParam", QueryParameterValue.range(datetimeRange)) - .addNamedParameter("timestampParam", QueryParameterValue.range(timestampRange)) - .build(); - TableResult result = bigquery.query(config); - assertEquals(1, Iterables.size(result.getValues())); - for (FieldValueList values : result.iterateAll()) { - assertEquals(dateRange.getStart(), values.get("date").getRangeValue().getStart()); - assertEquals(dateRange.getEnd(), values.get("date").getRangeValue().getEnd()); - assertEquals(datetimeRange.getStart(), values.get("datetime").getRangeValue().getStart()); - assertEquals(datetimeRange.getEnd(), values.get("datetime").getRangeValue().getEnd()); - // timestamps are returned as seconds since epoch - assertEquals( - 1388559600000000L, - values.get("timestamp").getRangeValue().getStart().getTimestampValue()); - assertEquals( - 1420095600000000L, values.get("timestamp").getRangeValue().getEnd().getTimestampValue()); - } + // Test Query Parameter by selecting for the bounded Range entry only. + String query = + String.format( + "SELECT name, date, datetime, timestamp\n" + + "FROM %s.%s\n" + + "WHERE date = @dateParam\n" + + "AND datetime = @datetimeParam\n" + + "AND timestamp = @timestampParam", + DATASET, tableName); + + QueryJobConfiguration config = + QueryJobConfiguration.newBuilder(query) + .setDefaultDataset(DatasetId.of(DATASET)) + .addNamedParameter( + "dateParam", QueryParameterValue.range(RANGE_TEST_VALUES_DATES.get("bounded"))) + .addNamedParameter( + "datetimeParam", + QueryParameterValue.range(RANGE_TEST_VALUES_DATETIME.get("bounded"))) + .addNamedParameter( + "timestampParam", + QueryParameterValue.range(RANGE_TEST_VALUES_TIMESTAMP.get("bounded"))) + .build(); + result = bigquery.query(config); - // Test listTableData. - Schema schema = result.getSchema(); - result = bigquery.listTableData(DATASET, tableName, schema); - assertEquals(1, Iterables.size(result.getValues())); - for (FieldValueList values : result.iterateAll()) { - assertEquals(dateRange.getStart(), values.get("date").getRangeValue().getStart()); - assertEquals(dateRange.getEnd(), values.get("date").getRangeValue().getEnd()); - assertEquals(datetimeRange.getStart(), values.get("datetime").getRangeValue().getStart()); - assertEquals(datetimeRange.getEnd(), values.get("datetime").getRangeValue().getEnd()); - // timestamps are returned as seconds since epoch - assertEquals( - 1388559600000000L, - values.get("timestamp").getRangeValue().getStart().getTimestampValue()); - assertEquals( - 1420095600000000L, values.get("timestamp").getRangeValue().getEnd().getTimestampValue()); + assertEquals(1, Iterables.size(result.getValues())); + for (FieldValueList values : result.iterateAll()) { + String name = values.get("name").getStringValue(); + assertEquals(RANGE_TEST_VALUES_DATES.get(name), values.get("date").getRangeValue()); + assertEquals(RANGE_TEST_VALUES_DATETIME.get(name), values.get("datetime").getRangeValue()); + assertEquals( + RANGE_TEST_VALUES_EXPECTED_TIMESTAMP.get(name), + values.get("timestamp").getRangeValue()); + } + } finally { + assertTrue(bigquery.delete(tableId)); } }