fix(import-export): Use query document count in export COMPASS-4537, COMPASS-4906 (#2307)

* temp commit

* Make an import and an export store to avoid double listening, show correct count, fix showing undefined for filter

* Code cleanup, better comments

* fix standalone running import/export

Author: Anemy

packages/compass-crud/src/stores/crud-store.js

@@ -652,11 +652,12 @@ const configureStore = (options = {}) => {
     },
 
     /**
-     * Open an import file dialog from compass-import-export-plugin.
+     * Open an export file dialog from compass-import-export-plugin.
      * Emits a global app registry event the plugin listens to.
      */
     openExportFileDialog() {
-      this.localAppRegistry.emit('open-export');
+      // Pass the doc count to the export modal so we can avoid re-counting.
+      this.localAppRegistry.emit('open-export', this.state.count);
     },
 
     /**

packages/compass-import-export/electron/renderer/components/import-export/import-export.jsx

@@ -11,7 +11,8 @@ class ImportExport extends Component {
 
   static propTypes = {
     appRegistry: PropTypes.object.isRequired,
-    store: PropTypes.object.isRequired
+    exportStore: PropTypes.object.isRequired,
+    importStore: PropTypes.object.isRequired
   };
 
   handleExportModalOpen = () => {
@@ -33,12 +34,17 @@ class ImportExport extends Component {
         <TextButton
           className="btn btn-default btn-sm"
           clickHandler={this.handleImportModalOpen}
-          text="Import" />
+          text="Import"
+        />
         <TextButton
           className="btn btn-default btn-sm"
           clickHandler={this.handleExportModalOpen}
-          text="Export" />
-        <Plugin store={this.props.store} />
+          text="Export"
+        />
+        <Plugin
+          exportStore={this.props.exportStore}
+          importStore={this.props.importStore}
+        />
       </div>
     );
   }

packages/compass-import-export/electron/renderer/index.js

@@ -10,7 +10,12 @@ import AppRegistry from 'hadron-app-registry';
 import { AppContainer } from 'react-hot-loader';
 import { activate } from '../../src/index.js';
 import ImportExportPlugin from './components/import-export';
-import configureStore, { setDataProvider } from '../../src/stores';
+import configureExportStore, {
+  setDataProvider as setExportDataProvider
+} from '../../src/stores/export-store';
+import configureImportStore, {
+  setDataProvider as setImportDataProvider
+} from '../../src/stores/import-store';
 import { activate as activateStats } from '@mongodb-js/compass-collection-stats';
 
 // Import global less file. Note: these styles WILL NOT be used in compass, as compass provides its own set
@@ -79,7 +84,11 @@ root.id = 'root';
 document.body.appendChild(root);
 
 const localAppRegistry = new AppRegistry();
-const store = configureStore({
+const exportStore = configureExportStore({
+  namespace: NS,
+  localAppRegistry: localAppRegistry
+});
+const importStore = configureImportStore({
   namespace: NS,
   localAppRegistry: localAppRegistry
 });
@@ -88,7 +97,11 @@ const store = configureStore({
 const render = (Component) => {
   ReactDOM.render(
     <AppContainer>
-      <Component store={store} appRegistry={localAppRegistry} />
+      <Component
+        exportStore={exportStore}
+        importStore={importStore}
+        appRegistry={localAppRegistry}
+      />
     </AppContainer>,
     document.getElementById('root')
   );
@@ -108,7 +121,8 @@ import DataService from 'mongodb-data-service';
 const dataService = new DataService(connection);
 
 dataService.connect((error, ds) => {
-  setDataProvider(store, error, ds);
+  setImportDataProvider(importStore, error, ds);
+  setExportDataProvider(exportStore, error, ds);
   onDataServiceConnected(localAppRegistry);
 });
 

packages/compass-import-export/src/index.js

@@ -1,15 +1,16 @@
 import Plugin from './plugin';
 import ImportPlugin from './import-plugin';
 import ExportPlugin from './export-plugin';
-import configureStore from './stores';
+import configureExportStore from './stores/export-store';
+import configureImportStore from './stores/import-store';
 
 /**
  * The import plugin.
  */
 const IMPORT_ROLE = {
   name: 'Import',
   component: ImportPlugin,
-  configureStore: configureStore,
+  configureStore: configureImportStore,
   configureActions: () => {},
   storeName: 'Import.Store'
 };
@@ -20,7 +21,7 @@ const IMPORT_ROLE = {
 const EXPORT_ROLE = {
   name: 'Export',
   component: ExportPlugin,
-  configureStore: configureStore,
+  configureStore: configureExportStore,
   configureActions: () => {},
   storeName: 'Export.Store'
 };
@@ -44,4 +45,11 @@ function deactivate(appRegistry) {
 }
 
 export default Plugin;
-export { activate, deactivate, ImportPlugin, ExportPlugin, configureStore };
+export {
+  activate,
+  deactivate,
+  ImportPlugin,
+  ExportPlugin,
+  configureExportStore,
+  configureImportStore
+};

packages/compass-import-export/src/modules/export.js

@@ -1,6 +1,7 @@
 /* eslint-disable valid-jsdoc */
 import fs from 'fs';
 import stream from 'stream';
+import { promisify } from 'util';
 
 import PROCESS_STATUS from '../constants/process-status';
 import EXPORT_STEP from '../constants/export-step';
@@ -322,16 +323,51 @@ export const changeExportStep = (status) => ({
   status: status
 });
 
+const fetchDocumentCount = async(dataService, ns, query) => {
+  // When there is no filter/limit/skip try to use the estimated count.
+  if (
+    (!query.filter || Object.keys(query.filter).length < 1)
+    && !query.limit
+    && !query.skip
+  ) {
+    try {
+      const runEstimatedDocumentCount = promisify(dataService.estimatedCount.bind(dataService));
+      const count = await runEstimatedDocumentCount(ns, {});
+
+      return count;
+    } catch (estimatedCountErr) {
+      // `estimatedDocumentCount` is currently unsupported for
+      // views and time-series collections, so we can fallback to a full
+      // count in these cases and ignore this error.
+    }
+  }
+
+  const runCount = promisify(dataService.count.bind(dataService));
+
+  const count = await runCount(
+    ns,
+    query.filter || {},
+    {
+      ...(query.limit ? { limit: query.limit } : {} ),
+      ...(query.skip ? { skip: query.skip } : {} )
+    }
+  );
+  return count;
+};
+
 /**
  * Open the export modal.
  *
+ * @param {number} [count] - optional pre supplied count to shortcut and
+ * avoid a possibly expensive re-count.
+ *
  * Counts the documents to be exported given the current query on modal open to
  * provide user with accurate export data
  *
  * @api public
  */
-export const openExport = () => {
-  return (dispatch, getState) => {
+export const openExport = (count) => {
+  return async(dispatch, getState) => {
     const {
       ns,
       exportData,
@@ -340,12 +376,16 @@ export const openExport = () => {
 
     const spec = exportData.query;
 
-    dataService.estimatedCount(ns, {query: spec.filter}, function(countErr, count) {
-      if (countErr) {
-        return onError(countErr);
-      }
-      dispatch(onModalOpen(count, spec));
-    });
+    if (count) {
+      return dispatch(onModalOpen(count, spec));
+    }
+
+    try {
+      const docCount = await fetchDocumentCount(dataService, ns, spec);
+      dispatch(onModalOpen(docCount, spec));
+    } catch (e) {
+      dispatch(onError(e));
+    }
   };
 };
 
@@ -389,7 +429,7 @@ export const sampleFields = () => {
  * @api public
  */
 export const startExport = () => {
-  return (dispatch, getState) => {
+  return async(dispatch, getState) => {
     const {
       ns,
       exportData,
@@ -400,87 +440,85 @@ export const startExport = () => {
       ? { filter: {} }
       : exportData.query;
 
+    const numDocsToExport = exportData.isFullCollection
+      ? await fetchDocumentCount(dataService, ns, spec)
+      : exportData.count;
+
     // filter out only the fields we want to include in our export data
     const projection = Object.fromEntries(
       Object.entries(exportData.fields)
         .filter((keyAndValue) => keyAndValue[1] === 1));
 
-    dataService.estimatedCount(ns, {query: spec.filter}, function(countErr, numDocsToExport) {
-      if (countErr) {
-        return onError(countErr);
-      }
+    debug('count says to expect %d docs in export', numDocsToExport);
+    const source = createReadableCollectionStream(dataService, ns, spec, projection);
 
-      debug('count says to expect %d docs in export', numDocsToExport);
-      const source = createReadableCollectionStream(dataService, ns, spec, projection);
+    const progress = createProgressStream({
+      objectMode: true,
+      length: numDocsToExport,
+      time: 250 /* ms */
+    });
 
-      const progress = createProgressStream({
-        objectMode: true,
-        length: numDocsToExport,
-        time: 250 /* ms */
-      });
+    progress.on('progress', function(info) {
+      dispatch(onProgress(info.percentage, info.transferred));
+    });
 
-      progress.on('progress', function(info) {
-        dispatch(onProgress(info.percentage, info.transferred));
-      });
+    // Pick the columns that are going to be matched by the projection,
+    // where some prefix the field (e.g. ['a', 'a.b', 'a.b.c'] for 'a.b.c')
+    // has an entry in the projection object.
+    const columns = Object.keys(exportData.allFields)
+      .filter(field => field.split('.').some(
+        (_part, index, parts) => projection[parts.slice(0, index + 1).join('.')]));
+    let formatter;
+    if (exportData.fileType === 'csv') {
+      formatter = createCSVFormatter({ columns });
+    } else {
+      formatter = createJSONFormatter();
+    }
 
-      // Pick the columns that are going to be matched by the projection,
-      // where some prefix the field (e.g. ['a', 'a.b', 'a.b.c'] for 'a.b.c')
-      // has an entry in the projection object.
-      const columns = Object.keys(exportData.allFields)
-        .filter(field => field.split('.').some(
-          (_part, index, parts) => projection[parts.slice(0, index + 1).join('.')]));
-      let formatter;
-      if (exportData.fileType === 'csv') {
-        formatter = createCSVFormatter({ columns });
-      } else {
-        formatter = createJSONFormatter();
-      }
+    const dest = fs.createWriteStream(exportData.fileName);
 
-      const dest = fs.createWriteStream(exportData.fileName);
+    debug('executing pipeline');
+    dispatch(onStarted(source, dest, numDocsToExport));
+    stream.pipeline(source, progress, formatter, dest, function(err) {
+      if (err) {
+        debug('error running export pipeline', err);
+        return dispatch(onError(err));
+      }
+      debug(
+        'done. %d docs exported to %s',
+        numDocsToExport,
+        exportData.fileName
+      );
+      dispatch(onFinished(numDocsToExport));
+      dispatch(
+        appRegistryEmit(
+          'export-finished',
+          numDocsToExport,
+          exportData.fileType
+        )
+      );
 
-      debug('executing pipeline');
-      dispatch(onStarted(source, dest, numDocsToExport));
-      stream.pipeline(source, progress, formatter, dest, function(err) {
-        if (err) {
-          debug('error running export pipeline', err);
-          return dispatch(onError(err));
-        }
-        debug(
-          'done. %d docs exported to %s',
+      /**
+       * TODO: lucas: For metrics:
+       *
+       * "resource": "Export",
+       * "action": "completed",
+       * "file_type": "<csv|json_array>",
+       * "num_docs": "<how many docs exported>",
+       * "full_collection": true|false
+       * "filter": true|false,
+       * "projection": true|false,
+       * "skip": true|false,
+       * "limit": true|false,
+       * "fields_selected": true|false
+       */
+      dispatch(
+        globalAppRegistryEmit(
+          'export-finished',
           numDocsToExport,
-          exportData.fileName
-        );
-        dispatch(onFinished(numDocsToExport));
-        dispatch(
-          appRegistryEmit(
-            'export-finished',
-            numDocsToExport,
-            exportData.fileType
-          )
-        );
-
-        /**
-         * TODO: lucas: For metrics:
-         *
-         * "resource": "Export",
-         * "action": "completed",
-         * "file_type": "<csv|json_array>",
-         * "num_docs": "<how many docs exported>",
-         * "full_collection": true|false
-         * "filter": true|false,
-         * "projection": true|false,
-         * "skip": true|false,
-         * "limit": true|false,
-         * "fields_selected": true|false
-         */
-        dispatch(
-          globalAppRegistryEmit(
-            'export-finished',
-            numDocsToExport,
-            exportData.fileType
-          )
-        );
-      });
+          exportData.fileType
+        )
+      );
     });
   };
 };