feat: Add Lucene/SQL search filter to Service Map page

packages/app/src/DBServiceMapPage.tsx

@@ -1,15 +1,37 @@
-import { useEffect, useMemo, useState } from 'react';
+import { useCallback, useEffect, useMemo, useState } from 'react';
 import dynamic from 'next/dynamic';
 import Head from 'next/head';
-import { parseAsInteger, useQueryState } from 'nuqs';
+import {
+  parseAsArrayOf,
+  parseAsInteger,
+  parseAsString,
+  parseAsStringEnum,
+  useQueryState,
+  useQueryStates,
+} from 'nuqs';
 import { useForm, useWatch } from 'react-hook-form';
+import { tcFromSource } from '@hyperdx/common-utils/dist/core/metadata';
 import { SourceKind, TTraceSource } from '@hyperdx/common-utils/dist/types';
-import { Box, Button, Group, Modal, Slider, Text } from '@mantine/core';
+import {
+  Box,
+  Button,
+  Flex,
+  Group,
+  Modal,
+  MultiSelect,
+  Slider,
+  Text,
+} from '@mantine/core';
 import { IconConnection } from '@tabler/icons-react';
 
 import EmptyState from '@/components/EmptyState';
+import SearchWhereInput, {
+  getStoredLanguage,
+} from '@/components/SearchInput/SearchWhereInput';
 import { IS_LOCAL_MODE } from '@/config';
+import { useGetKeyValues } from '@/hooks/useMetadata';
 import { withAppNav } from '@/layout';
+import { parseAsStringEncoded } from '@/utils/queryParsers';
 
 import OnboardingModal from './components/OnboardingModal';
 import ServiceMap from './components/ServiceMap/ServiceMap';
@@ -51,13 +73,22 @@ const defaultTimeRange = parseTimeQuery(DEFAULT_INTERVAL, false) as [
   Date,
 ];
 
+const searchQueryStateMap = {
+  where: parseAsStringEncoded,
+  whereLanguage: parseAsStringEnum<'sql' | 'lucene'>(['sql', 'lucene']),
+  services: parseAsArrayOf(parseAsString),
+};
+
 function DBServiceMapPage() {
   const brandName = useBrandDisplayName();
 
   const { data: sources } = useSources();
   const [sourceId, setSourceId] = useQueryState('source');
   const [isCreateSourceModalOpen, setIsCreateSourceModalOpen] = useState(false);
 
+  const [searchedConfig, setSearchedConfig] =
+    useQueryStates(searchQueryStateMap);
+
   const [displayedTimeInputValue, setDisplayedTimeInputValue] =
     useState(DEFAULT_INTERVAL);
 
@@ -78,9 +109,12 @@ function DBServiceMapPage() {
         ) ?? defaultSource)
       : defaultSource;
 
-  const { control } = useForm({
+  const { control, handleSubmit, setValue } = useForm({
     values: {
       source: source?.id,
+      where: searchedConfig.where ?? '',
+      whereLanguage:
+        searchedConfig.whereLanguage ?? getStoredLanguage() ?? 'lucene',
     },
   });
 
@@ -92,6 +126,59 @@ function DBServiceMapPage() {
     }
   }, [watchedSource, sourceId, setSourceId]);
 
