code: /*
   * Copyright (c) 2022 Airbyte, Inc., all rights reserved.
   */
  
  package io.airbyte.integrations.source.cockroachdb;
  
  import com.fasterxml.jackson.databind.JsonNode;
  import io.airbyte.commons.functional.CheckedConsumer;
  import io.airbyte.commons.functional.CheckedFunction;
  import io.airbyte.db.JdbcCompatibleSourceOperations;
  import io.airbyte.db.jdbc.JdbcDatabase;
  import java.sql.Connection;
  import java.sql.DatabaseMetaData;
  import java.sql.PreparedStatement;
  import java.sql.ResultSet;
  import java.sql.SQLException;
  import java.util.List;
  import java.util.stream.Stream;
  
  /**
   * This implementation uses non-streamed queries to CockroachDB. CockroachDB does not currently
   * support multiple active pgwire portals on the same session, which makes it impossible to
   * replicate tables that have over ~1000 rows using StreamingJdbcDatabase. See:
   * https://go.crdb.dev/issue-v/40195/v21.2 and in particular, the comment:
   * https://github.com/cockroachdb/cockroach/issues/40195?version=v21.2#issuecomment-870570351 The
   * same situation as kafka-connect applies to StreamingJdbcDatabase
   */
  public class CockroachJdbcDatabase
      extends JdbcDatabase {
  
    private final JdbcDatabase database;
  
    public CockroachJdbcDatabase(final JdbcDatabase database,
                                 final JdbcCompatibleSourceOperations<?> sourceOperations) {
      super(sourceOperations);
      this.database = database;
    }
  
    @Override
    public DatabaseMetaData getMetaData() throws SQLException {
      return database.getMetaData();
    }
  
    @Override
    public void execute(final CheckedConsumer<Connection, SQLException> query) throws SQLException {
      database.execute(query);
    }
  
    @Override
    public <T> List<T> bufferedResultSetQuery(final CheckedFunction<Connection, ResultSet, SQLException> query,
                                              final CheckedFunction<ResultSet, T, SQLException> recordTransform)
        throws SQLException {
      return database.bufferedResultSetQuery(query, recordTransform);
    }
  
    @Override
    public <T> Stream<T> unsafeResultSetQuery(final CheckedFunction<Connection, ResultSet, SQLException> query,
                                              final CheckedFunction<ResultSet, T, SQLException> recordTransform)
        throws SQLException {
      return database.unsafeResultSetQuery(query, recordTransform);
    }
  
    @Override
    public <T> Stream<T> unsafeQuery(final CheckedFunction<Connection, PreparedStatement, SQLException> statementCreator,
                                     final CheckedFunction<ResultSet, T, SQLException> recordTransform)
        throws SQLException {
      return database.unsafeQuery(statementCreator, recordTransform);
    }
  
    @Override
    public Stream<JsonNode> unsafeQuery(final String sql, final String... params) throws SQLException {
      return bufferedResultSetQuery(connection -> {
        final PreparedStatement statement = connection.prepareStatement(sql);
        int i = 1;
        for (final String param : params) {
          statement.setString(i, param);
          ++i;
        }
        return statement.executeQuery();
      }, sourceOperations::rowToJson).stream();
  
    }
  
  }
  