Merge pull request #2 from gadget-inc/allow-json-attributes-in-traces-2

Allow trace attributes to be stored as json objects

Author: udnay

provisioning/datasources/clickhouse.yml

@@ -19,4 +19,4 @@ datasources:
       port: 9000
       protocol: native
     secureJsonData:
-      # password: ""
+      password: ""

src/components/queryBuilder/views/TraceQueryBuilder.tsx

@@ -13,7 +13,7 @@ import useColumns from 'hooks/useColumns';
 import { BuilderOptionsReducerAction, setOptions, setOtelEnabled, setOtelVersion } from 'hooks/useBuilderOptionsState';
 import useIsNewQuery from 'hooks/useIsNewQuery';
 import { OtelVersionSelect } from '../OtelVersionSelect';
-import { useDefaultFilters, useOtelColumns, useTraceDefaultsOnMount } from './traceQueryBuilderHooks';
+import { useColumnTypes, useDefaultFilters, useOtelColumns, useTraceDefaultsOnMount } from './traceQueryBuilderHooks';
 import TraceIdInput from '../TraceIdInput';
 import { OrderByEditor, getOrderByOptions } from '../OrderByEditor';
 import { LimitEditor } from '../LimitEditor';
@@ -49,6 +49,7 @@ interface TraceQueryBuilderState {
   flattenNested?: boolean;
   traceEventsColumnPrefix?: string;
   traceLinksColumnPrefix?: string;
+  skipTraceAttributes?: boolean;
   traceId: string;
   orderBy: OrderBy[];
   limit: number;
@@ -94,6 +95,7 @@ export const TraceQueryBuilder = (props: TraceQueryBuilderProps) => {
       flattenNested: Boolean(builderOptions.meta?.flattenNested),
       traceEventsColumnPrefix: builderOptions.meta?.traceEventsColumnPrefix || '',
       traceLinksColumnPrefix: builderOptions.meta?.traceLinksColumnPrefix || '',
+      skipTraceAttributes: Boolean(builderOptions.meta?.skipTraceAttributes),
       traceId: builderOptions.meta?.traceId || '',
       orderBy: builderOptions.orderBy || [],
       limit: builderOptions.limit || 0,
@@ -135,6 +137,7 @@ export const TraceQueryBuilder = (props: TraceQueryBuilderProps) => {
           flattenNested: next.flattenNested,
           traceEventsColumnPrefix: next.traceEventsColumnPrefix,
           traceLinksColumnPrefix: next.traceLinksColumnPrefix,
+          skipTraceAttributes: next.skipTraceAttributes,
         },
       })
     );
