rocketchat: Import timezone-aware datetimes.

The bson library creates naive datetime objects by default.

Signed-off-by: Anders Kaseorg <anders@zulip.com>
This commit is contained in:
Anders Kaseorg
2022-12-25 14:53:26 -08:00
committed by Tim Abbott
parent 53238f0e8f
commit ce6f0e806a
2 changed files with 20 additions and 18 deletions

View File

@@ -979,6 +979,8 @@ def map_user_id_to_user(user_data_list: List[Dict[str, Any]]) -> Dict[str, Dict[
def rocketchat_data_to_dict(rocketchat_data_dir: str) -> Dict[str, Any]:
codec_options = bson.DEFAULT_CODEC_OPTIONS.with_options(tz_aware=True)
rocketchat_data: Dict[str, Any] = {}
rocketchat_data["instance"] = []
rocketchat_data["user"] = []
@@ -990,60 +992,60 @@ def rocketchat_data_to_dict(rocketchat_data_dir: str) -> Dict[str, Any]:
# Get instance
with open(os.path.join(rocketchat_data_dir, "instances.bson"), "rb") as fcache:
rocketchat_data["instance"] = bson.decode_all(fcache.read())
rocketchat_data["instance"] = bson.decode_all(fcache.read(), codec_options)
# Get user
with open(os.path.join(rocketchat_data_dir, "users.bson"), "rb") as fcache:
rocketchat_data["user"] = bson.decode_all(fcache.read())
rocketchat_data["user"] = bson.decode_all(fcache.read(), codec_options)
# Get avatar
with open(os.path.join(rocketchat_data_dir, "rocketchat_avatars.bson"), "rb") as fcache:
rocketchat_data["avatar"]["avatar"] = bson.decode_all(fcache.read())
rocketchat_data["avatar"]["avatar"] = bson.decode_all(fcache.read(), codec_options)
if rocketchat_data["avatar"]["avatar"]:
with open(
os.path.join(rocketchat_data_dir, "rocketchat_avatars.files.bson"), "rb"
) as fcache:
rocketchat_data["avatar"]["file"] = bson.decode_all(fcache.read())
rocketchat_data["avatar"]["file"] = bson.decode_all(fcache.read(), codec_options)
with open(
os.path.join(rocketchat_data_dir, "rocketchat_avatars.chunks.bson"), "rb"
) as fcache:
rocketchat_data["avatar"]["chunk"] = bson.decode_all(fcache.read())
rocketchat_data["avatar"]["chunk"] = bson.decode_all(fcache.read(), codec_options)
# Get room
with open(os.path.join(rocketchat_data_dir, "rocketchat_room.bson"), "rb") as fcache:
rocketchat_data["room"] = bson.decode_all(fcache.read())
rocketchat_data["room"] = bson.decode_all(fcache.read(), codec_options)
# Get messages
with open(os.path.join(rocketchat_data_dir, "rocketchat_message.bson"), "rb") as fcache:
rocketchat_data["message"] = bson.decode_all(fcache.read())
rocketchat_data["message"] = bson.decode_all(fcache.read(), codec_options)
# Get custom emoji
with open(os.path.join(rocketchat_data_dir, "rocketchat_custom_emoji.bson"), "rb") as fcache:
rocketchat_data["custom_emoji"]["emoji"] = bson.decode_all(fcache.read())
rocketchat_data["custom_emoji"]["emoji"] = bson.decode_all(fcache.read(), codec_options)
if rocketchat_data["custom_emoji"]["emoji"]:
with open(os.path.join(rocketchat_data_dir, "custom_emoji.files.bson"), "rb") as fcache:
rocketchat_data["custom_emoji"]["file"] = bson.decode_all(fcache.read())
rocketchat_data["custom_emoji"]["file"] = bson.decode_all(fcache.read(), codec_options)
with open(os.path.join(rocketchat_data_dir, "custom_emoji.chunks.bson"), "rb") as fcache:
rocketchat_data["custom_emoji"]["chunk"] = bson.decode_all(fcache.read())
rocketchat_data["custom_emoji"]["chunk"] = bson.decode_all(fcache.read(), codec_options)
# Get uploads
with open(os.path.join(rocketchat_data_dir, "rocketchat_uploads.bson"), "rb") as fcache:
rocketchat_data["upload"]["upload"] = bson.decode_all(fcache.read())
rocketchat_data["upload"]["upload"] = bson.decode_all(fcache.read(), codec_options)
if rocketchat_data["upload"]["upload"]:
with open(
os.path.join(rocketchat_data_dir, "rocketchat_uploads.files.bson"), "rb"
) as fcache:
rocketchat_data["upload"]["file"] = bson.decode_all(fcache.read())
rocketchat_data["upload"]["file"] = bson.decode_all(fcache.read(), codec_options)
with open(
os.path.join(rocketchat_data_dir, "rocketchat_uploads.chunks.bson"), "rb"
) as fcache:
rocketchat_data["upload"]["chunk"] = bson.decode_all(fcache.read())
rocketchat_data["upload"]["chunk"] = bson.decode_all(fcache.read(), codec_options)
return rocketchat_data

View File

@@ -1,5 +1,5 @@
import datetime
import os
from datetime import datetime, timezone
from typing import Any, Dict, List
import orjson
@@ -150,7 +150,7 @@ class RocketChatImporter(ZulipTestCase):
rocketchat_data["user"].append(
{
"_id": "s0m34ndmID",
"createdAt": datetime.datetime(2019, 11, 6, 0, 38, 42, 796000),
"createdAt": datetime(2019, 11, 6, 0, 38, 42, 796000, tzinfo=timezone.utc),
"type": "unknown",
"roles": ["unknown"],
"name": "Unknown user",
@@ -385,7 +385,7 @@ class RocketChatImporter(ZulipTestCase):
# Add a new channel with no user.
no_user_channel: Dict[str, Any] = {
"_id": "rand0mID",
"ts": datetime.datetime(2021, 7, 15, 10, 58, 23, 647000),
"ts": datetime(2021, 7, 15, 10, 58, 23, 647000, tzinfo=timezone.utc),
"t": "c",
"name": "no-user-channel",
}
@@ -672,7 +672,7 @@ class RocketChatImporter(ZulipTestCase):
{
"_id": "p4v37myxc6yLZ8AHh",
"t": "livechat_navigation_history",
"ts": datetime.datetime(2019, 11, 6, 0, 38, 42, 796000),
"ts": datetime(2019, 11, 6, 0, 38, 42, 796000, tzinfo=timezone.utc),
"msg": " - applewebdata://9124F033-BFEF-43C5-9215-DA369E4DA22D",
"u": {"_id": "rocket.cat", "username": "cat"},
"groupable": False,
@@ -687,7 +687,7 @@ class RocketChatImporter(ZulipTestCase):
},
"expireAt": 1575592722794.0,
"_hidden": True,
"_updatedAt": datetime.datetime(2019, 11, 6, 0, 38, 42, 796000),
"_updatedAt": datetime(2019, 11, 6, 0, 38, 42, 796000, tzinfo=timezone.utc),
}
)