diff --git a/src/main/java/at/favre/tools/rocketexporter/RocketExporter.java b/src/main/java/at/favre/tools/rocketexporter/RocketExporter.java index 3ccb80b..69dc81b 100644 --- a/src/main/java/at/favre/tools/rocketexporter/RocketExporter.java +++ b/src/main/java/at/favre/tools/rocketexporter/RocketExporter.java @@ -80,7 +80,7 @@ public interface RocketExporter { * @throws IOException on issues during the REST call * @throws TooManyRequestException if the server responds with 429, you need to throttle the requests */ - List exportPrivateGroupMessages(String roomName, String roomId, + SortedSet exportPrivateGroupMessages(SortedSet messages, String roomName, String roomId, int offset, int maxMessageCount, File out, ExportFormat exportFormat) throws IOException, TooManyRequestException; @@ -97,7 +97,7 @@ List exportPrivateGroupMessages(String roomName, String roomId, * @throws IOException on issues during the REST call * @throws TooManyRequestException if the server responds with 429, you need to throttle the requests */ - List exportChannelMessages(String channelName, String channelId, + SortedSet exportChannelMessages(SortedSet messages, String channelName, String channelId, int offset, int maxMessageCount, File out, ExportFormat exportFormat) throws IOException, TooManyRequestException; @@ -114,7 +114,7 @@ List exportChannelMessages(String channelName, String channelId, * @throws IOException on issues during the REST call * @throws TooManyRequestException if the server responds with 429, you need to throttle the requests */ - List exportDirectMessages(String dmName, String dmId, + SortedSet exportDirectMessages(SortedSet messages, String dmName, String dmId, int offset, int maxMessageCount, File out, ExportFormat exportFormat) throws IOException, TooManyRequestException; @@ -243,76 +243,81 @@ public List listDirectMessageChannels() throws IOExc } @Override - public List exportPrivateGroupMessages(String roomName, String roomId, + public SortedSet exportPrivateGroupMessages(SortedSet messages, String roomName, String roomId, int offset, int maxMessageCount, File out, ExportFormat exportFormat) throws IOException, TooManyRequestException { - return exportMessages(roomName, roomId, offset, maxMessageCount, ConversationType.GROUP, out, exportFormat); + return exportMessages(messages, roomName, roomId, offset, maxMessageCount, ConversationType.GROUP, out, exportFormat); } @Override - public List exportChannelMessages(String channelName, String channelId, + public SortedSet exportChannelMessages(SortedSet messages, String channelName, String channelId, int offset, int maxMessageCount, File out, ExportFormat exportFormat) throws IOException, TooManyRequestException { - return exportMessages(channelName, channelId, offset, maxMessageCount, ConversationType.CHANNEL, out, exportFormat); + return exportMessages(messages, channelName, channelId, offset, maxMessageCount, ConversationType.CHANNEL, out, exportFormat); } @Override - public List exportDirectMessages(String dmName, String dmId, + public SortedSet exportDirectMessages(SortedSet messages, String dmName, String dmId, int offset, int maxMessageCount, File out, ExportFormat exportFormat) throws IOException, TooManyRequestException { - return exportMessages(dmName, dmId, offset, maxMessageCount, ConversationType.DIRECT_MESSAGES, out, exportFormat); + return exportMessages(messages, dmName, dmId, offset, maxMessageCount, ConversationType.DIRECT_MESSAGES, out, exportFormat); } - private List exportMessages(String contextName, String id, + private SortedSet exportMessages(SortedSet messages, String contextName, String id, int offset, int maxMessageCount, ConversationType conversationType, File out, ExportFormat exportFormat) throws IOException, TooManyRequestException { checkAuthenticated(); - - Response response; - switch (conversationType) { - case GROUP: - response = getService().getAllMessagesFromGroup(authHeaders, id, offset, maxMessageCount).execute(); - break; - case CHANNEL: - response = getService().getAllMessagesFromChannels(authHeaders, id, offset, maxMessageCount).execute(); - break; - case DIRECT_MESSAGES: - response = getService().getAllMessagesFromDirectMessages(authHeaders, id, offset, maxMessageCount).execute(); - break; - default: - throw new IllegalStateException(); - } - - Map normalizedMessages = new HashMap<>(); - RocketChatMessageWrapperDto messagesBody; - - if (response.code() == 200 && (messagesBody = response.body()) != null) { - for (RocketChatMessageWrapperDto.Message message : messagesBody.getMessages()) { - Instant timestamp = Instant.parse(message.getTs()); - - normalizedMessages.put(timestamp.toEpochMilli(), - new Message( - message.getMsg(), - message.getU().getName(), - contextName, - timestamp - )); + do { + Response response; + switch (conversationType) { + case GROUP: + response = getService().getAllMessagesFromGroup(authHeaders, id, offset, maxMessageCount).execute(); + break; + case CHANNEL: + response = getService().getAllMessagesFromChannels(authHeaders, id, offset, maxMessageCount).execute(); + break; + case DIRECT_MESSAGES: + response = getService().getAllMessagesFromDirectMessages(authHeaders, id, offset, maxMessageCount).execute(); + break; + default: + throw new IllegalStateException(); } - } else if (response.code() == 429) { - throw new TooManyRequestException(response.body()); - } else { - throw new IllegalStateException("error response: " + response.code()); - } - List normalizedMessagesList = new ArrayList<>(normalizedMessages.values()); - normalizedMessagesList.sort(Comparator.comparingLong(m -> m.getTimestamp().toEpochMilli())); + Map normalizedMessages = new HashMap<>(); + RocketChatMessageWrapperDto messagesBody; + + if (response.code() == 200 && (messagesBody = response.body()) != null) { + for (RocketChatMessageWrapperDto.Message message : messagesBody.getMessages()) { + Instant timestamp = Instant.parse(message.getTs()); + + normalizedMessages.put(timestamp.toEpochMilli(), + new Message( + message.get_id(), + message.getMsg(), + message.getU().getName(), + contextName, + timestamp + )); + } + messages.addAll(normalizedMessages.values()); + if (messagesBody.getMessages().size() < 100 || messages.size() >= maxMessageCount) { + break; + } else { + offset += 100; + } + } else if (response.code() == 429) { + throw new TooManyRequestException(response.body(), offset, messages); + } else { + throw new IllegalStateException("error response: " + response.code()); + } + } while (true); exportFormat.export( - normalizedMessagesList, + new ArrayList<>(messages), new FileOutputStream(out)); - return normalizedMessagesList; + return messages; } private void checkAuthenticated() { diff --git a/src/main/java/at/favre/tools/rocketexporter/TooManyRequestException.java b/src/main/java/at/favre/tools/rocketexporter/TooManyRequestException.java index 5b0c4fd..3df7734 100644 --- a/src/main/java/at/favre/tools/rocketexporter/TooManyRequestException.java +++ b/src/main/java/at/favre/tools/rocketexporter/TooManyRequestException.java @@ -1,7 +1,25 @@ package at.favre.tools.rocketexporter; +import java.util.SortedSet; +import java.util.TreeSet; + +import at.favre.tools.rocketexporter.model.Message; +import lombok.Getter; + public class TooManyRequestException extends Exception { TooManyRequestException(Object body) { super(body != null ? body.toString() : ""); + offset = 0; + messages = new TreeSet<>(); + } + + @Getter + private final int offset; + @Getter + private final SortedSet messages; + TooManyRequestException(Object body, int offset, SortedSet messages) { + super(body != null ? body.toString() : ""); + this.offset = offset; + this.messages = messages; } } diff --git a/src/main/java/at/favre/tools/rocketexporter/cli/Export.java b/src/main/java/at/favre/tools/rocketexporter/cli/Export.java index 4aceb70..cc3184d 100644 --- a/src/main/java/at/favre/tools/rocketexporter/cli/Export.java +++ b/src/main/java/at/favre/tools/rocketexporter/cli/Export.java @@ -137,34 +137,38 @@ public void run() { toExport.add(allConversations.get(selection)); } + int offset = 0; + SortedSet messages = new TreeSet<>(); for (int i = 0; i < toExport.size(); i++) { Conversation selectedGroup = toExport.get(i); - final List messages; final ExportFormat format = new SlackCsvFormat(); - final int offset = 0; final int maxMsg = maxMessages; final File outFile = generateOutputFile(file, selectedGroup.getName(), type, format); try { switch (type) { case GROUP: - messages = exporter.exportPrivateGroupMessages(selectedGroup.getName(), selectedGroup.get_id(), offset, maxMsg, outFile, format); + messages = exporter.exportPrivateGroupMessages(messages, selectedGroup.getName(), selectedGroup.get_id(), offset, maxMsg, outFile, format); break; case CHANNEL: - messages = exporter.exportChannelMessages(selectedGroup.getName(), selectedGroup.get_id(), offset, maxMsg, outFile, format); + messages = exporter.exportChannelMessages(messages, selectedGroup.getName(), selectedGroup.get_id(), offset, maxMsg, outFile, format); break; case DIRECT_MESSAGES: - messages = exporter.exportDirectMessages(selectedGroup.getName(), selectedGroup.get_id(), offset, maxMsg, outFile, format); + messages = exporter.exportDirectMessages(messages, selectedGroup.getName(), selectedGroup.get_id(), offset, maxMsg, outFile, format); break; default: throw new IllegalStateException(); } out.println("Successfully exported " + messages.size() + " " + type.name + " messages to '" + outFile + "'"); + messages.clear(); + offset = 0; } catch (TooManyRequestException e) { out.println("Too many requests. Slowing down..."); Thread.sleep(5000); + offset = e.getOffset(); + messages = e.getMessages(); i--; } } diff --git a/src/main/java/at/favre/tools/rocketexporter/model/Message.java b/src/main/java/at/favre/tools/rocketexporter/model/Message.java index da366aa..9d1dfc4 100644 --- a/src/main/java/at/favre/tools/rocketexporter/model/Message.java +++ b/src/main/java/at/favre/tools/rocketexporter/model/Message.java @@ -7,10 +7,19 @@ @Data @AllArgsConstructor -public class Message { +public class Message implements Comparable { + private final String _id; private final String message; private final String username; private final String channel; private final Instant timestamp; + + @Override + public int compareTo(Message o) { + if (this.equals(o)) { + return 0; + } + return Long.valueOf(timestamp.toEpochMilli()).compareTo(o.getTimestamp().toEpochMilli()); + } } diff --git a/src/test/java/at/favre/tools/rocketexporter/RocketExporterTest.java b/src/test/java/at/favre/tools/rocketexporter/RocketExporterTest.java index 0105bde..8a5889b 100644 --- a/src/test/java/at/favre/tools/rocketexporter/RocketExporterTest.java +++ b/src/test/java/at/favre/tools/rocketexporter/RocketExporterTest.java @@ -12,6 +12,8 @@ import java.io.File; import java.net.URI; import java.util.List; +import java.util.SortedSet; +import java.util.TreeSet; import static com.github.tomakehurst.wiremock.client.WireMock.*; import static com.github.tomakehurst.wiremock.core.WireMockConfiguration.options; @@ -116,7 +118,7 @@ public void listDms() throws Exception { public void exportPrivateGroupMessages() throws Exception { login(); File tempFile = testFolder.newFile("out-test-group.csv"); - List msg = exporter.exportPrivateGroupMessages("roomName", "roomId", 0, 2000, tempFile, new SlackCsvFormat()); + SortedSet msg = exporter.exportPrivateGroupMessages(new TreeSet<>(), "roomName", "roomId", 0, 2000, tempFile, new SlackCsvFormat()); assertEquals(48, msg.size()); assertTrue(tempFile.exists() && tempFile.isFile() && tempFile.length() > 0); } @@ -125,7 +127,7 @@ public void exportPrivateGroupMessages() throws Exception { public void exportChannelMessages() throws Exception { login(); File tempFile = testFolder.newFile("out-test-channel.csv"); - List msg = exporter.exportChannelMessages("roomName", "roomId", 0, 2000, tempFile, new SlackCsvFormat()); + SortedSet msg = exporter.exportChannelMessages(new TreeSet<>(), "roomName", "roomId", 0, 2000, tempFile, new SlackCsvFormat()); assertEquals(3, msg.size()); assertTrue(tempFile.exists() && tempFile.isFile() && tempFile.length() > 0); } @@ -134,7 +136,7 @@ public void exportChannelMessages() throws Exception { public void exportDms() throws Exception { login(); File tempFile = testFolder.newFile("out-test-dm.csv"); - List msg = exporter.exportDirectMessages("roomName", "roomId", 0, 2000, tempFile, new SlackCsvFormat()); + SortedSet msg = exporter.exportDirectMessages(new TreeSet<>(), "roomName", "roomId", 0, 2000, tempFile, new SlackCsvFormat()); assertEquals(2, msg.size()); assertTrue(tempFile.exists() && tempFile.isFile() && tempFile.length() > 0); } diff --git a/src/test/java/at/favre/tools/rocketexporter/converter/SlackCsvFormatTest.java b/src/test/java/at/favre/tools/rocketexporter/converter/SlackCsvFormatTest.java index 7fc3738..edb6966 100644 --- a/src/test/java/at/favre/tools/rocketexporter/converter/SlackCsvFormatTest.java +++ b/src/test/java/at/favre/tools/rocketexporter/converter/SlackCsvFormatTest.java @@ -18,9 +18,9 @@ public void export() { ByteArrayOutputStream bout = new ByteArrayOutputStream(); exportFormat.export( List.of( - new Message("m1", "u1", "c1", EPOCH), - new Message("m2", "u2", "c3", EPOCH.plusSeconds(1)), - new Message(null, "u3", "c3", EPOCH) + new Message("a", "m1", "u1", "c1", EPOCH), + new Message("b", "m2", "u2", "c3", EPOCH.plusSeconds(1)), + new Message("c", null, "u3", "c3", EPOCH) ), bout);