Browse Source

[aws] MVP of S3 file download via URL

Note: get_key() fetches all the file currently
Romain Rigaux 4 years ago
parent
commit
b51ce37449

+ 13 - 7
desktop/libs/aws/src/aws/s3/s3connection.py

@@ -26,7 +26,6 @@ from boto.s3.bucketlistresultset import BucketListResultSet
 from boto.s3.prefix import Prefix
 
 from desktop.lib.raz.clients import S3RazClient
-from aws.s3.s3fs import S3FileSystemException
 
 
 LOG = logging.getLogger(__name__)
@@ -75,7 +74,6 @@ class RazUrlConnection(UrlConnection):
 class UrlKey(Key):
 
   def open_read(self, headers=None, query_args='', override_num_retries=None, response_headers=None):
-
     # Similar to Bucket.get_key()
     # data = self.resp.read(self.BufferSize)
     # For seek: headers={"Range": "bytes=%d-" % pos}
@@ -92,7 +90,7 @@ class UrlKey(Key):
     except BotoClientError as e:
       LOG.error(e)
       if tmp_url is None:
-        raise S3FileSystemException("Resource does not exist or permission missing : '%s'" % kwargs)
+        raise IOError("Resource does not exist or permission missing : '%s'" % kwargs)
 
     return tmp_url
 
@@ -108,8 +106,8 @@ class UrlBucket(Bucket):
 
 
   def get_key(self, key_name, headers=None, version_id=None, response_headers=None, validate=True):
-    # Note: in current FS API we get file even if we don't need the content, hence why it can be slow.
-    # To check if we should give a length in read() to mitigate.
+    # Note: in current FS API, this fetches the full file... hence why it can be very slow.
+    # To check if we should give a length in read() to mitigate and just get metadata: HEAD ?
     LOG.debug('key name: %s' % key_name)
     kwargs = {'bucket': self.name, 'key': key_name}
 
@@ -144,6 +142,14 @@ class UrlBucket(Bucket):
       k.handle_encryption_headers(response)
       k.handle_restore_headers(response)
       k.handle_addl_headers(response.getheaders())
+
+      class MockResponse():
+        def __init__(self, resp):
+          self.resp = resp
+        def read(self, size):
+          return self.resp.content
+
+      k.resp = MockResponse(response)
     else:
       # Currently needed as 404 on directories via stats_key()
       k = self.key_class(self, key_name)
@@ -181,7 +187,7 @@ class UrlBucket(Bucket):
     except BotoClientError as e:
       LOG.error(e)
       if tmp_url is None:
-        raise S3FileSystemException("Resource does not exist or permission missing : '%s'" % kwargs)
+        raise IOError("Resource does not exist or permission missing : '%s'" % kwargs)
 
     return tmp_url
 
@@ -206,7 +212,7 @@ class BotoUrlConnection(UrlConnection):
     except BotoClientError as e:
       LOG.error(e)
       if tmp_url is None:
-        raise S3FileSystemException("Resource does not exist or permission missing : '%s'" % kwargs)
+        raise IOError("Resource does not exist or permission missing : '%s'" % kwargs)
 
     return tmp_url
 

+ 4 - 4
desktop/libs/aws/src/aws/s3/s3connection_test.py

@@ -33,10 +33,10 @@ else:
 class TestBotoUrlConnection():
 
   def test_get_buckets(self):
-    with patch('aws.s3.s3connection.BotoUrlConnection._generate_url') as _generate_url:
+    with patch('aws.s3.s3connection.BotoUrlConnection.get_url_request') as get_url_request:
       with patch('aws.s3.s3connection.requests.get') as requests_get:
 
-        _generate_url.return_value = 'https://gethue-test.s3.amazonaws.com/?AWSAccessKeyId=AKIA23E77ZX2HVY76YGL' + \
+        get_url_request.return_value = 'https://gethue-test.s3.amazonaws.com/?AWSAccessKeyId=AKIA23E77ZX2HVY76YGL' + \
             '&Signature=3lhK%2BwtQ9Q2u5VDIqb4MEpoY3X4%3D&Expires=1617207304'
         requests_get.return_value = Mock(
           content=b'<?xml version="1.0" encoding="UTF-8"?>\n<ListAllMyBucketsResult '
@@ -55,11 +55,11 @@ class TestBotoUrlConnection():
 class TestRazUrlConnection():
 
   def test_get_buckets(self):
-    with patch('aws.s3.s3connection.RazUrlConnection._generate_url') as _generate_url:
+    with patch('aws.s3.s3connection.RazUrlConnection.get_url_request') as get_url_request:
       with patch('aws.s3.s3connection.requests.get') as requests_get:
 
         # TODO: update with potentially slightly different URL/headers
-        _generate_url.return_value = 'https://gethue-test.s3.amazonaws.com/?AWSAccessKeyId=AKIA23E77ZX2HVY76YGL' + \
+        get_url_request.return_value = 'https://gethue-test.s3.amazonaws.com/?AWSAccessKeyId=AKIA23E77ZX2HVY76YGL' + \
             '&Signature=3lhK%2BwtQ9Q2u5VDIqb4MEpoY3X4%3D&Expires=1617207304'
 
         requests_get.return_value = Mock(

+ 15 - 3
desktop/libs/aws/src/aws/s3/s3file.py

@@ -22,9 +22,11 @@ from boto.s3.keyfile import KeyFile
 
 from aws.conf import get_key_expiry
 from aws.s3 import translate_s3_error
+from aws.s3.s3connection import UrlKey
 
 DEFAULT_READ_SIZE = 1024 * 1024  # 1MB
 
+
 def open(key, mode='r'):
   if mode == 'r':
     return _ReadableS3File(key)
@@ -34,12 +36,22 @@ def open(key, mode='r'):
 
 class _ReadableS3File(KeyFile):
   def __init__(self, key):
-      key_copy = key.bucket.get_key(key.name, validate=False)
-      KeyFile.__init__(self, key_copy)
+    key_copy = key.bucket.get_key(key.name, validate=False)
+    KeyFile.__init__(self, key_copy)
+
+    self.is_raz = type(key) == UrlKey
+    if self.is_raz:
+      if self.key.resp:
+        self.content = self.key.resp.resp.content
+      else:
+        self.content = ''
 
   def read_url(self):
     return self.getkey().generate_url(get_key_expiry())
 
   @translate_s3_error
   def read(self, length=DEFAULT_READ_SIZE):
-    return KeyFile.read(self, length)
+    if self.is_raz:
+      return self.content
+    else:
+      return KeyFile.read(self, length)