Merge pull request #155 from peterrus/py3

Merge jimas14 improvements

tc_aws/aws/bucket.py

@@ -11,6 +11,7 @@
 
 
 class Bucket(object):
+    _client = None
     _instances = {}
 
     @staticmethod
@@ -35,8 +36,6 @@ def __init__(self, bucket, region, endpoint, max_retry=None):
         """
         self._bucket = bucket
 
-        session = aiobotocore.get_session()
-
         config = None
         if max_retry is not None:
             config = Config(
@@ -45,12 +44,27 @@ def __init__(self, bucket, region, endpoint, max_retry=None):
                 )
             )
 
-        self._client = session.create_client(
-            's3',
-            region_name=region,
-            endpoint_url=endpoint,
-            config=config
-        )
+        if self._client is None:
+            self._client = aiobotocore.get_session().create_client(
+                's3',
+                region_name=region,
+                endpoint_url=endpoint,
+                config=config
+            )
+
+    async def exists(self, path):
+        """
+        Checks if an object exists at a given path
+        :param string path: Path or 'key' to retrieve AWS object
+        """
+        try:
+            await self._client.head_object(
+                Bucket=self._bucket,
+                Key=self._clean_key(path),
+            )
+        except Exception:
+            return False
+        return True
 
     async def get(self, path):
         """

tc_aws/loaders/s3_loader.py

@@ -63,7 +63,8 @@ async def load(context, url):
         return result
 
     result.successful = True
-    result.buffer = await file_key['Body'].read()
+    async with file_key['Body'] as stream:
+        result.buffer = await stream.read()
 
     result.metadata.update(
         size=file_key['ContentLength'],

tc_aws/result_storages/s3_storage.py

@@ -60,7 +60,8 @@ async def get(self, path = None):
             return None
 
         result = ResultStorageResult()
-        result.buffer = await key['Body'].read()
+        async with key['Body'] as stream:
+            result.buffer = await stream.read()
         result.successful = True
 
         result.metadata = {

tc_aws/storages/s3_storage.py

@@ -109,7 +109,8 @@ async def get_crypto(self, path):
             logger.warn("[STORAGE] s3 key not found at %s" % crypto_path)
             return None
 
-        file_key = await file_key['Body'].read()
+        async with file_key['Body'] as stream:
+            file_key = await stream.read()
 
         return file_key.decode('utf-8')
 
@@ -129,7 +130,8 @@ async def get_detector_data(self, path):
         if not file_key or self.is_expired(file_key) or 'Body' not in file_key:
             return None
 
-        return loads(await file_key['Body'].read())
+        async with file_key['Body'] as stream:
+            return loads(await stream.read())
 
     async def get(self, path):
         """
@@ -142,23 +144,16 @@ async def get(self, path):
         except BotoCoreError:
             return None
 
-        return await file['Body'].read()
+        async with file['Body'] as stream:
+            return await stream.read()
 
     async def exists(self, path):
         """
         Tells if data exists at given path
         :param string path: Path to check
         """
         file_abspath = self._normalize_path(path)
-
-        try:
-            await self.storage.get(file_abspath)
-        except ClientError as err:
-            if err.response['Error']['Code'] == 'NoSuchKey':
-                return False
-            raise
-
-        return True
+        return await self.storage.exists(file_abspath)
 
     async def remove(self, path):
         """