refactor: Access a bucket by calling zerver.lib.uploads.get_bucket.

zerver/lib/export.py

@@ -15,7 +15,6 @@ import subprocess
 import tempfile
 from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union
 
-import boto3
 import orjson
 from boto3.resources.base import ServiceResource
 from django.apps import apps
@@ -29,6 +28,7 @@ from analytics.models import RealmCount, StreamCount, UserCount
 from scripts.lib.zulip_tools import overwrite_symlink
 from zerver.lib.avatar_hash import user_avatar_path_from_ids
 from zerver.lib.pysa import mark_sanitized
+from zerver.lib.upload import get_bucket
 from zerver.models import (
     AlertWord,
     Attachment,
@@ -1272,9 +1272,7 @@ def _save_s3_object_to_file(key: ServiceResource, output_dir: str, processing_av
 def export_files_from_s3(realm: Realm, bucket_name: str, output_dir: Path,
                          processing_avatars: bool=False, processing_emoji: bool=False,
                          processing_realm_icon_and_logo: bool=False) -> None:
-    session = boto3.Session(settings.S3_KEY, settings.S3_SECRET_KEY)
-    s3 = session.resource('s3')
-    bucket = s3.Bucket(bucket_name)
+    bucket = get_bucket(bucket_name)
     records = []
 
     logging.info("Downloading uploaded files from %s", bucket_name)

zerver/lib/import_realm.py

@@ -6,7 +6,6 @@ import secrets
 import shutil
 from typing import Any, Dict, Iterable, List, Optional, Tuple
 
-import boto3
 import orjson
 from bs4 import BeautifulSoup
 from django.conf import settings
@@ -32,7 +31,7 @@ from zerver.lib.message import get_last_message_id
 from zerver.lib.server_initialization import create_internal_realm, server_initialized
 from zerver.lib.streams import render_stream_description
 from zerver.lib.timestamp import datetime_to_timestamp
-from zerver.lib.upload import BadImageError, guess_type, sanitize_name
+from zerver.lib.upload import BadImageError, get_bucket, guess_type, sanitize_name
 from zerver.lib.utils import generate_api_key, process_list_in_batches
 from zerver.models import (
     AlertWord,
@@ -656,8 +655,7 @@ def import_uploads(realm: Realm, import_dir: Path, processes: int, processing_av
             bucket_name = settings.S3_AVATAR_BUCKET
         else:
             bucket_name = settings.S3_AUTH_UPLOADS_BUCKET
-        session = boto3.Session(settings.S3_KEY, settings.S3_SECRET_KEY)
-        bucket = session.resource('s3').Bucket(bucket_name)
+        bucket = get_bucket(bucket_name)
 
     count = 0
     for record in records:

zerver/lib/upload.py

@@ -267,9 +267,11 @@ class ZulipUploadBackend:
 
 ### S3
 
-def get_bucket(session: Session, bucket_name: str) -> ServiceResource:
+def get_bucket(bucket_name: str, session: Optional[Session]=None) -> ServiceResource:
     # See https://github.com/python/typeshed/issues/2706
     # for why this return type is a `ServiceResource`.
+    if session is None:
+        session = boto3.Session(settings.S3_KEY, settings.S3_SECRET_KEY)
     bucket = session.resource('s3').Bucket(bucket_name)
     return bucket
 
@@ -336,12 +338,12 @@ class S3UploadBackend(ZulipUploadBackend):
     def __init__(self) -> None:
         self.session = boto3.Session(settings.S3_KEY, settings.S3_SECRET_KEY)
 
-        self.avatar_bucket = get_bucket(self.session, settings.S3_AVATAR_BUCKET)
+        self.avatar_bucket = get_bucket(settings.S3_AVATAR_BUCKET, self.session)
         network_location = urllib.parse.urlparse(
             self.avatar_bucket.meta.client.meta.endpoint_url).netloc
         self.avatar_bucket_url = f"https://{self.avatar_bucket.name}.{network_location}"
 
-        self.uploads_bucket = get_bucket(self.session, settings.S3_AUTH_UPLOADS_BUCKET)
+        self.uploads_bucket = get_bucket(settings.S3_AUTH_UPLOADS_BUCKET, self.session)
 
     def delete_file_from_s3(self, path_id: str, bucket: ServiceResource) -> bool:
         key = bucket.Object(path_id)