feat(workspaces, databases, collections, collection): allow to open shell with initial state; add shell entrypoints to tab headers (#5999)

Author: gribnoysup

package-lock.json

@@ -51,10 +51,11 @@
     "@mongodb-js/compass-app-stores": "^7.21.0",
     "@mongodb-js/compass-components": "^1.27.0",
     "@mongodb-js/compass-connections": "^1.35.0",
-    "@mongodb-js/connection-info": "^0.5.0",
     "@mongodb-js/compass-logging": "^1.4.0",
     "@mongodb-js/compass-telemetry": "^1.1.0",
     "@mongodb-js/compass-workspaces": "^0.16.0",
+    "@mongodb-js/connection-info": "^0.5.0",
+    "@mongodb-js/mongodb-constants": "^0.10.2",
     "compass-preferences-model": "^2.24.0",
     "hadron-app-registry": "^9.2.0",
     "mongodb-collection-model": "^5.22.0",

packages/compass-collection/package.json

@@ -8,8 +8,9 @@ import {
 import { useConnectionInfo } from '@mongodb-js/compass-connections/provider';
 import { useOpenWorkspace } from '@mongodb-js/compass-workspaces/provider';
 import React from 'react';
-import { usePreference } from 'compass-preferences-model/provider';
+import { usePreferences } from 'compass-preferences-model/provider';
 import toNS from 'mongodb-ns';
+import { wrapField } from '@mongodb-js/mongodb-constants';
 
 const collectionHeaderActionsStyles = css({
   display: 'flex',
@@ -50,14 +51,41 @@ const CollectionHeaderActions: React.FunctionComponent<
   sourcePipeline,
 }: CollectionHeaderActionsProps) => {
   const { id: connectionId, atlasMetadata } = useConnectionInfo();
-  const { openCollectionWorkspace, openEditViewWorkspace } = useOpenWorkspace();
-  const preferencesReadOnly = usePreference('readOnly');
+  const { openCollectionWorkspace, openEditViewWorkspace, openShellWorkspace } =
+    useOpenWorkspace();
+  const {
+    readOnly: preferencesReadOnly,
+    enableShell,
+    enableNewMultipleConnectionSystem,
+  } = usePreferences([
+    'readOnly',
+    'enableShell',
+    'enableNewMultipleConnectionSystem',
+  ]);
+
+  const { database, collection } = toNS(namespace);
+
+  const showOpenShellButton = enableShell && enableNewMultipleConnectionSystem;
 
   return (
     <div
       className={collectionHeaderActionsStyles}
       data-testid="collection-header-actions"
     >
+      {showOpenShellButton && (
+        <Button
+          size="small"
+          onClick={() => {
+            openShellWorkspace(connectionId, {
+              initialEvaluate: `use ${database}`,
+              initialInput: `db[${wrapField(collection, true)}].find()`,
+            });
+          }}
+          leftGlyph={<Icon glyph="Shell"></Icon>}
+        >
+          Open MongoDB shell
+        </Button>
+      )}
       {atlasMetadata && (
         <Button
           data-testid="collection-header-visualize-your-data"

packages/compass-collection/src/components/collection-header-actions/collection-header-actions.tsx

@@ -5,11 +5,14 @@ import {
   createElement,
   createContext,
   useContext,
+  useRef,
+  useCallback,
 } from 'react';
 import { type AllPreferences } from './';
 import type { PreferencesAccess } from './preferences';
 import { ReadOnlyPreferenceAccess } from './read-only-preferences-access';
 import { createServiceLocator } from 'hadron-app-registry';
+import { pick } from 'lodash';
 
 const PreferencesContext = createContext<PreferencesAccess>(
   // Our context starts with our read-only preference access but we expect
@@ -44,11 +47,64 @@ export function usePreference<K extends keyof AllPreferences>(
   return value;
 }
 
+/**
+ * A version of usePreference hook that allows to get multiple preferences in a
+ * single object for brevity
+ *
+ * @example
+ * const {
+ *   enableShell,
+ *   readOnly,
+ * } = usePreferences(['enableShell', 'readOnly']);
+ *
+ * @param keys list of preferences keys to return as an object
+ */
+export function usePreferences<K extends (keyof AllPreferences)[]>(
+  keys: K
+): Pick<AllPreferences, K[number]> {
+  const preferences = usePreferencesContext();
+  const keysRef = useRef(keys);
+  keysRef.current = keys;
+  const [values, setValues] = useState(() => {
+    return pick(preferences.getPreferences(), keys);
+  });
+  const updateValue = useCallback((key: keyof AllPreferences, newValue) => {
+    setValues((values) => {
+      if (newValue === values[key]) {
+        return values;
+      }
+      return {
+        ...values,
+        [key]: newValue,
+      };
+    });
+  }, []);
+  useEffect(() => {
+    const unsubscribe = keysRef.current.map((key) => {
+      return preferences.onPreferenceValueChanged(key, (newValue) => {
+        updateValue(key, newValue);
+      });
+    });
+    return () => {
+      unsubscribe.forEach((fn) => {
+        fn();
+      });
+    };
+  }, [
+    // An easy way to depend on actual array values instead of the array itself
+    // and avoid extra effect calls if array defined inline
+    keysRef.current.join(''),
+    preferences,
+  ]);
+  return values;
+}
+
 type FirstArgument<F> = F extends (...args: [infer A, ...any]) => any
   ? A
   : F extends { new (...args: [infer A, ...any]): any }
   ? A
   : never;
+
 type OptionalOmit<T, K extends keyof T> = Omit<T, K> & Partial<Pick<T, K>>;
 
 /** Use as: const WrappedComponent = withPreferences(Component, ['enableMaps', 'trackUsageStatistics'], React); */

packages/compass-preferences-model/src/react.ts

@@ -1,19 +1,23 @@
 import { connect } from 'react-redux';
-import React, { useCallback, useRef, useState } from 'react';
-import { useTabState } from '@mongodb-js/compass-workspaces/provider';
+import React, { useCallback, useEffect, useRef, useState } from 'react';
+import {
+  useOnTabReplace,
+  useTabState,
+} from '@mongodb-js/compass-workspaces/provider';
 import {
   Banner,
   Link,
   css,
   getScrollbarStyles,
   palette,
+  rafraf,
   spacing,
 } from '@mongodb-js/compass-components';
 import type { WorkerRuntime } from '@mongosh/node-runtime-worker-thread';
 import ShellInfoModal from '../shell-info-modal';
 import ShellHeader from '../shell-header/shell-header';
 import { usePreference } from 'compass-preferences-model/provider';
-import { Shell } from '@mongosh/browser-repl';
+import { Shell as _Shell } from '@mongosh/browser-repl';
 import type { RootState } from '../../stores/store';
 import { selectRuntimeById, saveHistory } from '../../stores/store';
 
@@ -40,22 +44,70 @@ const compassShellContainerStyles = css({
   borderTop: `1px solid ${palette.gray.dark2}`,
 });
 
-type ShellProps = React.ComponentProps<typeof Shell>;
+type ShellProps = React.ComponentProps<typeof _Shell>;
 
 type ShellRef = Extract<Required<ShellProps>['ref'], { current: any }>;
 
+type ShellType = ShellRef['current'];
+
 type ShellOutputEntry = Required<ShellProps>['initialOutput'][number];
 
 type CompassShellProps = {
   runtime: WorkerRuntime | null;
   initialHistory: string[] | null;
   onHistoryChange: (history: string[]) => void;
+  initialEvaluate?: string | string[];
+  initialInput?: string;
 };
 
+function useInitialEval(initialEvaluate?: string | string[]) {
+  const [initialEvalApplied, setInitialEvalApplied] = useTabState(
+    'initialEvalApplied',
+    false
+  );
+  useEffect(() => {
+    setInitialEvalApplied(true);
+  }, [setInitialEvalApplied]);
+  return initialEvalApplied ? undefined : initialEvaluate;
+}
+
+const Shell = React.forwardRef<ShellType, ShellProps>(function Shell(
+  { initialEvaluate: _initialEvaluate, ...props },
+  ref
+) {
+  const shellRef = useRef<ShellType | null>(null);
+  const initialEvaluate = useInitialEval(_initialEvaluate);
+  const mergeRef = useCallback(
+    (shell: ShellType | null) => {
+      shellRef.current = shell;
+      if (typeof ref === 'function') {
+        ref(shell);
+      } else if (ref) {
+        ref.current = shell;
+      }
+    },
+    [ref]
+  );
+  useEffect(() => {
+    return rafraf(() => {
+      shellRef.current?.focusEditor();
+    });
+  }, []);
+  return (
+    <_Shell
+      ref={mergeRef}
+      initialEvaluate={initialEvaluate}
+      {...props}
+    ></_Shell>
+  );
+});
+
 const CompassShell: React.FC<CompassShellProps> = ({
   runtime,
   initialHistory,
   onHistoryChange,
+  initialEvaluate,
+  initialInput,
 }) => {
   const enableShell = usePreference('enableShell');
   const shellRef: ShellRef = useRef(null);
@@ -64,7 +116,16 @@ const CompassShell: React.FC<CompassShellProps> = ({
   const [shellOutput, setShellOutput] = useTabState<
     readonly ShellOutputEntry[]
   >('shellOutput', []);
-  const [shellInput, setShellInput] = useTabState('shellInput', '');
+  const [shellInput, setShellInput] = useTabState(
+    'shellInput',
+    initialInput ?? ''
+  );
+
+  useOnTabReplace(() => {
+    // Never allow to replace the shell tab to avoid destroying the runtime
+    // unnecessarily
+    return false;
+  });
 
   const showInfoModal = useCallback(() => {
     setInfoModalVisible(true);
@@ -78,7 +139,7 @@ const CompassShell: React.FC<CompassShellProps> = ({
     if (shellRef.current && window.getSelection()?.type !== 'Range') {
       shellRef.current.focusEditor();
     }
-  }, [shellRef]);
+  }, []);
 
   const updateShellOutput = useCallback(
     (output: readonly ShellOutputEntry[]) => {
@@ -95,6 +156,8 @@ const CompassShell: React.FC<CompassShellProps> = ({
     setIsOperationInProgress(false);
   }, []);
 
+  const canRenderShell = enableShell && initialHistory && runtime;
+
   if (!enableShell) {
     return (
       <div className={infoBannerContainerStyles}>
@@ -110,7 +173,7 @@ const CompassShell: React.FC<CompassShellProps> = ({
     );
   }
 
-  if (!runtime || !initialHistory) {
+  if (!canRenderShell) {
     return <div className={compassShellStyles} />;
   }
 
@@ -140,6 +203,7 @@ const CompassShell: React.FC<CompassShellProps> = ({
           <Shell
             ref={shellRef}
             runtime={runtime}
+            initialEvaluate={initialEvaluate}
             initialInput={shellInput}
             onInputChanged={setShellInput}
             initialOutput={shellOutput}
@@ -150,6 +214,8 @@ const CompassShell: React.FC<CompassShellProps> = ({
             }}
             onOperationStarted={notifyOperationStarted}
             onOperationEnd={notifyOperationEnd}
+            maxOutputLength={1000}
+            maxHistoryLength={1000}
           />
         </div>
       </div>

packages/compass-shell/src/components/compass-shell/tab-compass-shell.tsx

@@ -23,14 +23,19 @@ import { Theme, ThemeProvider } from '@mongodb-js/compass-components';
 
 const SHELL_THEME = { theme: Theme.Dark, enabled: true };
 
-export function ShellPlugin() {
+type ShellPluginProps = {
+  initialEvaluate?: string | string[];
+  initialInput?: string;
+};
+
+export function ShellPlugin(props: ShellPluginProps) {
   const multiConnectionsEnabled = usePreference(
     'enableNewMultipleConnectionSystem'
   );
   const ShellComponent = multiConnectionsEnabled ? TabShell : Shell;
   return (
     <ThemeProvider theme={SHELL_THEME}>
-      <ShellComponent />
+      <ShellComponent {...props} />
     </ThemeProvider>
   );
 }
@@ -48,7 +53,7 @@ export type ShellPluginExtraArgs = ShellPluginServices & {
 };
 
 export function onActivated(
-  _initialProps: unknown,
+  _initialProps: ShellPluginProps,
   services: ShellPluginServices,
   { addCleanup, cleanup }: ActivateHelpers
 ) {