Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -21,10 +21,12 @@

import com.google.api.gax.rpc.InvalidArgumentException;
import com.google.cloud.Date;
import com.google.cloud.bigtable.admin.v2.BigtableTableAdminClient;
import com.google.cloud.bigtable.admin.v2.models.CreateSchemaBundleRequest;
import com.google.cloud.bigtable.admin.v2.models.CreateTableRequest;
import com.google.cloud.bigtable.admin.v2.models.Table;
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.BoundStatement;
import com.google.cloud.bigtable.data.v2.models.sql.PreparedStatement;
import com.google.cloud.bigtable.data.v2.models.sql.ResultSet;
Expand All @@ -35,6 +37,7 @@
import com.google.cloud.bigtable.data.v2.test.SingerProto.Singer;
import com.google.cloud.bigtable.test_helpers.env.AbstractTestEnv;
import com.google.cloud.bigtable.test_helpers.env.EmulatorEnv;
import com.google.cloud.bigtable.test_helpers.env.PrefixGenerator;
import com.google.cloud.bigtable.test_helpers.env.TestEnvRule;
import com.google.protobuf.ByteString;
import com.google.protobuf.DescriptorProtos.FileDescriptorSet;
Expand All @@ -47,6 +50,7 @@
import java.util.UUID;
import java.util.logging.Level;
import java.util.logging.Logger;
import org.junit.AfterClass;
import org.junit.BeforeClass;
import org.junit.ClassRule;
import org.junit.Test;
Expand All @@ -58,7 +62,9 @@ public class ExecuteQueryIT {

@ClassRule public static TestEnvRule testEnvRule = new TestEnvRule();
private static BigtableDataClient dataClient;
private static TableId tableId;
private static BigtableTableAdminClient adminClient;
private static String tableId;
private static Table table;
private static String schemaBundleId;
private static String cf;
private static String uniquePrefix;
Expand All @@ -77,11 +83,22 @@ public static void setUpAll() throws IOException {
AbstractTestEnv.ConnectionMode.REQUIRE_DIRECT_PATH,
AbstractTestEnv.ConnectionMode.REQUIRE_DIRECT_PATH_IPV4);

tableId = testEnvRule.env().getTableId();
tableId = PrefixGenerator.newPrefix("ExecuteQueryIT");
dataClient = testEnvRule.env().getDataClient();
adminClient = testEnvRule.env().getTableAdminClient();
cf = testEnvRule.env().getFamilyId();
uniquePrefix = UUID.randomUUID() + "-execute-query-it-";
schemaBundleId = UUID.randomUUID() + "-bundle";
schemaBundleId = PrefixGenerator.newPrefix("ExecuteQueryIT#bundle");

table = adminClient.createTable(CreateTableRequest.of(tableId).addFamily(cf));
FileDescriptorSet fileDescriptorSet =
FileDescriptorSet.newBuilder()
.addFile(Singer.getDescriptor().getFile().toProto())
.addFile(Album.getDescriptor().getFile().toProto())
.build();
adminClient.createSchemaBundle(
CreateSchemaBundleRequest.of(tableId, schemaBundleId)
.setProtoSchema(fileDescriptorSet.toByteString()));

dataClient.mutateRow(
RowMutation.create(tableId, uniquePrefix + "a")
Expand All @@ -100,11 +117,19 @@ public static void setUpAll() throws IOException {
cf, ByteString.copyFromUtf8("qual2"), 10000, ByteString.copyFromUtf8("bval2")));
}

@AfterClass
public static void tearDownAll() {
if (table != null) {
// Deleting the table will also clean up all the schema bundles under it.
adminClient.deleteTable(tableId);
}
}

@Test
public void selectStar() {
PreparedStatement preparedStatement =
dataClient.prepareStatement(
"SELECT * FROM " + tableId.getTableId() + " WHERE _key LIKE '" + uniquePrefix + "%'",
"SELECT * FROM `" + tableId + "` WHERE _key LIKE '" + uniquePrefix + "%'",
new HashMap<>());
BoundStatement statement = preparedStatement.bind().build();
try (ResultSet rs = dataClient.executeQuery(statement)) {
Expand All @@ -131,7 +156,7 @@ public void withHistoryQuery() {
PreparedStatement preparedStatement =
dataClient.prepareStatement(
"SELECT * FROM `"
+ tableId.getTableId()
+ tableId
+ "`(with_history => true) WHERE _key LIKE '"
+ uniquePrefix
+ "%'",
Expand Down Expand Up @@ -166,7 +191,6 @@ public void withHistoryQuery() {
@SuppressWarnings("DoubleBraceInitialization")
@Test
public void allTypes() throws Exception {
createTestSchemaBundle();
Album album = Album.newBuilder().setTitle("Lover").build();

// For some reason the ExecuteQuery data path sometimes cannot resolve a newly-created schema
Expand All @@ -191,7 +215,7 @@ public void allTypes() throws Exception {
+ " `"
+ schemaBundleId
+ ".com.google.cloud.bigtable.data.v2.test.Genre`) as enumCol FROM `"
+ tableId.getTableId()
+ tableId
+ "` WHERE _key='"
+ uniquePrefix
+ "a' LIMIT 1",
Expand Down Expand Up @@ -256,8 +280,6 @@ public void allTypes() throws Exception {
assertThat(rs.getProtoEnum("enumCol", Genre::forNumber)).isEqualTo(Genre.JAZZ);
assertThat(rs.getProtoEnum(12, Genre::forNumber)).isEqualTo(Genre.JAZZ);
assertThat(rs.next()).isFalse();
} finally {
deleteTestSchemaBundle();
}
}

Expand Down Expand Up @@ -400,9 +422,9 @@ public void allQueryParamsTypes() {
public void testNullColumns() {
PreparedStatement preparedStatement =
dataClient.prepareStatement(
"SELECT cf['qual'] AS neverNull, cf['qual3'] AS maybeNull FROM "
+ tableId.getTableId()
+ " WHERE _key LIKE '"
"SELECT cf['qual'] AS neverNull, cf['qual3'] AS maybeNull FROM `"
+ tableId
+ "` WHERE _key LIKE '"
+ uniquePrefix
+ "%'",
new HashMap<>());
Expand All @@ -423,23 +445,4 @@ public void testNullColumns() {
assertThat(rs.next()).isFalse();
}
}

private static void deleteTestSchemaBundle() {
testEnvRule
.env()
.getTableAdminClient()
.deleteSchemaBundle(tableId.getTableId(), schemaBundleId);
}

private static void createTestSchemaBundle() throws Exception {
FileDescriptorSet fileDescriptorSet =
FileDescriptorSet.newBuilder()
.addFile(Singer.getDescriptor().getFile().toProto())
.addFile(Album.getDescriptor().getFile().toProto())
.build();
CreateSchemaBundleRequest request =
CreateSchemaBundleRequest.of(tableId.getTableId(), schemaBundleId)
.setProtoSchema(fileDescriptorSet.toByteString());
testEnvRule.env().getTableAdminClient().createSchemaBundle(request);
}
}
Loading