org.apache.calcite.avatica.AvaticaConnection Java Examples

@Test public void testRemoteStatementInsert() throws Exception {
  ConnectionSpec.getDatabaseLock().lock();
  try {
    final String t = AvaticaUtils.unique("TEST_TABLE2");
    AvaticaConnection conn = (AvaticaConnection) DriverManager.getConnection(url);
    Statement statement = conn.createStatement();
    final String create =
        String.format(Locale.ROOT, "create table if not exists %s ("
            + "  id int not null, msg varchar(255) not null)", t);
    int status = statement.executeUpdate(create);
    assertEquals(status, 0);

    statement = conn.createStatement();
    final String update =
        String.format(Locale.ROOT, "insert into %s values ('%d', '%s')",
            t, RANDOM.nextInt(Integer.MAX_VALUE), UUID.randomUUID());
    status = statement.executeUpdate(update);
    assertEquals(status, 1);
  } finally {
    ConnectionSpec.getDatabaseLock().unlock();
  }
}
 

@Test public void testUnicodeCharacters() throws Exception {
  ConnectionSpec.getDatabaseLock().lock();
  try (AvaticaConnection conn = (AvaticaConnection) DriverManager.getConnection(url)) {
    final AvaticaStatement statement = conn.createStatement();
    ResultSet rs = statement.executeQuery(
        "select * from (values ('您好', 'こんにちは', '안녕하세요'))");
    assertThat(rs.next(), is(true));
    assertEquals("您好", rs.getString(1));
    assertEquals("こんにちは", rs.getString(2));
    assertEquals("안녕하세요", rs.getString(3));
    rs.close();
    statement.close();
    conn.close();
  } finally {
    ConnectionSpec.getDatabaseLock().unlock();
  }
}
 

@Test public void testArrays() throws SQLException {
  ConnectionSpec.getDatabaseLock().lock();
  try (AvaticaConnection conn = (AvaticaConnection) DriverManager.getConnection(url);
       Statement stmt = conn.createStatement()) {
    ResultSet resultSet =
        stmt.executeQuery("select * from (values ('a', array['b', 'c']));");

    assertTrue(resultSet.next());
    assertEquals("a", resultSet.getString(1));
    Array arr = resultSet.getArray(2);
    assertTrue(arr instanceof ArrayImpl);
    Object[] values = (Object[]) ((ArrayImpl) arr).getArray();
    assertArrayEquals(new String[]{"b", "c"}, values);
  } finally {
    ConnectionSpec.getDatabaseLock().unlock();
  }
}
 

@Test public void testExceptionPropagation() throws Exception {
  final String sql = "SELECT * from EMP LIMIT FOOBARBAZ";
  ConnectionSpec.getDatabaseLock().lock();
  try (final AvaticaConnection conn = (AvaticaConnection) DriverManager.getConnection(url);
      final Statement stmt = conn.createStatement()) {
    try {
      // invalid SQL
      stmt.execute(sql);
      fail("Expected an AvaticaSqlException");
    } catch (AvaticaSqlException e) {
      assertEquals(ErrorResponse.UNKNOWN_ERROR_CODE, e.getErrorCode());
      assertEquals(ErrorResponse.UNKNOWN_SQL_STATE, e.getSQLState());
      assertTrue("Message should contain original SQL, was '" + e.getMessage() + "'",
          e.getMessage().contains(sql));
      assertEquals(1, e.getStackTraces().size());
      final String stacktrace = e.getStackTraces().get(0);
      final String substring = "unexpected token: FOOBARBAZ";
      assertTrue("Message should contain '" + substring + "', was '" + e.getMessage() + ",",
          stacktrace.contains(substring));
    }
  } finally {
    ConnectionSpec.getDatabaseLock().unlock();
  }
}
 

