use of org.skife.jdbi.v2.exceptions.ResultSetException 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);
}
}
Aggregations