Merge pull request #88 from uw-it-aca/hotfix/download-memory

Hotfix/download memory

sqlshare_rest/backend/base.py

@@ -21,6 +21,12 @@ def get_query_plan(self, sql, user):
         """
         return ""
 
+    def run_named_cursor_query(self, *args, **kwargs):
+        return self.run_query(*args, **kwargs)
+
+    def finish_named_cursor(self, *args, **kwargs):
+        pass
+
     def run_query(self, sql, user, params=None, return_cursor=False,
                   query=None):
         self._not_implemented("run_query")

sqlshare_rest/backend/pg.py

@@ -4,6 +4,8 @@
 from django.db import connection
 from django.conf import settings
 from logging import getLogger
+import time
+import random
 import re
 import tempfile
 
@@ -55,6 +57,32 @@ def remove_schema(self, schema):
         cursor.execute(sql)
         cursor.close()
 
+    def run_named_cursor_query(self, sql, user, params=None,
+                               return_cursor=False, query=None):
+        connection = self.get_connection_for_user(user)
+        connection.set_session(autocommit=False)
+
+        if query:
+            query.backend_terminate_data = "%s" % connection.get_backend_pid()
+            query.save()
+
+        cursor = connection.cursor('cursor-%s-%s' % (time.time(),
+                                                     random.random()))
+        cursor.execute(sql.encode('utf-8'), params)
+
+        if return_cursor:
+            return cursor
+        data = cursor.fetchall()
+        self.finish_named_cursor(user, cursor)
+        return data
+
+    def finish_named_cursor(self, user, cursor):
+        cursor.close()
+        connection = self.get_connection_for_user(user)
+        connection.commit()
+        connection.set_session(autocommit=True)
+        self.close_user_connection(user)
+
     def run_query(self, sql, user, params=None, return_cursor=False,
                   query=None):
         connection = self.get_connection_for_user(user)

sqlshare_rest/views/sql.py

@@ -37,10 +37,11 @@ def response_for_query(sql, user, download_name):
         backend = get_backend()
         if sql == "":
             raise Exception("Missing sql statement")
-        cursor = backend.run_query(sql, user, return_cursor=True)
+        cursor = backend.run_named_cursor_query(sql, user, return_cursor=True)
         disposition = 'attachment; filename="%s"' % download_name
         response = StreamingHttpResponse(stream_query(cursor, user),
                                          content_type='text/csv')
+
         response['Content-Disposition'] = disposition
         return response
     except Exception as ex:
@@ -52,6 +53,9 @@ def response_for_query(sql, user, download_name):
 
 
 def stream_query(cursor, user):
+    rows = cursor.fetchmany(100)
+
+    # Need to fetch columns after fetching a bit of data :(
     columns = frontend_description_from_cursor(cursor)
 
     names = ",".join(list(map(lambda x: csv_encode(x["name"]), columns)))
@@ -60,13 +64,14 @@ def stream_query(cursor, user):
     yield names
     yield "\n"
 
-    row = cursor.fetchone()
-    while row:
-        yield ",".join(list(map(lambda x: csv_encode("%s" % x), row)))
-        yield "\n"
-        row = cursor.fetchone()
+    while rows:
+        for row in rows:
+            yield ",".join(list(map(lambda x: csv_encode("%s" % x), row)))
+            yield "\n"
+        rows = cursor.fetchmany(100)
 
-    get_backend().close_user_connection(user)
+    backend = get_backend()
+    backend.finish_named_cursor(user, cursor)
 
 
 def csv_encode(value):