@Test public void testRemoteExecuteMaxRowCount() throws Exception {
  ConnectionSpec.getDatabaseLock().lock();
  try (AvaticaConnection conn = (AvaticaConnection) DriverManager.getConnection(url)) {
    final AvaticaStatement statement = conn.createStatement();
    prepareAndExecuteInternal(conn, statement,
        "select * from (values ('a', 1), ('b', 2))", 0);
    ResultSet rs = statement.getResultSet();
    int count = 0;
    while (rs.next()) {
      count++;
    }
    assertEquals("Check maxRowCount=0 and ResultSets is 0 row", count, 0);
    assertEquals("Check result set meta is still there",
        rs.getMetaData().getColumnCount(), 2);
    rs.close();
    statement.close();
    conn.close();
  } finally {
    ConnectionSpec.getDatabaseLock().unlock();
  }
}
 

private void checkLargeQuery(int n) throws Exception {
  try (AvaticaConnection conn = (AvaticaConnection) DriverManager.getConnection(url)) {
    final AvaticaStatement statement = conn.createStatement();
    final String frenchDisko = "It said human existence is pointless\n"
        + "As acts of rebellious solidarity\n"
        + "Can bring sense in this world\n"
        + "La resistance!\n";
    final String sql = "select '"
        + longString(frenchDisko, n)
        + "' as s from (values 'x')";
    prepareAndExecuteInternal(conn, statement, sql, -1);
    ResultSet rs = statement.getResultSet();
    int count = 0;
    while (rs.next()) {
      count++;
    }
    assertThat(count, is(1));
    rs.close();
    statement.close();
    conn.close();
  }
}
 

/**
 * Creates a {@link Service} with the given {@link AvaticaConnection} and configuration.
 *
 * @param connection The {@link AvaticaConnection} to use.
 * @param config Configuration properties
 * @return A Service implementation.
 */
Service createService(AvaticaConnection connection, ConnectionConfig config) {
  final Service.Factory metaFactory = config.factory();
  final Service service;
  if (metaFactory != null) {
    service = metaFactory.create(connection);
  } else if (config.url() != null) {
    final AvaticaHttpClient httpClient = getHttpClient(connection, config);
    final Serialization serializationType = getSerialization(config);

    LOG.debug("Instantiating {} service", serializationType);
    switch (serializationType) {
    case JSON:
      service = new RemoteService(httpClient);
      break;
    case PROTOBUF:
      service = new RemoteProtobufService(httpClient, new ProtobufTranslationImpl());
      break;
    default:
      throw new IllegalArgumentException("Unhandled serialization type: " + serializationType);
    }
  } else {
    service = new MockJsonService(Collections.<String, String>emptyMap());
  }
  return service;
}
 

/**
 * Creates the row type descriptor for the result of a DML operation, which is a single column named ROWCOUNT of
 * type BIGINT for INSERT; a single column named PLAN for EXPLAIN.
 *
 * @param kind Kind of node
 * @param typeFactory factory to use for creating type descriptor
 * @return created type
 */
public static RelDataType createDmlRowType(
    SqlKind kind,
    RelDataTypeFactory typeFactory) {
    switch (kind) {
        case INSERT:
        case DELETE:
        case UPDATE:
            return typeFactory.createStructType(
                ImmutableList.of(
                    Pair.of(AvaticaConnection.ROWCOUNT_COLUMN_NAME,
                        typeFactory.createSqlType(SqlTypeName.BIGINT))));
        case EXPLAIN:
            return typeFactory.createStructType(
                ImmutableList.of(
                    Pair.of(AvaticaConnection.PLAN_COLUMN_NAME,
                        typeFactory.createSqlType(
                            SqlTypeName.VARCHAR,
                            RelDataType.PRECISION_NOT_SPECIFIED))));
        default:
            throw Util.unexpected(kind);
    }
}
 

@Test public void testBigints() throws Exception {
  final String table = "TESTBIGINTS";
  ConnectionSpec.getDatabaseLock().lock();
  try (AvaticaConnection conn = (AvaticaConnection) DriverManager.getConnection(url);
      Statement stmt = conn.createStatement()) {
    assertFalse(stmt.execute("DROP TABLE IF EXISTS " + table));
    assertFalse(stmt.execute("CREATE TABLE " + table + " (id BIGINT)"));
    assertFalse(stmt.execute("INSERT INTO " + table + " values(10)"));
    ResultSet results = conn.getMetaData().getColumns(null, null, table, null);
    assertTrue(results.next());
    assertEquals(table, results.getString(3));
    // ordinal position
    assertEquals(1L, results.getLong(17));
  } finally {
    ConnectionSpec.getDatabaseLock().unlock();
  }
}
 

