Examples with SQLFirehoseDatabaseConnector org.apache.druid.metadata.SQLFirehoseDatabaseConnector used on opensource projects

Example 1 with SQLFirehoseDatabaseConnector

use of org.apache.druid.metadata.SQLFirehoseDatabaseConnector in project druid by druid-io.

the class SqlEntity method openCleanableFile.

/**
 * Executes a SQL query on the specified database and fetches the result into the given file.
 * The result file is deleted if the query execution or the file write fails.
 *
 * @param sql                          The SQL query to be executed
 * @param sqlFirehoseDatabaseConnector The database connector
 * @param objectMapper                 An object mapper, used for deserialization
 * @param foldCase                     A boolean flag used to enable or disabling case sensitivity while handling database column names
 *
 * @return A {@link InputEntity.CleanableFile} object that wraps the file containing the SQL results
 */
public static CleanableFile openCleanableFile(String sql, SQLFirehoseDatabaseConnector sqlFirehoseDatabaseConnector, ObjectMapper objectMapper, boolean foldCase, File tempFile) throws IOException {
    try (FileOutputStream fos = new FileOutputStream(tempFile);
        final JsonGenerator jg = objectMapper.getFactory().createGenerator(fos)) {
        // Execute the sql query and lazily retrieve the results into the file in json format.
        // foldCase is useful to handle differences in case sensitivity behavior across databases.
        sqlFirehoseDatabaseConnector.retryWithHandle((handle) -> {
            ResultIterator<Map<String, Object>> resultIterator = handle.createQuery(sql).map((index, r, ctx) -> {
                Map<String, Object> resultRow = foldCase ? new CaseFoldedMap() : new HashMap<>();
                ResultSetMetaData resultMetadata;
                try {
                    resultMetadata = r.getMetaData();
                } catch (SQLException e) {
                    throw new ResultSetException("Unable to obtain metadata from result set", e, ctx);
                }
                try {
                    for (int i = 1; i <= resultMetadata.getColumnCount(); i++) {
                        String key = resultMetadata.getColumnName(i);
                        String alias = resultMetadata.getColumnLabel(i);
                        Object value = r.getObject(i);
                        resultRow.put(alias != null ? alias : key, value);
                    }
                } catch (SQLException e) {
                    throw new ResultSetException("Unable to access specific metadata from " + "result set metadata", e, ctx);
                }
                return resultRow;
            }).iterator();
            jg.writeStartArray();
            while (resultIterator.hasNext()) {
                jg.writeObject(resultIterator.next());
            }
            jg.writeEndArray();
            jg.close();
            return null;
        }, (exception) -> sqlFirehoseDatabaseConnector.isTransientException(exception) && !(SQLMetadataStorageActionHandler.isStatementException(exception)));
        return new CleanableFile() {

            @Override
            public File file() {
                return tempFile;
            }

            @Override
            public void close() {
                if (!tempFile.delete()) {
                    LOG.warn("Failed to remove file[%s]", tempFile.getAbsolutePath());
                }
            }
        };
    } catch (Exception e) {
        if (!tempFile.delete()) {
            LOG.warn("Failed to remove file[%s]", tempFile.getAbsolutePath());
        }
        throw new IOException(e);
    }
}