warm search indexes on search intent

apps/site/components/__tests__/withSearch.test.jsx

@@ -0,0 +1,67 @@
+import assert from 'node:assert/strict';
+import { describe, it, mock } from 'node:test';
+
+const fakeDb = {};
+const createMock = mock.fn(() => fakeDb);
+const insertMultipleMock = mock.fn(async () => {});
+const searchMock = mock.fn(async (_db, options) => ({ options }));
+
+mock.module('@orama/orama', {
+  namedExports: {
+    create: createMock,
+    insertMultiple: insertMultipleMock,
+    search: searchMock,
+  },
+});
+
+const { addPrefixToDocs, createOramaClient } = await import(
+  '#site/components/withSearch'
+);
+
+describe('withSearch', () => {
+  it('adds the locale prefix to indexed document URLs', () => {
+    const db = {
+      docs: {
+        docs: {
+          one: { href: '/docs', title: 'Docs' },
+        },
+      },
+    };
+
+    assert.deepEqual(addPrefixToDocs(db, '/en'), {
+      docs: {
+        docs: {
+          one: { href: '/en/docs', title: 'Docs' },
+        },
+      },
+    });
+  });
+
+  it('warms the Orama indexes only once before searching', async () => {
+    createMock.mock.resetCalls();
+    insertMultipleMock.mock.resetCalls();
+    searchMock.mock.resetCalls();
+
+    global.fetch = mock.fn(async url => ({
+      json: async () => ({
+        docs: {
+          docs: {
+            [String(url)]: { href: '/result', title: 'Node.js' },
+          },
+        },
+      }),
+    }));
+
+    const { client, warmup } = createOramaClient({
+      '/docs': 'https://example.com/docs.json',
+      '/learn': 'https://example.com/learn.json',
+    });
+
+    await Promise.all([warmup(), warmup()]);
+    await client.search({ term: 'node' });
+
+    assert.equal(global.fetch.mock.callCount(), 2);
+    assert.equal(insertMultipleMock.mock.callCount(), 2);
+    assert.equal(searchMock.mock.callCount(), 1);
+  });
+});

apps/site/components/withSearch.tsx

@@ -3,7 +3,7 @@
 import SearchBox from '@node-core/ui-components/Common/Search';
 import { create, insertMultiple, search } from '@orama/orama';
 import { useTranslations } from 'next-intl';
-import { useMemo, useRef } from 'react';
+import { useMemo } from 'react';
 
 import { ORAMA_DB_URLS } from '#site/next.constants.mjs';
 
@@ -19,6 +19,13 @@ type SerializedOramaDb = {
   };
 };
 
+type OramaUrls = typeof ORAMA_DB_URLS;
+
+type OramaClient = {
+  client: OramaCloud;
+  warmup: () => Promise<void>;
+};
+
 export const addPrefixToDocs = <T extends SerializedOramaDb>(
   db: T,
   prefix: string
@@ -35,9 +42,12 @@ export const addPrefixToDocs = <T extends SerializedOramaDb>(
   };
 };
 