/**
 * Creates a {@link Service} with the given {@link AvaticaConnection} and configuration.
 *
 * @param connection The {@link AvaticaConnection} to use.
 * @param config Configuration properties
 * @return A Service implementation.
 */
Service createService(AvaticaConnection connection, ConnectionConfig config) {
  final Service.Factory metaFactory = config.factory();
  final Service service;
  if (metaFactory != null) {
    service = metaFactory.create(connection);
  } else if (config.url() != null) {
    final AvaticaHttpClient httpClient = getHttpClient(connection, config);
    final Serialization serializationType = getSerialization(config);

    LOG.debug("Instantiating {} service", serializationType);
    switch (serializationType) {
    case JSON:
      service = new RemoteService(httpClient);
      break;
    case PROTOBUF:
      service = new RemoteProtobufService(httpClient, new ProtobufTranslationImpl());
      break;
    default:
      throw new IllegalArgumentException("Unhandled serialization type: " + serializationType);
    }
  } else {
    service = new MockJsonService(Collections.<String, String>emptyMap());
  }
  return service;
}
 

@Test
public void testPrepareStatementQuery() {
    try {
        AvaticaConnection conn = getConnection();
        final AvaticaStatement statement = conn.createStatement();
        PreparedStatement preparedStatement = conn.prepareStatement("select * from (values (1, 'a'), (2, 'b'), "
            + "(3, 'c')) "
            + "where expr_col__0 = ? or expr_col__1 = ?");
        preparedStatement.setInt(1,1);
        preparedStatement.setString(2,"b");
        ResultSet rs = preparedStatement.executeQuery();
        while (rs.next()) {
            System.out.println(rs.getString(1));
            System.out.println(rs.getString(2));
        }
        close(rs, statement,conn);
    } catch (Exception e) {
        e.printStackTrace();
    }
}
 

@Test
public void testQueryMetaData() {
    try {
        AvaticaConnection conn = getConnection();
        final AvaticaStatement statement = conn.createStatement();
        ResultSet rs = statement.executeQuery("select * from (values (1, 'a'), (2, 'b'))");
        if (rs.getMetaData() != null) {
            System.out.println(rs.getMetaData().getColumnCount());
            System.out.println(rs.getMetaData().getColumnName(1));
            System.out.println(rs.getMetaData().getColumnName(2));
        }
        close(rs, statement,conn);
    } catch (Exception e) {
        e.printStackTrace();
    }
}
 

@Override
public DremioJdbc41PreparedStatement newPreparedStatement(AvaticaConnection connection,
                                                     StatementHandle h,
                                                     Meta.Signature signature,
                                                     int resultSetType,
                                                     int resultSetConcurrency,
                                                     int resultSetHoldability)
    throws SQLException {
  DremioConnectionImpl dremioConnection = (DremioConnectionImpl) connection;
  DremioClient client = dremioConnection.getClient();
  if (dremioConnection.getConfig().isServerPreparedStatementDisabled() || !client.getSupportedMethods().contains(ServerMethod.PREPARED_STATEMENT)) {
    // fallback to client side prepared statement
    return new DremioJdbc41PreparedStatement(dremioConnection, h, signature, null, resultSetType, resultSetConcurrency, resultSetHoldability);
  }
  return newServerPreparedStatement(dremioConnection, h, signature, resultSetType,
      resultSetConcurrency, resultSetHoldability);
}
 

private void checkLargeQuery(int n) throws Exception {
  try (AvaticaConnection conn = (AvaticaConnection) DriverManager.getConnection(url)) {
    final AvaticaStatement statement = conn.createStatement();
    final String frenchDisko = "It said human existence is pointless\n"
        + "As acts of rebellious solidarity\n"
        + "Can bring sense in this world\n"
        + "La resistance!\n";
    final String sql = "select '"
        + longString(frenchDisko, n)
        + "' as s from (values 'x')";
    prepareAndExecuteInternal(conn, statement, sql, -1);
    ResultSet rs = statement.getResultSet();
    int count = 0;
    while (rs.next()) {
      count++;
    }
    assertThat(count, is(1));
    rs.close();
    statement.close();
    conn.close();
  }
}
 

