Examples with PreparedStatementBuilder com.amazonaws.athena.connectors.jdbc.manager.PreparedStatementBuilder used on opensource projects

Example 1 with PreparedStatementBuilder

use of com.amazonaws.athena.connectors.jdbc.manager.PreparedStatementBuilder in project aws-athena-query-federation by awslabs.

the class OracleMetadataHandler method getPartitions.

/**
 * If it is a table with no partition, then data will be fetched with single split.
 * If it is a partitioned table, we are fetching the partition info and creating splits equals to the number of partitions
 * for parallel processing.
 * @param blockWriter
 * @param getTableLayoutRequest
 * @param queryStatusChecker
 */
@Override
public void getPartitions(final BlockWriter blockWriter, final GetTableLayoutRequest getTableLayoutRequest, QueryStatusChecker queryStatusChecker) {
    LOGGER.debug("{}: Schema {}, table {}", getTableLayoutRequest.getQueryId(), getTableLayoutRequest.getTableName().getSchemaName(), getTableLayoutRequest.getTableName().getTableName());
    try (Connection connection = getJdbcConnectionFactory().getConnection(getCredentialProvider())) {
        List<String> parameters = Arrays.asList(getTableLayoutRequest.getTableName().getTableName().toUpperCase());
        try (PreparedStatement preparedStatement = new PreparedStatementBuilder().withConnection(connection).withQuery(GET_PARTITIONS_QUERY).withParameters(parameters).build();
            ResultSet resultSet = preparedStatement.executeQuery()) {
            // Return a single partition if no partitions defined
            if (!resultSet.next()) {
                blockWriter.writeRows((Block block, int rowNum) -> {
                    block.setValue(BLOCK_PARTITION_COLUMN_NAME, rowNum, ALL_PARTITIONS);
                    LOGGER.info("Adding partition {}", ALL_PARTITIONS);
                    // we wrote 1 row so we return 1
                    return 1;
                });
            } else {
                do {
                    final String partitionName = resultSet.getString(PARTITION_COLUMN_NAME);
                    // 1. Returns all partitions of table, we are not supporting constraints push down to filter partitions.
                    // 2. This API is not paginated, we could use order by and limit clause with offsets here.
                    blockWriter.writeRows((Block block, int rowNum) -> {
                        block.setValue(BLOCK_PARTITION_COLUMN_NAME, rowNum, partitionName);
                        LOGGER.debug("Adding partition {}", partitionName);
                        // we wrote 1 row so we return 1
                        return 1;
                    });
                } while (resultSet.next() && queryStatusChecker.isQueryRunning());
            }
        }
    } catch (SQLException sqlException) {
        throw new RuntimeException(sqlException.getErrorCode() + ": " + sqlException.getMessage(), sqlException);
    }
}

Example 2 with PreparedStatementBuilder

use of com.amazonaws.athena.connectors.jdbc.manager.PreparedStatementBuilder in project aws-athena-query-federation by awslabs.

the class SaphanaMetadataHandler method getColumnDatatype.

/**
 * Logic to fetch column datatypes of table and view to handle data-type specific logic
 * @param connection
 * @param tableName
 * @return
 * @throws SQLException
 */
private ResultSet getColumnDatatype(Connection connection, TableName tableName) throws SQLException {
    List<String> statementParams = Arrays.asList(tableName.getTableName(), tableName.getSchemaName());
    LOGGER.debug("SaphanaMetadataHandler::getColumnDatatype statement params  {}", statementParams);
    PreparedStatement preparedStatement = new PreparedStatementBuilder().withConnection(connection).withQuery(DATA_TYPE_QUERY_FOR_TABLE).withParameters(statementParams).build();
    ResultSet dataTypeResultSet = preparedStatement.executeQuery();
    if (!dataTypeResultSet.isBeforeFirst()) {
        PreparedStatement preparedStatementForView = new PreparedStatementBuilder().withConnection(connection).withQuery(DATA_TYPE_QUERY_FOR_VIEW).withParameters(statementParams).build();
        ResultSet dataTypeResultSetView = preparedStatementForView.executeQuery();
        return dataTypeResultSetView;
    } else {
        LOGGER.debug("Metadata available for TABLE: " + tableName.getTableName());
        return dataTypeResultSet;
    }
}

