From 04e66f94fca69a73ce18893fbc1760a6f4c2c68a Mon Sep 17 00:00:00 2001
From: Ruslan Kuprieiev <kupruser@gmail.com>
Date: Sun, 20 Aug 2017 19:04:34 +0300
Subject: [PATCH] aws: use multipart to push changes, v2

Fixes #100

Signed-off-by: Ruslan Kuprieiev <kupruser@gmail.com>
---
 dvc/data_cloud.py | 35 ++++++++++++++++++++++++++++++++---
 1 file changed, 32 insertions(+), 3 deletions(-)

diff --git a/dvc/data_cloud.py b/dvc/data_cloud.py
index a194e5871..ea936e234 100644
--- a/dvc/data_cloud.py
+++ b/dvc/data_cloud.py
@@ -7,6 +7,7 @@ import tempfile
 import requests
 import configparser
 import filecmp
+import math
 
 from boto.s3.connection import S3Connection
 from google.cloud import storage as gc
@@ -418,9 +419,12 @@ class DataCloudAWS(DataCloudBase):
         return bucket
 
     def _cmp_checksum(self, key, fname):
-        md5_cloud = key.etag[1:-1]
+        md5_cloud = key.metadata.get('dvc-md5', None)
         md5_local = file_md5(fname)[0]
 
+        if md5_cloud == None:
+            md5_cloud = key.etag[1:-1]
+
         if md5_cloud == md5_local:
             return True
 
@@ -455,6 +459,32 @@ class DataCloudAWS(DataCloudBase):
 
         return data_item
 
+    def _push_multipart(self, key, fname):
+        # AWS doesn't provide easilly accessible md5 for multipart
+        # objects, so we have to store our own md5 sum to use later.
+        metadata = {'dvc-md5' : str(file_md5(fname)[0])}
+
+        mp = key.bucket.initiate_multipart_upload(key.name, metadata=metadata)
+
+        source_size = os.stat(fname).st_size
+        chunk_size = 5000*1024*1024
+        chunk_count = int(math.ceil(source_size / float(chunk_size)))
+
+        with open(fname, 'rb') as fp:
+            for i in range(chunk_count):
+                offset = i * chunk_size
+                left = source_size - offset
+                size = min([chunk_size, left])
+                part_num = i + 1
+
+                fp.seek(offset)
+                mp.upload_part_from_file(fp=fp, part_num=part_num, size=size, cb=create_cb(fname))
+
+        if len(mp.get_all_parts()) != chunk_count:
+            raise Exception("Couldn't upload all file parts")
+
+        mp.complete_upload()
+
     def push(self, data_item):
         """ push, aws version """
 
@@ -473,8 +503,7 @@ class DataCloudAWS(DataCloudBase):
         key = bucket.new_key(aws_key)
 
         try:
-            key.set_contents_from_filename(data_item.resolved_cache.relative,
-                                       cb=create_cb(data_item.resolved_cache.relative))
+            self._push_multipart(key, data_item.resolved_cache.relative)
         except Exception as exc:
             Logger.error('Failed to upload "{}": {}'.format(data_item.resolved_cache.relative, exc))
             return None