@Override
public QuarkStatement newStatement(AvaticaConnection connection,
                                   Meta.StatementHandle statementHandle,
                                   int resultSetType,
                                   int resultSetConcurrency,
                                   int resultSetHoldability) {
  return new QuarkJdbc41Statement((QuarkConnectionImpl) connection,
      statementHandle,
      resultSetType,
      resultSetConcurrency,
      resultSetHoldability);
}
 

@Test public void testQuery() throws Exception {
  ConnectionSpec.getDatabaseLock().lock();
  try (AvaticaConnection conn = (AvaticaConnection) DriverManager.getConnection(url);
      Statement statement = conn.createStatement()) {
    assertFalse(statement.execute("SET SCHEMA \"SCOTT\""));
    assertFalse(
        statement.execute(
            "CREATE TABLE \"FOO\"(\"KEY\" INTEGER NOT NULL, \"VALUE\" VARCHAR(10))"));
    assertFalse(statement.execute("SET TABLE \"FOO\" READONLY FALSE"));

    final int numRecords = 1000;
    for (int i = 0; i < numRecords; i++) {
      assertFalse(statement.execute("INSERT INTO \"FOO\" VALUES(" + i + ", '" + i + "')"));
    }

    // Make sure all the records are there that we expect
    ResultSet results = statement.executeQuery("SELECT count(KEY) FROM FOO");
    assertTrue(results.next());
    assertEquals(1000, results.getInt(1));
    assertFalse(results.next());

    results = statement.executeQuery("SELECT KEY, VALUE FROM FOO ORDER BY KEY ASC");
    for (int i = 0; i < numRecords; i++) {
      assertTrue(results.next());
      assertEquals(i, results.getInt(1));
      assertEquals(Integer.toString(i), results.getString(2));
    }
  } finally {
    ConnectionSpec.getDatabaseLock().unlock();
  }
}
 

@Override
public DremioStatementImpl newStatement(AvaticaConnection connection,
                                       StatementHandle h,
                                       int resultSetType,
                                       int resultSetConcurrency,
                                       int resultSetHoldability) {
  return new DremioStatementImpl((DremioConnectionImpl) connection,
                                h,
                                resultSetType,
                                resultSetConcurrency,
                                resultSetHoldability);
}
 

@Override public Meta createMeta(AvaticaConnection connection) {
  final ConnectionConfig config = connection.config();

  // Perform the login and launch the renewal thread if necessary
  final KerberosConnection kerberosUtil = createKerberosUtility(config);
  if (null != kerberosUtil) {
    kerberosUtil.login();
    connection.setKerberosConnection(kerberosUtil);
  }

  // Create a single Service and set it on the Connection instance
  final Service service = createService(connection, config);
  connection.setService(service);
  return new RemoteMeta(connection, service);
}
 

public AvaticaPreparedStatement newPreparedStatement(
    AvaticaConnection connection,
    Meta.StatementHandle h,
    Meta.Signature signature,
    int resultSetType,
    int resultSetConcurrency,
    int resultSetHoldability) throws SQLException {
  return new CalciteJdbc41PreparedStatement(
      (CalciteConnectionImpl) connection, h,
      (CalcitePrepare.CalciteSignature) signature, resultSetType,
      resultSetConcurrency, resultSetHoldability);
}
 