+  const sourceTableConnection = useMemo(() => tcFromSource(source), [source]);
+
+  const serviceNameKey = source?.serviceNameExpression ?? 'ServiceName';
+  const serviceNamesChartConfig = useMemo(
+    () =>
+      source
+        ? {
+            from: source.from,
+            connection: source.connection,
+            timestampValueExpression: source.timestampValueExpression,
+            where: '',
+            select: '',
+            dateRange: searchedTimeRange,
+          }
+        : undefined,
+    [source, searchedTimeRange],
+  );
+  const { data: serviceNameValues, isLoading: isServiceNamesLoading } =
+    useGetKeyValues(
+      {
+        chartConfig: serviceNamesChartConfig,
+        keys: [serviceNameKey],
+        disableRowLimit: true,
+        limit: 10000,
+      },
+      { enabled: !!source },
+    );
+  const serviceNameOptions = useMemo(
+    () =>
+      (serviceNameValues?.[0]?.value ?? [])
+        .map(v => String(v))
+        .sort((a, b) => a.localeCompare(b)),
+    [serviceNameValues],
+  );
+
+  const selectedServices = searchedConfig.services ?? [];
+  const setSelectedServices = useCallback(
+    (values: string[]) => {
+      setSearchedConfig(prev => ({
+        ...prev,
+        services: values.length > 0 ? values : null,
+      }));
+    },
+    [setSearchedConfig],
+  );
+
+  const onSubmit = useCallback(() => {
+    onSearch(displayedTimeInputValue);
+    handleSubmit(({ where, whereLanguage }) => {
+      setSearchedConfig(prev => ({ ...prev, where, whereLanguage }));
+    })();
+  }, [handleSubmit, setSearchedConfig, displayedTimeInputValue, onSearch]);
+
   const [samplingFactor, setSamplingFactor] = useQueryState(
     'samplingFactor',
     parseAsInteger.withDefault(10),
@@ -179,7 +266,7 @@ function DBServiceMapPage() {
       style={{ display: 'flex', flexDirection: 'column', height: '100vh' }}
     >
       {head}
-      <Group mb="md" justify="space-between">
+      <Group mb="xs" justify="space-between">
         <Group>
           <Text size="xl">Service Map</Text>
           <SourceSelectControlled
@@ -215,10 +302,51 @@ function DBServiceMapPage() {
           />
         </Group>
       </Group>
+      <Flex mb="sm" gap="sm" align="flex-start" wrap="wrap">
+        <MultiSelect
+          placeholder={
+            selectedServices.length === 0 ? 'All Services' : undefined
+          }
+          value={selectedServices}
+          data={serviceNameOptions}
+          onChange={setSelectedServices}
+          searchable
+          clearable
+          size="xs"
+          maxDropdownHeight={280}
+          disabled={isServiceNamesLoading}
+          variant="filled"
+          w={250}
+          limit={100}
+          data-testid="service-map-service-filter"
+        />
+        <SearchWhereInput
+          tableConnection={sourceTableConnection}
+          control={control}
+          name="where"
+          onSubmit={onSubmit}
+          onLanguageChange={lang =>
+            setValue('whereLanguage', lang, { shouldDirty: true })
+          }
+          enableHotkey
+          size="xs"
+          data-testid="service-map-search-input"
+          dateRange={searchedTimeRange}
+          sourceId={source?.id}
+          lucenePlaceholder="Filter spans w/ Lucene (ex. http.method:GET)"
+          sqlPlaceholder="SQL WHERE to filter spans (ex. Duration > 1000000)"
+          minWidth="min(500px, 100%)"
+        />
+      </Flex>
       <ServiceMap
         traceTableSource={source}
         dateRange={searchedTimeRange}
         samplingFactor={samplingFactor}
+        where={searchedConfig.where ?? undefined}
+        whereLanguage={searchedConfig.whereLanguage ?? undefined}
+        serviceNames={
+          selectedServices.length > 0 ? selectedServices : undefined
+        }
       />
     </Box>
   ) : null;

packages/app/src/components/ServiceMap/ServiceMap.tsx

@@ -246,6 +246,9 @@ interface ServiceMapProps {
   dateRange: [Date, Date];
   samplingFactor?: number;
   isSingleTrace?: boolean;
+  where?: string;
+  whereLanguage?: 'sql' | 'lucene';
+  serviceNames?: string[];
 }
 
 export default function ServiceMap({
@@ -254,6 +257,9 @@ export default function ServiceMap({
   dateRange,
   samplingFactor = 1,
   isSingleTrace,
+  where,
+  whereLanguage,
+  serviceNames,
 }: ServiceMapProps) {
   const {
     isLoading,
@@ -264,6 +270,9 @@ export default function ServiceMap({
     source: traceTableSource,
     dateRange,
     samplingFactor,
+    where,
+    whereLanguage,
+    serviceNames,
   });
 
   useEffect(() => {

packages/app/src/hooks/useServiceMap.tsx

@@ -22,12 +22,18 @@ async function getServiceMapQuery({
   traceId,
   metadata,
   samplingFactor,
+  where,
+  whereLanguage,
+  serviceNames,
 }: {
   source: TTraceSource;
   dateRange: [Date, Date];
   traceId?: string;
   metadata: Metadata;
   samplingFactor: number;
+  where?: string;
+  whereLanguage?: 'sql' | 'lucene';
+  serviceNames?: string[];
 }) {
   // Don't sample if we're looking for a specific trace
   const effectiveSamplingLevel = traceId ? 1 : samplingFactor;
@@ -37,6 +43,11 @@ async function getServiceMapQuery({
     connection: source.connection,
     dateRange,
     timestampValueExpression: source.timestampValueExpression,
+    ...(source.implicitColumnExpression != null
+      ? { implicitColumnExpression: source.implicitColumnExpression }
+      : {}),
+    where: where || '',
+    whereLanguage: (whereLanguage ?? 'lucene') as 'sql' | 'lucene',
     filters: [
       // Sample a subset of traces, for performance in the following join
       {
@@ -91,7 +102,6 @@ async function getServiceMapQuery({
             condition: `${source.spanKindExpression} IN ('Server', 'Consumer', 'SPAN_KIND_SERVER', 'SPAN_KIND_CONSUMER')`,
           },
         ],
-        where: '',
       },
       metadata,
       source.querySettings,
@@ -106,13 +116,22 @@ async function getServiceMapQuery({
             condition: `${source.spanKindExpression} IN ('Client', 'Producer', 'SPAN_KIND_CLIENT', 'SPAN_KIND_PRODUCER')`,
           },
         ],
-        where: '',
       },
       metadata,
       source.querySettings,
     ),
   ]);
 
+  const serviceNameInList = serviceNames?.length
+    ? { UNSAFE_RAW_SQL: serviceNames.map(s => SqlString.escape(s)).join(', ') }
+    : null;
+  const serviceNameFilter = serviceNameInList
+    ? chSql`AND (
+        ServerSpans.serviceName IN (${serviceNameInList})
+        OR ClientSpans.serviceName IN (${serviceNameInList})
+      )`
+    : chSql``;
+
   // Left join to support services which receive requests from clients that are not instrumented.
   // Ordering helps ensure stable graph layout.
   return chSql`
@@ -129,6 +148,7 @@ async function getServiceMapQuery({
         ON ServerSpans.traceId = ClientSpans.traceId
         AND ServerSpans.parentSpanId = ClientSpans.spanId
     WHERE (ClientSpans.serviceName IS NULL OR ServerSpans.serviceName != ClientSpans.serviceName)
+      ${serviceNameFilter}
     GROUP BY serverServiceName, serverStatusCode, clientServiceName
     ORDER BY serverServiceName, serverStatusCode, clientServiceName
   `;
@@ -246,24 +266,42 @@ export default function useServiceMap({
   dateRange,
   traceId,
   samplingFactor,
+  where,
+  whereLanguage,
+  serviceNames,
 }: {
   source: TTraceSource;
   dateRange: [Date, Date];
   traceId?: string;
   samplingFactor: number;
+  where?: string;
+  whereLanguage?: 'sql' | 'lucene';
+  serviceNames?: string[];
 }) {
   const client = useClickhouseClient();
   const metadata = useMetadataWithSettings();
 
   return useQuery({
-    queryKey: ['serviceMapData', traceId, source, dateRange, samplingFactor],
+    queryKey: [
+      'serviceMapData',
+      traceId,
+      source,
+      dateRange,
+      samplingFactor,
+      where,
+      whereLanguage,
+      serviceNames,
+    ],
     queryFn: async ({ signal }) => {
       const query = await getServiceMapQuery({
         source,
         dateRange,
         traceId,
         metadata,
         samplingFactor,
+        where,
+        whereLanguage,
+        serviceNames,
       });
 
       const data = await client