Search in sources :

Example 41 with CalciteConnection

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.jdbc.CalciteConnection in project beam by apache.

the class PubsubTableProviderIT method connect.

@SuppressWarnings("unchecked")
private CalciteConnection connect(PipelineOptions options, TableProvider... tableProviders) {
    // HACK: PipelineOptions should expose a prominent method to do this reliably
    // The actual options are in the "options" field of the converted map
    Map<String, String> argsMap = ((Map<String, Object>) MAPPER.convertValue(pipeline.getOptions(), Map.class).get("options")).entrySet().stream().filter((entry) -> {
        if (entry.getValue() instanceof List) {
            if (!((List) entry.getValue()).isEmpty()) {
                throw new IllegalArgumentException("Cannot encode list arguments");
            }
            // We can encode empty lists, just omit them.
            return false;
        }
        return true;
    }).collect(Collectors.toMap(Map.Entry::getKey, entry -> toArg(entry.getValue())));
    InMemoryMetaStore inMemoryMetaStore = new InMemoryMetaStore();
    for (TableProvider tableProvider : tableProviders) {
        inMemoryMetaStore.registerProvider(tableProvider);
    }
    JdbcConnection connection = JdbcDriver.connect(inMemoryMetaStore, options);
    connection.setPipelineOptionsMap(argsMap);
    return connection;
}
Also used : Arrays(java.util.Arrays) LoggerFactory(org.slf4j.LoggerFactory) TimeoutException(java.util.concurrent.TimeoutException) PubsubMessage(org.apache.beam.sdk.io.gcp.pubsub.PubsubMessage) Future(java.util.concurrent.Future) TestPubsub(org.apache.beam.sdk.io.gcp.pubsub.TestPubsub) ResultSet(java.sql.ResultSet) Map(java.util.Map) TestPubsubSignal(org.apache.beam.sdk.io.gcp.pubsub.TestPubsubSignal) Parameterized(org.junit.runners.Parameterized) ImmutableMap(org.apache.beam.vendor.calcite.v1_28_0.com.google.common.collect.ImmutableMap) GcpOptions(org.apache.beam.sdk.extensions.gcp.options.GcpOptions) Matchers.allOf(org.hamcrest.Matchers.allOf) Collection(java.util.Collection) SchemaCoder(org.apache.beam.sdk.schemas.SchemaCoder) Set(java.util.Set) FieldType(org.apache.beam.sdk.schemas.Schema.FieldType) SchemaIOTableProviderWrapper(org.apache.beam.sdk.extensions.sql.meta.provider.SchemaIOTableProviderWrapper) Collectors(java.util.stream.Collectors) StandardCharsets(java.nio.charset.StandardCharsets) Executors(java.util.concurrent.Executors) ImmutableSet(org.apache.beam.vendor.calcite.v1_28_0.com.google.common.collect.ImmutableSet) Serializable(java.io.Serializable) List(java.util.List) Matchers.equalTo(org.hamcrest.Matchers.equalTo) JdbcDriver(org.apache.beam.sdk.extensions.sql.impl.JdbcDriver) ReflectHelpers(org.apache.beam.sdk.util.common.ReflectHelpers) ImmutableList(org.apache.beam.vendor.calcite.v1_28_0.com.google.common.collect.ImmutableList) GenericRecordBuilder(org.apache.avro.generic.GenericRecordBuilder) JsonMatcher.jsonBytesLike(org.apache.beam.sdk.testing.JsonMatcher.jsonBytesLike) ByteArrayOutputStream(java.io.ByteArrayOutputStream) InMemoryMetaStore(org.apache.beam.sdk.extensions.sql.meta.store.InMemoryMetaStore) Duration(org.joda.time.Duration) RunWith(org.junit.runner.RunWith) Parameters(org.junit.runners.Parameterized.Parameters) HashMap(java.util.HashMap) Callable(java.util.concurrent.Callable) Matchers.hasProperty(org.hamcrest.Matchers.hasProperty) BeamSqlRelUtils(org.apache.beam.sdk.extensions.sql.impl.rel.BeamSqlRelUtils) TestPipeline(org.apache.beam.sdk.testing.TestPipeline) MatcherAssert.assertThat(org.hamcrest.MatcherAssert.assertThat) BeamSqlEnv(org.apache.beam.sdk.extensions.sql.impl.BeamSqlEnv) Row(org.apache.beam.sdk.values.Row) PipelineOptions(org.apache.beam.sdk.options.PipelineOptions) ExecutorService(java.util.concurrent.ExecutorService) AvroUtils(org.apache.beam.sdk.schemas.utils.AvroUtils) Matchers.hasEntry(org.hamcrest.Matchers.hasEntry) GenericRecord(org.apache.avro.generic.GenericRecord) Logger(org.slf4j.Logger) UTF_8(java.nio.charset.StandardCharsets.UTF_8) Parameter(org.junit.runners.Parameterized.Parameter) ObjectMapper(com.fasterxml.jackson.databind.ObjectMapper) JdbcConnection(org.apache.beam.sdk.extensions.sql.impl.JdbcConnection) JsonProcessingException(com.fasterxml.jackson.core.JsonProcessingException) IOException(java.io.IOException) Test(org.junit.Test) PCollection(org.apache.beam.sdk.values.PCollection) AvroCoder(org.apache.beam.sdk.coders.AvroCoder) Schema(org.apache.beam.sdk.schemas.Schema) ExecutionException(java.util.concurrent.ExecutionException) TimeUnit(java.util.concurrent.TimeUnit) PayloadMessages(org.apache.beam.sdk.extensions.protobuf.PayloadMessages) Rule(org.junit.Rule) Ignore(org.junit.Ignore) Matcher(org.hamcrest.Matcher) Instant(org.joda.time.Instant) Statement(java.sql.Statement) TableProvider(org.apache.beam.sdk.extensions.sql.meta.provider.TableProvider) CalciteConnection(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.jdbc.CalciteConnection) List(java.util.List) ImmutableList(org.apache.beam.vendor.calcite.v1_28_0.com.google.common.collect.ImmutableList) JdbcConnection(org.apache.beam.sdk.extensions.sql.impl.JdbcConnection) TableProvider(org.apache.beam.sdk.extensions.sql.meta.provider.TableProvider) Map(java.util.Map) ImmutableMap(org.apache.beam.vendor.calcite.v1_28_0.com.google.common.collect.ImmutableMap) HashMap(java.util.HashMap) InMemoryMetaStore(org.apache.beam.sdk.extensions.sql.meta.store.InMemoryMetaStore)

