From c52ddd0b868d42169cf62a4430759096f3b0b616 Mon Sep 17 00:00:00 2001
From: Vadim Ogievetsky <vadim@ogievetsky.com>
Date: Thu, 22 Feb 2024 14:07:04 -0800
Subject: [PATCH] make flattenSpec location adaptive (#15946)

---
 .../ingestion-spec/ingestion-spec.tsx         | 22 ++++++++++++
 web-console/src/utils/sampler.ts              |  3 +-
 .../views/load-data-view/load-data-view.tsx   | 36 +++++++++----------
 3 files changed, 40 insertions(+), 21 deletions(-)

diff --git a/web-console/src/druid-models/ingestion-spec/ingestion-spec.tsx b/web-console/src/druid-models/ingestion-spec/ingestion-spec.tsx
index eca148ada3b..92985d8158b 100644
--- a/web-console/src/druid-models/ingestion-spec/ingestion-spec.tsx
+++ b/web-console/src/druid-models/ingestion-spec/ingestion-spec.tsx
@@ -50,6 +50,7 @@ import {
   getDimensionSpecName,
   getDimensionSpecs,
 } from '../dimension-spec/dimension-spec';
+import type { FlattenSpec } from '../flatten-spec/flatten-spec';
 import type { IndexSpec } from '../index-spec/index-spec';
 import { summarizeIndexSpec } from '../index-spec/index-spec';
 import type { InputFormat } from '../input-format/input-format';
@@ -381,6 +382,27 @@ export function getPossibleSystemFieldsForSpec(spec: Partial<IngestionSpec>): st
   return getPossibleSystemFieldsForInputSource(inputSource);
 }
 
+export function getFlattenSpec(spec: Partial<IngestionSpec>): FlattenSpec | undefined {
+  const inputFormat: InputFormat | undefined = deepGet(spec, 'spec.ioConfig.inputFormat');
+  if (!inputFormat) return;
+  return (
+    (inputFormat.type === 'kafka'
+      ? inputFormat.valueFormat?.flattenSpec
+      : inputFormat.flattenSpec) || undefined
+  );
+}
+
+export function changeFlattenSpec(
+  spec: Partial<IngestionSpec>,
+  flattenSpec: FlattenSpec | undefined,
+): Partial<IngestionSpec> {
+  if (deepGet(spec, 'spec.ioConfig.inputFormat.type') === 'kafka') {
+    return deepSet(spec, 'spec.ioConfig.inputFormat.valueFormat.flattenSpec', flattenSpec);
+  } else {
+    return deepSet(spec, 'spec.ioConfig.inputFormat.flattenSpec', flattenSpec);
+  }
+}
+
 export function getPossibleSystemFieldsForInputSource(inputSource: InputSource): string[] {
   switch (inputSource.type) {
     case 's3':
diff --git a/web-console/src/utils/sampler.ts b/web-console/src/utils/sampler.ts
index 866b195d457..6c31f2bb2b6 100644
--- a/web-console/src/utils/sampler.ts
+++ b/web-console/src/utils/sampler.ts
@@ -36,6 +36,7 @@ import {
   DETECTION_TIMESTAMP_SPEC,
   getDimensionNamesFromTransforms,
   getDimensionSpecName,
+  getFlattenSpec,
   getSpecType,
   getTimestampSchema,
   isDruidSource,
@@ -157,7 +158,7 @@ export function applyCache(sampleSpec: SampleSpec, cacheRows: CacheRows) {
     data: cacheRows.map(r => JSONBig.stringify(r)).join('\n'),
   });
 
-  const flattenSpec = deepGet(sampleSpec, 'spec.ioConfig.inputFormat.flattenSpec');
+  const flattenSpec = getFlattenSpec(sampleSpec);
   let inputFormat: InputFormat = { type: 'json', keepNullColumns: true };
   if (flattenSpec) {
     inputFormat = deepSet(inputFormat, 'flattenSpec', flattenSpec);
diff --git a/web-console/src/views/load-data-view/load-data-view.tsx b/web-console/src/views/load-data-view/load-data-view.tsx
index ddc6c304711..ccc48a17459 100644
--- a/web-console/src/views/load-data-view/load-data-view.tsx
+++ b/web-console/src/views/load-data-view/load-data-view.tsx
@@ -76,6 +76,7 @@ import {
   adjustForceGuaranteedRollup,
   adjustId,
   BATCH_INPUT_FORMAT_FIELDS,
+  changeFlattenSpec,
   chooseByBestTimestamp,
   cleanSpec,
   computeFlattenPathsForData,
@@ -89,6 +90,7 @@ import {
   FLATTEN_FIELD_FIELDS,
   getArrayMode,
   getDimensionSpecName,
+  getFlattenSpec,
   getIngestionComboType,
   getIngestionImage,
   getIngestionTitle,
@@ -1459,8 +1461,7 @@ export class LoadDataView extends React.PureComponent<LoadDataViewProps, LoadDat
     const { columnFilter, specialColumnsOnly, parserQueryState, selectedFlattenField } = this.state;
     const spec = this.getEffectiveSpec();
     const inputFormat: InputFormat = deepGet(spec, 'spec.ioConfig.inputFormat') || EMPTY_OBJECT;
-    const flattenFields: FlattenField[] =
-      deepGet(spec, 'spec.ioConfig.inputFormat.flattenSpec.fields') || EMPTY_ARRAY;
+    const flattenFields: FlattenField[] = getFlattenSpec(spec)?.fields || EMPTY_ARRAY;
 
     const canHaveNestedData = inputFormatCanProduceNestedData(inputFormat);
 
@@ -1617,13 +1618,7 @@ export class LoadDataView extends React.PureComponent<LoadDataViewProps, LoadDat
                 icon={IconNames.LIGHTBULB}
                 text={`Auto add ${pluralIfNeeded(suggestedFlattenFields.length, 'flatten spec')}`}
                 onClick={() => {
-                  this.updateSpec(
-                    deepSet(
-                      spec,
-                      'spec.ioConfig.inputFormat.flattenSpec.fields',
-                      suggestedFlattenFields,
-                    ),
-                  );
+                  this.updateSpec(changeFlattenSpec(spec, { fields: suggestedFlattenFields }));
                 }}
               />
             </FormGroup>
@@ -1676,24 +1671,25 @@ export class LoadDataView extends React.PureComponent<LoadDataViewProps, LoadDat
           initValue={selectedFlattenField.value}
           onClose={this.resetSelected}
           onDirty={this.handleDirty}
-          onApply={flattenField =>
+          onApply={flattenField => {
+            const flattenSpec = getFlattenSpec(spec) || {};
             this.updateSpec(
-              deepSet(
+              changeFlattenSpec(
                 spec,
-                `spec.ioConfig.inputFormat.flattenSpec.fields.${selectedFlattenField.index}`,
-                flattenField,
+                deepSet(flattenSpec, `fields.${selectedFlattenField.index}`, flattenField),
               ),
-            )
-          }
+            );
+          }}
           showDelete={selectedFlattenField.index !== -1}
-          onDelete={() =>
+          onDelete={() => {
+            const flattenSpec = getFlattenSpec(spec) || {};
             this.updateSpec(
-              deepDelete(
+              changeFlattenSpec(
                 spec,
-                `spec.ioConfig.inputFormat.flattenSpec.fields.${selectedFlattenField.index}`,
+                deepDelete(flattenSpec, `fields.${selectedFlattenField.index}`),
               ),
-            )
-          }
+            );
+          }}
         />
       );
     } else {