@epcare/esm-report-builder 0.1.1-pre.39 → 0.1.1-pre.40

package/src/components/linelist/builder/observation-column-modal.component.tsx

@@ -15,6 +15,7 @@ import {
   Stack,
   FormGroup,
   Checkbox,
+  Button,
 } from '@carbon/react';
 import ConceptSelector from '../../shared/concept-selector.component';
 import type { ConceptSummary } from '../../../resources/concepts/concept-types';
@@ -30,6 +31,8 @@ type Props = {
   onAddColumn: (column: LinelistColumnDraft) => void;
   /** Existing columns to check for duplicates */
   existingColumns?: LinelistColumnDraft[];
+  /** Column to edit (if provided, modal is in edit mode) */
+  editingColumn?: LinelistColumnDraft | null;
 };
 
 type ColumnModifier = 'ANY' | 'FIRST' | 'MOST_RECENT' | 'FIRST_N' | 'MOST_RECENT_N';
@@ -57,6 +60,7 @@ const ObservationColumnModal: React.FC<Props> = ({
   onClose,
   onAddColumn,
   existingColumns = [],
+  editingColumn = null,
 }) => {
   const [conceptSelectorKey, setConceptSelectorKey] = useState(0);
   const [selectedConcept, setSelectedConcept] = useState<ConceptSummary | undefined>();
@@ -68,15 +72,6 @@ const ObservationColumnModal: React.FC<Props> = ({
   const [extraValues, setExtraValues] = useState<ExtraValue[]>([]);
   const [answerConcept, setAnswerConcept] = useState<ConceptSummary | undefined>();
 
-  /**
-   * Reset form state when modal opens/closes
-   */
-  useEffect(() => {
-    if (open) {
-      setConceptSelectorKey((k) => k + 1);
-    }
-  }, [open]);
-
   /**
    * Reset form state
    */
@@ -91,6 +86,79 @@ const ObservationColumnModal: React.FC<Props> = ({
     setAnswerConcept(undefined);
   }, []);
 
+  /**
+   * Reset form state when modal opens/closes
+   * Also populate form when editing existing column
+   */
+  useEffect(() => {
+    if (open) {
+      setConceptSelectorKey((k) => k + 1);
+      if (editingColumn) {
+        // Populate form with existing column data
+        const config = editingColumn.dataDefinitionConfig || {};
+
+        // Debug: Check for required fields
+        if (!config.conceptUuid) {
+          console.warn('⚠️ Missing conceptUuid in dataDefinitionConfig');
+        }
+        if (!config.conceptName && !editingColumn.name) {
+          console.warn('⚠️ Missing conceptName and column name');
+        }
+        if (!config.strategy) {
+          console.warn('⚠️ Missing strategy in dataDefinitionConfig');
+        }
+
+        // Map strategy back to columnModifier
+        let mappedColumnModifier: ColumnModifier = 'MOST_RECENT';
+        const strategy = config.strategy;
+
+        if (strategy === 'ALL_VALUES') {
+          mappedColumnModifier = 'ANY';
+        } else if (strategy === 'EARLIEST') {
+          mappedColumnModifier = 'FIRST';
+        } else if (strategy === 'LATEST') {
+          mappedColumnModifier = 'MOST_RECENT';
+        } else if (typeof strategy === 'string' && strategy.startsWith('FIRST_')) {
+          mappedColumnModifier = 'FIRST_N';
+        } else if (typeof strategy === 'string' && strategy.startsWith('LAST_')) {
+          mappedColumnModifier = 'MOST_RECENT_N';
+        }
+
+        // Extract modifier count from strategy if it's FIRST_N or LAST_N
+        let extractedModifierCount = 1;
+        if (typeof strategy === 'string' && strategy.startsWith('FIRST_')) {
+          const match = strategy.match(/^FIRST_(\d+)$/);
+          if (match) extractedModifierCount = parseInt(match[1]) || 1;
+        } else if (typeof strategy === 'string' && strategy.startsWith('LAST_')) {
+          const match = strategy.match(/^LAST_(\d+)$/);
+          if (match) extractedModifierCount = parseInt(match[1]) || 1;
+        }
+
+        const conceptForSelection = editingColumn.source?.conceptUuid ? {
+          uuid: editingColumn.source.conceptUuid,
+          display: config.conceptName || editingColumn.name,
+          datatype: { name: editingColumn.source?.fieldType || 'TEXT' },
+          conceptClass: undefined, // Not loaded from saved data
+        } as ConceptSummary : undefined;
+
+        setSelectedConcept(conceptForSelection);
+        setColumnName(editingColumn.name || '');
+        setColumnValue(editingColumn.description || '');
+        setColumnModifier(config.columnModifier || mappedColumnModifier);
+        setModifierCount(config.modifierCount || extractedModifierCount);
+        setReturnDisplay(config.returnDisplay !== false);
+        setExtraValues(config.extraValues || []);
+        setAnswerConcept(config.answerConceptUuid ? {
+          uuid: config.answerConceptUuid,
+          display: config.answerConceptName || '',
+        } as ConceptSummary : undefined);
+      } else {
+        // Reset for new column creation
+        resetForm();
+      }
+    }
+  }, [open, editingColumn, resetForm]);
+
   /**
    * Handle modal close
    */
@@ -133,9 +201,9 @@ const ObservationColumnModal: React.FC<Props> = ({
       return;
     }
 
-    // Check for duplicate column names
+    // Check for duplicate column names (exclude current column when editing)
     const duplicateName = existingColumns.some(col =>
-      col.name.toLowerCase() === columnName.toLowerCase()
+      col.id !== editingColumn?.id && col.name.toLowerCase() === columnName.toLowerCase()
     );
     if (duplicateName) {
       // TODO: Show error notification
@@ -174,13 +242,13 @@ const ObservationColumnModal: React.FC<Props> = ({
     }
 
     const newColumn: LinelistColumnDraft = {
-      id: `col-${now}-${selectedConcept.uuid}`,
+      id: editingColumn?.id || `col-${now}-${selectedConcept.uuid}`,
       name: columnName.trim(),
       description: columnValue.trim() || selectedConcept.display,
       source: {
         dataSourceUuid: 'observations',
         dataSourceName: 'Observations',
-        table: 'obs',
+        table: 'observations',
         field: selectedConcept.uuid,
         fieldType: selectedConcept.datatype?.name || 'TEXT',
         conceptUuid: selectedConcept.uuid,
@@ -197,18 +265,18 @@ const ObservationColumnModal: React.FC<Props> = ({
         ...(answerConcept && { answerConceptUuid: answerConcept.uuid }),
       },
       additionInfo: {
-        addedVia: 'VISUAL_SELECTOR',
-        addedAt: now,
-        orderAdded: existingColumns.length,
+        addedVia: editingColumn?.additionInfo?.addedVia || 'VISUAL_SELECTOR',
+        addedAt: editingColumn?.additionInfo?.addedAt || now,
+        orderAdded: editingColumn?.additionInfo?.orderAdded || existingColumns.length,
       },
-      display: {
+      display: editingColumn?.display || {
         width: 150,
         align: 'left',
         sortable: true,
         filterable: true,
         format: 'coded',
       },
-      sortOrder: existingColumns.length,
+      sortOrder: editingColumn?.sortOrder ?? existingColumns.length,
       repeatResolution: columnModifier === 'ANY' ? undefined : {
         strategy: repeatStrategy,
         orderBy: 'obs_datetime',
@@ -219,7 +287,7 @@ const ObservationColumnModal: React.FC<Props> = ({
 
     onAddColumn(newColumn);
     handleClose();
-  }, [selectedConcept, columnName, columnValue, columnModifier, modifierCount, returnDisplay, extraValues, answerConcept, existingColumns, onAddColumn, handleClose]);
+  }, [selectedConcept, columnName, columnValue, columnModifier, modifierCount, returnDisplay, extraValues, answerConcept, existingColumns, onAddColumn, handleClose, editingColumn]);
 
   /**
    * Check if form is valid
@@ -230,11 +298,11 @@ const ObservationColumnModal: React.FC<Props> = ({
     <Modal
       open={open}
       onRequestClose={handleClose}
-      modalHeading="Add Observation Column"
+      modalHeading={editingColumn ? "Edit Observation Column" : "Add Observation Column"}
       modalLabel="Data Catalogue"
       size="md"
       preventCloseOnClickOutside
-      primaryButtonText="Add Column"
+      primaryButtonText={editingColumn ? "Save Changes" : "Add Column"}
       primaryButtonDisabled={!isFormValid}
       onRequestSubmit={handleSubmit}
       secondaryButtonText="Cancel"
@@ -333,25 +401,42 @@ const ObservationColumnModal: React.FC<Props> = ({
                 </div>
 
                 {/* Answer Concept (for coded concepts with answers) */}
-                {selectedConcept.answers && selectedConcept.answers.length > 0 && (
-                  <Select
-                    id="observation-answer-concept"
-                    labelText="Filter by Answer (Optional)"
-                    value={answerConcept?.uuid || ''}
-                    onChange={(e) => {
-                      const answer = selectedConcept.answers?.find(a => a.uuid === e.target.value);
-                      setAnswerConcept(answer);
-                    }}
-                  >
-                    <SelectItem value="" text="All answers" />
-                    {selectedConcept.answers.map(answer => (
-                      <SelectItem
-                        key={answer.uuid}
-                        value={answer.uuid}
-                        text={answer.display}
-                      />
-                    ))}
-                  </Select>
+                {(selectedConcept.answers && selectedConcept.answers.length > 0 || answerConcept) && (
+                  <div>
+                    {answerConcept && (!selectedConcept.answers || selectedConcept.answers.length === 0) && (
+                      <div style={{ marginBottom: '0.5rem', fontSize: '0.875rem', opacity: 0.9 }}>
+                        <strong>Filter by Answer:</strong> {answerConcept.display || answerConcept.uuid}
+                        <Button
+                          kind="ghost"
+                          size="sm"
+                          onClick={() => setAnswerConcept(undefined)}
+                          style={{ marginLeft: '0.5rem' }}
+                        >
+                          Clear
+                        </Button>
+                      </div>
+                    )}
+                    {selectedConcept.answers && selectedConcept.answers.length > 0 && (
+                      <Select
+                        id="observation-answer-concept"
+                        labelText="Filter by Answer (Optional)"
+                        value={answerConcept?.uuid || ''}
+                        onChange={(e) => {
+                          const answer = selectedConcept.answers?.find(a => a.uuid === e.target.value);
+                          setAnswerConcept(answer);
+                        }}
+                      >
+                        <SelectItem value="" text="All answers" />
+                        {selectedConcept.answers.map(answer => (
+                          <SelectItem
+                            key={answer.uuid}
+                            value={answer.uuid}
+                            text={answer.display}
+                          />
+                        ))}
+                      </Select>
+                    )}
+                  </div>
                 )}
               </Stack>
             </FormGroup>

package/src/components/linelist/modals/linelist-report-builder-modal.component.tsx

@@ -68,6 +68,7 @@ import {
 import type { LinelistReportDto } from '../../../types/linelist-types';
 import { listReportCategories, type ReportCategoryDto } from '../../../resources/report-category/report-category.api';
 import { useETLTables } from '../../../hooks/theme';
+import { enhanceConfigForPreview } from '../../../utils/config-enhancer';
 
 import CohortSQLEditor from '../config/cohort-sql-editor.component';
 import ColumnSelector from '../builder/column-selector.component';
@@ -178,6 +179,38 @@ function extractTableAndFieldFromSql(sql: string): { table: string; field: strin
   return null;
 }
 
+/**
+ * Normalize a column to ensure it has required fields (id, sortOrder)
+ * This ensures backward compatibility with legacy configs
+ */
+function normalizeColumn(column: any, index: number, existingIds: Set<string>): string {
+  // Generate a unique ID if missing or duplicate
+  let columnId = column.id;
+
+  if (!columnId) {
+    // Use index + timestamp for uniqueness
+    const timestamp = Date.now().toString(36);
+    columnId = `col-${timestamp}-${index}`;
+    // If still duplicate, add random suffix
+    if (existingIds.has(columnId)) {
+      columnId = `col-${timestamp}-${index}-${Math.random().toString(36).substr(2, 5)}`;
+    }
+  } else if (existingIds.has(columnId)) {
+    // Duplicate ID - generate new one
+    const timestamp = Date.now().toString(36);
+    columnId = `col-${timestamp}-${index}-dup`;
+  }
+
+  existingIds.add(columnId);
+
+  // Ensure sortOrder is set on source if it exists
+  if (column && typeof column.sortOrder !== 'number') {
+    column.sortOrder = index;
+  }
+
+  return columnId;
+}
+
 /**
  * Convert LinelistReportDto to LinelistReportDraft for editing
  */
@@ -203,6 +236,7 @@ function reportDtoToDraft(report: LinelistReportDto): LinelistReportDraft {
   // Extract columns from dataset definitions
   const columns: LinelistColumnDraft[] = [];
   const dataSet = config.dataSetDefinitions?.[0];
+  const existingIds = new Set<string>(); // Track IDs to ensure uniqueness
 
   if (dataSet?.columns) {
     dataSet.columns.forEach((col, idx) => {
@@ -228,21 +262,58 @@ function reportDtoToDraft(report: LinelistReportDto): LinelistReportDraft {
         'CALCULATION': 'NUMBER',
         'PERSON_ADDRESS': 'TEXT',
         'SQL': 'TEXT',
+        'OBSERVATION': 'CODED',
+        'ENCOUNTER_DIAGNOSIS': 'CODED',
       };
       const fieldType = fieldTypeMap[dataDefType] || 'TEXT';
 
+      // Set data source name based on dataDefinitionType
+      let dataSourceNameToUse = dataSourceName;
+      let tableToUse = table;
+      let fieldToUse = field;
+
+      // Override source info for known column types
+      if (dataDefType === 'OBSERVATION') {
+        dataSourceNameToUse = 'Observations';
+        tableToUse = 'observations';
+        fieldToUse = dataDefConfig.conceptUuid || field;
+      } else if (dataDefType === 'ENCOUNTER_DIAGNOSIS') {
+        dataSourceNameToUse = 'Encounter Diagnoses';
+        tableToUse = 'encounter_diagnoses';
+        fieldToUse = dataDefConfig.conceptUuid || '*';
+      } else if (dataDefType === 'PERSON_NAME') {
+        dataSourceNameToUse = 'Person Names';
+        tableToUse = 'person_name';
+      } else if (dataDefType === 'PERSON_ADDRESS') {
+        dataSourceNameToUse = 'Person Address';
+        tableToUse = 'person_address';
+      } else if (dataDefType === 'IDENTIFIER') {
+        dataSourceNameToUse = 'Patient Identifiers';
+        tableToUse = 'patient_identifier';
+      } else if (dataDefType === 'PERSON_ATTRIBUTE') {
+        dataSourceNameToUse = 'Person Attributes';
+        tableToUse = 'person_attribute';
+      }
+
+      // Normalize ID and sortOrder
+      const columnId = normalizeColumn(col, idx, existingIds);
+
       columns.push({
-        id: `col-${idx}`,
+        id: columnId,
         name: col.name,
+        description: (col as any).description || '',
         dataDefinitionType: dataDefType,
         dataDefinitionConfig: dataDefConfig,
         config: dataDefConfig, // @deprecated
         source: {
-          dataSourceUuid,
-          dataSourceName,
-          table,
-          field,
+          dataSourceUuid: dataSourceUuid || (dataDefType === 'OBSERVATION' ? 'observations' : dataDefType === 'ENCOUNTER_DIAGNOSIS' ? 'encounter_diagnoses' : ''),
+          dataSourceName: dataSourceNameToUse,
+          table: tableToUse,
+          field: fieldToUse,
           fieldType,
+          conceptUuid: dataDefConfig.conceptUuid,
+          attributeTypeUuid: dataDefConfig.attributeTypeUuid,
+          identifierTypeUuid: dataDefConfig.identifierTypeUuid,
         },
         additionInfo: {
           addedVia: 'IMPORT',
@@ -254,6 +325,7 @@ function reportDtoToDraft(report: LinelistReportDto): LinelistReportDraft {
           align: 'left',
           sortable: true,
           filterable: true,
+          format: dataDefType === 'OBSERVATION' || dataDefType === 'ENCOUNTER_DIAGNOSIS' ? 'coded' : 'text',
         },
         sortOrder: idx,
         repeatResolution: (col as any).repeatResolution,
@@ -801,8 +873,6 @@ function BasicsPanel({
   categories,
   tables,
 }: BasicsPanelProps) {
-  console.log('BasicsPanel render, draft.name:', draft.name, 'draft.categoryUuid:', draft.categoryUuid);
-
   const rowGrainOptions: Array<{ value: LinelistRowGrain; label: string; description: string }> = [
     { value: 'PATIENT', label: 'Patient', description: 'One row per patient' },
     { value: 'ENCOUNTER', label: 'Encounter', description: 'One row per encounter/visit' },
@@ -1246,7 +1316,7 @@ function PreviewPanel({ draft }: PreviewPanelProps) {
       {showJson && (
         <div className={styles.sqlPreview}>
           <h4>Generated Configuration</h4>
-          <pre className={styles.jsonPreview}>{JSON.stringify(config, null, 2)}</pre>
+          <pre className={styles.jsonPreview}>{JSON.stringify(enhanceConfigForPreview(config, draft), null, 2)}</pre>
         </div>
       )}
     </Stack>

package/src/components/linelist/pages/linelist-report-editor-page.component.tsx

@@ -49,8 +49,6 @@ const LinelistReportEditorPage: React.FC<Props> = () => {
 
     getLinelistReport(reportId)
       .then((report) => {
-        console.error('🔴 LOADED REPORT DATA:', report?.name, report);
-        console.log('Loaded report data:', report);
         setInitialReport(report);
         setMode('edit');
       })

package/src/components/linelist/shared/json-preview.component.tsx

@@ -11,6 +11,7 @@ import { Code } from '@carbon/react/icons';
 import type { LinelistReportDraft } from '../../../types/linelist-types';
 import { draftToConfig } from '../../../types/linelist-types';
 import type { IndicatorDto } from '../../../resources/indicator/indicators.api';
+import { enhanceConfigForPreview } from '../../../utils/config-enhancer';
 import styles from './json-preview.scss';
 
 type Props = {
@@ -43,7 +44,8 @@ export default function JsonPreview({ draft, className, indicators = [] }: Props
    */
   const handleCopy = useCallback(() => {
     const config = draftToConfig(draft, indicatorMap);
-    navigator.clipboard.writeText(JSON.stringify(config, null, 2));
+    const enhancedConfig = enhanceConfigForPreview(config, draft);
+    navigator.clipboard.writeText(JSON.stringify(enhancedConfig, null, 2));
   }, [draft, indicatorMap]);
 
   /**
@@ -78,6 +80,7 @@ export default function JsonPreview({ draft, className, indicators = [] }: Props
   };
 
   const config = draftToConfig(draft, indicatorMap);
+  const enhancedConfig = enhanceConfigForPreview(config, draft);
 
   return (
     <div className={`${styles.preview} ${className || ''}`}>
@@ -100,7 +103,7 @@ export default function JsonPreview({ draft, className, indicators = [] }: Props
         feedback="Copied!"
         className={styles.code}
       >
-        {JSON.stringify(config, null, 2)}
+        {JSON.stringify(enhancedConfig, null, 2)}
       </CodeSnippet>
 
       <div className={styles.validation}>

package/src/components/report-sections/hooks/useSectionEditorState.ts

@@ -14,8 +14,10 @@ export function useSectionEditorState(args: {
     const { open, mode, initialSection, indicators, ageCategories } = args;
     const isEdit = mode === 'edit';
     React.useEffect(() => {
-        console.log('[hook] MOUNT');
-        return () => console.log('[hook] UNMOUNT');
+        // Component mount/unmount
+        return () => {
+            // Component unmount
+        };
     }, []);
 
     
@@ -196,11 +198,7 @@ export function useSectionEditorState(args: {
 
     const toggleIndicator = React.useCallback(
         (i: SectionIndicatorRef, checked: boolean) => {
-            console.log('[hook] toggleIndicator called', { id: i.id, checked });
-
             setSelected((prev) => {
-                console.log('[hook] setSelected(prev) BEFORE', prev.map((p) => p.id));
-
                 let next;
                 if (checked) {
                     if (prev.some((x) => x.id === i.id)) {
@@ -214,7 +212,6 @@ export function useSectionEditorState(args: {
                     next = kept.map((x, idx) => ({ ...x, sortOrder: idx + 1 }));
                 }
 
-                console.log('[hook] setSelected(prev) AFTER', next.map((p) => p.id));
                 return next;
             });
 

package/src/components/shared/concept-selector.component.tsx

@@ -66,6 +66,22 @@ const ConceptSelector: React.FC<Props> = ({
 
   const { loading, results, error } = useConceptSearch(conceptQuery);
 
+  /**
+   * Combine search results with the currently selected concept
+   * This ensures the selected concept appears in the list even if not in search results
+   */
+  const items = React.useMemo(() => {
+    const combined = [...results];
+    // If there's a selected concept and it's not already in the results, add it
+    if (selectedConcept) {
+      const exists = results.some(r => r.uuid === selectedConcept.uuid);
+      if (!exists) {
+        combined.unshift(selectedConcept); // Add at the beginning
+      }
+    }
+    return combined;
+  }, [results, selectedConcept]);
+
   /**
    * Handle search input change
    */
@@ -166,13 +182,14 @@ const ConceptSelector: React.FC<Props> = ({
             key={multiKey}
             id="concept-select"
             titleText="Select concept"
-            items={results}
+            items={items}
             itemToString={itemToString}
-            placeholder={results.length ? 'Select a concept…' : 'No results'}
+            placeholder={items.length ? 'Select a concept…' : 'No results'}
             helperText={selectedConcept ? 'Click to change selection' : 'Type at least 2 characters to search'}
             onChange={(e: any) => handleSelect((e.selectedItems ?? []) as ConceptSummary[])}
             disabled={disabled}
             selectionFeedback="top"
+            initialSelectedItems={selectedConcept ? [selectedConcept] : []}
           />
 
           {/* Helper Text */}

package/src/hooks/theme/useMambaTables.ts

@@ -32,8 +32,6 @@ export function useETLTables(enabled: boolean) {
 
         getSchemaTables(ac.signal).then((data) => {
                 // data is expected to be SchemaTable[] or something close
-
-            console.log("I am here trying to see if data is loaded")
                 const names = (data ?? [])
                     .map((t) => toTableName(t))
                     .map((x) => x.trim())

package/src/types/linelist-types.ts

@@ -1303,6 +1303,14 @@ export function draftToConfig(
       let table: string;
       if (col.dataDefinitionType === 'CALCULATION') {
         table = 'calculated';
+      } else if (col.dataDefinitionType === 'OBSERVATION') {
+        table = 'observations';
+      } else if (col.dataDefinitionType === 'ENCOUNTER_DIAGNOSIS') {
+        table = 'encounter_diagnoses';
+      } else if (col.dataDefinitionType === 'PERSON_ADDRESS') {
+        table = 'person_address';
+      } else if (col.dataDefinitionType === 'PERSON_NAME') {
+        table = 'person_name';
       } else {
         // For SQL columns, derive the table from the SQL expression
         table = extractSourceTable(col.dataDefinitionConfig?.sql);
@@ -1317,10 +1325,20 @@ export function draftToConfig(
           ? 'SQL'
           : dsUuid.startsWith('mamba_') ? 'ETL' : 'SQL';
 
+      // Map table names to user-friendly data source names
+      const dataSourceNames: Record<string, string> = {
+        'observations': 'Observations',
+        'encounter_diagnoses': 'Encounter Diagnoses',
+        'person_address': 'Person Address',
+        'person_name': 'Person Names',
+        'calculated': 'Calculated Fields',
+        'custom_sql': 'Custom SQL',
+      };
+
       if (!sourceMap.has(dsUuid)) {
         sourceMap.set(dsUuid, {
           uuid: dsUuid,
-          name: dsUuid,
+          name: dataSourceNames[dsUuid] || dsUuid,
           type: dsType as DataSourceConfig['type'],
           role: isPrimary ? 'PRIMARY' : 'SECONDARY',
           columns: {},

package/src/utils/config-enhancer.ts

@@ -0,0 +1,62 @@
+/**
+ * Config Enhancement Utilities
+ *
+ * Utilities for enhancing configuration objects with metadata
+ * for debugging and visibility in JSON previews.
+ */
+
+import type { LinelistReportDefinitionConfig } from '../types/linelist-types';
+import type { LinelistReportDraft } from '../types/linelist-types';
+
+/**
+ * Enhance config for preview by adding column metadata (id, position)
+ * This helps with debugging and visibility in the JSON preview
+ */
+export function enhanceConfigForPreview(
+  config: LinelistReportDefinitionConfig,
+  draft: LinelistReportDraft
+): LinelistReportDefinitionConfig {
+  const enhanced = { ...config };
+
+  // Add column metadata to dataSetDefinitions columns
+  if (enhanced.dataSetDefinitions && enhanced.dataSetDefinitions[0]) {
+    const dataSetDef = enhanced.dataSetDefinitions[0];
+    const columnsWithMetadata = dataSetDef.columns.map((col: any, idx: number) => {
+      const draftCol = draft.columns.find(c => c.name === col.name);
+      return {
+        ...col,
+        _metadata: {
+          id: draftCol?.id || `col-${idx}`,
+          position: draftCol?.sortOrder ?? idx,
+        },
+      };
+    });
+    (enhanced as any).dataSetDefinitions[0] = {
+      ...dataSetDef,
+      columns: columnsWithMetadata,
+    };
+  }
+
+  // Also add metadata to dataSources columns
+  if (enhanced.dataSources) {
+    (enhanced as any).dataSources = enhanced.dataSources.map((ds: any) => {
+      if (ds.columns) {
+        const columnsWithMetadata: any = {};
+        Object.entries(ds.columns).forEach(([name, col]: [string, any]) => {
+          const draftCol = draft.columns.find(c => c.name === name);
+          columnsWithMetadata[name] = {
+            ...col,
+            _metadata: {
+              id: draftCol?.id,
+              position: draftCol?.sortOrder,
+            },
+          };
+        });
+        return { ...ds, columns: columnsWithMetadata };
+      }
+      return ds;
+    });
+  }
+
+  return enhanced;
+}