diff --git a/zerver/lib/slack_data_to_zulip_data.py b/zerver/lib/slack_data_to_zulip_data.py index ae1b41afc7..2f20eb4610 100755 --- a/zerver/lib/slack_data_to_zulip_data.py +++ b/zerver/lib/slack_data_to_zulip_data.py @@ -16,7 +16,7 @@ from django.db import connection from django.utils.timezone import now as timezone_now from typing import Any, Dict, List, Tuple from zerver.models import UserProfile, Realm, Stream, UserMessage, \ - Subscription, Message, Recipient, DefaultStream + Subscription, Message, Recipient, DefaultStream, Attachment from zerver.forms import check_subdomain_available from zerver.lib.slack_message_conversion import convert_to_zulip_markdown, \ get_user_full_name @@ -492,13 +492,14 @@ def convert_slack_workspace_messages(slack_data_dir: str, users: List[ZerverFiel logging.info('######### IMPORTING MESSAGES STARTED #########\n') # To pre-compute the total number of messages and usermessages - total_messages, total_usermessages = get_total_messages_and_usermessages( + total_messages, total_usermessages, total_attachments = get_total_messages_and_attachments( realm['zerver_subscription'], added_recipient, all_messages) message_id_list = allocate_ids(Message, total_messages) usermessage_id_list = allocate_ids(UserMessage, total_usermessages) + attachment_id_list = allocate_ids(UserMessage, total_attachments) - id_list = [message_id_list, usermessage_id_list] + id_list = [message_id_list, usermessage_id_list, attachment_id_list] zerver_message, zerver_usermessage = channel_message_to_zerver_message( realm_id, users, added_users, added_recipient, all_messages, realm['zerver_subscription'], domain_name, id_list) @@ -524,29 +525,32 @@ def get_all_messages(slack_data_dir: str, added_channels: AddedChannelsT) -> Lis all_messages += messages return all_messages -def get_total_messages_and_usermessages(zerver_subscription: List[ZerverFieldsT], - added_recipient: AddedRecipientsT, - all_messages: List[ZerverFieldsT]) -> Tuple[int, int]: +def get_total_messages_and_attachments(zerver_subscription: List[ZerverFieldsT], + added_recipient: AddedRecipientsT, + all_messages: List[ZerverFieldsT]) -> Tuple[int, int, + int]: """ Returns: 1. message_id, which is total number of messages 2. usermessage_id, which is total number of usermessages + 3. attachment_id, which is total number of attachments """ - total_messages = 0 - total_usermessages = 0 + total_messages = total_usermessages = total_attachments = 0 for message in all_messages: if 'subtype' in message.keys(): subtype = message['subtype'] if subtype in ["channel_join", "channel_leave", "channel_name"]: continue + elif subtype == "file_share": + total_attachments += 1 for subscription in zerver_subscription: if subscription['recipient'] == added_recipient[message['channel_name']]: total_usermessages += 1 total_messages += 1 - return total_messages, total_usermessages + return total_messages, total_usermessages, total_attachments def channel_message_to_zerver_message(realm_id: int, users: List[ZerverFieldsT], added_users: AddedUsersT, @@ -561,11 +565,12 @@ def channel_message_to_zerver_message(realm_id: int, users: List[ZerverFieldsT], 1. zerver_message, which is a list of the messages 2. zerver_usermessage, which is a list of the usermessages """ - message_id_count = usermessage_id_count = 0 - message_id_list, usermessage_id_list = ids + message_id_count = usermessage_id_count = attachment_id_count = 0 + message_id_list, usermessage_id_list, attachment_id_list = ids zerver_message = [] zerver_usermessage = [] # type: List[ZerverFieldsT] uploads_list = [] # type: List[ZerverFieldsT] + zerver_attachment = [] # type: List[ZerverFieldsT] for message in all_messages: user = get_message_sending_user(message) @@ -579,6 +584,10 @@ def channel_message_to_zerver_message(realm_id: int, users: List[ZerverFieldsT], users, added_users) rendered_content = None + + recipient_id = added_recipient[message['channel_name']] + message_id = message_id_list[message_id_count] + if 'subtype' in message.keys(): subtype = message['subtype'] if subtype in ["channel_join", "channel_leave", "channel_name"]: @@ -600,8 +609,11 @@ def channel_message_to_zerver_message(realm_id: int, users: List[ZerverFieldsT], build_uploads(added_users[user], realm_id, file_user_email, fileinfo, s3_path, uploads_list) - recipient_id = added_recipient[message['channel_name']] - message_id = message_id_list[message_id_count] + attachment_id = attachment_id_list[attachment_id_count] + build_zerver_attachment(realm_id, message_id, attachment_id, added_users[user], + fileinfo, s3_path, zerver_attachment) + attachment_id_count += 1 + # construct message zulip_message = dict( sending_client=1, @@ -664,6 +676,21 @@ def build_uploads(user_id: int, realm_id: int, email: str, fileinfo: ZerverField size=fileinfo['size']) uploads_list.append(upload) +def build_zerver_attachment(realm_id: int, message_id: int, attachment_id: int, + user_id: int, fileinfo: ZerverFieldsT, s3_path: str, + zerver_attachment: List[ZerverFieldsT]) -> None: + attachment = dict( + owner=user_id, + messages=[message_id], + id=attachment_id, + size=fileinfo['size'], + create_time=fileinfo['created'], + is_realm_public=True, # is always true for stream message + path_id=s3_path, + realm=realm_id, + file_name=fileinfo['name']) + zerver_attachment.append(attachment) + def get_message_sending_user(message: ZerverFieldsT) -> str: try: user = message.get('user', message['file']['user']) diff --git a/zerver/tests/test_slack_importer.py b/zerver/tests/test_slack_importer.py index 0f88936868..33fe90133a 100644 --- a/zerver/tests/test_slack_importer.py +++ b/zerver/tests/test_slack_importer.py @@ -18,7 +18,7 @@ from zerver.lib.slack_data_to_zulip_data import ( build_subscription, channels_to_zerver_stream, slack_workspace_to_realm, - get_total_messages_and_usermessages, + get_total_messages_and_attachments, get_message_sending_user, build_zerver_usermessage, channel_message_to_zerver_message, @@ -346,10 +346,10 @@ class SlackImporter(ZulipTestCase): self.assertEqual(realm['zerver_userprofile'], []) self.assertEqual(realm['zerver_realm'], [{}]) - def test_get_total_messages_and_usermessages(self) -> None: + def test_get_total_messages_and_attachments(self) -> None: messages = [{"text": "<@U8VAHEVUY> has joined the channel", "subtype": "channel_join", "channel_name": "random"}, - {"text": "message", "channel_name": "random"}, + {"text": "message", "channel_name": "random", "subtype": "file_share"}, {"text": "random", "channel_name": "random"}, {"text": "test messsage", "channel_name": "general"}, {"text": "test message 2", "subtype": "channel_leave", "channel_name": "general"}, @@ -359,11 +359,11 @@ class SlackImporter(ZulipTestCase): added_recipient = {'random': 2, 'general': 4} zerver_subscription = [{'recipient': 2}, {'recipient': 4}, {'recipient': 2}] - total_messages, total_usermessages = get_total_messages_and_usermessages(zerver_subscription, - added_recipient, - messages) + total_messages, total_usermessages, total_attachments = get_total_messages_and_attachments( + zerver_subscription, added_recipient, messages) # subtype: channel_join, channel_leave are filtered out self.assertEqual(total_messages, 4) + self.assertEqual(total_attachments, 1) self.assertEqual(total_usermessages, 6) def test_get_message_sending_user(self) -> None: @@ -428,7 +428,7 @@ class SlackImporter(ZulipTestCase): "ts": "1433868669.000012", "channel_name": "general"}] # type: List[Dict[str, Any]] added_recipient = {'random': 2, 'general': 1} - ids = [[3, 4, 5, 6, 7], []] + ids = [[3, 4, 5, 6, 7], [], []] zerver_usermessage = [] # type: List[Dict[str, Any]] zerver_subscription = [] # type: List[Dict[str, Any]] @@ -467,8 +467,8 @@ class SlackImporter(ZulipTestCase): @mock.patch("zerver.lib.slack_data_to_zulip_data.channel_message_to_zerver_message") @mock.patch("zerver.lib.slack_data_to_zulip_data.allocate_ids") @mock.patch("zerver.lib.slack_data_to_zulip_data.get_all_messages") - @mock.patch("zerver.lib.slack_data_to_zulip_data.get_total_messages_and_usermessages", return_value=[2, 4]) - def test_convert_slack_workspace_messages(self, mock_get_total_messages_and_usermessages: mock.Mock, + @mock.patch("zerver.lib.slack_data_to_zulip_data.get_total_messages_and_attachments", return_value=[2, 4, 1]) + def test_convert_slack_workspace_messages(self, mock_get_total_messages_and_attachments: mock.Mock, mock_get_all_messages: mock.Mock, mock_allocate_ids: mock.Mock, mock_message: mock.Mock) -> None: added_channels = {'random': 1, 'general': 2}