From e883ab057fdf9aed14fa4c63bff6952690c5c81b Mon Sep 17 00:00:00 2001
From: Mateusz Mandera <mateusz.mandera@zulip.com>
Date: Sat, 19 Jun 2021 11:48:22 +0200
Subject: [PATCH] upload: Cache the boto client to improve performance.

Fixes #18915

This was very slow, causing performance issues. After investigating,
generate_presigned_url is the cheap part of this, but the
session.client() call is expensive - so that's what we should cache.

Before the change:
```
In [4]: t = time.time()
   ...: for i in range(250):
   ...:     x = u.get_public_upload_url("foo")
   ...: print(time.time()-t)
6.408717393875122
```

After:
```
In [4]: t = time.time()
   ...: for i in range(250):
   ...:     x = u.get_public_upload_url("foo")
   ...: print(time.time()-t)
0.48990607261657715
```

This is not good enough to avoid doing something ugly like replacing
generate_presigned_url with some manual URL manipulation, but it's a
helpful structure that we may find useful with further refactoring.
---
 zerver/lib/upload.py | 24 +++++++++++++++++-------
 1 file changed, 17 insertions(+), 7 deletions(-)

diff --git a/zerver/lib/upload.py b/zerver/lib/upload.py
index 0d1d9ff8ce..2b0dbebe78 100644
--- a/zerver/lib/upload.py
+++ b/zerver/lib/upload.py
@@ -386,6 +386,8 @@ class S3UploadBackend(ZulipUploadBackend):
         self.avatar_bucket = get_bucket(settings.S3_AVATAR_BUCKET, self.session)
         self.uploads_bucket = get_bucket(settings.S3_AUTH_UPLOADS_BUCKET, self.session)
 
+        self._boto_client = None
+
     def get_public_upload_url(
         self,
         key: str,
@@ -399,13 +401,7 @@ class S3UploadBackend(ZulipUploadBackend):
         # different URL format. Configuring no signature and providing
         # no access key makes `generate_presigned_url` just return the
         # normal public URL for a key.
-        config = Config(signature_version=botocore.UNSIGNED)
-        return self.session.client(
-            "s3",
-            region_name=settings.S3_REGION,
-            endpoint_url=settings.S3_ENDPOINT_URL,
-            config=config,
-        ).generate_presigned_url(
+        return self.get_boto_client().generate_presigned_url(
             ClientMethod="get_object",
             Params={
                 "Bucket": self.avatar_bucket.name,
@@ -414,6 +410,20 @@ class S3UploadBackend(ZulipUploadBackend):
             ExpiresIn=0,
         )
 
+    def get_boto_client(self) -> botocore.client.BaseClient:
+        """
+        Creating the client takes a long time so we need to cache it.
+        """
+        if self._boto_client is None:
+            config = Config(signature_version=botocore.UNSIGNED)
+            self._boto_client = self.session.client(
+                "s3",
+                region_name=settings.S3_REGION,
+                endpoint_url=settings.S3_ENDPOINT_URL,
+                config=config,
+            )
+        return self._boto_client
+
     def delete_file_from_s3(self, path_id: str, bucket: ServiceResource) -> bool:
         key = bucket.Object(path_id)