Examples with CalciteConnection org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.jdbc.CalciteConnection used on opensource projects

Example 36 with CalciteConnection

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

the class QueryRecord method query.

protected QueryResult query(final ProcessSession session, final FlowFile flowFile, final String sql, final ProcessContext context, final RecordReaderFactory recordParserFactory) throws SQLException {
    final Properties properties = new Properties();
    properties.put(CalciteConnectionProperty.LEX.camelName(), Lex.MYSQL_ANSI.name());
    Connection connection = null;
    ResultSet resultSet = null;
    Statement statement = null;
    try {
        connection = DriverManager.getConnection("jdbc:calcite:", properties);
        final CalciteConnection calciteConnection = connection.unwrap(CalciteConnection.class);
        final SchemaPlus rootSchema = calciteConnection.getRootSchema();
        final FlowFileTable<?, ?> flowFileTable = new FlowFileTable<>(session, flowFile, recordParserFactory, getLogger());
        rootSchema.add("FLOWFILE", flowFileTable);
        rootSchema.setCacheEnabled(false);
        statement = connection.createStatement();
        resultSet = statement.executeQuery(sql);
        final ResultSet rs = resultSet;
        final Statement stmt = statement;
        final Connection conn = connection;
        return new QueryResult() {

            @Override
            public void close() throws IOException {
                closeQuietly(rs, stmt, conn);
            }

            @Override
            public ResultSet getResultSet() {
                return rs;
            }

            @Override
            public int getRecordsRead() {
                return flowFileTable.getRecordsRead();
            }
        };
    } catch (final Exception e) {
        closeQuietly(resultSet, statement, connection);
        throw e;
    }
}

Example 37 with CalciteConnection

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

the class QueryRecord method buildCachedStatement.

private CachedStatement buildCachedStatement(final String sql, final Supplier<CalciteConnection> connectionSupplier, final ProcessSession session, final FlowFile flowFile, final RecordReaderFactory recordReaderFactory) throws SQLException {
    final CalciteConnection connection = connectionSupplier.get();
    final SchemaPlus rootSchema = connection.getRootSchema();
    final FlowFileTable<?, ?> flowFileTable = new FlowFileTable<>(session, flowFile, recordReaderFactory, getLogger());
    rootSchema.add("FLOWFILE", flowFileTable);
    rootSchema.setCacheEnabled(false);
    final PreparedStatement stmt = connection.prepareStatement(sql);
    return new CachedStatement(stmt, flowFileTable, connection);
}

Example 38 with CalciteConnection

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

the class SqlApi method checkAndAddSchema.

public boolean checkAndAddSchema(String dbName) throws Exception {
    synchronized (connection) {
        if (!engine.checkIfExists(dbName)) {
            return false;
        }
        CalciteConnection calciteConnection = connection.unwrap(CalciteConnection.class);
        String tdbName = dbName.toUpperCase();
        if (calciteConnection.getRootSchema().getSubSchema(tdbName) == null) {
            System.err.println("Adding DB to connection:" + dbName + "\t" + tdbName);
            calciteConnection.getRootSchema().add(tdbName, new SidewinderDatabaseSchema(engine, dbName));
        }
        return true;
    }
}

Example 39 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_defaultUserAgent.

/**
 * Tests that the userAgent is set in the pipeline options of the connection.
 */
@Test
public void testDriverManager_defaultUserAgent() throws Exception {
    Connection connection = DriverManager.getConnection(JdbcDriver.CONNECT_STRING_PREFIX);
    SchemaPlus rootSchema = ((CalciteConnection) connection).getRootSchema();
    BeamCalciteSchema beamSchema = (BeamCalciteSchema) CalciteSchema.from(rootSchema.getSubSchema("beam")).schema;
    Map<String, String> pipelineOptions = beamSchema.getPipelineOptions();
    assertThat(pipelineOptions.get("userAgent"), containsString("BeamSQL"));
}

Example 40 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_pipelineOptionsPlumbing.

/**
 * Tests that unknown pipeline options are passed verbatim from the JDBC URI.
 */
@Test
public void testDriverManager_pipelineOptionsPlumbing() throws Exception {
    Connection connection = DriverManager.getConnection(JdbcDriver.CONNECT_STRING_PREFIX + "beam.foo=baz;beam.foobizzle=mahshizzle;other=smother");
    SchemaPlus rootSchema = ((CalciteConnection) connection).getRootSchema();
    BeamCalciteSchema beamSchema = (BeamCalciteSchema) CalciteSchema.from(rootSchema.getSubSchema("beam")).schema;
    Map<String, String> pipelineOptions = beamSchema.getPipelineOptions();
    assertThat(pipelineOptions.get("foo"), equalTo("baz"));
    assertThat(pipelineOptions.get("foobizzle"), equalTo("mahshizzle"));
    assertThat(pipelineOptions.get("other"), nullValue());
}