Fix large downloads (#105)

tableauserverclient/server/endpoint/datasources_endpoint.py

@@ -6,6 +6,7 @@
 import logging
 import copy
 import cgi
+from contextlib import closing
 
 # The maximum size of a file that can be published in a single request is 64MB
 FILESIZE_LIMIT = 1024 * 1024 * 64   # 64MB
@@ -64,16 +65,18 @@ def download(self, datasource_id, filepath=None):
             error = "Datasource ID undefined."
             raise ValueError(error)
         url = "{0}/{1}/content".format(self.baseurl, datasource_id)
-        server_response = self.get_request(url)
-        _, params = cgi.parse_header(server_response.headers['Content-Disposition'])
-        filename = os.path.basename(params['filename'])
-        if filepath is None:
-            filepath = filename
-        elif os.path.isdir(filepath):
-            filepath = os.path.join(filepath, filename)
-
-        with open(filepath, 'wb') as f:
-            f.write(server_response.content)
+        with closing(self.get_request(url, parameters={'stream': True})) as server_response:
+            _, params = cgi.parse_header(server_response.headers['Content-Disposition'])
+            filename = os.path.basename(params['filename'])
+            if filepath is None:
+                filepath = filename
+            elif os.path.isdir(filepath):
+                filepath = os.path.join(filepath, filename)
+
+            with open(filepath, 'wb') as f:
+                for chunk in server_response.iter_content(1024):  # 1KB
+                    f.write(chunk)
+
         logger.info('Downloaded datasource to {0} (ID: {1})'.format(filepath, datasource_id))
         return os.path.abspath(filepath)
 

tableauserverclient/server/endpoint/endpoint.py

@@ -21,10 +21,11 @@ def _make_common_headers(auth_token, content_type):
 
         return headers
 
-    def _make_request(self, method, url, content=None, request_object=None, auth_token=None, content_type=None):
+    def _make_request(self, method, url, content=None, request_object=None,
+                      auth_token=None, content_type=None, parameters=None):
         if request_object is not None:
             url = request_object.apply_query_params(url)
-        parameters = {}
+        parameters = parameters or {}
         parameters.update(self.parent_srv.http_options)
         parameters['headers'] = Endpoint._make_common_headers(auth_token, content_type)
 
@@ -49,9 +50,9 @@ def _check_status(server_response):
     def get_unauthenticated_request(self, url, request_object=None):
         return self._make_request(self.parent_srv.session.get, url, request_object=request_object)
 
-    def get_request(self, url, request_object=None):
+    def get_request(self, url, request_object=None, parameters=None):
         return self._make_request(self.parent_srv.session.get, url, auth_token=self.parent_srv.auth_token,
-                                  request_object=request_object)
+                                  request_object=request_object, parameters=parameters)
 
     def delete_request(self, url):
         # We don't return anything for a delete

tableauserverclient/server/endpoint/workbooks_endpoint.py

@@ -7,6 +7,7 @@
 import logging
 import copy
 import cgi
+from contextlib import closing
 
 # The maximum size of a file that can be published in a single request is 64MB
 FILESIZE_LIMIT = 1024 * 1024 * 64   # 64MB
@@ -92,16 +93,18 @@ def download(self, workbook_id, filepath=None):
             error = "Workbook ID undefined."
             raise ValueError(error)
         url = "{0}/{1}/content".format(self.baseurl, workbook_id)
-        server_response = self.get_request(url)
-        _, params = cgi.parse_header(server_response.headers['Content-Disposition'])
-        filename = os.path.basename(params['filename'])
-        if filepath is None:
-            filepath = filename
-        elif os.path.isdir(filepath):
-            filepath = os.path.join(filepath, filename)
-
-        with open(filepath, 'wb') as f:
-            f.write(server_response.content)
+
+        with closing(self.get_request(url, parameters={"stream": True})) as server_response:
+            _, params = cgi.parse_header(server_response.headers['Content-Disposition'])
+            filename = os.path.basename(params['filename'])
+            if filepath is None:
+                filepath = filename
+            elif os.path.isdir(filepath):
+                filepath = os.path.join(filepath, filename)
+
+            with open(filepath, 'wb') as f:
+                for chunk in server_response.iter_content(1024):  # 1KB
+                    f.write(chunk)
         logger.info('Downloaded workbook to {0} (ID: {1})'.format(filepath, workbook_id))
         return os.path.abspath(filepath)