@Test public void testRemoteConnectionProperties() throws Exception {
  ConnectionSpec.getDatabaseLock().lock();
  try (AvaticaConnection conn = (AvaticaConnection) DriverManager.getConnection(url)) {
    String id = conn.id;
    final Map<String, ConnectionPropertiesImpl> m = ((RemoteMeta) getMeta(conn)).propsMap;
    assertFalse("remote connection map should start ignorant", m.containsKey(id));
    // force creating a connection object on the remote side.
    try (final Statement stmt = conn.createStatement()) {
      assertTrue("creating a statement starts a local object.", m.containsKey(id));
      assertTrue(stmt.execute("select count(1) from EMP"));
    }
    Connection remoteConn = getConnection(FullyRemoteJdbcMetaFactory.getInstance(), id);
    final boolean defaultRO = remoteConn.isReadOnly();
    final boolean defaultAutoCommit = remoteConn.getAutoCommit();
    final String defaultCatalog = remoteConn.getCatalog();
    final String defaultSchema = remoteConn.getSchema();
    conn.setReadOnly(!defaultRO);
    assertTrue("local changes dirty local state", m.get(id).isDirty());
    assertEquals("remote connection has not been touched", defaultRO, remoteConn.isReadOnly());
    conn.setAutoCommit(!defaultAutoCommit);
    assertEquals("remote connection has not been touched",
        defaultAutoCommit, remoteConn.getAutoCommit());

    // further interaction with the connection will force a sync
    try (final Statement stmt = conn.createStatement()) {
      assertEquals(!defaultAutoCommit, remoteConn.getAutoCommit());
      assertFalse("local values should be clean", m.get(id).isDirty());
    }
  } finally {
    ConnectionSpec.getDatabaseLock().unlock();
  }
}
 

@Test
public void testConnectionPropertiesSync() throws Exception {
  ConnectionSpec.getDatabaseLock().lock();
  try {
    AvaticaConnection conn = (AvaticaConnection) DriverManager.getConnection(url);
    conn.setAutoCommit(false);
    conn.setTransactionIsolation(Connection.TRANSACTION_REPEATABLE_READ);

    // sync connection properties
    conn.createStatement();
    Connection remoteConn = getConnection(
            AvaticaServersForTest.PropertyRemoteJdbcMetaFactory.getInstance(PROPERTIES), conn.id);

    assertFalse(remoteConn.getAutoCommit());
    assertEquals(remoteConn.getTransactionIsolation(),
            Connection.TRANSACTION_REPEATABLE_READ);

    // after 1s, remote connection expired and reopen
    Thread.sleep(1000);

    conn.createStatement();
    Connection remoteConn1 = getConnection(
            AvaticaServersForTest.PropertyRemoteJdbcMetaFactory.getInstance(PROPERTIES), conn.id);

    assertFalse(remoteConn1.getAutoCommit());
    assertEquals(remoteConn1.getTransactionIsolation(),
            Connection.TRANSACTION_REPEATABLE_READ);
  } finally {
    ConnectionSpec.getDatabaseLock().unlock();
  }
}
 

public CalciteJdbc41Statement newStatement(AvaticaConnection connection,
    Meta.StatementHandle h,
    int resultSetType,
    int resultSetConcurrency,
    int resultSetHoldability) {
  return new CalciteJdbc41Statement(
      (CalciteConnectionImpl) connection,
      h,
      resultSetType, resultSetConcurrency,
      resultSetHoldability);
}
 

@Override
public AvaticaPreparedStatement newPreparedStatement(AvaticaConnection connection,
                                                     Meta.StatementHandle h,
                                                     Meta.Signature signature,
                                                     int resultSetType,
                                                     int resultSetConcurrency,
                                                     int resultSetHoldability)
    throws SQLException {
  throw new SQLFeatureNotSupportedException();
}
 

/**
 * Creates a Quark connection for Avatica (in terms of Avatica types).
 * <p>
 * This implementation delegates to
 * {@link #newConnection(QuarkDriver, QuarkJdbcFactory, String, Properties, CalciteRootSchema,
 * JavaTypeFactory)}.
 * </p>
 */
@Override
public final AvaticaConnection newConnection(UnregisteredDriver driver,
                                             AvaticaFactory factory,
                                             String url,
                                             Properties info) throws SQLException {
  return newConnection((QuarkDriver) driver, (QuarkJdbcFactory) factory, url, info,
      null, null);
}
 

public final AvaticaConnection newConnection(
    UnregisteredDriver driver,
    AvaticaFactory factory,
    String url,
    Properties info) {
  return newConnection(driver, factory, url, info, null, null);
}
 

/**
 * Creates a Dremio connection for Avatica (in terms of Avatica types).
 * <p>
 *   This implementation delegates to
 *   {@link #newDremioConnection(DriverImpl, DremioFactory, String, Properties)}.
 * </p>
 */
