From 79b6f43d0ecc3984e7b41f3896dac1dfdac2437a Mon Sep 17 00:00:00 2001 From: Tim Abbott Date: Thu, 17 Oct 2024 11:09:44 -0700 Subject: [PATCH] rocketchat: Move bson_code_options to a global variable. This will make it a lot easier to only read files in when we actually need them. --- zerver/data_import/rocketchat.py | 36 +++++++++++++++++++------------- 1 file changed, 21 insertions(+), 15 deletions(-) diff --git a/zerver/data_import/rocketchat.py b/zerver/data_import/rocketchat.py index 906f0a1182..be0d6e113a 100644 --- a/zerver/data_import/rocketchat.py +++ b/zerver/data_import/rocketchat.py @@ -36,6 +36,8 @@ from zerver.lib.upload import sanitize_name from zerver.lib.utils import process_list_in_batches from zerver.models import Reaction, RealmEmoji, Recipient, UserProfile +bson_codec_options = bson.DEFAULT_CODEC_OPTIONS.with_options(tz_aware=True) + def make_realm( realm_id: int, realm_subdomain: str, domain_name: str, rc_instance: dict[str, Any] @@ -976,8 +978,6 @@ 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"] = [] @@ -989,60 +989,66 @@ 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(), codec_options) + rocketchat_data["instance"] = bson.decode_all(fcache.read(), bson_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(), codec_options) + rocketchat_data["user"] = bson.decode_all(fcache.read(), bson_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(), codec_options) + rocketchat_data["avatar"]["avatar"] = bson.decode_all(fcache.read(), bson_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(), codec_options) + rocketchat_data["avatar"]["file"] = bson.decode_all(fcache.read(), bson_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(), codec_options) + rocketchat_data["avatar"]["chunk"] = bson.decode_all(fcache.read(), bson_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(), codec_options) + rocketchat_data["room"] = bson.decode_all(fcache.read(), bson_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(), codec_options) + rocketchat_data["message"] = bson.decode_all(fcache.read(), bson_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(), codec_options) + rocketchat_data["custom_emoji"]["emoji"] = bson.decode_all( + fcache.read(), bson_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(), codec_options) + rocketchat_data["custom_emoji"]["file"] = bson.decode_all( + fcache.read(), bson_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(), codec_options) + rocketchat_data["custom_emoji"]["chunk"] = bson.decode_all( + fcache.read(), bson_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(), codec_options) + rocketchat_data["upload"]["upload"] = bson.decode_all(fcache.read(), bson_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(), codec_options) + rocketchat_data["upload"]["file"] = bson.decode_all(fcache.read(), bson_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(), codec_options) + rocketchat_data["upload"]["chunk"] = bson.decode_all(fcache.read(), bson_codec_options) return rocketchat_data