Example 3 with PreparedStatementBuilder

use of com.amazonaws.athena.connectors.jdbc.manager.PreparedStatementBuilder in project aws-athena-query-federation by awslabs.

the class SnowflakeMetadataHandler method checkForView.

/**
 * Check if the input table is a view and returns viewflag accordingly
 * @param getTableLayoutRequest
 * @return
 * @throws SQLException
 */
private boolean checkForView(GetTableLayoutRequest getTableLayoutRequest) throws SQLException {
    boolean viewFlag = false;
    List<String> viewparameters = Arrays.asList(getTableLayoutRequest.getTableName().getSchemaName(), getTableLayoutRequest.getTableName().getTableName());
    try (Connection connection = getJdbcConnectionFactory().getConnection(getCredentialProvider())) {
        try (PreparedStatement preparedStatement = new PreparedStatementBuilder().withConnection(connection).withQuery(VIEW_CHECK_QUERY).withParameters(viewparameters).build();
            ResultSet resultSet = preparedStatement.executeQuery()) {
            if (resultSet.next()) {
                viewFlag = true;
            }
            LOGGER.info("viewFlag: {}", viewFlag);
        } catch (SQLException sqlException) {
            LOGGER.info("Exception while querying view details for view {}", getTableLayoutRequest.getTableName().getTableName());
            throw new SQLException(sqlException.getErrorCode() + ": " + sqlException.getMessage(), sqlException);
        }
    }
    return viewFlag;
}

Example 4 with PreparedStatementBuilder

use of com.amazonaws.athena.connectors.jdbc.manager.PreparedStatementBuilder in project aws-athena-query-federation by awslabs.

the class SqlServerMetadataHandler method getPartitions.

/**
 * Check whether input table is a view or not. If it's a view, it will not have any partition info and
 * data will be fetched with single split.If it's a table with no partition, then data will be fetched with single split.
 * If it's a partitioned table, we are fetching the partition info and creating splits equals to the number of partitions
 * for parallel processing.
 * @param blockWriter
 * @param getTableLayoutRequest
 * @param queryStatusChecker
 * @throws Exception
 */
@Override
public void getPartitions(BlockWriter blockWriter, GetTableLayoutRequest getTableLayoutRequest, QueryStatusChecker queryStatusChecker) throws Exception {
    LOGGER.info("{}: Schema {}, table {}", getTableLayoutRequest.getQueryId(), getTableLayoutRequest.getTableName().getSchemaName(), getTableLayoutRequest.getTableName().getTableName());
    List<String> params = Arrays.asList(getTableLayoutRequest.getTableName().getTableName(), getTableLayoutRequest.getTableName().getSchemaName());
    // check whether the input table is a view or not
    String viewFlag = "N";
    try (Connection connection = getJdbcConnectionFactory().getConnection(getCredentialProvider());
        PreparedStatement preparedStatement = new PreparedStatementBuilder().withConnection(connection).withQuery(VIEW_CHECK_QUERY).withParameters(params).build();
        ResultSet resultSet = preparedStatement.executeQuery()) {
        if (resultSet.next()) {
            viewFlag = "VIEW".equalsIgnoreCase(resultSet.getString("TYPE_DESC")) ? "Y" : "N";
        }
        LOGGER.info("viewFlag: {}", viewFlag);
    } catch (SQLException sqlException) {
        throw new SQLException(sqlException.getErrorCode() + ": " + sqlException.getMessage(), sqlException);
    }
    List<String> parameters = Arrays.asList(getTableLayoutRequest.getTableName().getSchemaName() + "." + getTableLayoutRequest.getTableName().getTableName());
    try (Connection connection = getJdbcConnectionFactory().getConnection(getCredentialProvider())) {
        try (PreparedStatement preparedStatement = new PreparedStatementBuilder().withConnection(connection).withQuery(GET_PARTITIONS_QUERY).withParameters(parameters).build();
            PreparedStatement preparedStatement2 = new PreparedStatementBuilder().withConnection(connection).withQuery(ROW_COUNT_QUERY).withParameters(parameters).build();
            ResultSet resultSet = preparedStatement.executeQuery();
            ResultSet resultSet2 = preparedStatement2.executeQuery()) {
            // check whether the table have partitions or not using ROW_COUNT_QUERY
            if (resultSet2.next()) {
                rowCount = resultSet2.getInt("ROW_COUNT");
                LOGGER.info("rowCount: {}", rowCount);
            }
            // create a single split for view/non-partition table
            if ("Y".equals(viewFlag) || rowCount == 0) {
                LOGGER.debug("Getting as single Partition: ");
                blockWriter.writeRows((Block block, int rowNum) -> {
                    block.setValue(PARTITION_NUMBER, rowNum, ALL_PARTITIONS);
                    // we wrote 1 row so we return 1
                    return 1;
                });
            } else {
                LOGGER.debug("Getting data with diff Partitions: ");
                // get partition details from sql server meta data tables
                getPartitionFunction(params);
                // Include the first partition because it's not retrieved from GET_PARTITIONS_QUERY
                blockWriter.writeRows((Block block, int rowNum) -> {
                    block.setValue(PARTITION_NUMBER, rowNum, "1");
                    return 1;
                });
                if (resultSet.next()) {
                    do {
                        final String partitionNumber = resultSet.getString(PARTITION_NUMBER);
                        // 1. Returns all partitions of table, we are not supporting constraints push down to filter partitions.
                        // 2. This API is not paginated, we could use order by and limit clause with offsets here.
                        blockWriter.writeRows((Block block, int rowNum) -> {
                            block.setValue(PARTITION_NUMBER, rowNum, partitionNumber);
                            // we wrote 1 row so we return 1
                            return 1;
                        });
                    } while (resultSet.next());
                }
            }
        }
    } catch (SQLException sqlException) {
        throw new SQLException(sqlException.getErrorCode() + ": " + sqlException.getMessage(), sqlException);
    }
}