@Override
public final AvaticaConnection newConnection(UnregisteredDriver driver,
                                             AvaticaFactory factory,
                                             String url,
                                             Properties info) throws SQLException {
  return newConnection((DriverImpl) driver, (DremioFactory) factory, url, info);
}
 

public Service create(AvaticaConnection connection) {
  try {
    Connection localConnection = CalciteAssert.hr().connect();
    final Meta meta = CalciteConnectionImpl.TROJAN
        .getMeta((CalciteConnectionImpl) localConnection);
    return new LocalJsonService(new LocalService(meta));
  } catch (Exception e) {
    throw TestUtil.rethrow(e);
  }
}
 

@Test public void testRemoteConnectionProperties() throws Exception {
  ConnectionSpec.getDatabaseLock().lock();
  try (AvaticaConnection conn = (AvaticaConnection) DriverManager.getConnection(url)) {
    String id = conn.id;
    final Map<String, ConnectionPropertiesImpl> m = ((RemoteMeta) getMeta(conn)).propsMap;
    assertFalse("remote connection map should start ignorant", m.containsKey(id));
    // force creating a connection object on the remote side.
    try (final Statement stmt = conn.createStatement()) {
      assertTrue("creating a statement starts a local object.", m.containsKey(id));
      assertTrue(stmt.execute("select count(1) from EMP"));
    }
    Connection remoteConn = getConnection(FullyRemoteJdbcMetaFactory.getInstance(), id);
    final boolean defaultRO = remoteConn.isReadOnly();
    final boolean defaultAutoCommit = remoteConn.getAutoCommit();
    final String defaultCatalog = remoteConn.getCatalog();
    final String defaultSchema = remoteConn.getSchema();
    conn.setReadOnly(!defaultRO);
    assertTrue("local changes dirty local state", m.get(id).isDirty());
    assertEquals("remote connection has not been touched", defaultRO, remoteConn.isReadOnly());
    conn.setAutoCommit(!defaultAutoCommit);
    assertEquals("remote connection has not been touched",
        defaultAutoCommit, remoteConn.getAutoCommit());

    // further interaction with the connection will force a sync
    try (final Statement stmt = conn.createStatement()) {
      assertEquals(!defaultAutoCommit, remoteConn.getAutoCommit());
      assertFalse("local values should be clean", m.get(id).isDirty());
    }
  } finally {
    ConnectionSpec.getDatabaseLock().unlock();
  }
}
 

/** Test case for
 * <a href="https://issues.apache.org/jira/browse/CALCITE-1301">[CALCITE-1301]
 * Add cancel flag to AvaticaStatement</a>. */
@Test public void testCancel() throws Exception {
  ConnectionSpec.getDatabaseLock().lock();
  try (AvaticaConnection conn = (AvaticaConnection) DriverManager.getConnection(url)) {
    final AvaticaStatement statement = conn.createStatement();
    final String sql = "select * from (values ('a', 1), ('b', 2))";
    final ResultSet rs = statement.executeQuery(sql);
    int count = 0;
  loop:
    for (;;) {
      switch (count++) {
      case 0:
        assertThat(rs.next(), is(true));
        break;
      case 1:
        rs.getStatement().cancel();
        try {
          boolean x = rs.next();
          fail("expected exception, got " + x);
        } catch (SQLException e) {
          assertThat(e.getMessage(), is("Statement canceled"));
        }
        break loop;
      default:
        fail("count: " + count);
      }
    }
    assertThat(count, is(2));
    assertThat(statement.isClosed(), is(false));
    rs.close();
    assertThat(statement.isClosed(), is(false));
    statement.close();
    assertThat(statement.isClosed(), is(true));
    statement.close();
    assertThat(statement.isClosed(), is(true));
  } finally {
    ConnectionSpec.getDatabaseLock().unlock();
  }
}
 

private static Meta.ExecuteResult prepareAndExecuteInternal(AvaticaConnection conn,
    final AvaticaStatement statement, String sql, int maxRowCount) throws Exception {
  Method m =
      AvaticaConnection.class.getDeclaredMethod("prepareAndExecuteInternal",
          AvaticaStatement.class, String.class, long.class);
  m.setAccessible(true);
  return (Meta.ExecuteResult) m.invoke(conn, statement, sql, maxRowCount);
}