feat: add connection variables for raw SQL tiles and connection-scoped data source variables

.changeset/source-vars-5.md

@@ -0,0 +1,5 @@
+---
+'@hyperdx/app': minor
+---
+
+feat: add connection variables for raw SQL tiles and connection-scoped data source variables

packages/app/src/DBDashboardImportPage.tsx

@@ -20,6 +20,7 @@ import {
   type DashboardFilterQueryIssue,
   getAdhocFilterSourceKinds,
   isAdhocFilter,
+  isDataSourceFilter,
   isPrometheusLabelFilter,
   isQueryExpressionFilter,
   QUERY_EXPRESSION_FILTER_SOURCE_KINDS,
@@ -75,6 +76,7 @@ import {
 
 import SelectControlled from './components/SelectControlled';
 import { SourceMultiSelectControlled } from './components/SourceMultiSelect';
+import { useConnectionItems } from './hooks/useConnectionItems';
 import api from './api';
 import { usePageTitle } from './brand';
 import { useConnections } from './connection';
@@ -392,8 +394,8 @@ function resolveAppliesToSources(
 }
 
 /**
- * A data source variable reference (`$name`) is portable as written, so
- * import neither maps it nor asks the user to.
+ * A data source or connection variable reference (`$name`) is portable as
+ * written, so import neither maps it nor asks the user to.
  */
 const isVariableRef = (name: string | undefined) =>
   parseVariableReference(name) != null;
@@ -411,6 +413,15 @@ const isSourceBackedTemplateFilter = (
   (isQueryExpressionFilter(filter) || isPrometheusLabelFilter(filter)) &&
   !isVariableRef(filter.source);
 
+/** A data source variable limited to a connection that import must map. */
+const connectionBackedTemplateFilter = (filter: DashboardFilter | undefined) =>
+  filter &&
+  isDataSourceFilter(filter) &&
+  filter.connection &&
+  !isVariableRef(filter.connection)
+    ? filter
+    : undefined;
+
 /** Returns the filter if it names a source, or `undefined` if it doesn't. */
 const sourceBackedTemplateFilter = (filter: DashboardFilter | undefined) =>
   filter && isSourceBackedTemplateFilter(filter) ? filter : undefined;
@@ -473,6 +484,8 @@ export function buildMappingFormSchema(
       connectionMappings: z.array(MappingValueSchema),
       /** A list of filter source mappings, ordered by input filter index */
       filterSourceMappings: z.array(MappingValueSchema).optional(),
+      /** Connection mappings of data source variables, ordered by input filter index */
+      filterConnectionMappings: z.array(MappingValueSchema).optional(),
       /**
        * Per-filter applies-to source mappings. Each entry is an array of mapped
        * source IDs in the same order as the filter's `appliesToSourceIds` names
@@ -514,7 +527,11 @@ export function buildMappingFormSchema(
         }
         // A raw SQL tile always gets a connection row, and its connection is
         // not optional — leaving it unpicked used to throw on `connection!.id`.
-        if (isRawSqlSavedChartConfig(config) && !data.connectionMappings[idx]) {
+        if (
+          isRawSqlSavedChartConfig(config) &&
+          !isVariableRef(config.connection) &&
+          !data.connectionMappings[idx]
+        ) {
           ctx.addIssue({
             code: z.ZodIssueCode.custom,
             path: ['connectionMappings', idx],
@@ -530,6 +547,16 @@ export function buildMappingFormSchema(
         });
 
       input.filters?.forEach((filter, idx) => {
+        if (
+          connectionBackedTemplateFilter(filter) &&
+          !data.filterConnectionMappings?.[idx]
+        ) {
+          ctx.addIssue({
+            code: z.ZodIssueCode.custom,
+            path: ['filterConnectionMappings', idx],
+            message: 'Select a connection for this filter',
+          });
+        }
         if (isAdhocFilter(filter)) {
           const kinds = getAdhocFilterSourceKinds(filter.sourceType);
           const kindMessage =
@@ -601,6 +628,7 @@ export function Mapping({ input }: { input: DashboardTemplate }) {
   const router = useRouter();
   const { data: sources } = useSources();
   const { data: connections } = useConnections();
+  const connectionItems = useConnectionItems();
   const { data: dashboards } = useDashboards();
   const { data: existingTags } = api.useTags();
   const [dashboardId] = useQueryState('dashboardId', parseAsString);
@@ -620,6 +648,7 @@ export function Mapping({ input }: { input: DashboardTemplate }) {
         tileSourceMappings: input.tiles.map(() => ''),
         connectionMappings: input.tiles.map(() => ''),
         filterSourceMappings: input.filters?.map(() => '') ?? [],
+        filterConnectionMappings: input.filters?.map(() => '') ?? [],
         filterAppliesToSourceMappings: input.filters?.map(() => []) ?? [],
         filterAdhocSourceMappings: input.filters?.map(() => []) ?? [],
         onClickSourceMappings: input.tiles.map(() => ''),
@@ -650,13 +679,20 @@ export function Mapping({ input }: { input: DashboardTemplate }) {
       return match?.id || '';
     });
 
+    const findConnectionIdByName = (name: string) =>
+      connections.find(conn => conn.name.toLowerCase() === name.toLowerCase())
+        ?.id ?? '';
+
     const connectionMappings = input.tiles.map(tile => {
       const config = tile.config as SavedChartConfig;
-      if (!isRawSqlSavedChartConfig(config)) return '';
-      const match = connections.find(
-        conn => conn.name.toLowerCase() === config.connection.toLowerCase(),
-      );
-      return match?.id || '';
+      if (!isRawSqlSavedChartConfig(config) || isVariableRef(config.connection))
+        return '';
+      return findConnectionIdByName(config.connection);
+    });
+
+    const filterConnectionMappings = input.filters?.map(filter => {
+      const connection = connectionBackedTemplateFilter(filter)?.connection;
+      return connection ? findConnectionIdByName(connection) : '';
     });
 
     const filterSourceMappings = input.filters?.map(filter => {
@@ -716,6 +752,7 @@ export function Mapping({ input }: { input: DashboardTemplate }) {
 
     setValue('tileSourceMappings', tileSourceMappings);
     setValue('connectionMappings', connectionMappings);
+    setValue('filterConnectionMappings', filterConnectionMappings);
     setValue('filterSourceMappings', filterSourceMappings);
     multiMappingSlotsRef.current.clear();
     filterAppliesToSourceMappings?.forEach((mapping, idx) => {
@@ -1035,12 +1072,14 @@ export function Mapping({ input }: { input: DashboardTemplate }) {
         };
 
         if (isRawSqlSavedChartConfig(tile.config)) {
-          const connection = findConnection(data.connectionMappings[idx]);
+          const connection = isVariableRef(tile.config.connection)
+            ? tile.config.connection
+            : findConnection(data.connectionMappings[idx])!.id;
           return {
             ...tile,
             config: applyOnClick({
               ...tile.config,
-              connection: connection!.id,
+              connection,
               ...(source ? { source: source.id } : {}),
             }),
           };
@@ -1082,6 +1121,13 @@ export function Mapping({ input }: { input: DashboardTemplate }) {
             appliesToSourceIds,
           };
         }
+        if (connectionBackedTemplateFilter(filter)) {
+          return {
+            ...filter,
+            connection: findConnection(data.filterConnectionMappings?.[idx])
+              ?.id,
+          };
+        }
         const source = findSource(data.filterSourceMappings?.[idx]);
         return {
           ...filter,
@@ -1192,7 +1238,7 @@ export function Mapping({ input }: { input: DashboardTemplate }) {
                     </Table.Tr>
                   )}
                   {/** Mapping for the tile's connection, if it's a raw sql tile */}
-                  {isRawSql && (
+                  {isRawSql && !isVariableRef(config.connection) && (
                     <Table.Tr>
                       <Table.Td>{tile.config.name}</Table.Td>
                       <Table.Td>Data connection</Table.Td>
@@ -1201,10 +1247,7 @@ export function Mapping({ input }: { input: DashboardTemplate }) {
                         <SelectControlled
                           control={control}
                           name={`connectionMappings.${i}`}
-                          data={connections?.map(conn => ({
-                            value: conn.id,
-                            label: conn.name,
-                          }))}
+                          data={connectionItems}
                           placeholder="Select a connection"
                         />
                       </Table.Td>
@@ -1264,6 +1307,25 @@ export function Mapping({ input }: { input: DashboardTemplate }) {
                 : undefined;
               return (
                 <Fragment key={filter.id}>
+                  {connectionBackedTemplateFilter(filter) && (
+                    <Table.Tr>
+                      <Table.Td>{filter.name} (Filter)</Table.Td>
+                      <Table.Td>Data connection</Table.Td>
+                      <Table.Td>
+                        {connectionBackedTemplateFilter(filter)?.connection}
+                      </Table.Td>
+                      <Table.Td>
+                        <SelectControlled
+                          control={control}
+                          name={`filterConnectionMappings.${i}`}
+                          data={connectionItems}
+                          placeholder="Select a connection"
+                        />
+                      </Table.Td>
+                      <Table.Td />
+                      <Table.Td />
+                    </Table.Tr>
+                  )}
                   {isSourceBackedTemplateFilter(filter) && (
                     <Table.Tr>
                       <Table.Td>{filter.name} (Filter)</Table.Td>

packages/app/src/DashboardFilters.tsx

@@ -10,6 +10,7 @@ import {
   isFilterRequired,
   isFilterVariableEnabled,
   isQueryExpressionFilter,
+  isReferenceVariableFilter,
   isStaticListFilter,
 } from '@hyperdx/common-utils/dist/filters';
 import { ReferenceVariableValues } from '@hyperdx/common-utils/dist/referenceVariables';
@@ -30,7 +31,10 @@ import {
   getPendingVariablesTooltip,
 } from './components/DashboardFilterLabel';
 import { FilterLinkToggle } from './components/FilterLinkToggle';
-import { DataSourceDashboardFilter } from './components/ReferenceVariableDashboardFilter';
+import {
+  ConnectionDashboardFilter,
+  DataSourceDashboardFilter,
+} from './components/ReferenceVariableDashboardFilter';
 import { VirtualMultiSelect } from './components/VirtualMultiSelect/VirtualMultiSelect';
 import { useDashboardFilterValues } from './hooks/useDashboardFilterValues';
 
@@ -237,20 +241,28 @@ const DashboardFilters = ({
       {filters.map(filter => {
         const unresolvedVariables =
           unresolvedReferenceVariablesByFilterId?.get(filter.id) ?? [];
-        if (isDataSourceFilter(filter)) {
-          return (
+        if (isReferenceVariableFilter(filter)) {
+          const referenceProps = {
+            value: referenceVariableValues?.get(
+              getFilterVariableName(filter) ?? '',
+            ),
+            onChange: (value: string | undefined) =>
+              onSetFilterValue(filter.id, value ? [value] : []),
+            effect: getFilterEffect(filter),
+            requiredTooltip: getRequiredFilterTooltip(filter),
+          };
+          return isDataSourceFilter(filter) ? (
             <DataSourceDashboardFilter
               key={filter.id}
               filter={filter}
               pendingVariables={unresolvedVariables}
-              value={referenceVariableValues?.get(
-                getFilterVariableName(filter) ?? '',
-              )}
-              onChange={value =>
-                onSetFilterValue(filter.id, value ? [value] : [])
-              }
-              effect={getFilterEffect(filter)}
-              requiredTooltip={getRequiredFilterTooltip(filter)}
+              {...referenceProps}
+            />
+          ) : (
+            <ConnectionDashboardFilter
+              key={filter.id}
+              filter={filter}
+              {...referenceProps}
             />
           );
         }

packages/app/src/__tests__/DBDashboardImportPage.test.tsx

@@ -720,3 +720,67 @@ describe('Dashboard import - data source variables', () => {
     expect(payload.filters[2].sources).toEqual(['$logs', 'src-traces']);
   });
 });
+
+describe('Dashboard import - connection variables', () => {
+  const connectionVariableTemplate = DashboardTemplateSchema.parse({
+    version: '0.1.0',
+    name: 'Connection Variables',
+    tiles: [
+      {
+        id: 'tile-rawsql',
+        x: 0,
+        y: 0,
+        w: 6,
+        h: 3,
+        config: {
+          name: 'tile-rawsql',
+          displayType: 'table',
+          configType: 'sql',
+          connection: '$conn',
+          sqlTemplate: 'SELECT 1',
+        },
+      },
+    ],
+    filters: [
+      {
+        id: 'conn',
+        type: 'CONNECTION',
+        name: 'Cluster',
+        variableName: 'conn',
+        minSelections: 1,
+        maxSelections: 1,
+        isBroadcastEnabled: false,
+        isVariableEnabled: true,
+      },
+      {
+        id: 'ds',
+        type: 'DATA_SOURCE',
+        name: 'Logs source',
+        variableName: 'logs',
+        allowedSourceKinds: [SourceKind.Log],
+        connection: 'Default',
+        minSelections: 1,
+        maxSelections: 1,
+        isBroadcastEnabled: false,
+        isVariableEnabled: true,
+      },
+    ],
+  });
+
+  it('keeps connection references and maps a data source variable connection', async () => {
+    renderWithMantine(<Mapping input={connectionVariableTemplate} />);
+
+    const finish = await screen.findByRole('button', {
+      name: /finish import/i,
+    });
+    expect(screen.queryByText('$conn')).not.toBeInTheDocument();
+    await act(async () => {
+      fireEvent.click(finish);
+    });
+
+    await waitFor(() => expect(mockMutateAsync).toHaveBeenCalledTimes(1));
+    const payload = mockMutateAsync.mock.calls[0][0];
+    expect(payload.tiles[0].config.connection).toBe('$conn');
+    expect(payload.filters[1].connection).toBe('conn-default');
+  });
+});