Example 42 with CalciteConnection

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.jdbc.CalciteConnection 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();
}
Also used : Statement(java.sql.Statement) GenericRecordBuilder(org.apache.avro.generic.GenericRecordBuilder) PubsubMessage(org.apache.beam.sdk.io.gcp.pubsub.PubsubMessage) ExecutorService(java.util.concurrent.ExecutorService) ResultSet(java.sql.ResultSet) List(java.util.List) ImmutableList(org.apache.beam.vendor.calcite.v1_28_0.com.google.common.collect.ImmutableList) ExecutionException(java.util.concurrent.ExecutionException) CalciteConnection(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.jdbc.CalciteConnection) TimeoutException(java.util.concurrent.TimeoutException) Test(org.junit.Test)

Example 43 with CalciteConnection

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.jdbc.CalciteConnection in project beam by apache.

the class JdbcDriverTest method testDriverManager_setUserAgent.

/**
 * Tests that userAgent can be overridden on the querystring.
 */
@Test
public void testDriverManager_setUserAgent() throws Exception {
    Connection connection = DriverManager.getConnection(JdbcDriver.CONNECT_STRING_PREFIX + "beam.userAgent=Secret Agent");
    SchemaPlus rootSchema = ((CalciteConnection) connection).getRootSchema();
    BeamCalciteSchema beamSchema = (BeamCalciteSchema) CalciteSchema.from(rootSchema.getSubSchema("beam")).schema;
    Map<String, String> pipelineOptions = beamSchema.getPipelineOptions();
    assertThat(pipelineOptions.get("userAgent"), equalTo("Secret Agent"));
}
Also used : Connection(java.sql.Connection) CalciteConnection(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.jdbc.CalciteConnection) SchemaPlus(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.schema.SchemaPlus) Matchers.containsString(org.hamcrest.Matchers.containsString) CalciteConnection(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.jdbc.CalciteConnection) Test(org.junit.Test)

