Autocomplete database names for tsh proxy db

packages/teleterm/src/ui/QuickInput/QuickInputList/QuickInputList.tsx

@@ -16,10 +16,10 @@ limitations under the License.
 
 import React, { useEffect, useRef } from 'react';
 import styled from 'styled-components';
-import { Box, Flex, Label } from 'design';
+import { Box, Flex, Label, Text } from 'design';
 import { makeLabelTag } from 'teleport/components/formatters';
 import * as types from 'teleterm/ui/services/quickInput/types';
-import { Cli, Server, Person } from 'design/Icon';
+import { Cli, Server, Person, Database } from 'design/Icon';
 
 const QuickInputList = React.forwardRef<HTMLElement, Props>((props, ref) => {
   const activeItemRef = useRef<HTMLDivElement>();
@@ -114,6 +114,34 @@ function ServerItem(props: { item: types.SuggestionServer }) {
   );
 }
 
+function DatabaseItem(props: { item: types.SuggestionDatabase }) {
+  const db = props.item.data;
+  const $labels = db.labelsList.map((label, index) => (
+    <Label mr="1" key={index} kind="secondary">
+      {makeLabelTag(label)}
+    </Label>
+  ));
+
+  return (
+    <Flex alignItems="center" p={1} minWidth="300px">
+      <SquareIconBackground color="#4DB2F0">
+        <Database fontSize="10px" />
+      </SquareIconBackground>
+      <Flex flexDirection="column" ml={1} flex={1}>
+        <Flex justifyContent="space-between" alignItems="center">
+          <Box mr={2}>{db.name}</Box>
+          <Box mr={2}>
+            <Text typography="body2" fontSize={0}>
+              {db.type}/{db.protocol}
+            </Text>
+          </Box>
+        </Flex>
+        <Box>{$labels}</Box>
+      </Flex>
+    </Flex>
+  );
+}
+
 function UnknownItem(props: { item: types.Suggestion }) {
   const { kind } = props.item;
   return <div>unknown kind: {kind} </div>;
@@ -163,6 +191,7 @@ const ComponentMap: Record<
   ['suggestion.cmd']: CmdItem,
   ['suggestion.ssh-login']: SshLoginItem,
   ['suggestion.server']: ServerItem,
+  ['suggestion.database']: DatabaseItem,
 };
 
 type Props = {

packages/teleterm/src/ui/QuickInput/useQuickInput.ts

@@ -23,11 +23,8 @@ import {
 } from 'teleterm/ui/services/quickInput/types';
 
 export default function useQuickInput() {
-  const {
-    quickInputService,
-    workspacesService,
-    commandLauncher,
-  } = useAppContext();
+  const { quickInputService, workspacesService, commandLauncher } =
+    useAppContext();
   workspacesService.useState();
   const documentsService =
     workspacesService.getActiveWorkspaceDocumentService();

packages/teleterm/src/ui/services/quickInput/quickInputService.test.ts

@@ -175,6 +175,54 @@ test('getAutocompleteResult returns correct result for an SSH login suggestion w
   });
 });
 
+test('getAutocompleteResult returns correct result for a database name suggestion', () => {
+  mockCommandLauncherAutocompleteCommands(CommandLauncherMock, [
+    {
+      name: 'autocomplete.tsh-proxy-db',
+      displayName: 'tsh proxy db',
+      description: '',
+      run: () => {},
+    },
+  ]);
+  jest
+    .spyOn(WorkspacesServiceMock.prototype, 'getActiveWorkspace')
+    .mockImplementation(() => ({
+      localClusterUri: 'test_uri',
+      documents: [],
+      location: '',
+    }));
+  jest
+    .spyOn(ClustersServiceMock.prototype, 'searchDbs')
+    .mockImplementation(() => {
+      return [
+        {
+          hostname: 'foobar',
+          uri: '',
+          name: '',
+          desc: '',
+          protocol: '',
+          type: '',
+          addr: '',
+          labelsList: null,
+        },
+      ];
+    });
+  const quickInputService = new QuickInputService(
+    new CommandLauncherMock(undefined),
+    new ClustersServiceMock(undefined),
+    new WorkspacesServiceMock(undefined, undefined, undefined)
+  );
+
+  const autocompleteResult =
+    quickInputService.getAutocompleteResult('tsh proxy db foo');
+  expect(autocompleteResult.kind).toBe('autocomplete.partial-match');
+  expect((autocompleteResult as AutocompletePartialMatch).targetToken).toEqual({
+    value: 'foo',
+    startIndex: 13,
+  });
+  expect(autocompleteResult.command).toEqual({ kind: 'command.unknown' });
+});
+
 test("getAutocompleteResult doesn't return any suggestions if the only suggestion completely matches the target token", () => {
   jest.mock('./quickPickers');
   const QuickCommandPickerMock = pickers.QuickCommandPicker as jest.MockedClass<

packages/teleterm/src/ui/services/quickInput/quickInputService.ts

@@ -50,14 +50,18 @@ export class QuickInputService extends Store<State> {
       workspacesService,
       clustersService
     );
+    const databasePicker = new pickers.QuickDatabasePicker(
+      workspacesService,
+      clustersService
+    );
 
     this.quickCommandPicker.registerPickerForCommand(
       'tsh ssh',
       new pickers.QuickTshSshPicker(sshLoginPicker, serverPicker)
     );
     this.quickCommandPicker.registerPickerForCommand(
       'tsh proxy db',
-      new pickers.QuickTshProxyDbPicker()
+      new pickers.QuickTshProxyDbPicker(databasePicker)
     );
   }
 

packages/teleterm/src/ui/services/quickInput/quickPickers.ts

@@ -23,6 +23,7 @@ import {
   SuggestionCmd,
   SuggestionServer,
   SuggestionSshLogin,
+  SuggestionDatabase,
   AutocompleteResult,
 } from './types';
 
@@ -238,11 +239,46 @@ export class QuickTshSshPicker implements QuickInputPicker {
   }
 }
 
