Added support to download binary data from result grid. #4011

Author: pravesh-sharma

docs/en_US/preferences.rst

@@ -444,11 +444,17 @@ Use the fields on the *File Downloads* panel to manage file downloads related pr
 
 * When the *Automatically open downloaded files?* switch is set to *True*
   the downloaded file will automatically open in the system's default
-  application associated with that file type.
+  application associated with that file type. **Note:** This option is applicable and
+  visible only in desktop mode.
+
+* When the *Enable binary data download?* switch is set to *True*,
+  binary data can be downloaded from the result grid. Default is set to *False*
+  to prevent excessive memory usage on the server.
 
 * When the *Prompt for the download location?* switch is set to *True*
   a prompt will appear after clicking the download button, allowing you
-  to choose the download location.
+  to choose the download location. **Note:** This option is applicable and
+  visible only in desktop mode.
 
 **Note:** File Downloads related settings are applicable and visible only in desktop mode.
 

web/pgadmin/misc/__init__.py

@@ -166,6 +166,18 @@ def register_preferences(self):
                 )
             )
 
+        self.preference.register(
+            'file_downloads', 'enable_binary_data_download',
+            gettext("Enable binary data download?"),
+            'boolean', False,
+            category_label=PREF_LABEL_FILE_DOWNLOADS,
+            help_str=gettext(
+                'If set to True, binary data can be downloaded '
+                'from the result grid. The default is False to '
+                'prevent excessive memory usage on the server.'
+            )
+        )
+
     def get_exposed_url_endpoints(self):
         """
         Returns:

web/pgadmin/tools/sqleditor/__init__.py

@@ -14,6 +14,7 @@
 import secrets
 from urllib.parse import unquote
 from threading import Lock
+from io import BytesIO
 import threading
 import math
 
@@ -23,7 +24,8 @@
 
 from config import PG_DEFAULT_DRIVER, ALLOW_SAVE_PASSWORD
 from werkzeug.user_agent import UserAgent
-from flask import Response, url_for, render_template, session, current_app
+from flask import Response, url_for, render_template, session, current_app, \
+    send_file
 from flask import request
 from flask_babel import gettext
 from pgadmin.tools.sqleditor.utils.query_tool_connection_check \
@@ -70,6 +72,8 @@
 from pgadmin.browser.server_groups.servers.utils import \
     convert_connection_parameter, get_db_disp_restriction
 from pgadmin.misc.workspaces import check_and_delete_adhoc_server
+from pgadmin.utils.driver.psycopg3.typecast import \
+    register_binary_data_typecasters, register_binary_typecasters
 
 MODULE_NAME = 'sqleditor'
 TRANSACTION_STATUS_CHECK_FAILED = gettext("Transaction status check failed.")
@@ -147,6 +151,7 @@ def get_exposed_url_endpoints(self):
             'sqleditor.server_cursor',
             'sqleditor.nlq_chat_stream',
             'sqleditor.explain_analyze_stream',
+            'sqleditor.download_binary_data',
         ]
 
     def on_logout(self):
@@ -2182,6 +2187,66 @@ def start_query_download_tool(trans_id):
         return internal_server_error(errormsg=err_msg)
 
 
+@blueprint.route(
+    '/download_binary_data/<int:trans_id>',
+    methods=["POST"], endpoint='download_binary_data'
+)
+@pga_login_required
+def download_binary_data(trans_id):
+    """
+    This method is used to download binary data.
+    """
+
+    (status, error_msg, conn, trans_obj,
+     session_obj) = check_transaction_status(trans_id)
+
+    if isinstance(error_msg, Response):
+        return error_msg
+    if error_msg == ERROR_MSG_TRANS_ID_NOT_FOUND:
+        return make_json_response(
+            success=0,
+            errormsg=error_msg,
+            info='DATAGRID_TRANSACTION_REQUIRED',
+            status=404
+        )
+
+    if not status or conn is None or trans_obj is None or \
+            session_obj is None:
+        return internal_server_error(
+            errormsg=TRANSACTION_STATUS_CHECK_FAILED
+        )
+
+    cur = conn._Connection__async_cursor
+    if cur is None:
+        return internal_server_error(
+            errormsg=gettext('No active result cursor.')
+        )
+
+    data = request.values if request.values else request.get_json(silent=True)
+    if data is None:
+        return make_json_response(
+            status=410,
+            success=0,
+            errormsg=gettext(
+                "Could not find the required parameters (rowpos, colpos)."
+            )
+        )
+
+    binary_data = conn.download_binary_data(cur, data)
+
+    if binary_data is None:
+        return bad_request(
+            errormsg=gettext('The selected cell contains NULL.')
+        )
+
+    return send_file(
+        BytesIO(binary_data),
+        as_attachment=True,
+        download_name='binary_data',
+        mimetype='application/octet-stream'
+    )
+
+
 @blueprint.route(
     '/status/<int:trans_id>',
     methods=["GET"],

web/pgadmin/tools/sqleditor/static/js/components/QueryToolConstants.js

@@ -30,6 +30,7 @@ export const QUERY_TOOL_EVENTS = {
   TRIGGER_SELECT_ALL: 'TRIGGER_SELECT_ALL',
   TRIGGER_SAVE_QUERY_TOOL_DATA: 'TRIGGER_SAVE_QUERY_TOOL_DATA',
   TRIGGER_GET_QUERY_CONTENT: 'TRIGGER_GET_QUERY_CONTENT',
+  TRIGGER_SAVE_BINARY_DATA: 'TRIGGER_SAVE_BINARY_DATA',
 
   COPY_DATA: 'COPY_DATA',
   SET_LIMIT_VALUE: 'SET_LIMIT_VALUE',

web/pgadmin/tools/sqleditor/static/js/components/QueryToolDataGrid/Formatters.jsx

@@ -6,12 +6,18 @@
 // This software is released under the PostgreSQL Licence
 //
 //////////////////////////////////////////////////////////////
+import { useContext } from 'react';
 import { styled } from '@mui/material/styles';
 import _ from 'lodash';
 import PropTypes from 'prop-types';
+import gettext from 'sources/gettext';
 import CustomPropTypes from '../../../../../../static/js/custom_prop_types';
 import usePreferences from '../../../../../../preferences/static/js/store';
-
+import GetAppRoundedIcon from '@mui/icons-material/GetAppRounded';
+import { PgIconButton } from '../../../../../../static/js/components/Buttons';
+import { QUERY_TOOL_EVENTS } from '../QueryToolConstants';
+import { QueryToolEventsContext } from '../QueryToolComponent';
+import { DataGridExtrasContext } from './index';
 
 const StyledNullAndDefaultFormatter = styled(NullAndDefaultFormatter)(({theme}) => ({
   '& .Formatters-disabledCell': {
@@ -68,12 +74,20 @@ export function NumberFormatter({row, column}) {
 }
 NumberFormatter.propTypes = FormatterPropTypes;
 
-export function BinaryFormatter({row, column}) {
+export function BinaryFormatter({row, column, ...props}) {
   let value = row[column.key];
+  const eventBus = useContext(QueryToolEventsContext);
+  const dataGridExtras = useContext(DataGridExtrasContext);
+  const downloadBinaryData = usePreferences().getPreferences('misc', 'enable_binary_data_download').value;
+
+  const absoluteRowPos = (dataGridExtras?.startRowNum ?? 1) - 1 + props.rowIdx;
 
   return (
     <StyledNullAndDefaultFormatter value={value} column={column}>
-      <span className='Formatters-disabledCell'>[{value}]</span>
+      <span className='Formatters-disabledCell'>[{value}]</span>&nbsp;&nbsp;
+      {downloadBinaryData &&
+        <PgIconButton size="xs" title={gettext('Download binary data')} icon={<GetAppRoundedIcon />}
+          onClick={()=>eventBus.fireEvent(QUERY_TOOL_EVENTS.TRIGGER_SAVE_BINARY_DATA, absoluteRowPos, column.pos)}/>}
     </StyledNullAndDefaultFormatter>
   );
 }

web/pgadmin/tools/sqleditor/static/js/components/sections/ResultSet.jsx

@@ -493,6 +493,23 @@ export class ResultSetUtils {
     }
   }
 
+  async saveBinaryResultsToFile(fileName, rowPos, colPos, onProgress) {
+    try {
+      await DownloadUtils.downloadFileStream({
+        url: url_for('sqleditor.download_binary_data', {
+          'trans_id': this.transId,
+        }),
+        options: {
+          method: 'POST',
+          body: JSON.stringify({filename: fileName, rowpos: rowPos, colpos: colPos})
+        }}, fileName, 'application/octet-stream', onProgress);
+      this.eventBus.fireEvent(QUERY_TOOL_EVENTS.TRIGGER_SAVE_RESULTS_END);
+    } catch (error) {
+      this.eventBus.fireEvent(QUERY_TOOL_EVENTS.TRIGGER_SAVE_RESULTS_END);
+      this.eventBus.fireEvent(QUERY_TOOL_EVENTS.HANDLE_API_ERROR, error);
+    }
+  }
+
   includeFilter(reqData) {
     return this.api.post(
       url_for('sqleditor.inclusive_filter', {
@@ -1048,6 +1065,15 @@ export function ResultSet() {
       setLoaderText('');
     });
 
+    eventBus.registerListener(QUERY_TOOL_EVENTS.TRIGGER_SAVE_BINARY_DATA, async (rowPos, colPos)=>{
+      let fileName = 'data-' + new Date().getTime();
+      setLoaderText(gettext('Downloading results...'));
+      await rsu.current.saveBinaryResultsToFile(fileName, rowPos, colPos, (p)=>{
+        setLoaderText(gettext('Downloading results(%s)...', p));
+      });
+      setLoaderText('');
+    });
+
     eventBus.registerListener(QUERY_TOOL_EVENTS.TRIGGER_SET_LIMIT, async (limit)=>{
       setLoaderText(gettext('Setting the limit on the result...'));
       try {

web/pgadmin/utils/driver/psycopg3/connection.py

@@ -30,12 +30,13 @@
 from pgadmin.model import User
 from pgadmin.utils.exception import ConnectionLost, CryptKeyMissing
 from pgadmin.utils import get_complete_file_path
+from pgadmin.utils.ajax import internal_server_error
 from ..abstract import BaseConnection
 from .cursor import DictCursor, AsyncDictCursor, AsyncDictServerCursor
-from .typecast import register_global_typecasters,\
-    register_string_typecasters, register_binary_typecasters, \
-    register_array_to_string_typecasters, ALL_JSON_TYPES, \
-    register_numeric_typecasters
+from .typecast import register_binary_data_typecasters,\
+    register_global_typecasters, register_string_typecasters,\
+    register_binary_typecasters, register_array_to_string_typecasters,\
+    register_numeric_typecasters, ALL_JSON_TYPES
 from .encoding import get_encoding, configure_driver_encodings
 from pgadmin.utils import csv_lib as csv
 from pgadmin.utils.master_password import get_crypt_key
@@ -1913,3 +1914,43 @@ def mogrify(self, query, parameters):
                     return _cur.mogrify(query, parameters)
             else:
                 return query
+
+    def download_binary_data(self, cur, params):
+        """
+        This function will return the binary data for the given query.
+        :param cur: cursor object
+        :param params: row/col params
+        :return:
+        """
+        try:
+            register_binary_data_typecasters(cur)
+            row_pos = int(params['rowpos'])
+            col_pos = int(params['colpos'])
+            if row_pos < 0 or col_pos < 0:
+                raise ValueError
+
+            # Save the current cursor position
+            saved_pos = cur.rownumber if cur.rownumber is not None else 0
+
+            try:
+                # Scroll to the requested row and fetch it
+                cur.scroll(row_pos, mode='absolute')
+                row = cur.fetchone()
+            finally:
+                # Always restore the cursor position
+                cur.scroll(saved_pos, mode='absolute')
+
+            if row is None or col_pos >= len(row):
+                return internal_server_error(
+                    errormsg=gettext('Requested cell is out of range.')
+                )
+            return row[col_pos]
+        except (ValueError, IndexError, TypeError) as e:
+            current_app.logger.error(e)
+            return internal_server_error(
+                errormsg='Invalid row/column position.'
+            )
+        finally:
+            # Always restore the original typecasters
+            # (works on connection or cursor)
+            register_binary_typecasters(cur)

web/pgadmin/utils/driver/psycopg3/typecast.py

@@ -212,6 +212,21 @@ def register_array_to_string_typecasters(connection=None):
                                                 TextLoaderpgAdmin)
 
 
+def register_binary_data_typecasters(cur):
+    # Register type caster to fetch original binary data for bytea type.
+    cur.adapters.register_loader(17,
+                                 ByteaDataLoader)
+
+    cur.adapters.register_loader(1001,
+                                 ByteaDataLoader)
+
+    cur.adapters.register_loader(17,
+                                 ByteaBinaryDataLoader)
+
+    cur.adapters.register_loader(1001,
+                                 ByteaBinaryDataLoader)
+
+
 class InetLoader(InetLoader):
     def load(self, data):
         if isinstance(data, memoryview):
@@ -240,6 +255,24 @@ def load(self, data):
         return 'binary data' if data is not None else None
 
 
+class ByteaDataLoader(Loader):
+    # Loads the actual binary data.
+    def load(self, data):
+        if data is None:
+            return None
+        raw = bytes(data) if isinstance(data, memoryview) else data
+        if isinstance(raw, bytes) and raw.startswith(b'\\x'):
+            return bytes.fromhex(raw[2:].decode())
+        return raw
+
+
+class ByteaBinaryDataLoader(Loader):
+    format = _pq_Format.BINARY
+
+    def load(self, data):
+        return data if data is not None else None
+
+
 class TextLoaderpgAdmin(TextLoader):
     def load(self, data):
         postgres_encoding, python_encoding = get_encoding(