Skip to content

Commit c88b31e

Browse files
jackdingilianjackdingilian
authored
test: Add ExecuteQuery integration tests for SQL (#2285)
Change-Id: I5d7cdfcdc7f7a5db60b88278c13e6fe1358f0920
1 parent ceebef5 commit c88b31e

File tree

1 file changed

+348
-0
lines changed

1 file changed

+348
-0
lines changed

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

Lines changed: 348 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,348 @@
1+
/*
2+
* Copyright 2024 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package com.google.cloud.bigtable.data.v2.it;
17+
18+
import static com.google.common.truth.Truth.assertThat;
19+
import static com.google.common.truth.TruthJUnit.assume;
20+
import static org.junit.Assert.assertThrows;
21+
22+
import com.google.cloud.Date;
23+
import com.google.cloud.bigtable.data.v2.BigtableDataClient;
24+
import com.google.cloud.bigtable.data.v2.models.RowMutation;
25+
import com.google.cloud.bigtable.data.v2.models.TableId;
26+
import com.google.cloud.bigtable.data.v2.models.sql.ResultSet;
27+
import com.google.cloud.bigtable.data.v2.models.sql.SqlType;
28+
import com.google.cloud.bigtable.data.v2.models.sql.Statement;
29+
import com.google.cloud.bigtable.data.v2.models.sql.Struct;
30+
import com.google.cloud.bigtable.test_helpers.env.EmulatorEnv;
31+
import com.google.cloud.bigtable.test_helpers.env.TestEnvRule;
32+
import com.google.protobuf.ByteString;
33+
import java.io.IOException;
34+
import java.util.Arrays;
35+
import java.util.HashMap;
36+
import java.util.List;
37+
import java.util.Map;
38+
import java.util.UUID;
39+
import org.junit.BeforeClass;
40+
import org.junit.ClassRule;
41+
import org.junit.Test;
42+
import org.junit.runner.RunWith;
43+
import org.junit.runners.JUnit4;
44+
import org.threeten.bp.Instant;
45+
46+
@RunWith(JUnit4.class)
47+
public class ExecuteQueryIT {
48+
49+
@ClassRule public static TestEnvRule testEnvRule = new TestEnvRule();
50+
private static BigtableDataClient dataClient;
51+
private static String tableId;
52+
private static String cf;
53+
private static String uniquePrefix;
54+
55+
@BeforeClass
56+
public static void setUpAll() throws IOException {
57+
assume()
58+
.withMessage("ExecuteQuery is not supported on Emulator")
59+
.that(testEnvRule.env())
60+
.isNotInstanceOf(EmulatorEnv.class);
61+
62+
tableId = testEnvRule.env().getTableId();
63+
dataClient = testEnvRule.env().getDataClient();
64+
cf = testEnvRule.env().getFamilyId();
65+
uniquePrefix = UUID.randomUUID() + "-execute-query-it-";
66+
67+
dataClient.mutateRow(
68+
RowMutation.create(TableId.of(tableId), uniquePrefix + "a")
69+
.setCell(cf, ByteString.copyFromUtf8("qual"), 1000, ByteString.copyFromUtf8("old"))
70+
.setCell(cf, ByteString.copyFromUtf8("qual2"), 1000, ByteString.copyFromUtf8("old2")));
71+
// Overwrite the previous values. Used for testing with_history
72+
dataClient.mutateRow(
73+
RowMutation.create(TableId.of(tableId), uniquePrefix + "a")
74+
.setCell(cf, ByteString.copyFromUtf8("qual"), 10000, ByteString.copyFromUtf8("val"))
75+
.setCell(cf, ByteString.copyFromUtf8("qual2"), 10000, ByteString.copyFromUtf8("val2"))
76+
.setCell(cf, ByteString.copyFromUtf8("qual3"), 10000, ByteString.copyFromUtf8("val3")));
77+
dataClient.mutateRow(
78+
RowMutation.create(TableId.of(tableId), uniquePrefix + "b")
79+
.setCell(cf, ByteString.copyFromUtf8("qual"), 10000, ByteString.copyFromUtf8("bval"))
80+
.setCell(
81+
cf, ByteString.copyFromUtf8("qual2"), 10000, ByteString.copyFromUtf8("bval2")));
82+
}
83+
84+
@Test
85+
public void selectStar() {
86+
try (ResultSet rs =
87+
dataClient.executeQuery(
88+
Statement.of(
89+
"SELECT * FROM " + tableId + " WHERE _key LIKE '" + uniquePrefix + "%'"))) {
90+
assertThat(rs.next()).isTrue();
91+
assertThat(rs.getBytes("_key")).isEqualTo(ByteString.copyFromUtf8(uniquePrefix + "a"));
92+
assertThat(
93+
rs.getMap(cf, SqlType.mapOf(SqlType.bytes(), SqlType.bytes()))
94+
.get(ByteString.copyFromUtf8("qual")))
95+
.isEqualTo(ByteString.copyFromUtf8("val"));
96+
97+
assertThat(rs.next()).isTrue();
98+
assertThat(rs.getBytes("_key")).isEqualTo(ByteString.copyFromUtf8(uniquePrefix + "b"));
99+
assertThat(
100+
rs.getMap(cf, SqlType.mapOf(SqlType.bytes(), SqlType.bytes()))
101+
.get(ByteString.copyFromUtf8("qual")))
102+
.isEqualTo(ByteString.copyFromUtf8("bval"));
103+
104+
assertThat(rs.next()).isFalse();
105+
}
106+
}
107+
108+
@Test
109+
public void withHistoryQuery() {
110+
try (ResultSet rs =
111+
dataClient.executeQuery(
112+
Statement.of(
113+
"SELECT * FROM `"
114+
+ tableId
115+
+ "`(with_history => true) WHERE _key LIKE '"
116+
+ uniquePrefix
117+
+ "%'"))) {
118+
119+
assertThat(rs.next()).isTrue();
120+
assertThat(rs.getBytes("_key")).isEqualTo(ByteString.copyFromUtf8(uniquePrefix + "a"));
121+
Map<ByteString, List<Struct>> rowACf = rs.getMap(cf, SqlType.historicalMap());
122+
List<Struct> rowAQual = rowACf.get(ByteString.copyFromUtf8("qual"));
123+
assertThat(rowAQual.size()).isEqualTo(2);
124+
Struct rowAQual_0 = rowAQual.get(0);
125+
assertThat(rowAQual_0.getBytes("value")).isEqualTo(ByteString.copyFromUtf8("val"));
126+
// timestamp in micros above so we divide by 1000
127+
assertThat(rowAQual_0.getTimestamp("timestamp")).isEqualTo(Instant.ofEpochMilli(10));
128+
Struct rowAQual_1 = rowAQual.get(1);
129+
assertThat(rowAQual_1.getBytes("value")).isEqualTo(ByteString.copyFromUtf8("old"));
130+
assertThat(rowAQual_1.getTimestamp("timestamp")).isEqualTo(Instant.ofEpochMilli(1));
131+
132+
assertThat(rs.next()).isTrue();
133+
assertThat(rs.getBytes("_key")).isEqualTo(ByteString.copyFromUtf8(uniquePrefix + "b"));
134+
Map<ByteString, List<Struct>> rowBCf = rs.getMap(cf, SqlType.historicalMap());
135+
List<Struct> rowBQual = rowBCf.get(ByteString.copyFromUtf8("qual"));
136+
assertThat(rowBQual.size()).isEqualTo(1);
137+
Struct rowBQual_0 = rowBQual.get(0);
138+
assertThat(rowBQual_0.getBytes("value")).isEqualTo(ByteString.copyFromUtf8("bval"));
139+
140+
assertThat(rs.next()).isFalse();
141+
}
142+
}
143+
144+
@Test
145+
public void allTypes() {
146+
try (ResultSet rs =
147+
dataClient.executeQuery(
148+
Statement.of(
149+
"SELECT 'stringVal' AS strCol, b'foo' as bytesCol, 1 AS intCol, CAST(1.2 AS FLOAT32) as f32Col, "
150+
+ "CAST(1.3 AS FLOAT64) as f64Col, true as boolCol, TIMESTAMP_FROM_UNIX_MILLIS(1000) AS tsCol, "
151+
+ "DATE(2024, 06, 01) as dateCol, STRUCT(1 as a, \"foo\" as b) AS structCol, [1,2,3] AS arrCol, "
152+
+ cf
153+
+ " as mapCol FROM `"
154+
+ tableId
155+
+ "` WHERE _key='"
156+
+ uniquePrefix
157+
+ "a' LIMIT 1"))) {
158+
159+
assertThat(rs.next()).isTrue();
160+
assertThat(rs.getString("strCol")).isEqualTo("stringVal");
161+
assertThat(rs.getString(0)).isEqualTo("stringVal");
162+
assertThat(rs.getBytes("bytesCol")).isEqualTo(ByteString.copyFromUtf8("foo"));
163+
assertThat(rs.getBytes(1)).isEqualTo(ByteString.copyFromUtf8("foo"));
164+
assertThat(rs.getLong("intCol")).isEqualTo(1L);
165+
assertThat(rs.getLong(2)).isEqualTo(1L);
166+
assertThat(rs.getFloat("f32Col")).isEqualTo(1.2f);
167+
assertThat(rs.getFloat(3)).isEqualTo(1.2f);
168+
assertThat(rs.getDouble("f64Col")).isEqualTo(1.3d);
169+
assertThat(rs.getDouble(4)).isEqualTo(1.3d);
170+
assertThat(rs.getBoolean("boolCol")).isTrue();
171+
assertThat(rs.getBoolean(5)).isTrue();
172+
assertThat(rs.getTimestamp("tsCol")).isEqualTo(Instant.ofEpochMilli(1000));
173+
assertThat(rs.getTimestamp(6)).isEqualTo(Instant.ofEpochMilli(1000));
174+
assertThat(rs.getDate("dateCol")).isEqualTo(Date.fromYearMonthDay(2024, 6, 1));
175+
assertThat(rs.getDate(7)).isEqualTo(Date.fromYearMonthDay(2024, 6, 1));
176+
assertThat(rs.getStruct("structCol").getLong("a")).isEqualTo(1);
177+
assertThat(rs.getStruct("structCol").getString("b")).isEqualTo("foo");
178+
assertThat(rs.getStruct(8).getLong("a")).isEqualTo(1);
179+
assertThat(rs.getStruct(8).getString("b")).isEqualTo("foo");
180+
assertThat(rs.getList("arrCol", SqlType.arrayOf(SqlType.int64())))
181+
.isEqualTo(Arrays.asList(1L, 2L, 3L));
182+
assertThat(rs.getList(9, SqlType.arrayOf(SqlType.int64())))
183+
.isEqualTo(Arrays.asList(1L, 2L, 3L));
184+
assertThat(rs.getMap("mapCol", SqlType.mapOf(SqlType.bytes(), SqlType.bytes())))
185+
.isEqualTo(
186+
new HashMap<ByteString, ByteString>() {
187+
{
188+
put(ByteString.copyFromUtf8("qual"), ByteString.copyFromUtf8("val"));
189+
put(ByteString.copyFromUtf8("qual2"), ByteString.copyFromUtf8("val2"));
190+
put(ByteString.copyFromUtf8("qual3"), ByteString.copyFromUtf8("val3"));
191+
}
192+
});
193+
assertThat(rs.getMap(10, SqlType.mapOf(SqlType.bytes(), SqlType.bytes())))
194+
.isEqualTo(
195+
new HashMap<ByteString, ByteString>() {
196+
{
197+
put(ByteString.copyFromUtf8("qual"), ByteString.copyFromUtf8("val"));
198+
put(ByteString.copyFromUtf8("qual2"), ByteString.copyFromUtf8("val2"));
199+
put(ByteString.copyFromUtf8("qual3"), ByteString.copyFromUtf8("val3"));
200+
}
201+
});
202+
203+
assertThat(rs.next()).isFalse();
204+
}
205+
}
206+
207+
@Test
208+
public void allQueryParamsTypes() {
209+
ResultSet rs =
210+
dataClient.executeQuery(
211+
Statement.newBuilder(
212+
"SELECT @stringParam AS strCol, @bytesParam as bytesCol, @int64Param AS intCol, "
213+
+ "@doubleParam AS doubleCol, @floatParam AS floatCol, @boolParam AS boolCol, "
214+
+ "@tsParam AS tsCol, @dateParam AS dateCol, @byteArrayParam AS byteArrayCol, "
215+
+ "@stringArrayParam AS stringArrayCol, @intArrayParam AS intArrayCol, "
216+
+ "@floatArrayParam AS floatArrayCol, @doubleArrayParam AS doubleArrayCol, "
217+
+ "@boolArrayParam AS boolArrayCol, @tsArrayParam AS tsArrayCol, "
218+
+ "@dateArrayParam AS dateArrayCol")
219+
.setStringParam("stringParam", "stringVal")
220+
.setBytesParam("bytesParam", ByteString.copyFromUtf8("foo"))
221+
.setLongParam("int64Param", 1L)
222+
.setDoubleParam("doubleParam", 1.3d)
223+
.setFloatParam("floatParam", 1.4f)
224+
.setBooleanParam("boolParam", true)
225+
.setTimestampParam("tsParam", Instant.ofEpochMilli(1000))
226+
.setDateParam("dateParam", Date.fromYearMonthDay(2024, 6, 1))
227+
.setListParam(
228+
"byteArrayParam",
229+
Arrays.asList(
230+
ByteString.copyFromUtf8("foo"), null, ByteString.copyFromUtf8("bar")),
231+
SqlType.arrayOf(SqlType.bytes()))
232+
.setListParam(
233+
"stringArrayParam",
234+
Arrays.asList("foo", null, "bar"),
235+
SqlType.arrayOf(SqlType.string()))
236+
.setListParam(
237+
"intArrayParam", Arrays.asList(1L, null, 2L), SqlType.arrayOf(SqlType.int64()))
238+
.setListParam(
239+
"floatArrayParam",
240+
Arrays.asList(1.2f, null, 1.3f),
241+
SqlType.arrayOf(SqlType.float32()))
242+
.setListParam(
243+
"doubleArrayParam",
244+
Arrays.asList(1.4d, null, 1.5d),
245+
SqlType.arrayOf(SqlType.float64()))
246+
.setListParam(
247+
"boolArrayParam",
248+
Arrays.asList(true, null, false),
249+
SqlType.arrayOf(SqlType.bool()))
250+
.setListParam(
251+
"tsArrayParam",
252+
Arrays.asList(
253+
Instant.ofEpochSecond(1000, 1000), null, Instant.ofEpochSecond(2000, 2000)),
254+
SqlType.arrayOf(SqlType.timestamp()))
255+
.setListParam(
256+
"dateArrayParam",
257+
Arrays.asList(
258+
Date.fromYearMonthDay(2024, 8, 1), null, Date.fromYearMonthDay(2024, 8, 2)),
259+
SqlType.arrayOf(SqlType.date()))
260+
.build());
261+
262+
assertThat(rs.next()).isTrue();
263+
assertThat(rs.getString("strCol")).isEqualTo("stringVal");
264+
assertThat(rs.getString(0)).isEqualTo("stringVal");
265+
assertThat(rs.getBytes("bytesCol")).isEqualTo(ByteString.copyFromUtf8("foo"));
266+
assertThat(rs.getBytes(1)).isEqualTo(ByteString.copyFromUtf8("foo"));
267+
assertThat(rs.getLong("intCol")).isEqualTo(1L);
268+
assertThat(rs.getLong(2)).isEqualTo(1L);
269+
assertThat(rs.getDouble("doubleCol")).isEqualTo(1.3d);
270+
assertThat(rs.getDouble(3)).isEqualTo(1.3d);
271+
assertThat(rs.getFloat("floatCol")).isEqualTo(1.4f);
272+
assertThat(rs.getFloat(4)).isEqualTo(1.4f);
273+
assertThat(rs.getBoolean("boolCol")).isTrue();
274+
assertThat(rs.getBoolean(5)).isTrue();
275+
assertThat(rs.getTimestamp("tsCol")).isEqualTo(Instant.ofEpochMilli(1000));
276+
assertThat(rs.getTimestamp(6)).isEqualTo(Instant.ofEpochMilli(1000));
277+
assertThat(rs.getDate("dateCol")).isEqualTo(Date.fromYearMonthDay(2024, 6, 1));
278+
assertThat(rs.getDate(7)).isEqualTo(Date.fromYearMonthDay(2024, 6, 1));
279+
assertThat(rs.getList("byteArrayCol", SqlType.arrayOf(SqlType.bytes())))
280+
.isEqualTo(
281+
Arrays.asList(ByteString.copyFromUtf8("foo"), null, ByteString.copyFromUtf8("bar")));
282+
assertThat(rs.getList(8, SqlType.arrayOf(SqlType.bytes())))
283+
.isEqualTo(
284+
Arrays.asList(ByteString.copyFromUtf8("foo"), null, ByteString.copyFromUtf8("bar")));
285+
assertThat(rs.getList("stringArrayCol", SqlType.arrayOf(SqlType.string())))
286+
.isEqualTo(Arrays.asList("foo", null, "bar"));
287+
assertThat(rs.getList(9, SqlType.arrayOf(SqlType.string())))
288+
.isEqualTo(Arrays.asList("foo", null, "bar"));
289+
assertThat(rs.getList("intArrayCol", SqlType.arrayOf(SqlType.int64())))
290+
.isEqualTo(Arrays.asList(1L, null, 2L));
291+
assertThat(rs.getList(10, SqlType.arrayOf(SqlType.int64())))
292+
.isEqualTo(Arrays.asList(1L, null, 2L));
293+
assertThat(rs.getList("floatArrayCol", SqlType.arrayOf(SqlType.float32())))
294+
.isEqualTo(Arrays.asList(1.2f, null, 1.3f));
295+
assertThat(rs.getList(11, SqlType.arrayOf(SqlType.float32())))
296+
.isEqualTo(Arrays.asList(1.2f, null, 1.3f));
297+
assertThat(rs.getList("doubleArrayCol", SqlType.arrayOf(SqlType.float64())))
298+
.isEqualTo(Arrays.asList(1.4d, null, 1.5d));
299+
assertThat(rs.getList(12, SqlType.arrayOf(SqlType.float64())))
300+
.isEqualTo(Arrays.asList(1.4d, null, 1.5d));
301+
assertThat(rs.getList("boolArrayCol", SqlType.arrayOf(SqlType.bool())))
302+
.isEqualTo(Arrays.asList(true, null, false));
303+
assertThat(rs.getList(13, SqlType.arrayOf(SqlType.bool())))
304+
.isEqualTo(Arrays.asList(true, null, false));
305+
assertThat(rs.getList("tsArrayCol", SqlType.arrayOf(SqlType.timestamp())))
306+
.isEqualTo(
307+
Arrays.asList(
308+
Instant.ofEpochSecond(1000, 1000), null, Instant.ofEpochSecond(2000, 2000)));
309+
assertThat(rs.getList(14, SqlType.arrayOf(SqlType.timestamp())))
310+
.isEqualTo(
311+
Arrays.asList(
312+
Instant.ofEpochSecond(1000, 1000), null, Instant.ofEpochSecond(2000, 2000)));
313+
assertThat(rs.getList("dateArrayCol", SqlType.arrayOf(SqlType.date())))
314+
.isEqualTo(
315+
Arrays.asList(
316+
Date.fromYearMonthDay(2024, 8, 1), null, Date.fromYearMonthDay(2024, 8, 2)));
317+
assertThat(rs.getList(15, SqlType.arrayOf(SqlType.date())))
318+
.isEqualTo(
319+
Arrays.asList(
320+
Date.fromYearMonthDay(2024, 8, 1), null, Date.fromYearMonthDay(2024, 8, 2)));
321+
}
322+
323+
@Test
324+
public void testNullColumns() {
325+
try (ResultSet rs =
326+
dataClient.executeQuery(
327+
Statement.of(
328+
"SELECT cf['qual'] AS neverNull, cf['qual3'] AS maybeNull FROM "
329+
+ tableId
330+
+ " WHERE _key LIKE '"
331+
+ uniquePrefix
332+
+ "%'"))) {
333+
assertThat(rs.next()).isTrue();
334+
assertThat(rs.getBytes("neverNull")).isEqualTo(ByteString.copyFromUtf8("val"));
335+
// qual3 is set in row A but not row B
336+
assertThat(rs.isNull("maybeNull")).isFalse();
337+
assertThat(rs.isNull(1)).isFalse();
338+
assertThat(rs.getBytes("maybeNull")).isEqualTo(ByteString.copyFromUtf8("val3"));
339+
assertThat(rs.next()).isTrue();
340+
assertThat(rs.getBytes("neverNull")).isEqualTo(ByteString.copyFromUtf8("bval"));
341+
assertThat(rs.isNull("maybeNull")).isTrue();
342+
assertThat(rs.isNull(1)).isTrue();
343+
assertThrows(NullPointerException.class, () -> rs.getBytes("maybeNull"));
344+
assertThrows(NullPointerException.class, () -> rs.getBytes(1));
345+
assertThat(rs.next()).isFalse();
346+
}
347+
}
348+
}

0 commit comments

Comments
 (0)