Optimize storage performance (#225)

* rework doanload all archives

* fix typo

* add exists method for storage

* avoid unnecessary upload

Author: iberryful

src/api/handlers/job_api.py

@@ -558,7 +558,8 @@ def post(self):
 
             g.release_db()
 
-            storage.upload_output(stream, key)
+            if not storage.exists(key):
+                storage.upload_output(stream, key)
 
             for c in clusters:
                 stream.seek(0)

src/api/handlers/project.py

@@ -298,9 +298,9 @@ def post(self, project_id, build_id):
             abort(400, 'Project is not of type "upload"')
 
         key = '%s.zip' % build_id
-
-        stream = request.files['project.zip'].stream
-        storage.upload_project(stream, key)
+        if not storage.exists(key):
+            stream = request.files['project.zip'].stream
+            storage.upload_project(stream, key)
 
         return OK('successfully uploaded data')
 
@@ -403,4 +403,4 @@ def post(self, project_id):
 
             g.db.commit()
 
-            return OK('successfully started build', data=data)
+            return OK('successfully started build', data=data)

src/api/handlers/projects/jobs.py

@@ -9,7 +9,7 @@
 
 import requests
 
-from flask import g, abort, Response, send_file, request, after_this_request
+from flask import g, abort, Response, send_file, request, redirect
 from flask_restplus import Resource, fields
 
 from pyinfraboxutils import get_logger, get_env
@@ -394,9 +394,8 @@ def get(self, project_id, job_id):
         job_cluster = result['cluster_name']
         key = '%s/%s' % (job_id, filename)
 
-        if os.environ['INFRABOX_CLUSTER_NAME'] == job_cluster:
-            f = storage.download_archive(key)
-        else:
+        f = storage.download_archive(key)
+        if not f and os.environ['INFRABOX_CLUSTER_NAME'] != job_cluster:
             c = g.db.execute_one_dict('''
                 SELECT *
                 FROM cluster
@@ -450,53 +449,7 @@ def get(self, project_id, job_id):
         '''
         Returns all archives
         '''
-        result = g.db.execute_one_dict('''
-                    SELECT archive
-                    FROM job
-                    WHERE   id = %s
-                        AND project_id = %s
-                ''', [job_id, project_id])
-
-        if not result or not result['archive']:
-            abort(404)
-
-        base_path = os.path.join('/tmp', str(uuid.uuid4()))
-        archive_dir = os.path.join(base_path, 'archive')
-        os.mkdir(base_path)
-        os.mkdir(archive_dir)
-
-        @after_this_request
-        def _remove_file(response):
-            if os.path.exists(base_path):
-                shutil.rmtree(base_path)
-            return response
-
-        for item in result['archive']:
-            filename = item['filename']
-            url = "%s/api/v1/projects/%s/jobs/%s/archive/download?filename=%s" % \
-                  (get_env('INFRABOX_ROOT_URL'), project_id, job_id, filename)
-            try:
-                token = encode_user_token(g.token['user']['id'])
-            except Exception:
-                #public project has no token here.
-                token = ""
-            headers = {'Authorization': 'bearer ' + token}
-
-            r = requests.get(url, headers=headers, timeout=120, verify=False)
-            if r.status_code != 200:
-                continue
-
-            with open(os.path.join(archive_dir, os.path.basename(filename)), 'w') as f:
-                f.write(r.content)
-
-        if not os.listdir(archive_dir):
-            abort(404)
-
-        tar_file = os.path.join(base_path, 'archive_%s' % job_id +'.tar.gz')
-        with tarfile.open(tar_file, mode='w:gz') as archive:
-            archive.add(archive_dir, arcname='archive')
-
-        return send_file(tar_file, as_attachment=True, attachment_filename=os.path.basename(tar_file))
+        return redirect("/api/v1/projects/%s/jobs/%s/archive/download?filename=all_archives.tar.gz" %(project_id, job_id))
 
 @ns.route('/<job_id>/console')
 @api.response(403, 'Not Authorized')

src/job/job.py

@@ -12,6 +12,7 @@
 import traceback
 import urllib3
 import yaml
+import tarfile
 
 from pyinfrabox.infrabox import validate_json
 from pyinfrabox.docker_compose import create_from
@@ -318,13 +319,15 @@ def convert_coverage_result(self, f):
 
     def upload_archive(self):
         c = self.console
+        archive_exists = False
+        testresult_exists = False
 
         if os.path.exists(self.infrabox_archive_dir):
             files = self.get_files_in_dir(self.infrabox_archive_dir)
 
             if files:
                 c.collect("Uploading /infrabox/upload/archive", show=True)
-
+                archive_exists = True
                 for f in files:
                     c.collect("%s" % f, show=True)
                     self.post_file_to_api_server("/archive", f, filename=f.replace(self.infrabox_upload_dir, ''))
@@ -333,10 +336,19 @@ def upload_archive(self):
             files = self.get_files_in_dir(self.infrabox_testresult_dir)
 
             if files:
-
+                testresult_exists = True
                 for f in files:
                     c.collect("%s" % f, show=True)
 
+        tar_file = os.path.join(self.infrabox_upload_dir, 'all_archives' + '.tar.gz')
+        with tarfile.open(tar_file, mode='w:gz') as archive:
+            if archive_exists:
+                archive.add(self.infrabox_archive_dir, arcname='archive')
+            if testresult_exists:
+                archive.add(self.infrabox_testresult_dir, arcname='testresult')
+
+        self.post_file_to_api_server("/archive", tar_file)
+
 
     def upload_coverage_results(self):
         if not os.path.exists(self.infrabox_coverage_dir):

src/pyinfraboxutils/storage.py

@@ -3,6 +3,7 @@
 import uuid
 
 import boto3
+from botocore.errorfactory import ClientError
 from google.cloud import storage as gcs
 from flask import after_this_request
 from flask import _app_ctx_stack as stack
@@ -69,6 +70,9 @@ def download_cache(self, key):
     def delete_cache(self, key):
         return self._delete('cache/%s' % key)
 
+    def exists(self, key):
+        return
+
 class S3(Storage):
     def __init__(self):
         super(Storage, self).__init__()
@@ -85,6 +89,14 @@ def __init__(self):
         self.bucket = get_env('INFRABOX_STORAGE_S3_BUCKET')
         self.create_buckets()
 
+    def exists(self, key):
+        client = self._get_client()
+        try:
+            client.head_object(Bucket=self.bucket, Key=key)
+        except ClientError:
+            return False
+        return True
+
     def _upload(self, stream, key):
         client = self._get_client()
         client.put_object(Body=stream,
@@ -146,6 +158,12 @@ def _delete(self, key):
         except:
             pass
 
+    def exists(self, key):
+        client = gcs.Client()
+        bucket = client.get_bucket(self.bucket)
+        blob = bucket.blob(key)
+        return blob.exists()
+
     def _upload(self, stream, key):
         client = gcs.Client()
         bucket = client.get_bucket(self.bucket)
@@ -173,6 +191,10 @@ def __init__(self):
         super(Storage, self).__init__()
         self.container = 'infrabox'
 
+    def exists(self, key):
+        client = self._get_client()
+        return client.exists(container_name=self.container, blob_name=key)
+
     def _upload(self, stream, key):
         client = self._get_client()
         if not client.exists(container_name=self.container):
@@ -217,6 +239,14 @@ def __init__(self):
         self.project_name = get_env('INFRABOX_STORAGE_SWIFT_PROJECT_NAME')
         self.project_domain_name = get_env('INFRABOX_STORAGE_SWIFT_PROJECT_DOMAIN_NAME')
 
+    def exists(self, key):
+        client = self._get_client()
+        try:
+            client.head_object(self.container, key)
+        except ClientException:
+            return False
+        return True
+
     def _upload(self, stream, key):
         client = self._get_client()
         try: