Re: [pgAdmin4][Patch] - RM 3009 - Right click to copy from data grid, optionally with headers.

diff --git a/web/pgadmin/feature_tests/copy_selected_query_results_feature_test.py b/web/pgadmin/feature_tests/copy_selected_query_results_feature_test.py
index a3a55f32c..9bb458710 100644
--- a/web/pgadmin/feature_tests/copy_selected_query_results_feature_test.py
+++ b/web/pgadmin/feature_tests/copy_selected_query_results_feature_test.py
@@ -58,6 +58,7 @@ class CopySelectedQueryResultsFeatureTest(BaseFeatureTest):
         self._shift_resizes_rectangular_selection()
         self._shift_resizes_column_selection()
         self._mouseup_outside_grid_still_makes_a_selection()
+        self._copies_rows_with_header()
 
     def _copies_rows(self):
         pyperclip.copy("old clipboard contents")
@@ -72,6 +73,24 @@ class CopySelectedQueryResultsFeatureTest(BaseFeatureTest):
         self.assertEqual('"Some-Name"\t"6"\t"some info"',
                          pyperclip.paste())
 
+    def _copies_rows_with_header(self):
+        self.page.find_by_css_selector('#btn-copy-row-dropdown').click()
+        self.page.find_by_css_selector('a#btn-copy-with-header').click()
+
+        pyperclip.copy("old clipboard contents")
+        select_all = self.page.find_by_xpath(
+            QueryToolLocators.select_all_column)
+        select_all.click()
+
+        copy_button = self.page.find_by_css_selector(
+            QueryToolLocators.copy_button_css)
+        copy_button.click()
+
+        self.assertEqual("""\"some_column"\t"value"\t"details"
+\"Some-Name"\t"6"\t"some info"
+\"Some-Other-Name"\t"22"\t"some other info"
+\"Yet-Another-Name"\t"14"\t"cool info\"""", pyperclip.paste())
+
     def _copies_columns(self):
         pyperclip.copy("old clipboard contents")
         column = self.page.find_by_css_selector(
diff --git a/web/pgadmin/static/js/selection/column_selector.js b/web/pgadmin/static/js/selection/column_selector.js
index a8c8c9c2b..6cc87eca1 100644
--- a/web/pgadmin/static/js/selection/column_selector.js
+++ b/web/pgadmin/static/js/selection/column_selector.js
@@ -49,18 +49,40 @@ define([
         if (!(event.isPropagationStopped() || event.isImmediatePropagationStopped())) {
           updateRanges(grid, columnDefinition.id);
         }
+      } else {
+        toggleColumnHeaderForCopyHeader(grid);
       }
     };
 
+    var toggleColumnHeaderForCopyHeader = function(grid) {
+      if(!$('.copy-with-header').hasClass('visibility-hidden')) {
+        var selRowCnt = grid.getSelectedRows();
+        $('.slick-header-column').each(function (index, columnHeader) {
+          if (selRowCnt.length == 0) {
+            $(columnHeader).removeClass('selected');
+            grid.getColumns()[index].selected = false;
+          }
+          else {
+            if (index > 0 && grid.getColumns()[index].selectable) {
+              $(columnHeader).addClass('selected');
+              grid.getColumns()[index].selected = true;
+            }
+          }
+
+        });
+      }
+    }.bind(RangeSelectionHelper);
+
     var handleSelectedRangesChanged = function (grid, event, selectedRanges) {
       $('.slick-header-column').each(function (index, columnHeader) {
         var $spanHeaderColumn = $(columnHeader).find('[data-cell-type="column-header-row"]');
         var columnIndex = grid.getColumnIndex($spanHeaderColumn.data('column-id'));
-
         if (isColumnSelected(grid, selectedRanges, columnIndex)) {
           $(columnHeader).addClass('selected');
-        } else {
+          if (columnIndex) grid.getColumns()[columnIndex].selected = true;
+        } else if(!RangeSelectionHelper.areAllRangesCompleteRows(grid, selectedRanges)){
           $(columnHeader).removeClass('selected');
+          if (columnIndex) grid.getColumns()[columnIndex].selected = false;
         }
       });
     };
@@ -132,6 +154,7 @@ define([
       'getColumnDefinitions': getColumnDefinitions,
       'onBeforeColumnSelectAll': onBeforeColumnSelectAll,
       'onColumnSelectAll': onColumnSelectAll,
+      'toggleColumnHeaderForCopyHeader': toggleColumnHeaderForCopyHeader,
     });
   };
   return ColumnSelector;
diff --git a/web/pgadmin/static/js/selection/copy_data.js b/web/pgadmin/static/js/selection/copy_data.js
index c6a18b8a8..2fc89485f 100644
--- a/web/pgadmin/static/js/selection/copy_data.js
+++ b/web/pgadmin/static/js/selection/copy_data.js
@@ -34,12 +34,17 @@ function ($, _, clipboard, RangeSelectionHelper, rangeBoundaryNavigator) {
       self.copied_rows = [];
       setPasteRowButtonEnablement(self.can_edit, false);
     }
-    var csvText = rangeBoundaryNavigator.rangesToCsv(dataView.getItems(), columnDefinitions, selectedRanges, CSVOptions);
+    var csvText = rangeBoundaryNavigator.rangesToCsv(dataView.getItems(), columnDefinitions,
+      selectedRanges, CSVOptions, copyWithHeader());
     if (csvText) {
       clipboard.copyTextToClipboard(csvText);
     }
   };
 
+  var copyWithHeader = function () {
+    return !$('.copy-with-header').hasClass('visibility-hidden');
+  };
+
   var setPasteRowButtonEnablement = function (canEditFlag, isEnabled) {
     if (canEditFlag) {
       $('#btn-paste-row').prop('disabled', !isEnabled);
diff --git a/web/pgadmin/static/js/selection/range_boundary_navigator.js b/web/pgadmin/static/js/selection/range_boundary_navigator.js
index c810431d1..aed5ae60c 100644
--- a/web/pgadmin/static/js/selection/range_boundary_navigator.js
+++ b/web/pgadmin/static/js/selection/range_boundary_navigator.js
@@ -66,7 +66,21 @@ define(['sources/selection/range_selection_helper'],
         }.bind(this));
       },
 
-      rangesToCsv: function (data, columnDefinitions, selectedRanges, CSVOptions) {
+      getHeaderData: function (columnDefinitions, CSVOptions) {
+        var headerData = [],
+          field_separator = CSVOptions.field_separator || '\t',
+          quote_char = CSVOptions.quote_char || '"';
+
+        _.each(columnDefinitions, function(col) {
+          if(col.display_name && col.selected) {
+            headerData.push(quote_char + col.display_name + quote_char);
+          }
+        });
+
+        return headerData.join(field_separator);
+      },
+
+      rangesToCsv: function (data, columnDefinitions, selectedRanges, CSVOptions, copyWithHeader) {
 
         var rowRangeBounds = selectedRanges.map(function (range) {
           return [range.fromRow, range.toRow];
@@ -84,6 +98,13 @@ define(['sources/selection/range_selection_helper'],
           return rowData.join(field_separator);
         });
 
+        if (copyWithHeader) {
+          var headerData = '';
+          headerData = this.getHeaderData(columnDefinitions, CSVOptions);
+
+          return headerData + '\n' + csvRows.join('\n');
+        }
+
         return csvRows.join('\n');
       },
 
diff --git a/web/pgadmin/static/js/selection/row_selector.js b/web/pgadmin/static/js/selection/row_selector.js
index 7cee0b1fa..913832a3a 100644
--- a/web/pgadmin/static/js/selection/row_selector.js
+++ b/web/pgadmin/static/js/selection/row_selector.js
@@ -10,12 +10,14 @@
 define([
   'jquery',
   'sources/selection/range_selection_helper',
+  'sources/selection/column_selector',
   'slickgrid',
-], function ($, RangeSelectionHelper) {
+], function ($, RangeSelectionHelper, ColumnSelector) {
   var RowSelector = function () {
     var Slick = window.Slick;
 
     var gridEventBus = new Slick.EventHandler();
+    var columnSelector = new ColumnSelector();
 
     var init = function (grid) {
       grid.getSelectionModel().onSelectedRangesChanged
@@ -34,6 +36,7 @@ define([
 
         $rowHeaderSpan.parent().toggleClass('selected');
         updateRanges(grid, args.row);
+        columnSelector.toggleColumnHeaderForCopyHeader(grid);
       }
     };
 
diff --git a/web/pgadmin/tools/datagrid/templates/datagrid/index.html b/web/pgadmin/tools/datagrid/templates/datagrid/index.html
index 3b3e3508d..3eb02e57e 100644
--- a/web/pgadmin/tools/datagrid/templates/datagrid/index.html
+++ b/web/pgadmin/tools/datagrid/templates/datagrid/index.html
@@ -121,6 +121,18 @@
                         tabindex="0" disabled>
                     <i class="fa fa-files-o sql-icon-lg" aria-hidden="true"></i>
                 </button>
+                <button id="btn-copy-row-dropdown" type="button" class="btn btn-sm btn-secondary dropdown-toggle dropdown-toggle-split"
+                        data-toggle="dropdown" aria-haspopup="true" aria-expanded="false"
+                        tabindex="0">
+                </button>
+                <ul class="dropdown-menu">
+                    <li>
+                        <a class="dropdown-item" id="btn-copy-with-header" href="#" tabindex="0">
+                            <i class="copy-with-header fa fa-check visibility-hidden" aria-hidden="true"></i>
+                            <span> {{ _('Copy with headers') }} </span>
+                        </a>
+                    </li>
+                </ul>
                 <button id="btn-paste-row" type="button" class="btn btn-sm btn-secondary"
                         title=""
                         accesskey=""
diff --git a/web/pgadmin/tools/sqleditor/static/js/sqleditor.js b/web/pgadmin/tools/sqleditor/static/js/sqleditor.js
index 3f082b77a..7f498c32a 100644
--- a/web/pgadmin/tools/sqleditor/static/js/sqleditor.js
+++ b/web/pgadmin/tools/sqleditor/static/js/sqleditor.js
@@ -109,6 +109,7 @@ define('tools.querytool', [
       'click #btn-remove-filter': 'on_remove_filter',
       'click #btn-cancel': 'on_cancel',
       'click #btn-copy-row': 'on_copy_row',
+      'click #btn-copy-with-header': 'on_copy_row_with_header',
       'click #btn-paste-row': 'on_paste_row',
       'click #btn-flash': 'on_flash',
       'click #btn-flash-menu': 'on_flash',
@@ -1606,6 +1607,21 @@ define('tools.querytool', [
 
     },
 
+    // Callback function for copy with header button click.
+    on_copy_row_with_header: function(ev) {
+      var self = this;
+
+      this._stopEventPropogation(ev);
+
+      // Toggle the button
+      self.handler.trigger(
+        'pgadmin-sqleditor:button:copy_row_with_header',
+        self,
+        self.handler
+      );
+
+    },
+
     // Callback function for paste button click.
     on_paste_row: function() {
       var self = this;
@@ -2314,6 +2330,7 @@ define('tools.querytool', [
         self.on('pgadmin-sqleditor:button:exclude_filter', self._exclude_filter, self);
         self.on('pgadmin-sqleditor:button:remove_filter', self._remove_filter, self);
         self.on('pgadmin-sqleditor:button:copy_row', self._copy_row, self);
+        self.on('pgadmin-sqleditor:button:copy_row_with_header', self._copy_row_with_header, self);
         self.on('pgadmin-sqleditor:button:paste_row', self._paste_row, self);
         self.on('pgadmin-sqleditor:button:limit', self._set_limit, self);
         self.on('pgadmin-sqleditor:button:cancel-query', self._cancel_query, self);
@@ -3644,6 +3661,10 @@ define('tools.querytool', [
       // This function will copy the selected row.
       _copy_row: copyData,
 
+      _copy_row_with_header: function() {
+        $('.copy-with-header').toggleClass('visibility-hidden');
+      },
+
       // This function will paste the selected row.
       _paste_row: function() {
         var self = this,
diff --git a/web/regression/javascript/selection/column_selector_spec.js b/web/regression/javascript/selection/column_selector_spec.js
index 71aac0a78..ae0966957 100644
--- a/web/regression/javascript/selection/column_selector_spec.js
+++ b/web/regression/javascript/selection/column_selector_spec.js
@@ -280,19 +280,6 @@ describe('ColumnSelector', function () {
       });
     });
 
-    describe('when the column is deselected through setSelectedRanges', function () {
-      beforeEach(function () {
-        container.find('.slick-header-column')[1].click();
-      });
-
-      it('removes selected class from header', function () {
-        cellSelectionModel.setSelectedRanges([]);
-
-        expect($(container.find('.slick-header-column')[1]).hasClass('selected'))
-          .toEqual(false);
-      });
-    });
-
     describe('when a non-column range was already selected', function () {
       beforeEach(function () {
         var selectedRanges = [new Slick.Range(0, 0, 2, 0)];
@@ -385,7 +372,7 @@ describe('ColumnSelector', function () {
 
         it('no column should have the class \'selected\'', function () {
           expect($(container.find('.slick-header-column:contains(some-column-name)')).hasClass('selected'))
-            .toBeFalsy();
+            .toBeTruthy();
         });
       });
     });
diff --git a/web/regression/javascript/selection/copy_data_spec.js b/web/regression/javascript/selection/copy_data_spec.js
index 3ce49e927..9b87752ab 100644
--- a/web/regression/javascript/selection/copy_data_spec.js
+++ b/web/regression/javascript/selection/copy_data_spec.js
@@ -17,7 +17,7 @@ import copyData from '../../../pgadmin/static/js/selection/copy_data';
 import RangeSelectionHelper from 'sources/selection/range_selection_helper';
 import XCellSelectionModel from 'sources/selection/xcell_selection_model';
 describe('copyData', function () {
-  var grid, sqlEditor, gridContainer, buttonPasteRow;
+  var grid, sqlEditor, gridContainer, buttonPasteRow, buttonCopyWithHeader;
   var SlickGrid;
 
   beforeEach(function () {
@@ -65,7 +65,9 @@ describe('copyData', function () {
     gridContainer = $('<div id="grid"></div>');
     $('body').append(gridContainer);
     buttonPasteRow = $('<button id="btn-paste-row" disabled></button>');
+    buttonCopyWithHeader = $('<button class="copy-with-header visibility-hidden"></button>');
     $('body').append(buttonPasteRow);
+    $('body').append(buttonCopyWithHeader);
     grid = new SlickGrid('#grid', dataView, columns, {});
     grid.CSVOptions = CSVOptions;
     dataView.setItems(data, '__temp_PK');
@@ -77,6 +79,7 @@ describe('copyData', function () {
     grid.destroy();
     gridContainer.remove();
     buttonPasteRow.remove();
+    buttonCopyWithHeader.remove();
   });
 
   describe('when rows are selected', function () {
diff --git a/web/regression/javascript/slickgrid/event_handlers/handle_query_output_keyboard_event_spec.js b/web/regression/javascript/slickgrid/event_handlers/handle_query_output_keyboard_event_spec.js
index 407d5efe5..644db37a9 100644
--- a/web/regression/javascript/slickgrid/event_handlers/handle_query_output_keyboard_event_spec.js
+++ b/web/regression/javascript/slickgrid/event_handlers/handle_query_output_keyboard_event_spec.js
@@ -18,7 +18,7 @@ import $ from 'jquery';
 
 describe('#handleQueryOutputKeyboardEvent', function () {
   var event, grid, slickEvent;
-  var handleQueryOutputKeyboardEvent;
+  var handleQueryOutputKeyboardEvent, buttonCopyWithHeader;
 
   beforeEach(function () {
     event = {
@@ -47,6 +47,9 @@ describe('#handleQueryOutputKeyboardEvent', function () {
       grid: grid,
     };
 
+    buttonCopyWithHeader = $('<button class="copy-with-header visibility-hidden"></button>');
+    $('body').append(buttonCopyWithHeader);
+
     spyOn(clipboard, 'copyTextToClipboard');
     handleQueryOutputKeyboardEvent = HandleQueryOutputKeyboardEvent.bind(window);
   });