From ffaff5a27e3d84157d526da7c8b939eae1e866ac Mon Sep 17 00:00:00 2001
From: Joseph Schorr <josephschorr@users.noreply.github.com>
Date: Tue, 17 Oct 2017 13:09:37 -0400
Subject: [PATCH] Add caching of get_auth call in Swift

Should make calls significantly faster for our Swift customers

Fixes https://jira.prod.coreos.systems/browse/QS-39
---
 storage/swift.py           | 22 ++++++++++++++++------
 storage/test/test_swift.py | 25 ++++++++++++++++++++++---
 2 files changed, 38 insertions(+), 9 deletions(-)

diff --git a/storage/swift.py b/storage/swift.py
index ea794c72e..4af54c6a7 100644
--- a/storage/swift.py
+++ b/storage/swift.py
@@ -8,6 +8,8 @@ import hmac
 import string
 import logging
 import json
+
+from cachetools import lru_cache
 from _pyio import BufferedReader
 from uuid import uuid4
 
@@ -147,6 +149,14 @@ class SwiftStorage(BaseStorage):
       logger.exception('Could not head object at path %s: %s', path, ex)
       return None
 
+  @lru_cache(maxsize=1)
+  def _get_root_storage_url(self):
+    """ Returns the root storage URL for this Swift storage. Note that since this requires a call
+        to Swift, we cache the result of this function call.
+    """
+    storage_url, _ = self._get_connection().get_auth()
+    return storage_url
+
   def get_direct_download_url(self, object_path, request_ip=None, expires_in=60, requires_cors=False, head=False):
     if requires_cors:
       return None
@@ -155,17 +165,17 @@ class SwiftStorage(BaseStorage):
     if not self._temp_url_key:
       return None
 
-    # Retrieve the auth details for the connection.
+    # Retrieve the root storage URL for the connection.
     try:
-      object_url_value, _ = self._get_connection().get_auth()
+      root_storage_url = self._get_root_storage_url()
     except ClientException:
       logger.exception('Got client exception when trying to load Swift auth')
       return None
 
-    object_url = urlparse(object_url_value)
-    scheme = object_url.scheme
-    path = object_url.path.rstrip('/')
-    hostname = object_url.netloc
+    parsed_storage_url = urlparse(root_storage_url)
+    scheme = parsed_storage_url.scheme
+    path = parsed_storage_url.path.rstrip('/')
+    hostname = parsed_storage_url.netloc
 
     object_path = self._normalize_path(object_path)
 
diff --git a/storage/test/test_swift.py b/storage/test/test_swift.py
index 50d6a360c..0667a0934 100644
--- a/storage/test/test_swift.py
+++ b/storage/test/test_swift.py
@@ -8,6 +8,7 @@ from mock import MagicMock
 
 from storage import StorageContext
 from storage.swift import SwiftStorage
+from swiftclient.client import ClientException
 
 base_args = {
   'context': StorageContext('nyc', None, None, None, None),
@@ -27,18 +28,26 @@ class MockSwiftStorage(SwiftStorage):
     return self._connection
 
 class FakeSwiftStorage(SwiftStorage):
-  def __init__(self, fail_checksum=False, connection=None, *args, **kwargs):
+  def __init__(self, fail_checksum=False,connection=None, *args, **kwargs):
     super(FakeSwiftStorage, self).__init__(*args, **kwargs)
-    self._connection = connection or FakeSwift(fail_checksum=fail_checksum)
+    self._connection = connection or FakeSwift(fail_checksum=fail_checksum,
+                                               temp_url_key=kwargs.get('temp_url_key'))
 
   def _get_connection(self):
     return self._connection
 
 
 class FakeSwift(object):
-  def __init__(self, fail_checksum=False):
+  def __init__(self, fail_checksum=False, temp_url_key=None):
     self.containers = defaultdict(dict)
     self.fail_checksum = fail_checksum
+    self.temp_url_key = temp_url_key
+
+  def get_auth(self):
+    if self.temp_url_key == 'exception':
+      raise ClientException('I failed!')
+      
+    return 'http://fake/swift', None
 
   def head_object(self, container, path):
     return self.containers[container].get(path)
@@ -258,3 +267,13 @@ def test_empty_chunks_queued_for_deletion():
 
   found2 = chunk_cleanup_queue.get()
   assert found2 is None
+
+@pytest.mark.parametrize('temp_url_key, expects_url', [
+  (None, False),
+  ('foobarbaz', True),
+  ('exception', False),
+])
+def test_get_direct_download_url(temp_url_key, expects_url):
+  swift = FakeSwiftStorage(temp_url_key=temp_url_key, **base_args)
+  swift.put_content('somepath', 'hello world!')
+  assert (swift.get_direct_download_url('somepath') is not None) == expects_url