test: Add ExecuteQuery integration tests for SQL (#2285)

Change-Id: I5d7cdfcdc7f7a5db60b88278c13e6fe1358f0920

google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/it/ExecuteQueryIT.java

@@ -0,0 +1,348 @@
+/*
+ * Copyright 2024 Google LLC
+ *
+ * Licensed 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
+ *
+ *     https://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 com.google.cloud.bigtable.data.v2.it;
+
+import static com.google.common.truth.Truth.assertThat;
+import static com.google.common.truth.TruthJUnit.assume;
+import static org.junit.Assert.assertThrows;
+
+import com.google.cloud.Date;
+import com.google.cloud.bigtable.data.v2.BigtableDataClient;
+import com.google.cloud.bigtable.data.v2.models.RowMutation;
+import com.google.cloud.bigtable.data.v2.models.TableId;
+import com.google.cloud.bigtable.data.v2.models.sql.ResultSet;
+import com.google.cloud.bigtable.data.v2.models.sql.SqlType;
+import com.google.cloud.bigtable.data.v2.models.sql.Statement;
+import com.google.cloud.bigtable.data.v2.models.sql.Struct;
+import com.google.cloud.bigtable.test_helpers.env.EmulatorEnv;
+import com.google.cloud.bigtable.test_helpers.env.TestEnvRule;
+import com.google.protobuf.ByteString;
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.UUID;
+import org.junit.BeforeClass;
+import org.junit.ClassRule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.threeten.bp.Instant;
+
+@RunWith(JUnit4.class)
+public class ExecuteQueryIT {
+
+  @ClassRule public static TestEnvRule testEnvRule = new TestEnvRule();
+  private static BigtableDataClient dataClient;
+  private static String tableId;
+  private static String cf;
+  private static String uniquePrefix;
+
+  @BeforeClass
+  public static void setUpAll() throws IOException {
+    assume()
+        .withMessage("ExecuteQuery is not supported on Emulator")
+        .that(testEnvRule.env())
+        .isNotInstanceOf(EmulatorEnv.class);
+
+    tableId = testEnvRule.env().getTableId();
+    dataClient = testEnvRule.env().getDataClient();
+    cf = testEnvRule.env().getFamilyId();
+    uniquePrefix = UUID.randomUUID() + "-execute-query-it-";
+
+    dataClient.mutateRow(
+        RowMutation.create(TableId.of(tableId), uniquePrefix + "a")
+            .setCell(cf, ByteString.copyFromUtf8("qual"), 1000, ByteString.copyFromUtf8("old"))
+            .setCell(cf, ByteString.copyFromUtf8("qual2"), 1000, ByteString.copyFromUtf8("old2")));
+    // Overwrite the previous values. Used for testing with_history
+    dataClient.mutateRow(
+        RowMutation.create(TableId.of(tableId), uniquePrefix + "a")
+            .setCell(cf, ByteString.copyFromUtf8("qual"), 10000, ByteString.copyFromUtf8("val"))
+            .setCell(cf, ByteString.copyFromUtf8("qual2"), 10000, ByteString.copyFromUtf8("val2"))
+            .setCell(cf, ByteString.copyFromUtf8("qual3"), 10000, ByteString.copyFromUtf8("val3")));
+    dataClient.mutateRow(
+        RowMutation.create(TableId.of(tableId), uniquePrefix + "b")
+            .setCell(cf, ByteString.copyFromUtf8("qual"), 10000, ByteString.copyFromUtf8("bval"))
+            .setCell(
+                cf, ByteString.copyFromUtf8("qual2"), 10000, ByteString.copyFromUtf8("bval2")));
+  }
+
+  @Test
+  public void selectStar() {
+    try (ResultSet rs =
+        dataClient.executeQuery(
+            Statement.of(
+                "SELECT * FROM " + tableId + " WHERE _key LIKE '" + uniquePrefix + "%'"))) {
+      assertThat(rs.next()).isTrue();
+      assertThat(rs.getBytes("_key")).isEqualTo(ByteString.copyFromUtf8(uniquePrefix + "a"));
+      assertThat(
+              rs.getMap(cf, SqlType.mapOf(SqlType.bytes(), SqlType.bytes()))
+                  .get(ByteString.copyFromUtf8("qual")))
+          .isEqualTo(ByteString.copyFromUtf8("val"));
+
+      assertThat(rs.next()).isTrue();
+      assertThat(rs.getBytes("_key")).isEqualTo(ByteString.copyFromUtf8(uniquePrefix + "b"));
+      assertThat(
+              rs.getMap(cf, SqlType.mapOf(SqlType.bytes(), SqlType.bytes()))
+                  .get(ByteString.copyFromUtf8("qual")))
+          .isEqualTo(ByteString.copyFromUtf8("bval"));
+
+      assertThat(rs.next()).isFalse();
+    }
+  }
+
+  @Test
+  public void withHistoryQuery() {
+    try (ResultSet rs =
+        dataClient.executeQuery(
+            Statement.of(
+                "SELECT * FROM `"
+                    + tableId
+                    + "`(with_history => true) WHERE _key LIKE '"
+                    + uniquePrefix
+                    + "%'"))) {
+
+      assertThat(rs.next()).isTrue();
+      assertThat(rs.getBytes("_key")).isEqualTo(ByteString.copyFromUtf8(uniquePrefix + "a"));
+      Map<ByteString, List<Struct>> rowACf = rs.getMap(cf, SqlType.historicalMap());
+      List<Struct> rowAQual = rowACf.get(ByteString.copyFromUtf8("qual"));
+      assertThat(rowAQual.size()).isEqualTo(2);
+      Struct rowAQual_0 = rowAQual.get(0);
+      assertThat(rowAQual_0.getBytes("value")).isEqualTo(ByteString.copyFromUtf8("val"));
+      // timestamp in micros above so we divide by 1000
+      assertThat(rowAQual_0.getTimestamp("timestamp")).isEqualTo(Instant.ofEpochMilli(10));
+      Struct rowAQual_1 = rowAQual.get(1);
+      assertThat(rowAQual_1.getBytes("value")).isEqualTo(ByteString.copyFromUtf8("old"));
+      assertThat(rowAQual_1.getTimestamp("timestamp")).isEqualTo(Instant.ofEpochMilli(1));
+
+      assertThat(rs.next()).isTrue();
+      assertThat(rs.getBytes("_key")).isEqualTo(ByteString.copyFromUtf8(uniquePrefix + "b"));
+      Map<ByteString, List<Struct>> rowBCf = rs.getMap(cf, SqlType.historicalMap());
+      List<Struct> rowBQual = rowBCf.get(ByteString.copyFromUtf8("qual"));
+      assertThat(rowBQual.size()).isEqualTo(1);
+      Struct rowBQual_0 = rowBQual.get(0);
+      assertThat(rowBQual_0.getBytes("value")).isEqualTo(ByteString.copyFromUtf8("bval"));
+
+      assertThat(rs.next()).isFalse();
+    }
+  }
+
+  @Test
+  public void allTypes() {
+    try (ResultSet rs =
+        dataClient.executeQuery(
+            Statement.of(
+                "SELECT 'stringVal' AS strCol, b'foo' as bytesCol, 1 AS intCol, CAST(1.2 AS FLOAT32) as f32Col, "
+                    + "CAST(1.3 AS FLOAT64) as f64Col, true as boolCol, TIMESTAMP_FROM_UNIX_MILLIS(1000) AS tsCol, "
+                    + "DATE(2024, 06, 01) as dateCol, STRUCT(1 as a, \"foo\" as b) AS structCol, [1,2,3] AS arrCol, "
+                    + cf
+                    + " as mapCol FROM `"
+                    + tableId
+                    + "` WHERE _key='"
+                    + uniquePrefix
+                    + "a' LIMIT 1"))) {
+
+      assertThat(rs.next()).isTrue();
+      assertThat(rs.getString("strCol")).isEqualTo("stringVal");
+      assertThat(rs.getString(0)).isEqualTo("stringVal");
+      assertThat(rs.getBytes("bytesCol")).isEqualTo(ByteString.copyFromUtf8("foo"));
+      assertThat(rs.getBytes(1)).isEqualTo(ByteString.copyFromUtf8("foo"));
+      assertThat(rs.getLong("intCol")).isEqualTo(1L);
+      assertThat(rs.getLong(2)).isEqualTo(1L);
+      assertThat(rs.getFloat("f32Col")).isEqualTo(1.2f);
+      assertThat(rs.getFloat(3)).isEqualTo(1.2f);
+      assertThat(rs.getDouble("f64Col")).isEqualTo(1.3d);
+      assertThat(rs.getDouble(4)).isEqualTo(1.3d);
+      assertThat(rs.getBoolean("boolCol")).isTrue();
+      assertThat(rs.getBoolean(5)).isTrue();
+      assertThat(rs.getTimestamp("tsCol")).isEqualTo(Instant.ofEpochMilli(1000));
+      assertThat(rs.getTimestamp(6)).isEqualTo(Instant.ofEpochMilli(1000));
+      assertThat(rs.getDate("dateCol")).isEqualTo(Date.fromYearMonthDay(2024, 6, 1));
+      assertThat(rs.getDate(7)).isEqualTo(Date.fromYearMonthDay(2024, 6, 1));
+      assertThat(rs.getStruct("structCol").getLong("a")).isEqualTo(1);
+      assertThat(rs.getStruct("structCol").getString("b")).isEqualTo("foo");
+      assertThat(rs.getStruct(8).getLong("a")).isEqualTo(1);
+      assertThat(rs.getStruct(8).getString("b")).isEqualTo("foo");
+      assertThat(rs.getList("arrCol", SqlType.arrayOf(SqlType.int64())))
+          .isEqualTo(Arrays.asList(1L, 2L, 3L));
+      assertThat(rs.getList(9, SqlType.arrayOf(SqlType.int64())))
+          .isEqualTo(Arrays.asList(1L, 2L, 3L));
+      assertThat(rs.getMap("mapCol", SqlType.mapOf(SqlType.bytes(), SqlType.bytes())))
+          .isEqualTo(
+              new HashMap<ByteString, ByteString>() {
+                {
+                  put(ByteString.copyFromUtf8("qual"), ByteString.copyFromUtf8("val"));
+                  put(ByteString.copyFromUtf8("qual2"), ByteString.copyFromUtf8("val2"));
+                  put(ByteString.copyFromUtf8("qual3"), ByteString.copyFromUtf8("val3"));
+                }
+              });
+      assertThat(rs.getMap(10, SqlType.mapOf(SqlType.bytes(), SqlType.bytes())))
+          .isEqualTo(
+              new HashMap<ByteString, ByteString>() {
+                {
+                  put(ByteString.copyFromUtf8("qual"), ByteString.copyFromUtf8("val"));
+                  put(ByteString.copyFromUtf8("qual2"), ByteString.copyFromUtf8("val2"));
+                  put(ByteString.copyFromUtf8("qual3"), ByteString.copyFromUtf8("val3"));
+                }
+              });
+
+      assertThat(rs.next()).isFalse();
+    }
+  }
+
+  @Test
+  public void allQueryParamsTypes() {
+    ResultSet rs =
+        dataClient.executeQuery(
+            Statement.newBuilder(
+                    "SELECT @stringParam AS strCol, @bytesParam as bytesCol, @int64Param AS intCol, "
+                        + "@doubleParam AS doubleCol, @floatParam AS floatCol, @boolParam AS boolCol, "
+                        + "@tsParam AS tsCol, @dateParam AS dateCol, @byteArrayParam AS byteArrayCol, "
+                        + "@stringArrayParam AS stringArrayCol, @intArrayParam AS intArrayCol, "
+                        + "@floatArrayParam AS floatArrayCol, @doubleArrayParam AS doubleArrayCol, "
+                        + "@boolArrayParam AS boolArrayCol, @tsArrayParam AS tsArrayCol, "
+                        + "@dateArrayParam AS dateArrayCol")
+                .setStringParam("stringParam", "stringVal")
+                .setBytesParam("bytesParam", ByteString.copyFromUtf8("foo"))
+                .setLongParam("int64Param", 1L)
+                .setDoubleParam("doubleParam", 1.3d)
+                .setFloatParam("floatParam", 1.4f)
+                .setBooleanParam("boolParam", true)
+                .setTimestampParam("tsParam", Instant.ofEpochMilli(1000))
+                .setDateParam("dateParam", Date.fromYearMonthDay(2024, 6, 1))
+                .setListParam(
+                    "byteArrayParam",
+                    Arrays.asList(
+                        ByteString.copyFromUtf8("foo"), null, ByteString.copyFromUtf8("bar")),
+                    SqlType.arrayOf(SqlType.bytes()))
+                .setListParam(
+                    "stringArrayParam",
+                    Arrays.asList("foo", null, "bar"),
+                    SqlType.arrayOf(SqlType.string()))
+                .setListParam(
+                    "intArrayParam", Arrays.asList(1L, null, 2L), SqlType.arrayOf(SqlType.int64()))
+                .setListParam(
+                    "floatArrayParam",
+                    Arrays.asList(1.2f, null, 1.3f),
+                    SqlType.arrayOf(SqlType.float32()))
+                .setListParam(
+                    "doubleArrayParam",
+                    Arrays.asList(1.4d, null, 1.5d),
+                    SqlType.arrayOf(SqlType.float64()))
+                .setListParam(
+                    "boolArrayParam",
+                    Arrays.asList(true, null, false),
+                    SqlType.arrayOf(SqlType.bool()))
+                .setListParam(
+                    "tsArrayParam",
+                    Arrays.asList(
+                        Instant.ofEpochSecond(1000, 1000), null, Instant.ofEpochSecond(2000, 2000)),
+                    SqlType.arrayOf(SqlType.timestamp()))
+                .setListParam(
+                    "dateArrayParam",
+                    Arrays.asList(
+                        Date.fromYearMonthDay(2024, 8, 1), null, Date.fromYearMonthDay(2024, 8, 2)),
+                    SqlType.arrayOf(SqlType.date()))
+                .build());
+
+    assertThat(rs.next()).isTrue();
+    assertThat(rs.getString("strCol")).isEqualTo("stringVal");
+    assertThat(rs.getString(0)).isEqualTo("stringVal");
+    assertThat(rs.getBytes("bytesCol")).isEqualTo(ByteString.copyFromUtf8("foo"));
+    assertThat(rs.getBytes(1)).isEqualTo(ByteString.copyFromUtf8("foo"));
+    assertThat(rs.getLong("intCol")).isEqualTo(1L);
+    assertThat(rs.getLong(2)).isEqualTo(1L);
+    assertThat(rs.getDouble("doubleCol")).isEqualTo(1.3d);
+    assertThat(rs.getDouble(3)).isEqualTo(1.3d);
+    assertThat(rs.getFloat("floatCol")).isEqualTo(1.4f);
+    assertThat(rs.getFloat(4)).isEqualTo(1.4f);
+    assertThat(rs.getBoolean("boolCol")).isTrue();
+    assertThat(rs.getBoolean(5)).isTrue();
+    assertThat(rs.getTimestamp("tsCol")).isEqualTo(Instant.ofEpochMilli(1000));
+    assertThat(rs.getTimestamp(6)).isEqualTo(Instant.ofEpochMilli(1000));
+    assertThat(rs.getDate("dateCol")).isEqualTo(Date.fromYearMonthDay(2024, 6, 1));
+    assertThat(rs.getDate(7)).isEqualTo(Date.fromYearMonthDay(2024, 6, 1));
+    assertThat(rs.getList("byteArrayCol", SqlType.arrayOf(SqlType.bytes())))
+        .isEqualTo(
+            Arrays.asList(ByteString.copyFromUtf8("foo"), null, ByteString.copyFromUtf8("bar")));
+    assertThat(rs.getList(8, SqlType.arrayOf(SqlType.bytes())))
+        .isEqualTo(
+            Arrays.asList(ByteString.copyFromUtf8("foo"), null, ByteString.copyFromUtf8("bar")));
+    assertThat(rs.getList("stringArrayCol", SqlType.arrayOf(SqlType.string())))
+        .isEqualTo(Arrays.asList("foo", null, "bar"));
+    assertThat(rs.getList(9, SqlType.arrayOf(SqlType.string())))
+        .isEqualTo(Arrays.asList("foo", null, "bar"));
+    assertThat(rs.getList("intArrayCol", SqlType.arrayOf(SqlType.int64())))
+        .isEqualTo(Arrays.asList(1L, null, 2L));
+    assertThat(rs.getList(10, SqlType.arrayOf(SqlType.int64())))
+        .isEqualTo(Arrays.asList(1L, null, 2L));
+    assertThat(rs.getList("floatArrayCol", SqlType.arrayOf(SqlType.float32())))
+        .isEqualTo(Arrays.asList(1.2f, null, 1.3f));
+    assertThat(rs.getList(11, SqlType.arrayOf(SqlType.float32())))
+        .isEqualTo(Arrays.asList(1.2f, null, 1.3f));
+    assertThat(rs.getList("doubleArrayCol", SqlType.arrayOf(SqlType.float64())))
+        .isEqualTo(Arrays.asList(1.4d, null, 1.5d));
+    assertThat(rs.getList(12, SqlType.arrayOf(SqlType.float64())))
+        .isEqualTo(Arrays.asList(1.4d, null, 1.5d));
+    assertThat(rs.getList("boolArrayCol", SqlType.arrayOf(SqlType.bool())))
+        .isEqualTo(Arrays.asList(true, null, false));
+    assertThat(rs.getList(13, SqlType.arrayOf(SqlType.bool())))
+        .isEqualTo(Arrays.asList(true, null, false));
+    assertThat(rs.getList("tsArrayCol", SqlType.arrayOf(SqlType.timestamp())))
+        .isEqualTo(
+            Arrays.asList(
+                Instant.ofEpochSecond(1000, 1000), null, Instant.ofEpochSecond(2000, 2000)));
+    assertThat(rs.getList(14, SqlType.arrayOf(SqlType.timestamp())))
+        .isEqualTo(
+            Arrays.asList(
+                Instant.ofEpochSecond(1000, 1000), null, Instant.ofEpochSecond(2000, 2000)));
+    assertThat(rs.getList("dateArrayCol", SqlType.arrayOf(SqlType.date())))
+        .isEqualTo(
+            Arrays.asList(
+                Date.fromYearMonthDay(2024, 8, 1), null, Date.fromYearMonthDay(2024, 8, 2)));
+    assertThat(rs.getList(15, SqlType.arrayOf(SqlType.date())))
+        .isEqualTo(
+            Arrays.asList(
+                Date.fromYearMonthDay(2024, 8, 1), null, Date.fromYearMonthDay(2024, 8, 2)));
+  }
+
+  @Test
+  public void testNullColumns() {
+    try (ResultSet rs =
+        dataClient.executeQuery(
+            Statement.of(
+                "SELECT cf['qual'] AS neverNull, cf['qual3'] AS maybeNull FROM "
+                    + tableId
+                    + " WHERE _key LIKE '"
+                    + uniquePrefix
+                    + "%'"))) {
+      assertThat(rs.next()).isTrue();
+      assertThat(rs.getBytes("neverNull")).isEqualTo(ByteString.copyFromUtf8("val"));
+      // qual3 is set in row A but not row B
+      assertThat(rs.isNull("maybeNull")).isFalse();
+      assertThat(rs.isNull(1)).isFalse();
+      assertThat(rs.getBytes("maybeNull")).isEqualTo(ByteString.copyFromUtf8("val3"));
+      assertThat(rs.next()).isTrue();
+      assertThat(rs.getBytes("neverNull")).isEqualTo(ByteString.copyFromUtf8("bval"));
+      assertThat(rs.isNull("maybeNull")).isTrue();
+      assertThat(rs.isNull(1)).isTrue();
+      assertThrows(NullPointerException.class, () -> rs.getBytes("maybeNull"));
+      assertThrows(NullPointerException.class, () -> rs.getBytes(1));
+      assertThat(rs.next()).isFalse();
+    }
+  }
+}