Example 44 with CalciteConnection

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.jdbc.CalciteConnection in project beam by apache.

the class JdbcDriverTest method testInternalConnect_setDirectRunner.

@Test
public void testInternalConnect_setDirectRunner() throws Exception {
    CalciteConnection connection = JdbcDriver.connect(BOUNDED_TABLE, PipelineOptionsFactory.create());
    Statement statement = connection.createStatement();
    assertEquals(0, statement.executeUpdate("SET runner = direct"));
    assertTrue(statement.execute("SELECT * FROM test"));
}
Also used : Statement(java.sql.Statement) CalciteConnection(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.jdbc.CalciteConnection) Test(org.junit.Test)

Example 45 with CalciteConnection

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.jdbc.CalciteConnection in project beam by apache.

the class JdbcDriverTest method testInternalConnect_unbounded_limit.

@Test
public void testInternalConnect_unbounded_limit() throws Exception {
    ReadOnlyTableProvider tableProvider = new ReadOnlyTableProvider("test", ImmutableMap.of("test", TestUnboundedTable.of(Schema.FieldType.INT32, "order_id", Schema.FieldType.INT32, "site_id", Schema.FieldType.INT32, "price", Schema.FieldType.DATETIME, "order_time").timestampColumnIndex(3).addRows(Duration.ZERO, 1, 1, 1, FIRST_DATE, 1, 2, 6, FIRST_DATE)));
    CalciteConnection connection = JdbcDriver.connect(tableProvider, PipelineOptionsFactory.create());
    Statement statement = connection.createStatement();
    ResultSet resultSet1 = statement.executeQuery("SELECT * FROM test LIMIT 1");
    assertTrue(resultSet1.next());
    assertFalse(resultSet1.next());
    ResultSet resultSet2 = statement.executeQuery("SELECT * FROM test LIMIT 2");
    assertTrue(resultSet2.next());
    assertTrue(resultSet2.next());
    assertFalse(resultSet2.next());
}
Also used : Statement(java.sql.Statement) ReadOnlyTableProvider(org.apache.beam.sdk.extensions.sql.meta.provider.ReadOnlyTableProvider) ResultSet(java.sql.ResultSet) CalciteConnection(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.jdbc.CalciteConnection) Test(org.junit.Test)

Aggregations

CalciteConnection (org.apache.calcite.jdbc.CalciteConnection)65 Test (org.junit.Test)52 Connection (java.sql.Connection)51 ResultSet (java.sql.ResultSet)42 SchemaPlus (org.apache.calcite.schema.SchemaPlus)42 Statement (java.sql.Statement)32 PreparedStatement (java.sql.PreparedStatement)24 AbstractSchema (org.apache.calcite.schema.impl.AbstractSchema)22 Properties (java.util.Properties)17 ReflectiveSchema (org.apache.calcite.adapter.java.ReflectiveSchema)17 SQLException (java.sql.SQLException)16 AvaticaConnection (org.apache.calcite.avatica.AvaticaConnection)12 CalciteConnection (org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.jdbc.CalciteConnection)11 CoreMatchers.containsString (org.hamcrest.CoreMatchers.containsString)11 TableFunction (org.apache.calcite.schema.TableFunction)10 AvaticaStatement (org.apache.calcite.avatica.AvaticaStatement)9 org.hsqldb.jdbcDriver (org.hsqldb.jdbcDriver)5 AssertThat (org.apache.calcite.test.CalciteAssert.AssertThat)4 IOException (java.io.IOException)3 ResultSetMetaData (java.sql.ResultSetMetaData)3