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
 
 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,49 @@ 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)
+
+    cur = conn._Connection__async_cursor
+    register_binary_data_typecasters(cur)
+    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
+        )
+
+    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 parameter (query)."
+            )
+        )
+    col_pos = data['colpos']
+    cur.scroll(int(data['rowpos']))
+    binary_data = cur.fetchone()
+    binary_data = binary_data[col_pos]
+
+    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,17 @@
 // 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';
 
 const StyledNullAndDefaultFormatter = styled(NullAndDefaultFormatter)(({theme}) => ({
   '& .Formatters-disabledCell': {
@@ -70,10 +75,14 @@ NumberFormatter.propTypes = FormatterPropTypes;
 
 export function BinaryFormatter({row, column}) {
   let value = row[column.key];
-
+  const eventBus = useContext(QueryToolEventsContext);
+  const downloadBinaryData = usePreferences().getPreferences('misc', 'enable_binary_data_download').value;
   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, row.__temp_PK, 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', {
@@ -1038,6 +1055,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/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,32 @@ 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:
+            if isinstance(data, memoryview):
+                data = bytes(data).decode()
+                if data.startswith('\\x'):
+                    data = data[2:]
+                try:
+                    return bytes.fromhex(data)
+                except ValueError:
+                    # In case of error while converting hex to bytes, return
+                    # original data.
+                    return data
+            else:
+                return data
+        return data if data is not None else None
+
+
+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(