-// TODO: Implement the rest of this class.
 export class QuickTshProxyDbPicker implements QuickInputPicker {
-  constructor() {}
+  private totalDbNameRegex = /^\S+$/i;
+
+  constructor(private databasePicker: QuickDatabasePicker) {}
+
+  getAutocompleteResult(
+    rawInput: string,
+    startIndex: number
+  ): AutocompleteResult {
+    // We can safely ignore any whitespace at the start. However, `startIndex` needs to account for
+    // any removed whitespace.
+    const input = rawInput.trimStart();
+    if (input === '') {
+      // input is empty, so rawInput must include only whitespace.
+      // Add length of the whitespace to startIndex.
+      startIndex += rawInput.length;
+    } else {
+      startIndex += rawInput.indexOf(input);
+    }
+
+    // Show autocomplete only after at least one space after `tsh proxy db`.
+    if (rawInput !== '' && input === '') {
+      return {
+        ...this.databasePicker.getAutocompleteResult('', startIndex),
+        command: { kind: 'command.unknown' },
+      };
+    }
+
+    const dbNameMatch = input.match(this.totalDbNameRegex);
+
+    if (dbNameMatch) {
+      return {
+        ...this.databasePicker.getAutocompleteResult(
+          dbNameMatch[0],
+          startIndex
+        ),
+        command: { kind: 'command.unknown' },
+      };
+    }
 
-  getAutocompleteResult(input: string): AutocompleteResult {
     return {
       kind: 'autocomplete.no-match',
       command: { kind: 'command.unknown' },
@@ -335,3 +371,40 @@ export class QuickServerPicker implements QuickInputPicker {
     };
   }
 }
+
+export class QuickDatabasePicker implements QuickInputPicker {
+  constructor(
+    private workspacesService: WorkspacesService,
+    private clustersService: ClustersService
+  ) {}
+
+  private filterDatabases(input: string): SuggestionDatabase[] {
+    const localClusterUri =
+      this.workspacesService.getActiveWorkspace()?.localClusterUri;
+    if (!localClusterUri) {
+      return [];
+    }
+    const databases = this.clustersService.searchDbs(localClusterUri, {
+      search: input,
+    });
+
+    return databases.map(database => ({
+      kind: 'suggestion.database' as const,
+      token: database.name,
+      data: database,
+    }));
+  }
+
+  getAutocompleteResult(input: string, startIndex: number): AutocompleteResult {
+    const suggestions = this.filterDatabases(input);
+    return {
+      kind: 'autocomplete.partial-match',
+      suggestions,
+      command: { kind: 'command.unknown' },
+      targetToken: {
+        startIndex,
+        value: input,
+      },
+    };
+  }
+}

packages/teleterm/src/ui/services/quickInput/types.ts

@@ -19,7 +19,16 @@ export type SuggestionSshLogin = SuggestionBase<
 
 export type SuggestionServer = SuggestionBase<'suggestion.server', tsh.Server>;
 
-export type Suggestion = SuggestionCmd | SuggestionSshLogin | SuggestionServer;
+export type SuggestionDatabase = SuggestionBase<
+  'suggestion.database',
+  tsh.Database
+>;
+
+export type Suggestion =
+  | SuggestionCmd
+  | SuggestionSshLogin
+  | SuggestionServer
+  | SuggestionDatabase;
 
 export type QuickInputPicker = {
   getAutocompleteResult(input: string, startIndex: number): AutocompleteResult;