-const loadOrama = async (db: AnyOrama): Promise<void> => {
+const loadOrama = async (
+  db: AnyOrama,
+  urls: OramaUrls = ORAMA_DB_URLS
+): Promise<void> => {
   const indexes = await Promise.all(
-    Object.entries(ORAMA_DB_URLS).map(async ([key, url]) => {
+    Object.entries(urls).map(async ([key, url]) => {
       const response = await fetch(url);
       const fetchedDb = (await response.json()) as SerializedOramaDb;
 
@@ -50,36 +60,43 @@ const loadOrama = async (db: AnyOrama): Promise<void> => {
   }
 };
 
-export const useOrama = () => {
-  const loadPromiseRef = useRef<Promise<void> | null>(null);
-
-  return useMemo(() => {
-    const db = create({
-      schema: {
-        title: 'string',
-        description: 'string',
-        href: 'string',
-        siteSection: 'string',
-      },
-    });
-
-    // @ts-expect-error We are overriding a method, an error is expected.
-    db.search = async options => {
-      await (loadPromiseRef.current ??= loadOrama(db));
-      return search(db, options);
-    };
-
-    return db;
-  }, []) as unknown as OramaCloud;
+export const createOramaClient = (
+  urls: OramaUrls = ORAMA_DB_URLS
+): OramaClient => {
+  const db = create({
+    schema: {
+      title: 'string',
+      description: 'string',
+      href: 'string',
+      siteSection: 'string',
+    },
+  });
+
+  let loadPromise: Promise<void> | null = null;
+  const warmup = () => (loadPromise ??= loadOrama(db, urls));
+
+  // @ts-expect-error We are overriding a method, an error is expected.
+  db.search = async options => {
+    await warmup();
+    return search(db, options);
+  };
+
+  return {
+    client: db as unknown as OramaCloud,
+    warmup,
+  };
 };
 
+export const useOrama = () => useMemo(() => createOramaClient(), []);
+
 const WithSearch: FC = () => {
   const t = useTranslations();
-  const client = useOrama();
+  const { client, warmup } = useOrama();
 
   return (
     <SearchBox
       client={client}
+      onWarmup={warmup}
       closeShortcutLabel={t('components.search.keyboardShortcuts.close')}
       navigateShortcutLabel={t('components.search.keyboardShortcuts.navigate')}
       noResultsTitle={t('components.search.noResultsFoundFor')}

packages/ui-components/src/Common/Search/Modal/__tests__/index.test.jsx

@@ -0,0 +1,46 @@
+import assert from 'node:assert/strict';
+import { describe, it } from 'node:test';
+
+import { fireEvent, render, screen } from '@testing-library/react';
+import userEvent from '@testing-library/user-event';
+
+import SearchModal from '..';
+
+describe('SearchModal', () => {
+  it('warms search when the trigger receives hover or focus intent', async t => {
+    const warmup = t.mock.fn();
+
+    render(
+      <SearchModal client={{}} placeholder="Search docs" onWarmup={warmup}>
+        <div>Results</div>
+      </SearchModal>
+    );
+
+    const trigger = screen.getByRole('button', { name: /search docs/i });
+
+    await userEvent.hover(trigger);
+    fireEvent.focus(trigger);
+
+    assert.equal(warmup.mock.callCount(), 2);
+  });
+
+  it('warms search when the keyboard shortcut is pressed', () => {
+    let warmupCalls = 0;
+
+    render(
+      <SearchModal
+        client={{}}
+        placeholder="Search docs"
+        onWarmup={() => {
+          warmupCalls += 1;
+        }}
+      >
+        <div>Results</div>
+      </SearchModal>
+    );
+
+    fireEvent.keyDown(document, { key: 'k', ctrlKey: true });
+
+    assert.equal(warmupCalls, 1);
+  });
+});

packages/ui-components/src/Common/Search/Modal/index.tsx

@@ -1,5 +1,6 @@
 import { MagnifyingGlassIcon } from '@heroicons/react/24/solid';
 import { SearchRoot, Modal } from '@orama/ui/components';
+import { useCallback, useEffect } from 'react';
 
 import SearchInput from '#ui/Common/Search/Input';
 
@@ -10,47 +11,88 @@ import styles from './index.module.css';
 
 type SearchModalProps = {
   client: OramaCloud | null;
+  onWarmup?: () => Promise<void> | void;
   placeholder: string;
 } & ComponentProps<typeof SearchInput>;
 
 const SearchModal: FC<PropsWithChildren<SearchModalProps>> = ({
   children,
   client,
+  onWarmup,
   placeholder,
-}) => (
-  <div className={styles.searchboxContainer}>
-    <Modal.Root>
-      <Modal.Trigger
-        type="button"
-        disabled={!client}
-        enableCmdK
-        className={styles.searchButton}
-      >
-        <div className={styles.searchButtonContent}>
-          <MagnifyingGlassIcon />
-          {placeholder}
-        </div>
-
-        <span className={styles.searchButtonShortcut}>⌘ K</span>
-      </Modal.Trigger>
-
-      <Modal.Wrapper
-        closeOnOutsideClick
-        closeOnEscape
-        className={styles.modalWrapper}
-      >
-        <SearchRoot client={client}>
-          <Modal.Inner className={styles.modalInner}>
-            <Modal.Content className={styles.modalContent}>
-              <SearchInput placeholder={placeholder} ariaLabel={placeholder} />
-              <Modal.Close className={styles.modalCloseButton} />
-              {children}
-            </Modal.Content>
-          </Modal.Inner>
-        </SearchRoot>
-      </Modal.Wrapper>
-    </Modal.Root>
-  </div>
-);
+}) => {
+  const warmup = useCallback(() => {
+    if (!onWarmup) {
+      return;
+    }
+
+    try {
+      const result = onWarmup();
+      Promise.resolve(result).catch((error) => {
+        console.error('Search warmup failed', error);
+      });
+    } catch (error) {
+      console.error('Search warmup failed', error);
+    }
+  }, [onWarmup]);
+
+  useEffect(() => {
+    if (!onWarmup) {
+      return;
+    }
+
+    const handleGlobalKeyDown = (event: KeyboardEvent) => {
+      if ((event.metaKey || event.ctrlKey) && event.key === 'k') {
+        warmup();
+      }
+    };
+
+    document.addEventListener('keydown', handleGlobalKeyDown);
+
+    return () => document.removeEventListener('keydown', handleGlobalKeyDown);
+  }, [onWarmup, warmup]);
+
+  return (
+    <div className={styles.searchboxContainer}>
+      <Modal.Root>
+        <Modal.Trigger
+          type="button"
+          disabled={!client}
+          enableCmdK
+          className={styles.searchButton}
+          onFocus={warmup}
+          onMouseEnter={warmup}
+          onPointerDown={warmup}
+        >
+          <div className={styles.searchButtonContent}>
+            <MagnifyingGlassIcon />
+            {placeholder}
+          </div>
+
+          <span className={styles.searchButtonShortcut}>⌘ K</span>
+        </Modal.Trigger>
+
+        <Modal.Wrapper
+          closeOnOutsideClick
+          closeOnEscape
+          className={styles.modalWrapper}
+        >
+          <SearchRoot client={client}>
+            <Modal.Inner className={styles.modalInner}>
+              <Modal.Content className={styles.modalContent}>
+                <SearchInput
+                  placeholder={placeholder}
+                  ariaLabel={placeholder}
+                />
+                <Modal.Close className={styles.modalCloseButton} />
+                {children}
+              </Modal.Content>
+            </Modal.Inner>
+          </SearchRoot>
+        </Modal.Wrapper>
+      </Modal.Root>
+    </div>
+  );
+};
 
 export default SearchModal;

packages/ui-components/src/Common/Search/index.tsx

@@ -17,6 +17,7 @@ type SearchBoxProps = {
   closeShortcutLabel?: string;
   navigateShortcutLabel?: string;
   noResultsTitle?: string;
+  onWarmup?: () => Promise<void> | void;
   placeholder?: string;
   selectShortcutLabel?: string;
 };
@@ -27,9 +28,10 @@ const SearchBox: React.FC<SearchBoxProps> = ({
   noResultsTitle = 'No results found for',
   closeShortcutLabel = 'to close',
   navigateShortcutLabel = 'to navigate',
+  onWarmup,
   selectShortcutLabel = 'to select',
 }) => (
-  <SearchModal client={client} placeholder={placeholder}>
+  <SearchModal client={client} placeholder={placeholder} onWarmup={onWarmup}>
     <div className={styles.searchResultsContainer}>
       <SearchResults
         noResultsTitle={noResultsTitle}