diff --git a/bin/run-search-job-result-dumper.sh b/bin/run-search-job-result-dumper.sh index 5eed129..e0f87b1 100755 --- a/bin/run-search-job-result-dumper.sh +++ b/bin/run-search-job-result-dumper.sh @@ -1 +1 @@ -java -cp "target/sumo-java-client-1.1-SNAPSHOT-jar-with-dependencies.jar" -server -Xmx128m com.sumologic.client.searchjob.SearchJobResultDumper $* +java -cp "target/sumo-java-client-2.4-SNAPSHOT-jar-with-dependencies.jar" -server -Xmx128m com.sumologic.client.searchjob.SearchJobResultDumper $* diff --git a/src/main/java/com/sumologic/client/searchjob/SearchJobResultDumper.java b/src/main/java/com/sumologic/client/searchjob/SearchJobResultDumper.java index d5cd3ad..512de08 100644 --- a/src/main/java/com/sumologic/client/searchjob/SearchJobResultDumper.java +++ b/src/main/java/com/sumologic/client/searchjob/SearchJobResultDumper.java @@ -1,34 +1,10 @@ package com.sumologic.client.searchjob; -import java.io.BufferedReader; -import java.io.BufferedWriter; -import java.io.FileInputStream; -import java.io.FileOutputStream; -import java.io.IOException; -import java.io.InputStream; -import java.io.InputStreamReader; -import java.io.OutputStreamWriter; -import java.net.URL; -import java.util.ArrayList; -import java.util.Collections; -import java.util.Date; -import java.util.List; -import java.util.Map; -import java.util.TimeZone; -import java.util.concurrent.atomic.AtomicBoolean; - import au.com.bytecode.opencsv.CSVWriter; +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.core.type.TypeReference; import com.fasterxml.jackson.databind.ObjectMapper; -import org.apache.commons.cli.CommandLine; -import org.apache.commons.cli.CommandLineParser; -import org.apache.commons.cli.GnuParser; -import org.apache.commons.cli.HelpFormatter; -import org.apache.commons.cli.OptionBuilder; -import org.apache.commons.cli.Options; -import org.apache.commons.cli.ParseException; -import org.joda.time.format.DateTimeFormatter; -import org.joda.time.format.ISODateTimeFormat; - +import com.fasterxml.jackson.databind.util.ClassUtil; import com.sumologic.client.Credentials; import com.sumologic.client.SumoLogicClient; import com.sumologic.client.model.LogMessage; @@ -36,6 +12,14 @@ import com.sumologic.client.searchjob.model.GetRecordsForSearchJobResponse; import com.sumologic.client.searchjob.model.GetSearchJobStatusResponse; import com.sumologic.client.searchjob.model.SearchJobRecord; +import org.apache.commons.cli.*; +import org.joda.time.format.DateTimeFormatter; +import org.joda.time.format.ISODateTimeFormat; + +import java.io.*; +import java.net.URL; +import java.util.*; +import java.util.concurrent.atomic.AtomicBoolean; /** * A small but useful tool that executes a search job and dumps the results to @@ -101,6 +85,19 @@ public static void main(String[] args) throws Exception { // How many times to retry if the query fails. int retry = 1; + // If outputting JSON, which field to flatten, or: add to the output as columns; + // this assumes that the lifted field contains valid JSON. + String flattenJson = null; + + // Whether to skip arrays when flattening JSON. + boolean skipArrays = false; + + // Whether to turn arrays into a JSON string when flattening JSON. + boolean arraysAsJson = false; + + // How many messages to grab in each request. + int messagesPerRequest = 1000; + // Create the command line options. Options options = createOptions(); @@ -184,6 +181,30 @@ public static void main(String[] args) throws Exception { if (commandLine.hasOption("json")) { outputFormat = OutputFormat.JSON; } + if (commandLine.hasOption("flatten")) { + if (commandLine.hasOption("json")) { + flattenJson = commandLine.getOptionValue("flatten"); + } else { + throw new ParseException("--json required if --flatten is specified"); + } + } + if (commandLine.hasOption("skip-arrays")) { + if (commandLine.hasOption("flatten")) { + skipArrays = true; + } else { + throw new ParseException("--flatten required if --skip-arrays is specified"); + } + } + if (commandLine.hasOption("arrays-as-json")) { + if (commandLine.hasOption("flatten")) { + if (commandLine.hasOption("skip-arrays")) { + throw new ParseException("--skip-arrays and --arrays-as-json cannot be specified together"); + } + arraysAsJson = true; + } else { + throw new ParseException("--flatten required if --skip-arrays is specified"); + } + } if (commandLine.hasOption("aggregates")) { dumpAggregates = true; } @@ -201,6 +222,11 @@ public static void main(String[] args) throws Exception { retry = Integer.parseInt(retryValue); } + if (commandLine.hasOption("messages-per-request")) { + String messagesPerRequestValue = commandLine.getOptionValue("messages-per-request"); + messagesPerRequest = Integer.parseInt(messagesPerRequestValue); + } + } catch (ParseException exp) { System.err.println(exp.getMessage()); @@ -280,7 +306,11 @@ public static void main(String[] args) throws Exception { "" + chunkEndMillis, timezone, retry, - lastEndFile); + lastEndFile, + flattenJson, + skipArrays, + arraysAsJson, + messagesPerRequest); if (failure) { break; } @@ -422,6 +452,22 @@ private static Options createOptions() { .withArgName("json") .withDescription("Format the output as JSON") .create()); + options.addOption( + OptionBuilder.withLongOpt("flatten") + .withArgName("flatten") + .withDescription("Name of the JSON field in the result to flatten into columns in the output") + .hasArg() + .create()); + options.addOption( + OptionBuilder.withLongOpt("skip-arrays") + .withArgName("skip-arrays") + .withDescription("When flattening JSON in the output, skip arrays") + .create()); + options.addOption( + OptionBuilder.withLongOpt("arrays-as-json") + .withArgName("arrays-as-json") + .withDescription("Turns arrays into a JSON string") + .create()); options.addOption( OptionBuilder.withLongOpt("last-end-file") .withArgName("last-end-file") @@ -434,6 +480,12 @@ private static Options createOptions() { .withDescription("Number of times to retry a query in case of an error") .hasArg() .create("r")); + options.addOption( + OptionBuilder.withLongOpt("messages-per-request") + .withArgName("messages-per-request") + .withDescription("Number of messages to fetch per request") + .hasArg() + .create()); return options; } @@ -476,7 +528,11 @@ private static boolean executeSearchJobWithRetry(CSVWriter csvWriter, String endTimestamp, String timeZone, int retry, - String lastEndFile) { + String lastEndFile, + String jsonFieldToFlatten, + boolean skipArrays, + boolean arraysAsJson, + int messagesPerRequest) { int triesLeft = retry; int attempt = 1; @@ -499,7 +555,11 @@ private static boolean executeSearchJobWithRetry(CSVWriter csvWriter, endTimestamp, timeZone, attempt, - lastEndFile); + lastEndFile, + jsonFieldToFlatten, + skipArrays, + arraysAsJson, + messagesPerRequest); if (failure) { System.err.println(String.format( @@ -525,7 +585,11 @@ private static boolean executeSearch(CSVWriter csvWriter, String endTimestamp, String timeZone, int attempt, - String lastEndFile) { + String lastEndFile, + String jsonFieldToFlatten, + boolean skipArrays, + boolean arraysAsJson, + int messagesPerRequest) { // Create the search job. String searchJobId = sumoClient.createSearchJob( @@ -535,7 +599,7 @@ private static boolean executeSearch(CSVWriter csvWriter, timeZone); System.err.printf("[%s] %s - Search job ID: '%s', attempt: '%d'\n", - new Date(), prefix, searchJobId, attempt); + new Date(), prefix, searchJobId, attempt); try { @@ -593,7 +657,11 @@ private static boolean executeSearch(CSVWriter csvWriter, sumoClient, searchJobId, offset, - messageCount); + messageCount, + jsonFieldToFlatten, + skipArrays, + arraysAsJson, + messagesPerRequest); } // Wait if necessary. @@ -669,7 +737,11 @@ private static int getMessages(CSVWriter csvWriter, SumoLogicClient sumoClient, String searchJobId, int messageOffset, - int messageCount) { + int messageCount, + String jsonFieldToFlatten, + boolean skipArrays, + boolean arraysAsJson, + int messagesPerRequest) { int messageLength = 0; while ((messageLength = messageCount - messageOffset) > 0) { @@ -694,7 +766,7 @@ private static int getMessages(CSVWriter csvWriter, } } - messageLength = Math.min(messageLength, 1000); + messageLength = Math.min(messageLength, messagesPerRequest); if (messageLength > 0) { System.err.printf( "[%s] %s - Search job ID: '%s', messages: '%s', getting offset: '%d', length: '%d'\n", @@ -705,6 +777,7 @@ private static int getMessages(CSVWriter csvWriter, searchJobId, messageOffset, messageLength); messageOffset += messageLength; + Map hasJson = new HashMap(); try { List messages = getMessagesForSearchJobResponse.getMessages(); for (LogMessage message : messages) { @@ -725,7 +798,41 @@ private static int getMessages(CSVWriter csvWriter, // Write as JSON. if (outputFormat == OutputFormat.JSON) { - String json = objectMapper.writeValueAsString(fields); + Map jsonFields = new HashMap(); + for (int i = 0; i < fieldNames.size(); i++) { + String fieldName = fieldNames.get(i); + String fieldValue = fields.get(fieldName); + + // Replace with JSON if possible. + Boolean fieldHasJson = hasJson.get(fieldName); + if (fieldHasJson == null || fieldHasJson) { + TypeReference> typeRef = + new TypeReference>() { + }; + try { + HashMap jsonValue = objectMapper.readValue(fieldValue, typeRef); + if (jsonFieldToFlatten != null && fieldName.equals(jsonFieldToFlatten)) { + addToJsonFields(objectMapper, + jsonFields, + jsonValue, + fieldName + "_", + skipArrays, + arraysAsJson); + } else { + jsonFields.put(fieldName, jsonValue); + } + hasJson.put(fieldName, true); + } catch (JsonProcessingException jpe) { + hasJson.put(fieldName, false); + jsonFields.put(fieldName, fieldValue); + } + } else { + jsonFields.put(fieldName, fieldValue); + } + } + + String json = objectMapper.writeValueAsString( + new TreeMap((Map) jsonFields)); System.out.println(json); } } @@ -738,6 +845,69 @@ private static int getMessages(CSVWriter csvWriter, return messageOffset; } + private static void addToJsonFields(ObjectMapper objectMapper, + Map jsonFields, + Map jsonValue, + String prefix, + boolean skipArrays, + boolean arraysAsJson) { + for (Map.Entry entry : jsonValue.entrySet()) { + String fieldName = entry.getKey(); + Object fieldValue = entry.getValue(); + if (ClassUtil.isCollectionMapOrArray(fieldValue.getClass())) { + if (fieldValue instanceof Map) { + String fieldNamePrefix = (prefix == null) + ? fieldName + "_" + : prefix + fieldName + "_"; + addToJsonFields(objectMapper, + jsonFields, + ((Map) fieldValue), + fieldNamePrefix, + skipArrays, + arraysAsJson); + } else { + if (!skipArrays) { + if (!arraysAsJson) { + List valueList = (List) fieldValue; + for (int i = 0; i < valueList.size(); i++) { + String fieldNamePrefix = (prefix == null) + ? fieldName + "_" + i + "_" + : prefix + fieldName + "_" + i + "_"; + addToJsonFields(objectMapper, + jsonFields, + ((Map) valueList.get(i)), + fieldNamePrefix, + skipArrays, + arraysAsJson); + } + } else { + try { + String value = objectMapper.writeValueAsString(fieldValue); + addToJsonFieldsWithPrefix(jsonFields, prefix, fieldName, value); + } catch (IOException ioe) { + throw new RuntimeException(ioe); + } + } + } + } + } else { + // Primitive + addToJsonFieldsWithPrefix(jsonFields, prefix, fieldName, fieldValue); + } + } + } + + private static void addToJsonFieldsWithPrefix(Map jsonFields, + String prefix, + String fieldName, + Object fieldValue) { + if (prefix != null) { + jsonFields.put(prefix + fieldName, fieldValue); + } else { + jsonFields.put(fieldName, fieldValue); + } + } + private static int getRecords(CSVWriter csvWriter, AtomicBoolean headerWritten, ObjectMapper objectMapper, diff --git a/src/main/java/com/sumologic/client/util/HttpUtils.java b/src/main/java/com/sumologic/client/util/HttpUtils.java index c870107..92536e9 100644 --- a/src/main/java/com/sumologic/client/util/HttpUtils.java +++ b/src/main/java/com/sumologic/client/util/HttpUtils.java @@ -8,10 +8,12 @@ import com.sumologic.client.model.HttpPostRequest; import com.sumologic.client.model.HttpPutRequest; import org.apache.http.HttpEntity; +import org.apache.http.client.AuthCache; import org.apache.http.client.CookieStore; import org.apache.http.client.CredentialsProvider; import org.apache.http.client.config.RequestConfig; import org.apache.http.client.methods.*; +import org.apache.http.client.protocol.HttpClientContext; import org.apache.http.client.utils.URIBuilder; import org.apache.http.entity.StringEntity; import org.apache.http.impl.client.*; @@ -31,6 +33,8 @@ public class HttpUtils { private final CookieStore cookieStore = new BasicCookieStore(); + private final AuthCache authCache = new BasicAuthCache(); + // Public HTTP request methods public Response @@ -168,13 +172,6 @@ public static Map toRequestHeaders(String... parts) { // Private methods - private CloseableHttpClient getHttpClient(ConnectionConfig config) { - CredentialsProvider provider = new BasicCredentialsProvider(); - provider.setCredentials(config.getAuthScope(), config.getUsernamePasswordCredentials()); - return HttpClients.custom().setDefaultCookieStore(cookieStore) - .setDefaultCredentialsProvider(provider) - .build(); - } private static String getEndpointURI(String endpoint) { return "/" + UrlParameters.API_SERVICE + @@ -189,20 +186,30 @@ private static String getEndpointURI(String endpoint) { } private Response - doRequest(ConnectionConfig config, int timeout, HttpUriRequest method, Map requestHeaders, + doRequest(ConnectionConfig config, int timeout, HttpUriRequest uriRequest, Map requestHeaders, Request request, ResponseHandler handler, int expectedStatusCode) { // Set headers for (Map.Entry header : requestHeaders.entrySet()) { - method.setHeader(header.getKey(), header.getValue()); + uriRequest.setHeader(header.getKey(), header.getValue()); } - CloseableHttpClient httpClient = getHttpClient(config); + CredentialsProvider credsProvider = new BasicCredentialsProvider(); + credsProvider.setCredentials(config.getAuthScope(), config.getUsernamePasswordCredentials()); + CloseableHttpClient httpClient = HttpClients.custom() + .setDefaultCookieStore(cookieStore) + .build(); + + // NOTE(stefan, 2017-08-21): Pass in a long-lived authCache so that on subsequent calls we don't have to make + // two requests. + HttpClientContext context = HttpClientContext.create(); + context.setCredentialsProvider(credsProvider); + context.setAuthCache(authCache); InputStream httpStream = null; CloseableHttpResponse httpResponse = null; try { - httpResponse = httpClient.execute(method); + httpResponse = httpClient.execute(uriRequest, context); HttpEntity entity = httpResponse.getEntity(); httpStream = entity.getContent(); @@ -223,10 +230,10 @@ private static String getEndpointURI(String endpoint) { String json = writer.toString(); if (JacksonUtils.isValidJson(json)) - throw new SumoServerException(method.getURI().toString(), writer.toString()); + throw new SumoServerException(uriRequest.getURI().toString(), writer.toString()); else throw new SumoServerException( - method.getURI().toString(), + uriRequest.getURI().toString(), httpResponse.getStatusLine().getStatusCode()); } } @@ -254,8 +261,8 @@ private static String getEndpointURI(String endpoint) { } } - if (method != null) { - try { method.abort();} catch (Exception ex) {} + if (uriRequest != null) { + try { uriRequest.abort();} catch (Exception ex) {} } try {