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.
This commit is contained in:
Tim Abbott 2024-10-17 11:09:44 -07:00
parent b93b05eddd
commit 79b6f43d0e
1 changed files with 21 additions and 15 deletions

View File

@ -36,6 +36,8 @@ from zerver.lib.upload import sanitize_name
from zerver.lib.utils import process_list_in_batches from zerver.lib.utils import process_list_in_batches
from zerver.models import Reaction, RealmEmoji, Recipient, UserProfile from zerver.models import Reaction, RealmEmoji, Recipient, UserProfile
bson_codec_options = bson.DEFAULT_CODEC_OPTIONS.with_options(tz_aware=True)
def make_realm( def make_realm(
realm_id: int, realm_subdomain: str, domain_name: str, rc_instance: dict[str, Any] 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]: 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: dict[str, Any] = {}
rocketchat_data["instance"] = [] rocketchat_data["instance"] = []
rocketchat_data["user"] = [] rocketchat_data["user"] = []
@ -989,60 +989,66 @@ def rocketchat_data_to_dict(rocketchat_data_dir: str) -> dict[str, Any]:
# Get instance # Get instance
with open(os.path.join(rocketchat_data_dir, "instances.bson"), "rb") as fcache: 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 # Get user
with open(os.path.join(rocketchat_data_dir, "users.bson"), "rb") as fcache: 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 # Get avatar
with open(os.path.join(rocketchat_data_dir, "rocketchat_avatars.bson"), "rb") as fcache: 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"]: if rocketchat_data["avatar"]["avatar"]:
with open( with open(
os.path.join(rocketchat_data_dir, "rocketchat_avatars.files.bson"), "rb" os.path.join(rocketchat_data_dir, "rocketchat_avatars.files.bson"), "rb"
) as fcache: ) 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( with open(
os.path.join(rocketchat_data_dir, "rocketchat_avatars.chunks.bson"), "rb" os.path.join(rocketchat_data_dir, "rocketchat_avatars.chunks.bson"), "rb"
) as fcache: ) 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 # Get room
with open(os.path.join(rocketchat_data_dir, "rocketchat_room.bson"), "rb") as fcache: 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 # Get messages
with open(os.path.join(rocketchat_data_dir, "rocketchat_message.bson"), "rb") as fcache: 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 # Get custom emoji
with open(os.path.join(rocketchat_data_dir, "rocketchat_custom_emoji.bson"), "rb") as fcache: 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"]: if rocketchat_data["custom_emoji"]["emoji"]:
with open(os.path.join(rocketchat_data_dir, "custom_emoji.files.bson"), "rb") as fcache: 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: 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 # Get uploads
with open(os.path.join(rocketchat_data_dir, "rocketchat_uploads.bson"), "rb") as fcache: 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"]: if rocketchat_data["upload"]["upload"]:
with open( with open(
os.path.join(rocketchat_data_dir, "rocketchat_uploads.files.bson"), "rb" os.path.join(rocketchat_data_dir, "rocketchat_uploads.files.bson"), "rb"
) as fcache: ) 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( with open(
os.path.join(rocketchat_data_dir, "rocketchat_uploads.chunks.bson"), "rb" os.path.join(rocketchat_data_dir, "rocketchat_uploads.chunks.bson"), "rb"
) as fcache: ) 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 return rocketchat_data