zulip/zerver/data_import/hipchat.py

886 lines
27 KiB
Python
Executable File

import base64
import dateutil
import glob
import logging
import os
import re
import shutil
import subprocess
import ujson
from typing import Any, Callable, Dict, List, Optional, Set
from django.conf import settings
from django.forms.models import model_to_dict
from django.utils.timezone import now as timezone_now
from zerver.lib.utils import (
process_list_in_batches,
)
from zerver.models import (
RealmEmoji,
Recipient,
UserProfile,
)
from zerver.data_import.import_util import (
build_message,
build_realm,
build_realm_emoji,
build_recipients,
build_stream,
build_personal_subscriptions,
build_public_stream_subscriptions,
build_private_stream_subscriptions,
build_user_message,
build_user_profile,
build_zerver_realm,
create_converted_data_files,
make_subscriber_map,
write_avatar_png,
)
from zerver.data_import.hipchat_attachment import AttachmentHandler
from zerver.data_import.hipchat_user import UserHandler
from zerver.data_import.hipchat_subscriber import SubscriberHandler
from zerver.data_import.sequencer import NEXT_ID, IdMapper
# stubs
ZerverFieldsT = Dict[str, Any]
def str_date_to_float(date_str: str) -> float:
'''
Dates look like this:
"2018-08-08T14:23:54Z 626267"
'''
parts = date_str.split(' ')
time_str = parts[0].replace('T', ' ')
date_time = dateutil.parser.parse(time_str)
timestamp = date_time.timestamp()
if len(parts) == 2:
microseconds = int(parts[1])
timestamp += microseconds / 1000000.0
return timestamp
def untar_input_file(tar_file: str) -> str:
data_dir = tar_file.replace('.tar', '')
data_dir = os.path.abspath(data_dir)
if os.path.exists(data_dir):
logging.info('input data was already untarred to %s, we will use it' % (data_dir,))
return data_dir
os.makedirs(data_dir)
subprocess.check_call(['tar', '-xf', tar_file, '-C', data_dir])
logging.info('input data was untarred to %s' % (data_dir,))
return data_dir
def read_user_data(data_dir: str) -> List[ZerverFieldsT]:
fn = 'users.json'
data_file = os.path.join(data_dir, fn)
with open(data_file, "r") as fp:
return ujson.load(fp)
def convert_user_data(user_handler: UserHandler,
slim_mode: bool,
user_id_mapper: IdMapper,
raw_data: List[ZerverFieldsT],
realm_id: int) -> None:
flat_data = [
d['User']
for d in raw_data
]
if slim_mode:
flat_data = [
d for d in flat_data
if not d['is_deleted']
]
def process(in_dict: ZerverFieldsT) -> ZerverFieldsT:
delivery_email = in_dict['email']
email = in_dict['email']
full_name = in_dict['name']
id = user_id_mapper.get(in_dict['id'])
is_realm_admin = in_dict['account_type'] == 'admin'
is_guest = in_dict['account_type'] == 'guest'
is_mirror_dummy = False
short_name = in_dict['mention_name']
timezone = in_dict['timezone']
date_joined = int(timezone_now().timestamp())
is_active = not in_dict['is_deleted']
if not email:
if is_guest:
# Hipchat guest users don't have emails, so
# we just fake them.
email = 'guest-{id}@example.com'.format(id=id)
delivery_email = email
else:
# Hipchat sometimes doesn't export an email for deactivated users.
assert not is_active
email = delivery_email = "deactivated-{id}@example.com".format(id=id)
# unmapped fields:
# title - Developer, Project Manager, etc.
# rooms - no good sample data
# created - we just use "now"
# roles - we just use account_type
if in_dict.get('avatar'):
avatar_source = 'U'
else:
avatar_source = 'G'
return build_user_profile(
avatar_source=avatar_source,
date_joined=date_joined,
delivery_email=delivery_email,
email=email,
full_name=full_name,
id=id,
is_active=is_active,
is_realm_admin=is_realm_admin,
is_guest=is_guest,
is_mirror_dummy=is_mirror_dummy,
realm_id=realm_id,
short_name=short_name,
timezone=timezone,
)
for raw_item in flat_data:
user = process(raw_item)
user_handler.add_user(user)
def convert_avatar_data(avatar_folder: str,
raw_data: List[ZerverFieldsT],
user_id_mapper: IdMapper,
realm_id: int) -> List[ZerverFieldsT]:
'''
This code is pretty specific to how Hipchat sends us data.
They give us the avatar payloads in base64 in users.json.
We process avatars in our own pass of that data, rather
than doing it while we're getting other user data. I
chose to keep this separate, as otherwise you have a lot
of extraneous data getting passed around.
This code has MAJOR SIDE EFFECTS--namely writing a bunch
of files to the avatars directory.
'''
avatar_records = []
for d in raw_data:
raw_user = d['User']
avatar_payload = raw_user.get('avatar')
if not avatar_payload:
continue
bits = base64.b64decode(avatar_payload)
raw_user_id = raw_user['id']
if not user_id_mapper.has(raw_user_id):
continue
user_id = user_id_mapper.get(raw_user_id)
metadata = write_avatar_png(
avatar_folder=avatar_folder,
realm_id=realm_id,
user_id=user_id,
bits=bits,
)
avatar_records.append(metadata)
return avatar_records
def read_room_data(data_dir: str) -> List[ZerverFieldsT]:
fn = 'rooms.json'
data_file = os.path.join(data_dir, fn)
with open(data_file) as f:
data = ujson.load(f)
return data
def convert_room_data(raw_data: List[ZerverFieldsT],
subscriber_handler: SubscriberHandler,
stream_id_mapper: IdMapper,
user_id_mapper: IdMapper,
realm_id: int) -> List[ZerverFieldsT]:
flat_data = [
d['Room']
for d in raw_data
]
def get_invite_only(v: str) -> bool:
if v == 'public':
return False
elif v == 'private':
return True
else:
raise Exception('unexpected value')
streams = []
for in_dict in flat_data:
now = int(timezone_now().timestamp())
stream_id = stream_id_mapper.get(in_dict['id'])
invite_only = get_invite_only(in_dict['privacy'])
stream = build_stream(
date_created=now,
realm_id=realm_id,
name=in_dict['name'],
description=in_dict['topic'],
stream_id=stream_id,
deactivated=in_dict['is_archived'],
invite_only=invite_only,
)
if invite_only:
users = {
user_id_mapper.get(key)
for key in in_dict['members']
if user_id_mapper.has(key)
} # type: Set[int]
if user_id_mapper.has(in_dict['owner']):
owner = user_id_mapper.get(in_dict['owner'])
users.add(owner)
if not users:
continue
subscriber_handler.set_info(
stream_id=stream_id,
users=users,
)
# unmapped fields:
# guest_access_url: no Zulip equivalent
# created: we just use "now"
# participants: no good sample data
streams.append(stream)
return streams
def make_realm(realm_id: int) -> ZerverFieldsT:
NOW = float(timezone_now().timestamp())
domain_name = settings.EXTERNAL_HOST
realm_subdomain = ""
zerver_realm = build_zerver_realm(realm_id, realm_subdomain, NOW, 'HipChat')
realm = build_realm(zerver_realm, realm_id, domain_name)
# We may override these later.
realm['zerver_defaultstream'] = []
return realm
def write_avatar_data(raw_user_data: List[ZerverFieldsT],
output_dir: str,
user_id_mapper: IdMapper,
realm_id: int) -> None:
avatar_folder = os.path.join(output_dir, 'avatars')
avatar_realm_folder = os.path.join(avatar_folder, str(realm_id))
os.makedirs(avatar_realm_folder, exist_ok=True)
avatar_records = convert_avatar_data(
avatar_folder=avatar_folder,
raw_data=raw_user_data,
user_id_mapper=user_id_mapper,
realm_id=realm_id,
)
create_converted_data_files(avatar_records, output_dir, '/avatars/records.json')
def write_emoticon_data(realm_id: int,
data_dir: str,
output_dir: str) -> List[ZerverFieldsT]:
'''
This function does most of the work for processing emoticons, the bulk
of which is copying files. We also write a json file with metadata.
Finally, we return a list of RealmEmoji dicts to our caller.
In our data_dir we have a pretty simple setup:
emoticons.json - has very simple metadata on emojis:
{
"Emoticon": {
"id": 9875487,
"path": "emoticons/yasss.jpg",
"shortcut": "yasss"
}
},
{
"Emoticon": {
"id": 718017,
"path": "emoticons/yayyyyy.gif",
"shortcut": "yayyyyy"
}
}
emoticons/ - contains a bunch of image files:
slytherinsnake.gif
spanishinquisition.jpg
sparkle.png
spiderman.gif
stableparrot.gif
stalkerparrot.gif
supergirl.png
superman.png
We move all the relevant files to Zulip's more nested
directory structure.
'''
logging.info('Starting to process emoticons')
fn = 'emoticons.json'
data_file = os.path.join(data_dir, fn)
with open(data_file) as f:
data = ujson.load(f)
flat_data = [
dict(
path=d['Emoticon']['path'],
name=d['Emoticon']['shortcut'],
)
for d in data
]
emoji_folder = os.path.join(output_dir, 'emoji')
os.makedirs(emoji_folder, exist_ok=True)
def process(data: ZerverFieldsT) -> ZerverFieldsT:
source_sub_path = data['path']
source_fn = os.path.basename(source_sub_path)
source_path = os.path.join(data_dir, source_sub_path)
# Use our template from RealmEmoji
# PATH_ID_TEMPLATE = "{realm_id}/emoji/images/{emoji_file_name}"
target_fn = source_fn
target_sub_path = RealmEmoji.PATH_ID_TEMPLATE.format(
realm_id=realm_id,
emoji_file_name=target_fn,
)
target_path = os.path.join(emoji_folder, target_sub_path)
os.makedirs(os.path.dirname(target_path), exist_ok=True)
source_path = os.path.abspath(source_path)
target_path = os.path.abspath(target_path)
shutil.copyfile(source_path, target_path)
return dict(
path=target_path,
s3_path=target_path,
file_name=target_fn,
realm_id=realm_id,
name=data['name'],
)
emoji_records = list(map(process, flat_data))
create_converted_data_files(emoji_records, output_dir, '/emoji/records.json')
realmemoji = [
build_realm_emoji(
realm_id=realm_id,
name=rec['name'],
id=NEXT_ID('realmemoji'),
file_name=rec['file_name'],
)
for rec in emoji_records
]
logging.info('Done processing emoticons')
return realmemoji
def write_message_data(realm_id: int,
slim_mode: bool,
message_key: str,
zerver_recipient: List[ZerverFieldsT],
subscriber_map: Dict[int, Set[int]],
data_dir: str,
output_dir: str,
masking_content: bool,
stream_id_mapper: IdMapper,
user_id_mapper: IdMapper,
user_handler: UserHandler,
attachment_handler: AttachmentHandler) -> None:
stream_id_to_recipient_id = {
d['type_id']: d['id']
for d in zerver_recipient
if d['type'] == Recipient.STREAM
}
user_id_to_recipient_id = {
d['type_id']: d['id']
for d in zerver_recipient
if d['type'] == Recipient.PERSONAL
}
def get_stream_recipient_id(raw_message: ZerverFieldsT) -> int:
fn_id = raw_message['fn_id']
stream_id = stream_id_mapper.get(fn_id)
recipient_id = stream_id_to_recipient_id[stream_id]
return recipient_id
def get_pm_recipient_id(raw_message: ZerverFieldsT) -> int:
raw_user_id = raw_message['receiver_id']
assert(raw_user_id)
user_id = user_id_mapper.get(raw_user_id)
recipient_id = user_id_to_recipient_id[user_id]
return recipient_id
if message_key in ['UserMessage', 'NotificationMessage']:
is_pm_data = False
dir_glob = os.path.join(data_dir, 'rooms', '*', 'history.json')
get_recipient_id = get_stream_recipient_id
get_files_dir = lambda fn_id: os.path.join(data_dir, 'rooms', str(fn_id), 'files')
elif message_key == 'PrivateUserMessage':
is_pm_data = True
dir_glob = os.path.join(data_dir, 'users', '*', 'history.json')
get_recipient_id = get_pm_recipient_id
get_files_dir = lambda fn_id: os.path.join(data_dir, 'users', 'files')
else:
raise Exception('programming error: invalid message_key: ' + message_key)
history_files = glob.glob(dir_glob)
for fn in history_files:
dir = os.path.dirname(fn)
fn_id = os.path.basename(dir)
files_dir = get_files_dir(fn_id)
process_message_file(
realm_id=realm_id,
slim_mode=slim_mode,
fn=fn,
fn_id=fn_id,
files_dir=files_dir,
get_recipient_id=get_recipient_id,
message_key=message_key,
subscriber_map=subscriber_map,
data_dir=data_dir,
output_dir=output_dir,
is_pm_data=is_pm_data,
masking_content=masking_content,
user_id_mapper=user_id_mapper,
user_handler=user_handler,
attachment_handler=attachment_handler,
)
def get_hipchat_sender_id(realm_id: int,
slim_mode: bool,
message_dict: Dict[str, Any],
user_id_mapper: IdMapper,
user_handler: UserHandler) -> Optional[int]:
'''
The HipChat export is inconsistent in how it renders
senders, and sometimes we don't even get an id.
'''
if isinstance(message_dict['sender'], str):
if slim_mode:
return None
# Some Hipchat instances just give us a person's
# name in the sender field for NotificationMessage.
# We turn them into a mirror user.
mirror_user = user_handler.get_mirror_user(
realm_id=realm_id,
name=message_dict['sender'],
)
sender_id = mirror_user['id']
return sender_id
raw_sender_id = message_dict['sender']['id']
if raw_sender_id == 0:
if slim_mode:
return None
mirror_user = user_handler.get_mirror_user(
realm_id=realm_id,
name=message_dict['sender']['name']
)
sender_id = mirror_user['id']
return sender_id
if not user_id_mapper.has(raw_sender_id):
if slim_mode:
return None
mirror_user = user_handler.get_mirror_user(
realm_id=realm_id,
name=message_dict['sender']['id']
)
sender_id = mirror_user['id']
return sender_id
# HAPPY PATH: Hipchat just gave us an ordinary
# sender_id.
sender_id = user_id_mapper.get(raw_sender_id)
return sender_id
def process_message_file(realm_id: int,
slim_mode: bool,
fn: str,
fn_id: str,
files_dir: str,
get_recipient_id: Callable[[ZerverFieldsT], int],
message_key: str,
subscriber_map: Dict[int, Set[int]],
data_dir: str,
output_dir: str,
is_pm_data: bool,
masking_content: bool,
user_id_mapper: IdMapper,
user_handler: UserHandler,
attachment_handler: AttachmentHandler) -> None:
def get_raw_messages(fn: str) -> List[ZerverFieldsT]:
with open(fn) as f:
data = ujson.load(f)
flat_data = [
d[message_key]
for d in data
if message_key in d
]
def get_raw_message(d: Dict[str, Any]) -> Optional[ZerverFieldsT]:
sender_id = get_hipchat_sender_id(
realm_id=realm_id,
slim_mode=slim_mode,
message_dict=d,
user_id_mapper=user_id_mapper,
user_handler=user_handler,
)
if sender_id is None:
return None
if is_pm_data:
if sender_id != fn_id:
# PMs are in multiple places in the Hipchat export,
# and we only use the copy from the sender
return None
content = d['message']
if masking_content:
content = re.sub('[a-z]', 'x', content)
content = re.sub('[A-Z]', 'X', content)
return dict(
fn_id=fn_id,
sender_id=sender_id,
receiver_id=d.get('receiver', {}).get('id'),
content=content,
mention_user_ids=d.get('mentions', []),
pub_date=str_date_to_float(d['timestamp']),
attachment=d.get('attachment'),
files_dir=files_dir,
)
raw_messages = []
for d in flat_data:
raw_message = get_raw_message(d)
if raw_message is not None:
raw_messages.append(raw_message)
return raw_messages
raw_messages = get_raw_messages(fn)
def process_batch(lst: List[Any]) -> None:
process_raw_message_batch(
realm_id=realm_id,
raw_messages=lst,
subscriber_map=subscriber_map,
user_id_mapper=user_id_mapper,
user_handler=user_handler,
attachment_handler=attachment_handler,
get_recipient_id=get_recipient_id,
is_pm_data=is_pm_data,
output_dir=output_dir,
)
chunk_size = 1000
process_list_in_batches(
lst=raw_messages,
chunk_size=chunk_size,
process_batch=process_batch,
)
def process_raw_message_batch(realm_id: int,
raw_messages: List[Dict[str, Any]],
subscriber_map: Dict[int, Set[int]],
user_id_mapper: IdMapper,
user_handler: UserHandler,
attachment_handler: AttachmentHandler,
get_recipient_id: Callable[[ZerverFieldsT], int],
is_pm_data: bool,
output_dir: str) -> None:
def fix_mentions(content: str,
mention_user_ids: Set[int]) -> str:
for user_id in mention_user_ids:
user = user_handler.get_user(user_id=user_id)
hipchat_mention = '@{short_name}'.format(**user)
zulip_mention = '@**{full_name}**'.format(**user)
content = content.replace(hipchat_mention, zulip_mention)
content = content.replace('@here', '@**all**')
return content
mention_map = dict() # type: Dict[int, Set[int]]
zerver_message = []
for raw_message in raw_messages:
# One side effect here:
message_id = NEXT_ID('message')
mention_user_ids = {
user_id_mapper.get(id)
for id in set(raw_message['mention_user_ids'])
if user_id_mapper.has(id)
}
mention_map[message_id] = mention_user_ids
content = fix_mentions(
content=raw_message['content'],
mention_user_ids=mention_user_ids,
)
if len(content) > 10000:
logging.info('skipping too-long message of length %s' % (len(content),))
continue
pub_date = raw_message['pub_date']
try:
recipient_id = get_recipient_id(raw_message)
except KeyError:
logging.debug("Could not find recipient_id for a message, skipping.")
continue
rendered_content = None
if is_pm_data:
topic_name = ''
else:
topic_name = 'imported from hipchat'
user_id = raw_message['sender_id']
# Another side effect:
extra_content = attachment_handler.handle_message_data(
realm_id=realm_id,
message_id=message_id,
sender_id=user_id,
attachment=raw_message['attachment'],
files_dir=raw_message['files_dir'],
)
if extra_content:
has_attachment = True
content += '\n' + extra_content
else:
has_attachment = False
message = build_message(
content=content,
message_id=message_id,
pub_date=pub_date,
recipient_id=recipient_id,
rendered_content=rendered_content,
topic_name=topic_name,
user_id=user_id,
has_attachment=has_attachment,
)
zerver_message.append(message)
zerver_usermessage = make_user_messages(
zerver_message=zerver_message,
subscriber_map=subscriber_map,
is_pm_data=is_pm_data,
mention_map=mention_map,
)
message_json = dict(
zerver_message=zerver_message,
zerver_usermessage=zerver_usermessage,
)
dump_file_id = NEXT_ID('dump_file_id')
message_file = "/messages-%06d.json" % (dump_file_id,)
create_converted_data_files(message_json, output_dir, message_file)
def make_user_messages(zerver_message: List[ZerverFieldsT],
subscriber_map: Dict[int, Set[int]],
is_pm_data: bool,
mention_map: Dict[int, Set[int]]) -> List[ZerverFieldsT]:
zerver_usermessage = []
for message in zerver_message:
message_id = message['id']
recipient_id = message['recipient']
sender_id = message['sender']
mention_user_ids = mention_map[message_id]
subscriber_ids = subscriber_map.get(recipient_id, set())
user_ids = subscriber_ids | {sender_id}
for user_id in user_ids:
is_mentioned = user_id in mention_user_ids
user_message = build_user_message(
user_id=user_id,
message_id=message_id,
is_private=is_pm_data,
is_mentioned=is_mentioned,
)
zerver_usermessage.append(user_message)
return zerver_usermessage
def do_convert_data(input_tar_file: str,
output_dir: str,
masking_content: bool) -> None:
input_data_dir = untar_input_file(input_tar_file)
attachment_handler = AttachmentHandler()
user_handler = UserHandler()
subscriber_handler = SubscriberHandler()
user_id_mapper = IdMapper()
stream_id_mapper = IdMapper()
realm_id = 0
realm = make_realm(realm_id=realm_id)
slim_mode = True
# users.json -> UserProfile
raw_user_data = read_user_data(data_dir=input_data_dir)
convert_user_data(
user_handler=user_handler,
slim_mode=slim_mode,
user_id_mapper=user_id_mapper,
raw_data=raw_user_data,
realm_id=realm_id,
)
normal_users = user_handler.get_normal_users()
# Don't write zerver_userprofile here, because we
# may add more users later.
# streams.json -> Stream
raw_stream_data = read_room_data(data_dir=input_data_dir)
zerver_stream = convert_room_data(
raw_data=raw_stream_data,
subscriber_handler=subscriber_handler,
stream_id_mapper=stream_id_mapper,
user_id_mapper=user_id_mapper,
realm_id=realm_id,
)
realm['zerver_stream'] = zerver_stream
zerver_recipient = build_recipients(
zerver_userprofile=normal_users,
zerver_stream=zerver_stream,
)
realm['zerver_recipient'] = zerver_recipient
if slim_mode:
public_stream_subscriptions = [] # type: List[ZerverFieldsT]
else:
public_stream_subscriptions = build_public_stream_subscriptions(
zerver_userprofile=normal_users,
zerver_recipient=zerver_recipient,
zerver_stream=zerver_stream,
)
private_stream_subscriptions = build_private_stream_subscriptions(
get_users=subscriber_handler.get_users,
zerver_recipient=zerver_recipient,
zerver_stream=zerver_stream,
)
personal_subscriptions = build_personal_subscriptions(
zerver_recipient=zerver_recipient,
)
zerver_subscription = \
public_stream_subscriptions + \
personal_subscriptions + \
private_stream_subscriptions
realm['zerver_subscription'] = zerver_subscription
zerver_realmemoji = write_emoticon_data(
realm_id=realm_id,
data_dir=input_data_dir,
output_dir=output_dir,
)
realm['zerver_realmemoji'] = zerver_realmemoji
subscriber_map = make_subscriber_map(
zerver_subscription=zerver_subscription,
)
logging.info('Start importing message data')
for message_key in ['UserMessage',
'NotificationMessage',
'PrivateUserMessage']:
write_message_data(
realm_id=realm_id,
slim_mode=slim_mode,
message_key=message_key,
zerver_recipient=zerver_recipient,
subscriber_map=subscriber_map,
data_dir=input_data_dir,
output_dir=output_dir,
masking_content=masking_content,
stream_id_mapper=stream_id_mapper,
user_id_mapper=user_id_mapper,
user_handler=user_handler,
attachment_handler=attachment_handler,
)
# Order is important here...don't write users until
# we process everything else, since we may introduce
# mirror users when processing messages.
realm['zerver_userprofile'] = user_handler.get_all_users()
realm['sort_by_date'] = True
create_converted_data_files(realm, output_dir, '/realm.json')
logging.info('Start importing avatar data')
write_avatar_data(
raw_user_data=raw_user_data,
output_dir=output_dir,
user_id_mapper=user_id_mapper,
realm_id=realm_id,
)
attachment_handler.write_info(
output_dir=output_dir,
realm_id=realm_id,
)
logging.info('Start making tarball')
subprocess.check_call(["tar", "-czf", output_dir + '.tar.gz', output_dir, '-P'])
logging.info('Done making tarball')