Skip to content

Commit

Permalink
FIX: Only able to export 100 messages at a time
Browse files Browse the repository at this point in the history 
fix for patrickfav#18
  • Loading branch information
@bvfalcon
bvfalcon committed Aug 22, 2023
1 parent 0ee6fda commit 9fcc5f7
Show file tree
Hide file tree
Showing 6 changed files with 99 additions and 61 deletions.
103 changes: 54 additions & 49 deletions src/main/java/at/favre/tools/rocketexporter/RocketExporter.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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<Message> exportPrivateGroupMessages(String roomName, String roomId,
SortedSet<Message> exportPrivateGroupMessages(SortedSet<Message> messages, String roomName, String roomId,
int offset, int maxMessageCount,
File out, ExportFormat exportFormat) throws IOException, TooManyRequestException;

Expand All @@ -97,7 +97,7 @@ List<Message> 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<Message> exportChannelMessages(String channelName, String channelId,
SortedSet<Message> exportChannelMessages(SortedSet<Message> messages, String channelName, String channelId,
int offset, int maxMessageCount,
File out, ExportFormat exportFormat) throws IOException, TooManyRequestException;

Expand All @@ -114,7 +114,7 @@ List<Message> 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<Message> exportDirectMessages(String dmName, String dmId,
SortedSet<Message> exportDirectMessages(SortedSet<Message> messages, String dmName, String dmId,
int offset, int maxMessageCount,
File out, ExportFormat exportFormat) throws IOException, TooManyRequestException;

Expand Down Expand Up @@ -243,76 +243,81 @@ public List<RocketChatDm.DirectMessage> listDirectMessageChannels() throws IOExc
}

@Override
public List<Message> exportPrivateGroupMessages(String roomName, String roomId,
public SortedSet<Message> exportPrivateGroupMessages(SortedSet<Message> 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<Message> exportChannelMessages(String channelName, String channelId,
public SortedSet<Message> exportChannelMessages(SortedSet<Message> 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<Message> exportDirectMessages(String dmName, String dmId,
public SortedSet<Message> exportDirectMessages(SortedSet<Message> 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<Message> exportMessages(String contextName, String id,
private SortedSet<Message> exportMessages(SortedSet<Message> messages, String contextName, String id,
int offset, int maxMessageCount,
ConversationType conversationType, File out, ExportFormat exportFormat) throws IOException, TooManyRequestException {
checkAuthenticated();


Response<RocketChatMessageWrapperDto> 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<Long, Message> 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<RocketChatMessageWrapperDto> 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<Message> normalizedMessagesList = new ArrayList<>(normalizedMessages.values());
normalizedMessagesList.sort(Comparator.comparingLong(m -> m.getTimestamp().toEpochMilli()));
Map<Long, Message> 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() {
Expand Down
18 changes: 18 additions & 0 deletions src/main/java/at/favre/tools/rocketexporter/TooManyRequestException.java
View file
Original file line number Diff line number Diff line change
@@ -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<Message> messages;
TooManyRequestException(Object body, int offset, SortedSet<Message> messages) {
super(body != null ? body.toString() : "");
this.offset = offset;
this.messages = messages;
}
}
14 changes: 9 additions & 5 deletions src/main/java/at/favre/tools/rocketexporter/cli/Export.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -137,34 +137,38 @@ public void run() {
toExport.add(allConversations.get(selection));
}

int offset = 0;
SortedSet<Message> messages = new TreeSet<>();
for (int i = 0; i < toExport.size(); i++) {
Conversation selectedGroup = toExport.get(i);

final List<Message> 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--;
}
}
Expand Down
11 changes: 10 additions & 1 deletion src/main/java/at/favre/tools/rocketexporter/model/Message.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,10 +7,19 @@

@Data
@AllArgsConstructor
public class Message {
public class Message implements Comparable<Message> {
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());
}
}

8 changes: 5 additions & 3 deletions src/test/java/at/favre/tools/rocketexporter/RocketExporterTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -116,7 +118,7 @@ public void listDms() throws Exception {
public void exportPrivateGroupMessages() throws Exception {
login();
File tempFile = testFolder.newFile("out-test-group.csv");
List<Message> msg = exporter.exportPrivateGroupMessages("roomName", "roomId", 0, 2000, tempFile, new SlackCsvFormat());
SortedSet<Message> msg = exporter.exportPrivateGroupMessages(new TreeSet<>(), "roomName", "roomId", 0, 2000, tempFile, new SlackCsvFormat());
assertEquals(48, msg.size());
assertTrue(tempFile.exists() && tempFile.isFile() && tempFile.length() > 0);
}
Expand All @@ -125,7 +127,7 @@ public void exportPrivateGroupMessages() throws Exception {
public void exportChannelMessages() throws Exception {
login();
File tempFile = testFolder.newFile("out-test-channel.csv");
List<Message> msg = exporter.exportChannelMessages("roomName", "roomId", 0, 2000, tempFile, new SlackCsvFormat());
SortedSet<Message> msg = exporter.exportChannelMessages(new TreeSet<>(), "roomName", "roomId", 0, 2000, tempFile, new SlackCsvFormat());
assertEquals(3, msg.size());
assertTrue(tempFile.exists() && tempFile.isFile() && tempFile.length() > 0);
}
Expand All @@ -134,7 +136,7 @@ public void exportChannelMessages() throws Exception {
public void exportDms() throws Exception {
login();
File tempFile = testFolder.newFile("out-test-dm.csv");
List<Message> msg = exporter.exportDirectMessages("roomName", "roomId", 0, 2000, tempFile, new SlackCsvFormat());
SortedSet<Message> msg = exporter.exportDirectMessages(new TreeSet<>(), "roomName", "roomId", 0, 2000, tempFile, new SlackCsvFormat());
assertEquals(2, msg.size());
assertTrue(tempFile.exists() && tempFile.isFile() && tempFile.length() > 0);
}
Expand Down
6 changes: 3 additions & 3 deletions src/test/java/at/favre/tools/rocketexporter/converter/SlackCsvFormatTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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);

Expand Down

0 comments on commit 9fcc5f7

Please sign in to comment.