@@ -143,6 +146,7 @@ export const TraceQueryBuilder = (props: TraceQueryBuilderProps) => {
   useTraceDefaultsOnMount(datasource, isNewQuery, builderOptions, builderOptionsDispatch);
   useOtelColumns(builderState.otelEnabled, builderState.otelVersion, builderOptionsDispatch);
   useDefaultFilters(builderOptions.table, builderState.isTraceIdMode, isNewQuery, builderOptionsDispatch);
+  useColumnTypes(allColumns, builderOptions.columns, builderOptions.meta?.useJsonAttributes, builderOptionsDispatch);
 
   const configWarning = showConfigWarning && (
     <Alert title="" severity="warning" buttonContent="Close" onRemove={() => setConfigWarningOpen(false)}>
@@ -384,6 +388,15 @@ export const TraceQueryBuilder = (props: TraceQueryBuilderProps) => {
             wide
           />
         </div>
+        <div className="gf-form">
+          <Switch
+            label={labels.columns.skipTraceAttributes.label}
+            tooltip={labels.columns.skipTraceAttributes.tooltip}
+            value={Boolean(builderState.skipTraceAttributes)}
+            onChange={onOptionChange('skipTraceAttributes')}
+            wide
+          />
+        </div>
         <div className="gf-form">
           <LabeledInput
             disabled={builderState.otelEnabled}

src/components/queryBuilder/views/traceQueryBuilderHooks.ts

@@ -12,6 +12,7 @@ import {
   QueryBuilderOptions,
   SelectedColumn,
   StringFilter,
+  TableColumn,
 } from 'types/queryBuilder';
 import { BuilderOptionsReducerAction, setOptions } from 'hooks/useBuilderOptionsState';
 
@@ -115,6 +116,70 @@ export const useOtelColumns = (
   }, [otelEnabled, otelVersion, builderOptionsDispatch]);
 };
 
+/**
+ * Auto-populates column types from the table schema and detects JSON attribute columns.
+ * When JSON columns are detected for TraceTags or TraceServiceTags, automatically sets useJsonAttributes.
+ */
+export const useColumnTypes = (
+  allColumns: readonly TableColumn[],
+  columns: SelectedColumn[] | undefined,
+  useJsonAttributes: boolean | undefined,
+  builderOptionsDispatch: React.Dispatch<BuilderOptionsReducerAction>
+) => {
+  const didPopulateTypes = useRef<boolean>(false);
+
+  useEffect(() => {
+    // Wait until we have both columns and schema
+    if (!columns?.length || !allColumns?.length) {
+      didPopulateTypes.current = false;
+      return;
+    }
+
+    // Check if any columns are missing types
+    const columnsNeedTypes = columns.some((c) => !c.type);
+    if (!columnsNeedTypes && didPopulateTypes.current) {
+      return;
+    }
+
+    // Populate types from schema
+    let hasChanges = false;
+    const updatedColumns = columns.map((col) => {
+      if (col.type) {
+        return col;
+      }
+      const schemaCol = allColumns.find((c) => c.name === col.name);
+      if (schemaCol?.type) {
+        hasChanges = true;
+        return { ...col, type: schemaCol.type };
+      }
+      return col;
+    });
+
+    // Auto-detect JSON attributes
+    let detectedJsonAttributes = true;
+    if (!useJsonAttributes) {
+      const tagsCol = updatedColumns.find((c) => c.hint === ColumnHint.TraceTags);
+      const serviceTagsCol = updatedColumns.find((c) => c.hint === ColumnHint.TraceServiceTags);
+      
+      detectedJsonAttributes = 
+        tagsCol?.type?.toLowerCase().startsWith('json') ||
+        serviceTagsCol?.type?.toLowerCase().startsWith('json') ||
+        detectedJsonAttributes;
+    }
+
+    if (hasChanges || detectedJsonAttributes) {
+      builderOptionsDispatch(
+        setOptions({
+          columns: updatedColumns,
+          meta: detectedJsonAttributes ? { useJsonAttributes: true } : undefined,
+        })
+      );
+    }
+
+    didPopulateTypes.current = true;
+  }, [allColumns, columns, useJsonAttributes, builderOptionsDispatch]);
+};
+
 // Apply default filters on table change
 export const useDefaultFilters = (
   table: string,

src/data/sqlGenerator.ts

@@ -147,19 +147,20 @@ const generateTraceIdQuery = (options: QueryBuilderOptions): string => {
     selectParts.push(getTraceDurationSelectSql(escapeIdentifier(traceDurationTime.name), timeUnit));
   }
 
-  // TODO: for tags and serviceTags, consider the column type. They might not require mapping, they could already be JSON.
-  const traceTags = getColumnByHint(options, ColumnHint.TraceTags);
-  if (traceTags !== undefined) {
-    selectParts.push(
-      `arrayMap(key -> map('key', key, 'value',${escapeIdentifier(traceTags.name)}[key]), mapKeys(${escapeIdentifier(traceTags.name)})) as tags`
-    );
-  }
+  const useJsonAttributes = options.meta?.useJsonAttributes;
+  const skipTraceAttributes = options.meta?.skipTraceAttributes;
+
+  // Skip expensive attribute columns when skipTraceAttributes is enabled (for performance)
+  if (!skipTraceAttributes) {
+    const traceTags = getColumnByHint(options, ColumnHint.TraceTags);
+    if (traceTags !== undefined) {
+      selectParts.push(getAttributesSelectSql(escapeIdentifier(traceTags.name), traceTags.type, 'tags', useJsonAttributes));
+    }
 
-  const traceServiceTags = getColumnByHint(options, ColumnHint.TraceServiceTags);
-  if (traceServiceTags !== undefined) {
-    selectParts.push(
-      `arrayMap(key -> map('key', key, 'value',${escapeIdentifier(traceServiceTags.name)}[key]), mapKeys(${escapeIdentifier(traceServiceTags.name)})) as serviceTags`
-    );
+    const traceServiceTags = getColumnByHint(options, ColumnHint.TraceServiceTags);
+    if (traceServiceTags !== undefined) {
+      selectParts.push(getAttributesSelectSql(escapeIdentifier(traceServiceTags.name), traceServiceTags.type, 'serviceTags', useJsonAttributes));
+    }
   }
 
   const traceStatusCode = getColumnByHint(options, ColumnHint.TraceStatusCode);
@@ -171,49 +172,54 @@ const generateTraceIdQuery = (options: QueryBuilderOptions): string => {
 
   const flattenNested = Boolean(options.meta?.flattenNested);
 
-  const traceEventsPrefix = options.meta?.traceEventsColumnPrefix || '';
-  if (traceEventsPrefix !== '') {
-    if (flattenNested) {
-      selectParts.push(
-        [
-          `arrayMap(event -> tuple(multiply(toFloat64(event.Timestamp), 1000),`,
-          `arrayConcat(arrayMap(key -> map('key', key, 'value', event.Attributes[key]),`,
-          `mapKeys(event.Attributes)), [map('key', 'message', 'value', event.Name)]))::Tuple(timestamp Float64, fields Array(Map(String, String))),`,
-          `${escapeIdentifier(traceEventsPrefix)}) as logs`,
-        ].join(' ')
-      );
-    } else {
-      selectParts.push(
-        [
-          `arrayMap((name, timestamp, attributes) -> tuple(name, toString(toUnixTimestamp64Milli(timestamp)),`,
-          `arrayMap( key -> map('key', key, 'value', attributes[key]),`,
-          `mapKeys(attributes)))::Tuple(name String, timestamp String, fields Array(Map(String, String))),`,
-          `${escapeIdentifier(traceEventsPrefix)}.Name, ${escapeIdentifier(traceEventsPrefix)}.Timestamp,`,
-          `${escapeIdentifier(traceEventsPrefix)}.Attributes) AS logs`,
-        ].join(' ')
-      );
+  // Skip Events and Links when skipTraceAttributes is enabled (they contain nested attributes)
+  if (!skipTraceAttributes) {
+    const traceEventsPrefix = options.meta?.traceEventsColumnPrefix || '';
+    if (traceEventsPrefix !== '') {
+      const eventAttrsSql = getNestedAttributesSql('event.Attributes', useJsonAttributes);
+      const attrsSql = getNestedAttributesSql('attributes', useJsonAttributes);
+      
+      if (flattenNested) {
+        selectParts.push(
+          [
+            `arrayMap(event -> tuple(multiply(toFloat64(event.Timestamp), 1000),`,
+            `arrayConcat(${eventAttrsSql}, [map('key', 'message', 'value', event.Name)]))::Tuple(timestamp Float64, fields Array(Map(String, String))),`,
+            `${escapeIdentifier(traceEventsPrefix)}) as logs`,
+          ].join(' ')
+        );
+      } else {
+        selectParts.push(
+          [
+            `arrayMap((name, timestamp, attributes) -> tuple(name, toString(toUnixTimestamp64Milli(timestamp)),`,
+            `${attrsSql})::Tuple(name String, timestamp String, fields Array(Map(String, String))),`,
+            `${escapeIdentifier(traceEventsPrefix)}.Name, ${escapeIdentifier(traceEventsPrefix)}.Timestamp,`,
+            `${escapeIdentifier(traceEventsPrefix)}.Attributes) AS logs`,
+          ].join(' ')
+        );
+      }
     }
-  }
 
-  const traceLinksPrefix = options.meta?.traceLinksColumnPrefix || '';
-  if (traceLinksPrefix !== '') {
-    if (flattenNested) {
-      selectParts.push(
-        [
-          `arrayMap(link -> tuple(link.TraceId, link.SpanId, arrayMap(key -> map('key', key, 'value', link.Attributes[key]),`,
-          `mapKeys(link.Attributes)))::Tuple(traceID String, spanID String, tags Array(Map(String, String))),`,
-          `${escapeIdentifier(traceLinksPrefix)}) AS references`,
-        ].join(' ')
-      );
-    } else {
-      selectParts.push(
-        [
-          `arrayMap((traceID, spanID, attributes) -> tuple(traceID, spanID, arrayMap(key -> map('key', key, 'value', attributes[key]),`,
-          `mapKeys(attributes)))::Tuple(traceID String, spanID String, tags Array(Map(String, String))),`,
-          `${escapeIdentifier(traceLinksPrefix)}.TraceId, ${escapeIdentifier(traceLinksPrefix)}.SpanId,`,
-          `${escapeIdentifier(traceLinksPrefix)}.Attributes) AS references`,
-        ].join(' ')
-      );
+    const traceLinksPrefix = options.meta?.traceLinksColumnPrefix || '';
+    if (traceLinksPrefix !== '') {
+      const linkAttrsSql = getNestedAttributesSql('link.Attributes', useJsonAttributes);
+      const linkAttrsSqlNonFlat = getNestedAttributesSql('attributes', useJsonAttributes);
+      
+      if (flattenNested) {
+        selectParts.push(
+          [
+            `arrayMap(link -> tuple(link.TraceId, link.SpanId, ${linkAttrsSql})::Tuple(traceID String, spanID String, tags Array(Map(String, String))),`,
+            `${escapeIdentifier(traceLinksPrefix)}) AS references`,
+          ].join(' ')
+        );
+      } else {
+        selectParts.push(
+          [
+            `arrayMap((traceID, spanID, attributes) -> tuple(traceID, spanID, ${linkAttrsSqlNonFlat})::Tuple(traceID String, spanID String, tags Array(Map(String, String))),`,
+            `${escapeIdentifier(traceLinksPrefix)}.TraceId, ${escapeIdentifier(traceLinksPrefix)}.SpanId,`,
+            `${escapeIdentifier(traceLinksPrefix)}.Attributes) AS references`,
+          ].join(' ')
+        );
+      }
     }
   }
 
@@ -657,7 +663,7 @@ const getTableIdentifier = (database: string, table: string): string => {
   return `${escapeIdentifier(database)}${sep}${escapeIdentifier(table)}`;
 };
 
-const escapeIdentifier = (id: string): string => {
+export const escapeIdentifier = (id: string): string => {
   return id ? `"${id}"` : '';
 };
 
@@ -669,6 +675,38 @@ const escapeValue = (value: string): string => {
   return `'${value}'`;
 };
 
+/**
+ * Returns the SELECT SQL for attributes columns (TraceTags/TraceServiceTags).
+ * Handles both Map and JSON column types.
+ * For Map types: uses arrayMap with mapKeys to convert to array of key/value maps
+ * For JSON types: return the JSON column as is
+ */
+export const getAttributesSelectSql = (columnIdentifier: string, columnType: string | undefined, alias: string, useJsonAttributes?: boolean): string => {
+  const isJsonType = useJsonAttributes || columnType?.toLowerCase().startsWith('json');
+  
+  if (isJsonType) {
+    // For native JSON columns (ClickHouse 25.x+), convert to string first then extract keys/values.
+    // Native JSON can be used directly in Grafana.
+    return `toString(${columnIdentifier}) as ${alias}`;
+  }
+  
+  // Default: Map type - convert to array of {key, value} maps for Grafana trace panel
+  return `arrayMap(key -> map('key', key, 'value',${columnIdentifier}[key]), mapKeys(${columnIdentifier})) as ${alias}`;
+};
+
+/**
+ * Returns the SQL fragment for converting attributes to array of key/value maps.
+ * Used within nested structures like Events and Links.
+ * @param attributesExpr - The expression to access attributes (e.g., "event.Attributes" or "attributes")
+ */
+const getNestedAttributesSql = (attributesExpr: string, useJsonAttributes?: boolean): string => {
+  if (useJsonAttributes) {
+    // For native JSON columns (ClickHouse 25.x+), convert to string first then extract keys/values
+    return `toString(${attributesExpr})`;
+  }
+  return `arrayMap(key -> map('key', key, 'value', ${attributesExpr}[key]), mapKeys(${attributesExpr}))`;
+};
+
 /**
  * Returns the SELECT column for trace duration.
  * Time unit is used to convert the value to milliseconds, as is required by Grafana's Trace panel.