fix incorrect split for select query

packages/app/src/DBSearchPage.tsx

@@ -26,7 +26,7 @@ import {
   DisplayType,
   Filter,
 } from '@hyperdx/common-utils/dist/types';
-import { splitAndTrimCSV } from '@hyperdx/common-utils/dist/utils';
+import { splitAndTrimWithBracket } from '@hyperdx/common-utils/dist/utils';
 import {
   ActionIcon,
   Box,
@@ -835,7 +835,7 @@ function DBSearchPage() {
     };
   }, [chartConfig, searchedTimeRange]);
 
-  const displayedColumns = splitAndTrimCSV(
+  const displayedColumns = splitAndTrimWithBracket(
     dbSqlRowTableConfig?.select ??
       searchedSource?.defaultTableSelectExpression ??
       '',

packages/app/src/components/DBRowTable.tsx

@@ -25,7 +25,7 @@ import {
   ChartConfigWithDateRange,
   SelectList,
 } from '@hyperdx/common-utils/dist/types';
-import { Granularity, splitAndTrimCSV } from '@hyperdx/common-utils/dist/utils';
+import { splitAndTrimWithBracket } from '@hyperdx/common-utils/dist/utils';
 import { Box, Code, Flex, Text } from '@mantine/core';
 import { FetchNextPageOptions } from '@tanstack/react-query';
 import {
@@ -810,10 +810,10 @@ function appendSelectWithPrimaryAndPartitionKey(
     .map(k => extractColumnReference(k.trim()))
     .filter((k): k is string => k != null && k.length > 0);
   const primaryKeyArr =
-    primaryKeys.trim() !== '' ? splitAndTrimCSV(primaryKeys) : [];
+    primaryKeys.trim() !== '' ? splitAndTrimWithBracket(primaryKeys) : [];
   const allKeys = [...partitionKeyArr, ...primaryKeyArr];
   if (typeof select === 'string') {
-    const selectSplit = splitAndTrimCSV(select);
+    const selectSplit = splitAndTrimWithBracket(select);
     const selectColumns = new Set(selectSplit);
     const additionalKeys = allKeys.filter(k => !selectColumns.has(k));
     return {

packages/app/src/source.ts

@@ -8,7 +8,10 @@ import {
   JSDataType,
 } from '@hyperdx/common-utils/dist/clickhouse';
 import { MetricsDataType, TSource } from '@hyperdx/common-utils/dist/types';
-import { hashCode, splitAndTrimCSV } from '@hyperdx/common-utils/dist/utils';
+import {
+  hashCode,
+  splitAndTrimWithBracket,
+} from '@hyperdx/common-utils/dist/utils';
 import { useMutation, useQuery, useQueryClient } from '@tanstack/react-query';
 
 import { hdxServer } from '@/api';
@@ -43,7 +46,7 @@ function getLocalSources(): TSource[] {
 // If a user specifies a timestampValueExpression with multiple columns,
 // this will return the first one. We'll want to refine this over time
 export function getFirstTimestampValueExpression(valueExpression: string) {
-  return splitAndTrimCSV(valueExpression)[0];
+  return splitAndTrimWithBracket(valueExpression)[0];
 }
 
 export function getSpanEventBody(eventModel: TSource) {
@@ -65,7 +68,7 @@ export function getEventBody(eventModel: TSource) {
       : undefined) ??
     eventModel.implicitColumnExpression; //??
   // (eventModel.kind === 'log' ? 'Body' : 'SpanName')
-  const multiExpr = splitAndTrimCSV(expression ?? '');
+  const multiExpr = splitAndTrimWithBracket(expression ?? '');
   return multiExpr.length === 1 ? expression : multiExpr[0]; // TODO: check if we want to show multiple columns
 }
 
@@ -215,7 +218,7 @@ export async function inferTableSourceConfig({
       connectionId,
     })
   ).primary_key;
-  const keys = splitAndTrimCSV(primaryKeys);
+  const keys = splitAndTrimWithBracket(primaryKeys);
 
   const isOtelLogSchema = hasAllColumns(columns, [
     'Timestamp',

packages/common-utils/src/__tests__/utils.test.ts

@@ -1,4 +1,4 @@
-import { splitAndTrimCSV } from '../utils';
+import { splitAndTrimCSV, splitAndTrimWithBracket } from '../utils';
 
 describe('utils', () => {
   describe('splitAndTrimCSV', () => {
@@ -26,4 +26,143 @@ describe('utils', () => {
       expect(splitAndTrimCSV(',,  ,,')).toEqual([]);
     });
   });
+
+  describe('splitAndTrimWithBracket', () => {
+    it('should split a simple comma-separated string', () => {
+      const input = 'column1, column2, column3';
+      const expected = ['column1', 'column2', 'column3'];
+      expect(splitAndTrimWithBracket(input)).toEqual(expected);
+    });
+
+    it('should handle function calls with commas in parameters', () => {
+      const input =
+        "Timestamp, ServiceName, JSONExtractString(Body, 'c'), JSONExtractString(Body, 'msg')";
+      const expected = [
+        'Timestamp',
+        'ServiceName',
+        "JSONExtractString(Body, 'c')",
+        "JSONExtractString(Body, 'msg')",
+      ];
+      expect(splitAndTrimWithBracket(input)).toEqual(expected);
+    });
+
+    it('should handle nested function calls', () => {
+      const input = 'col1, func1(a, b), col2, func2(c, func3(d, e)), col3';
+      const expected = [
+        'col1',
+        'func1(a, b)',
+        'col2',
+        'func2(c, func3(d, e))',
+        'col3',
+      ];
+      expect(splitAndTrimWithBracket(input)).toEqual(expected);
+    });
+
+    it('should handle square brackets in column expressions', () => {
+      const input = "col1, array[1, 2, 3], jsonb_path_query(data, '$[*]')";
+      const expected = [
+        'col1',
+        'array[1, 2, 3]',
+        "jsonb_path_query(data, '$[*]')",
+      ];
+      expect(splitAndTrimWithBracket(input)).toEqual(expected);
+    });
+
+    it('should handle mixed parentheses and square brackets', () => {
+      const input = "col1, func(array[1, 2], obj['key']), col2['nested'][0]";
+      const expected = [
+        'col1',
+        "func(array[1, 2], obj['key'])",
+        "col2['nested'][0]",
+      ];
+      expect(splitAndTrimWithBracket(input)).toEqual(expected);
+    });
+
+    it('should trim whitespace from resulting columns', () => {
+      const input = '  col1  ,   func(a, b)  ,  col2  ';
+      const expected = ['col1', 'func(a, b)', 'col2'];
+      expect(splitAndTrimWithBracket(input)).toEqual(expected);
+    });
+
+    it('should handle empty input', () => {
+      expect(splitAndTrimWithBracket('')).toEqual([]);
+    });
+
+    it('should handle input with only spaces', () => {
+      expect(splitAndTrimWithBracket('   ')).toEqual([]);
+    });
+
+    it('should skip empty elements', () => {
+      const input = 'col1,,col2, ,col3';
+      const expected = ['col1', 'col2', 'col3'];
+      expect(splitAndTrimWithBracket(input)).toEqual(expected);
+    });
+
+    it('should handle quoted strings with commas', () => {
+      const input = "col1, concat('Hello, World!'), col2";
+      const expected = ['col1', "concat('Hello, World!')", 'col2'];
+      expect(splitAndTrimWithBracket(input)).toEqual(expected);
+    });
+
+    it('should handle double quoted strings with commas', () => {
+      const input = 'col1, "quoted, string", col3';
+      const expected = ['col1', '"quoted, string"', 'col3'];
+      expect(splitAndTrimWithBracket(input)).toEqual(expected);
+    });
+
+    it('should handle single quoted strings with commas', () => {
+      const input = `col1, 'quoted, string', col3`;
+      const expected = ['col1', `'quoted, string'`, 'col3'];
+      expect(splitAndTrimWithBracket(input)).toEqual(expected);
+    });
+
+    it('should handle mixed quotes with commas', () => {
+      const input = `col1, "double, quoted", col2, 'single, quoted', col3`;
+      const expected = [
+        'col1',
+        `"double, quoted"`,
+        'col2',
+        `'single, quoted'`,
+        'col3',
+      ];
+      expect(splitAndTrimWithBracket(input)).toEqual(expected);
+    });
+
+    it('should handle quotes inside function calls', () => {
+      const input = 'col1, func("text with , comma", \'another, text\'), col2';
+      const expected = [
+        'col1',
+        'func("text with , comma", \'another, text\')',
+        'col2',
+      ];
+      expect(splitAndTrimWithBracket(input)).toEqual(expected);
+    });
+
+    it('should handle brackets inside quoted strings', () => {
+      const input =
+        'col1, "string with (brackets, inside)", col2, \'string with [brackets, inside]\', col3';
+      const expected = [
+        'col1',
+        '"string with (brackets, inside)"',
+        'col2',
+        "'string with [brackets, inside]'",
+        'col3',
+      ];
+      expect(splitAndTrimWithBracket(input)).toEqual(expected);
+    });
+
+    it('should handle real-world SQL column list example', () => {
+      const input =
+        "Timestamp, ServiceName, JSONExtractString(Body, 'c'), JSONExtractString(Body, 'msg'), Timestamp, \"foo, bar\"";
+      const expected = [
+        'Timestamp',
+        'ServiceName',
+        "JSONExtractString(Body, 'c')",
+        "JSONExtractString(Body, 'msg')",
+        'Timestamp',
+        '"foo, bar"',
+      ];
+      expect(splitAndTrimWithBracket(input)).toEqual(expected);
+    });
+  });
 });

packages/common-utils/src/queryParser.ts

@@ -3,7 +3,7 @@ import SqlString from 'sqlstring';
 
 import { convertCHTypeToPrimitiveJSType, JSDataType } from '@/clickhouse';
 import { Metadata } from '@/metadata';
-import { splitAndTrimCSV } from '@/utils';
+import { splitAndTrimWithBracket } from '@/utils';
 
 function encodeSpecialTokens(query: string): string {
   return query
@@ -550,7 +550,9 @@ export class CustomSchemaSQLSerializerV2 extends SQLSerializer {
         );
       }
 
-      const expressions = splitAndTrimCSV(this.implicitColumnExpression);
+      const expressions = splitAndTrimWithBracket(
+        this.implicitColumnExpression,
+      );
 
       return {
         column:

packages/common-utils/src/renderChartConfig.ts

@@ -25,7 +25,7 @@ import {
 import {
   convertDateRangeToGranularityString,
   getFirstTimestampValueExpression,
-  splitAndTrimCSV,
+  splitAndTrimWithBracket,
 } from '@/utils';
 
 // FIXME: SQLParser.ColumnRef is incomplete
@@ -467,7 +467,7 @@ async function timeFilterExpr({
   with?: ChartConfigWithDateRange['with'];
   includedDataInterval?: string;
 }) {
-  const valueExpressions = splitAndTrimCSV(timestampValueExpression);
+  const valueExpressions = splitAndTrimWithBracket(timestampValueExpression);
   const startTime = dateRange[0].getTime();
   const endTime = dateRange[1].getTime();
 

packages/common-utils/src/utils.ts

@@ -18,10 +18,61 @@ export function splitAndTrimCSV(input: string): string[] {
     .filter(column => column.length > 0);
 }
 
+// Replace splitAndTrimCSV, should remove splitAndTrimCSV later
+export function splitAndTrimWithBracket(input: string): string[] {
+  let parenCount: number = 0;
+  let squareCount: number = 0;
+  let inSingleQuote: boolean = false;
+  let inDoubleQuote: boolean = false;
+
+  const res: string[] = [];
+  let cur: string = '';
+  for (const c of input + ',') {
+    if (c === '"' && !inSingleQuote) {
+      inDoubleQuote = !inDoubleQuote;
+      cur += c;
+      continue;
+    }
+
+    if (c === "'" && !inDoubleQuote) {
+      inSingleQuote = !inSingleQuote;
+      cur += c;
+      continue;
+    }
+    // Only count brackets when not in quotes
+    if (!inSingleQuote && !inDoubleQuote) {
+      if (c === '(') {
+        parenCount++;
+      } else if (c === ')') {
+        parenCount--;
+      } else if (c === '[') {
+        squareCount++;
+      } else if (c === ']') {
+        squareCount--;
+      }
+    }
+
+    if (
+      c === ',' &&
+      parenCount === 0 &&
+      squareCount === 0 &&
+      !inSingleQuote &&
+      !inDoubleQuote
+    ) {
+      const trimString = cur.trim();
+      if (trimString) res.push(trimString);
+      cur = '';
+    } else {
+      cur += c;
+    }
+  }
+  return res;
+}
+
 // If a user specifies a timestampValueExpression with multiple columns,
 // this will return the first one. We'll want to refine this over time
 export function getFirstTimestampValueExpression(valueExpression: string) {
-  return splitAndTrimCSV(valueExpression)[0];
+  return splitAndTrimWithBracket(valueExpression)[0];
 }
 
 export enum Granularity {