Example 5 with PreparedStatementBuilder

use of com.amazonaws.athena.connectors.jdbc.manager.PreparedStatementBuilder in project aws-athena-query-federation by awslabs.

the class MySqlMetadataHandler method getPartitions.

@Override
public void getPartitions(final BlockWriter blockWriter, final GetTableLayoutRequest getTableLayoutRequest, QueryStatusChecker queryStatusChecker) {
    LOGGER.info("{}: Schema {}, table {}", getTableLayoutRequest.getQueryId(), getTableLayoutRequest.getTableName().getSchemaName(), getTableLayoutRequest.getTableName().getTableName());
    try (Connection connection = getJdbcConnectionFactory().getConnection(getCredentialProvider())) {
        final String escape = connection.getMetaData().getSearchStringEscape();
        List<String> parameters = Arrays.asList(getTableLayoutRequest.getTableName().getTableName(), getTableLayoutRequest.getTableName().getSchemaName());
        try (PreparedStatement preparedStatement = new PreparedStatementBuilder().withConnection(connection).withQuery(GET_PARTITIONS_QUERY).withParameters(parameters).build();
            ResultSet resultSet = preparedStatement.executeQuery()) {
            // Return a single partition if no partitions defined
            if (!resultSet.next()) {
                blockWriter.writeRows((Block block, int rowNum) -> {
                    block.setValue(BLOCK_PARTITION_COLUMN_NAME, rowNum, ALL_PARTITIONS);
                    LOGGER.info("Adding partition {}", ALL_PARTITIONS);
                    // we wrote 1 row so we return 1
                    return 1;
                });
            } else {
                do {
                    final String partitionName = resultSet.getString(PARTITION_COLUMN_NAME);
                    // 1. Returns all partitions of table, we are not supporting constraints push down to filter partitions.
                    // 2. This API is not paginated, we could use order by and limit clause with offsets here.
                    blockWriter.writeRows((Block block, int rowNum) -> {
                        block.setValue(BLOCK_PARTITION_COLUMN_NAME, rowNum, partitionName);
                        LOGGER.info("Adding partition {}", partitionName);
                        // we wrote 1 row so we return 1
                        return 1;
                    });
                } while (resultSet.next() && queryStatusChecker.isQueryRunning());
            }
        }
    } catch (SQLException sqlException) {
        throw new RuntimeException(sqlException.getErrorCode() + ": " + sqlException.getMessage(), sqlException);
    }
}