Fix large downloads (#105) (#111)

Large responses were being read into memory. For most calls that's fine, but download could cause the python process to go out of memory due to holding large workbooks or datasources all in memory before writing to disk.

Requests has a feature called `iter_content` which when used in combination with `stream=True` on a request will download only the headers, allow us to determine the filename, and then read through the response body in chunks. I picked a size of 1024 bytes, since that's what most of the internet appears to use and I noticed little perf difference between a 1024 byte chunk size and a 1MB chunk size.

This is all enabled by exposing the `parameters` argument to `requests.get` by pluming it through our wrapper functions. All tests pass, and manual testing showed the memory problem went away.

Author: t8y8

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)