diff --git a/web/pgadmin/tools/datagrid/templates/datagrid/index.html b/web/pgadmin/tools/datagrid/templates/datagrid/index.html index f2f608588..67cf107f8 100644 --- a/web/pgadmin/tools/datagrid/templates/datagrid/index.html +++ b/web/pgadmin/tools/datagrid/templates/datagrid/index.html @@ -182,6 +182,7 @@
+ {% endblock %} diff --git a/web/pgadmin/tools/sqleditor/__init__.py b/web/pgadmin/tools/sqleditor/__init__.py index acaca18b4..cfc791338 100644 --- a/web/pgadmin/tools/sqleditor/__init__.py +++ b/web/pgadmin/tools/sqleditor/__init__.py @@ -1215,3 +1215,58 @@ def save_file(): 'status': True, } ) + + +@blueprint.route('/query_tool/download/', methods=["GET"]) +@login_required +def start_query_download_tool(trans_id): + sync_conn = None + status, error_msg, conn, trans_obj, session_obj = check_transaction_status(trans_id) + + if status and conn is not None \ + and trans_obj is not None and session_obj is not None: + + data = request.args if request.args else None + try: + if data and 'query' in data: + sql = data['query'] + conn_id = str(random.randint(1, 9999999)) + sync_conn = conn.manager.connection( + did=trans_obj.did, + conn_id=conn_id, + auto_reconnect=False, + async=False + ) + + sync_conn.connect(autocommit=False) + + # This returns generator of records. + status, gen = sync_conn.execute_on_server_as_csv(sql, records=2000) + + if not status: + conn.manager.release(conn_id=conn_id, did=trans_obj.did) + return internal_server_error(errormsg=str(gen)) + + def cleanup(): + conn.manager.connections[sync_conn.conn_id]._release() + del conn.manager.connections[sync_conn.conn_id] + + r = Response(gen(), mimetype='text/csv') + + if 'filename' in data and data['filename'] != "": + filename = data['filename'] + else: + import time + filename = str(int(time.time())) + ".csv" + + r.headers["Content-Disposition"] = "attachment;filename={0}".format(filename) + + r.call_on_close(cleanup) + + return r + + except Exception as e: + conn.manager.release(conn_id=conn_id, did=trans_obj.did) + return internal_server_error(errormsg=str(e)) + else: + return internal_server_error(errormsg=gettext("Transaction status check failed.")) diff --git a/web/pgadmin/tools/sqleditor/templates/sqleditor/js/sqleditor.js b/web/pgadmin/tools/sqleditor/templates/sqleditor/js/sqleditor.js index 036b07d19..5544135b6 100644 --- a/web/pgadmin/tools/sqleditor/templates/sqleditor/js/sqleditor.js +++ b/web/pgadmin/tools/sqleditor/templates/sqleditor/js/sqleditor.js @@ -2602,76 +2602,63 @@ define( // This function will download the grid data as CSV file. _download: function() { - var self = this; - var coll = self.collection.fullCollection === undefined ? self.collection : self.collection.fullCollection; + var self = this, + selected_code = self.gridView.query_tool_obj.getSelection(), + sql = ""; - if (self.columns != undefined && - coll != undefined && - coll.length > 0) - { - var csv_col = _.indexBy(self.columns, 'name'), - labels = _.pluck(self.columns, 'label'), - keys = _.pluck(self.columns, 'name'); + if (selected_code.length > 0) + sql = selected_code; + else + sql = self.gridView.query_tool_obj.getValue(); - // Fetch the items from fullCollection and convert it as csv format - var csv = keys.join(',') + '\n'; - csv += coll.map(function(item) { - return _.map(keys, function(key) { - var cell = csv_col [key].cell, - // suppose you want to preserve custom formatters - formatter = cell.prototype && cell.prototype.formatter; + // If it is an empty query, do nothing. + if (sql.length <= 0) return; - return formatter && formatter.fromRaw ? - formatter.fromRaw(item.get(key), item) : item.get(key); - }).join(','); - }).join('\n'); - - // Download the file. - var encodedUri = encodeURI('data:text/csv&charset=utf-8&filename=download.csv&value=' + csv), - link = document.createElement('a'); - link.setAttribute('href', encodedUri); - - /* If download is from view data then file name should be - * the object name for which data is to be displayed. - */ - if (!self.is_query_tool) { - $.ajax({ - url: "{{ url_for('sqleditor.index') }}" + "object/get/" + self.transId, - method: 'GET', - success: function(res) { - if (res.data.status) { - filename = res.data.result + '.csv'; - link.setAttribute('download', filename); - link.click(); - } - }, - error: function(e) { - if (e.readyState == 0) { - alertify.alert('Get Object Name Error', - '{{ _('Not connected to the server or the connection to the server has been closed.') }}' - ); - return; - } - - var msg = e.responseText; - if (e.responseJSON != undefined && - e.responseJSON.errormsg != undefined) - msg = e.responseJSON.errormsg; - - alertify.alert('Get Object Name Error', msg); + /* If download is from view data then file name should be + * the object name for which data is to be displayed. + */ + if (!self.is_query_tool) { + $.ajax({ + url: "{{ url_for('sqleditor.index') }}" + "object/get/" + self.transId, + method: 'GET', + success: function(res) { + if (res.data.status) { + filename = res.data.result + '.csv'; + self._trigger_csv_download(sql, filename); + } + }, + error: function(e) { + if (e.readyState == 0) { + alertify.alert('Get Object Name Error', + '{{ _('Not connected to the server or the connection to the server has been closed.') }}' + ); + return; } - }); - } - else { - var cur_time = new Date(); - var filename = 'data-' + cur_time.getTime() + '.csv'; - link.setAttribute('download', filename); - link.click(); - } - } - else { - alertify.alert('Download Data', 'No data is available to download'); - } + + var msg = e.responseText; + if (e.responseJSON != undefined && + e.responseJSON.errormsg != undefined) + msg = e.responseJSON.errormsg; + + alertify.alert('Get Object Name Error', msg); + } + }); + } + else { + var cur_time = new Date(); + var filename = 'data-' + cur_time.getTime() + '.csv'; + self._trigger_csv_download(sql, filename); + } + + }, + // Trigger query result download to csv. + _trigger_csv_download: function(query, filename) { + var self = this, + link = $(this.container).find("#download-csv"), + url = "{{ url_for('sqleditor.index') }}" + "query_tool/download/" + self.transId; + + url +="?" + $.param({query:query, filename:filename}); + link.attr("src", url); }, _auto_rollback: function() { diff --git a/web/pgadmin/utils/driver/psycopg2/__init__.py b/web/pgadmin/utils/driver/psycopg2/__init__.py index ec0a8fce5..7b901d403 100644 --- a/web/pgadmin/utils/driver/psycopg2/__init__.py +++ b/web/pgadmin/utils/driver/psycopg2/__init__.py @@ -17,6 +17,8 @@ import datetime import os import random import select +import sys +import csv import psycopg2 import psycopg2.extras @@ -32,6 +34,11 @@ from .keywords import ScanKeyword from ..abstract import BaseDriver, BaseConnection from .cursor import DictCursor +if sys.version_info < (3,): + from StringIO import StringIO +else: + from io import StringIO + _ = gettext ASYNC_WAIT_TIMEOUT = 0.01 # in seconds or 10 milliseconds @@ -284,7 +291,10 @@ Failed to connect to the database server(#{server_id}) for connection ({conn_id} # autocommit flag does not work with asynchronous connections. # By default asynchronous connection runs in autocommit mode. if self.async == 0: - self.conn.autocommit = True + if 'autocommit' in kwargs and kwargs['autocommit'] == False: + self.conn.autocommit = False + else: + self.conn.autocommit = True register_date_typecasters(self.conn) status, res = self.execute_scalar(""" @@ -384,11 +394,12 @@ WHERE return True, None - def __cursor(self): + def __cursor(self, server_cursor=False): cur = getattr(g, str(self.manager.sid) + '#' + self.conn_id, None) if self.connected() and cur and not cur.closed: - return True, cur + if not server_cursor or (server_cursor and cur.name): + return True, cur if not self.connected(): status = False @@ -419,7 +430,13 @@ Attempt to reconnect failed with the error: return False, msg try: - cur = self.conn.cursor(cursor_factory=DictCursor) + if server_cursor: + # Providing name to cursor will create server side cursor. + cursor_name = "CURSOR:{0}".format(self.conn_id) + cur = self.conn.cursor(name=cursor_name, + cursor_factory=DictCursor) + else: + cur = self.conn.cursor(cursor_factory=DictCursor) except psycopg2.Error as pe: errmsg = gettext(""" Failed to create cursor for psycopg2 connection with error message for the \ @@ -471,6 +488,76 @@ Attempt to reconnect it failed with the error: if self.async == 1: self._wait(cur.connection) + + def execute_on_server_as_csv(self, query, params=None, formatted_exception_msg=False, records=2000): + status, cur = self.__cursor(server_cursor=True) + self.row_count = 0 + + if not status: + return False, str(cur) + query_id = random.randint(1, 9999999) + + current_app.logger.log(25, + "Execute (with server cursor) for server #{server_id} - {conn_id} (Query-id: {query_id}):\n{query}".format( + server_id=self.manager.sid, + conn_id=self.conn_id, + query=query, + query_id=query_id + ) + ) + try: + self.__internal_blocking_execute(cur, query, params) + except psycopg2.Error as pe: + cur.close() + errmsg = self._formatted_exception_msg(pe, formatted_exception_msg) + current_app.logger.error( + "Failed to execute query ((with server cursor) for the server #{server_id} - {conn_id} (Query-id: {query_id}):\nError Message:{errmsg}".format( + server_id=self.manager.sid, + conn_id=self.conn_id, + query=query, + errmsg=errmsg, + query_id=query_id + ) + ) + return False, errmsg + + def gen(): + + results = cur.fetchmany(records) + if not results: + if not cur.closed: + cur.close() + return + + header = [c.to_dict()['name'] for c in cur.ordered_description()] + + res_io = StringIO() + + csv_writer = csv.DictWriter( + res_io, fieldnames=header, delimiter=str(','), quoting=csv.QUOTE_NONNUMERIC + ) + csv_writer.writeheader() + csv_writer.writerows(results) + + yield res_io.getvalue().strip(str('\r\n')) + + while True: + results = cur.fetchmany(records) + + if not results: + if not cur.closed: + cur.close() + break + res_io = StringIO() + + csv_writer = csv.DictWriter( + res_io, fieldnames=header, delimiter=str(','), quoting=csv.QUOTE_NONNUMERIC + ) + csv_writer.writerows(results) + yield res_io.getvalue().strip(str('\r\n')) + + return True, gen + def execute_scalar(self, query, params=None, formatted_exception_msg=False): status, cur = self.__cursor() self.row_count = 0 @@ -1151,7 +1238,8 @@ class ServerManager(object): raise Exception("Information is not available.") def connection( - self, database=None, conn_id=None, auto_reconnect=True, did=None + self, database=None, conn_id=None, auto_reconnect=True, did=None, + async=None ): msg_active_conn = gettext( "Server has no active connection. Please connect to the server." @@ -1197,7 +1285,10 @@ WHERE db.oid = {0}""".format(did)) if my_id in self.connections: return self.connections[my_id] else: - async = 1 if conn_id is not None else 0 + if async is None: + async = 1 if conn_id is not None else 0 + else: + async = 1 if async is True else 0 self.connections[my_id] = Connection( self, my_id, database, auto_reconnect, async )