use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.schema.Table in project beam by apache.
the class PubsubTableProviderIT method testSQLLimit.
@Test
@SuppressWarnings({ "unchecked", "rawtypes" })
public void testSQLLimit() throws Exception {
String createTableString = String.format("CREATE EXTERNAL TABLE message (\n" + "event_timestamp TIMESTAMP, \n" + "attributes MAP<VARCHAR, VARCHAR>, \n" + "payload ROW< \n" + " id INTEGER, \n" + " name VARCHAR \n" + " > \n" + ") \n" + "TYPE '%s' \n" + "LOCATION '%s' \n" + "TBLPROPERTIES " + " '{ " + " %s" + " \"timestampAttributeKey\" : \"ts\", " + " \"deadLetterQueue\" : \"%s\", " + " \"protoClass\" : \"%s\" " + " }'", tableProvider.getTableType(), eventsTopic.topicPath(), payloadFormatParam(), dlqTopic.topicPath(), PayloadMessages.SimpleMessage.class.getName());
List<PubsubMessage> messages = ImmutableList.of(objectsProvider.messageIdName(ts(1), 3, "foo"), objectsProvider.messageIdName(ts(2), 5, "bar"), objectsProvider.messageIdName(ts(3), 7, "baz"), objectsProvider.messageIdName(ts(4), 9, "ba2"), objectsProvider.messageIdName(ts(5), 10, "ba3"), objectsProvider.messageIdName(ts(6), 13, "ba4"), objectsProvider.messageIdName(ts(7), 15, "ba5"));
// We need the default options on the schema to include the project passed in for the
// integration test
CalciteConnection connection = connect(pipeline.getOptions(), new PubsubTableProvider());
Statement statement = connection.createStatement();
statement.execute(createTableString);
// Because Pubsub only allow new subscription receives message after the subscription is
// created, eventsTopic.publish(messages) can only be called after statement.executeQuery.
// However, because statement.executeQuery is a blocking call, it has to be put into a
// separate thread to execute.
ExecutorService pool = Executors.newFixedThreadPool(1);
Future<List<String>> queryResult = pool.submit((Callable) () -> {
ResultSet resultSet = statement.executeQuery("SELECT message.payload.id FROM message LIMIT 3");
ImmutableList.Builder<String> result = ImmutableList.builder();
while (resultSet.next()) {
result.add(resultSet.getString(1));
}
return result.build();
});
try {
eventsTopic.assertSubscriptionEventuallyCreated(pipeline.getOptions().as(GcpOptions.class).getProject(), Duration.standardMinutes(5));
} catch (AssertionError assertionError) {
// Check if the forked thread had an exception.
try {
queryResult.get(0, TimeUnit.SECONDS);
} catch (TimeoutException e) {
// Nothing went wrong on the forked thread, but a subscription still wasn't created.
} catch (ExecutionException e) {
// up to the user.
throw new AssertionError("Exception occurred in statement.executeQuery thread", e);
}
// Just re-throw the timeout assertion.
throw assertionError;
}
eventsTopic.publish(messages);
assertThat(queryResult.get(2, TimeUnit.MINUTES).size(), equalTo(3));
pool.shutdown();
}
use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.schema.Table in project beam by apache.
the class PubsubTableProviderIT method testSQLSelectsPayloadContentFlat.
@Test
public void testSQLSelectsPayloadContentFlat() throws Exception {
String createTableString = String.format("CREATE EXTERNAL TABLE message (\n" + "event_timestamp TIMESTAMP, \n" + "id INTEGER, \n" + "name VARCHAR \n" + ") \n" + "TYPE '%s' \n" + "LOCATION '%s' \n" + "TBLPROPERTIES " + " '{ " + " %s" + " \"protoClass\" : \"%s\", " + " \"timestampAttributeKey\" : \"ts\" " + " }'", tableProvider.getTableType(), eventsTopic.topicPath(), payloadFormatParam(), PayloadMessages.SimpleMessage.class.getName());
String queryString = "SELECT message.id, message.name from message";
// Initialize SQL environment and create the pubsub table
BeamSqlEnv sqlEnv = BeamSqlEnv.inMemory(new PubsubTableProvider());
sqlEnv.executeDdl(createTableString);
// Apply the PTransform to query the pubsub topic
PCollection<Row> queryOutput = query(sqlEnv, pipeline, queryString);
// Observe the query results and send success signal after seeing the expected messages
queryOutput.apply("waitForSuccess", resultSignal.signalSuccessWhen(SchemaCoder.of(PAYLOAD_SCHEMA), observedRows -> observedRows.equals(ImmutableSet.of(row(PAYLOAD_SCHEMA, 3, "foo"), row(PAYLOAD_SCHEMA, 5, "bar"), row(PAYLOAD_SCHEMA, 7, "baz")))));
// Start the pipeline
pipeline.run();
// Block until a subscription for this topic exists
eventsTopic.assertSubscriptionEventuallyCreated(pipeline.getOptions().as(GcpOptions.class).getProject(), Duration.standardMinutes(5));
// Start publishing the messages when main pipeline is started and signaling topic is ready
eventsTopic.publish(ImmutableList.of(objectsProvider.messageIdName(ts(1), 3, "foo"), objectsProvider.messageIdName(ts(2), 5, "bar"), objectsProvider.messageIdName(ts(3), 7, "baz")));
// Poll the signaling topic for success message
resultSignal.waitForSuccess(timeout);
}
use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.schema.Table in project beam by apache.
the class TableResolutionTest method testResolveNestedWithDots.
/**
* Unit test for resolving a table with dots in the subschema names and the table name.
*/
@Test
public void testResolveNestedWithDots() {
String subSchema = "fake.schema";
String tableName = "fake.table";
when(mockSchemaPlus.getSubSchema(subSchema)).thenReturn(innerSchemaPlus);
when(innerSchemaPlus.getTable(tableName)).thenReturn(mockTable);
Table table = TableResolution.resolveCalciteTable(mockSchemaPlus, ImmutableList.of(subSchema, tableName));
assertThat(table, Matchers.is(mockTable));
}
use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.schema.Table in project beam by apache.
the class TableResolutionTest method testResolveFlat.
/**
* Unit test for resolving a table with no hierarchy.
*/
@Test
public void testResolveFlat() {
String tableName = "fake_table";
when(mockSchemaPlus.getTable(tableName)).thenReturn(mockTable);
Table table = TableResolution.resolveCalciteTable(mockSchemaPlus, ImmutableList.of(tableName));
assertThat(table, Matchers.is(mockTable));
}
use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.schema.Table in project beam by apache.
the class TableResolutionTest method testMissingTableInSubschema.
/**
* Unit test for resolving a table with some hierarchy and the table is missing.
*/
@Test
public void testMissingTableInSubschema() {
String subSchema = "fake_schema";
String tableName = "fake_table";
when(mockSchemaPlus.getSubSchema(subSchema)).thenReturn(innerSchemaPlus);
when(innerSchemaPlus.getTable(tableName)).thenReturn(null);
Table table = TableResolution.resolveCalciteTable(mockSchemaPlus, ImmutableList.of(subSchema, tableName));
assertThat(table, Matchers.nullValue());
}
Aggregations