fix(ui): expand icon after updating column details (#22790)

* fix(ui): expand icon after updating column details

* add tests and fix container schema

* fix tests

(cherry picked from commit 07c3acfdbc21ba6577709301f8dd64e85fa2207c)

openmetadata-ui/src/main/resources/ui/playwright/e2e/Features/Container.spec.ts

@@ -15,6 +15,10 @@ import { CONTAINER_CHILDREN } from '../../constant/contianer';
 import { ContainerClass } from '../../support/entity/ContainerClass';
 import { performAdminLogin } from '../../utils/admin';
 import { redirectToHomePage } from '../../utils/common';
+import {
+  assignTagToChildren,
+  removeTagsFromChildren,
+} from '../../utils/entity';
 import { test } from '../fixtures/pages';
 
 const container = new ContainerClass();
@@ -116,4 +120,36 @@ test.describe('Container entity specific tests ', () => {
       'Page 1 of 2'
     );
   });
+
+  test('expand / collapse should not appear after updating nested fields for container', async ({
+    page,
+  }) => {
+    await page.goto('/container/s3_storage_sample.departments.finance');
+
+    await page.waitForLoadState('networkidle');
+    await page.waitForSelector('[data-testid="loader"]', {
+      state: 'detached',
+    });
+
+    await assignTagToChildren({
+      page,
+      tag: 'PersonalData.Personal',
+      rowId: 'budget_executor',
+      entityEndpoint: 'containers',
+    });
+
+    // Should not show expand icon for non-nested columns
+    expect(
+      page
+        .locator('[data-row-key="budget_executor"]')
+        .getByTestId('expand-icon')
+    ).not.toBeVisible();
+
+    await removeTagsFromChildren({
+      page,
+      tags: ['PersonalData.Personal'],
+      rowId: 'budget_executor',
+      entityEndpoint: 'containers',
+    });
+  });
 });

openmetadata-ui/src/main/resources/ui/playwright/e2e/Features/Dashboards.spec.ts

@@ -14,14 +14,18 @@ import { expect } from '@playwright/test';
 import { DashboardServiceClass } from '../../support/entity/service/DashboardServiceClass';
 import { performAdminLogin } from '../../utils/admin';
 import { redirectToHomePage } from '../../utils/common';
-import { generateEntityChildren } from '../../utils/entity';
+import {
+  assignTagToChildren,
+  generateEntityChildren,
+  removeTagsFromChildren,
+} from '../../utils/entity';
 import { test } from '../fixtures/pages';
 
 const dashboardEntity = new DashboardServiceClass();
 
-test.slow(true);
-
 test.describe('Dashboards', () => {
+  test.slow(true);
+
   test.beforeAll('Setup pre-requests', async ({ browser }) => {
     const { apiContext, afterAction } = await performAdminLogin(browser);
 
@@ -79,3 +83,39 @@ test.describe('Dashboards', () => {
     );
   });
 });
+
+test.describe('Data Model', () => {
+  test('expand / collapse should not appear after updating nested fields for dashboardDataModels', async ({
+    page,
+  }) => {
+    await page.goto(
+      '/dashboardDataModel/sample_superset.model.big_analytics_data_model_with_nested_columns'
+    );
+
+    await page.waitForLoadState('networkidle');
+    await page.waitForSelector('[data-testid="loader"]', {
+      state: 'detached',
+    });
+
+    await assignTagToChildren({
+      page,
+      tag: 'PersonalData.Personal',
+      rowId: 'revenue_metrics_0031',
+      entityEndpoint: 'dashboard/datamodels',
+    });
+
+    // Should not show expand icon for non-nested columns
+    expect(
+      page
+        .locator('[data-row-key="revenue_metrics_0031"]')
+        .getByTestId('expand-icon')
+    ).not.toBeVisible();
+
+    await removeTagsFromChildren({
+      page,
+      tags: ['PersonalData.Personal'],
+      rowId: 'revenue_metrics_0031',
+      entityEndpoint: 'dashboard/datamodels',
+    });
+  });
+});

openmetadata-ui/src/main/resources/ui/playwright/e2e/Features/Table.spec.ts

@@ -15,14 +15,16 @@ import { SidebarItem } from '../../constant/sidebar';
 import { TableClass } from '../../support/entity/TableClass';
 import { performAdminLogin } from '../../utils/admin';
 import { redirectToHomePage } from '../../utils/common';
-import { getFirstRowColumnLink } from '../../utils/entity';
+import {
+  assignTagToChildren,
+  getFirstRowColumnLink,
+  removeTagsFromChildren,
+} from '../../utils/entity';
 import { sidebarClick } from '../../utils/sidebar';
 import { test } from '../fixtures/pages';
 
 const table1 = new TableClass();
 
-test.slow(true);
-
 test.describe('Table pagination sorting search scenarios ', () => {
   test.beforeAll('Setup pre-requests', async ({ browser }) => {
     test.slow(true);
@@ -438,4 +440,42 @@ test.describe('Table & Data Model columns table pagination', () => {
       page.locator('[data-row-key="shop_id"]').getByTestId('expand-icon')
     ).not.toBeVisible();
   });
+
+  test('expand / collapse should not appear after updating nested fields table', async ({
+    page,
+  }) => {
+    await page.goto(
+      '/table/sample_data.ecommerce_db.shopify.performance_test_table'
+    );
+
+    await page.waitForLoadState('networkidle');
+    await page.waitForSelector('[data-testid="loader"]', {
+      state: 'detached',
+    });
+
+    await assignTagToChildren({
+      page,
+      tag: 'PersonalData.Personal',
+      rowId:
+        'sample_data.ecommerce_db.shopify.performance_test_table.test_col_0044',
+      entityEndpoint: 'tables',
+    });
+
+    // Should not show expand icon for non-nested columns
+    expect(
+      page
+        .locator(
+          '[data-row-key="sample_data.ecommerce_db.shopify.performance_test_table.test_col_0044"]'
+        )
+        .getByTestId('expand-icon')
+    ).not.toBeVisible();
+
+    await removeTagsFromChildren({
+      page,
+      tags: ['PersonalData.Personal'],
+      rowId:
+        'sample_data.ecommerce_db.shopify.performance_test_table.test_col_0044',
+      entityEndpoint: 'tables',
+    });
+  });
 });

openmetadata-ui/src/main/resources/ui/src/components/Container/ContainerDataModel/ContainerDataModel.test.tsx

@@ -114,6 +114,7 @@ jest.mock('../../../utils/TableUtils', () => ({
       'glossary',
     ]),
   handleUpdateTableColumnSelections: jest.fn(),
+  pruneEmptyChildren: jest.fn().mockImplementation((value) => value),
 }));
 
 jest.mock('../../../utils/TableTags/TableTags.utils', () => ({

openmetadata-ui/src/main/resources/ui/src/components/Container/ContainerDataModel/ContainerDataModel.tsx

@@ -42,7 +42,10 @@ import {
   getAllTags,
   searchTagInData,
 } from '../../../utils/TableTags/TableTags.utils';
-import { getTableExpandableConfig } from '../../../utils/TableUtils';
+import {
+  getTableExpandableConfig,
+  pruneEmptyChildren,
+} from '../../../utils/TableUtils';
 import { EntityAttachmentProvider } from '../../common/EntityDescription/EntityAttachmentProvider/EntityAttachmentProvider';
 import ErrorPlaceHolder from '../../common/ErrorWithPlaceholder/ErrorPlaceHolder';
 import Table from '../../common/Table/Table';
@@ -66,6 +69,8 @@ const ContainerDataModel: FC<ContainerDataModelProps> = ({
   const [editContainerColumnDescription, setEditContainerColumnDescription] =
     useState<Column>();
 
+  const schema = pruneEmptyChildren(dataModel?.columns ?? []);
+
   const handleFieldTagsChange = useCallback(
     async (selectedTags: EntityTags[], editColumnTag: Column) => {
       if (selectedTags && editColumnTag) {
@@ -99,13 +104,13 @@ const ContainerDataModel: FC<ContainerDataModelProps> = ({
   };
 
   const tagFilter = useMemo(() => {
-    const tags = getAllTags(dataModel?.columns ?? []);
+    const tags = getAllTags(schema);
 
     return groupBy(uniqBy(tags, 'value'), (tag) => tag.source) as Record<
       TagSource,
       TagFilterOptions[]
     >;
-  }, [dataModel?.columns]);
+  }, [schema]);
 
   const columns: ColumnsType<Column> = useMemo(
     () => [
@@ -226,7 +231,7 @@ const ContainerDataModel: FC<ContainerDataModelProps> = ({
     ]
   );
 
-  if (isEmpty(dataModel?.columns)) {
+  if (isEmpty(schema)) {
     return <ErrorPlaceHolder className="border-default border-radius-sm" />;
   }
 
@@ -236,7 +241,7 @@ const ContainerDataModel: FC<ContainerDataModelProps> = ({
         className="align-table-filter-left"
         columns={columns}
         data-testid="container-data-model-table"
-        dataSource={dataModel?.columns}
+        dataSource={schema}
         defaultVisibleColumns={DEFAULT_CONTAINER_DATA_MODEL_VISIBLE_COLUMNS}
         expandable={{
           ...getTableExpandableConfig<Column>(),

openmetadata-ui/src/main/resources/ui/src/components/Container/ContainerVersion/ContainerVersion.component.tsx

@@ -35,6 +35,7 @@ import {
   getEntityVersionTags,
 } from '../../../utils/EntityVersionUtils';
 import { getVersionPath } from '../../../utils/RouterUtils';
+import { pruneEmptyChildren } from '../../../utils/TableUtils';
 import { useRequiredParams } from '../../../utils/useRequiredParams';
 import { CustomPropertyTable } from '../../common/CustomPropertyTable/CustomPropertyTable';
 import DescriptionV1 from '../../common/EntityDescription/DescriptionV1';
@@ -88,7 +89,9 @@ const ContainerVersion: React.FC<ContainerVersionProp> = ({
     );
 
   const columns = useMemo(() => {
-    const colList = cloneDeep(currentVersionData.dataModel?.columns);
+    const colList = cloneDeep(
+      pruneEmptyChildren(currentVersionData.dataModel?.columns ?? [])
+    );
 
     return getColumnsDataWithVersionChanges<Column>(
       changeDescription,

openmetadata-ui/src/main/resources/ui/src/components/Dashboard/DataModel/DataModels/ModelTab/ModelTab.component.tsx

@@ -12,7 +12,7 @@
  */
 import { Typography } from 'antd';
 import { ColumnsType } from 'antd/lib/table';
-import { groupBy, omit, uniqBy } from 'lodash';
+import { groupBy, isEmpty, omit, uniqBy } from 'lodash';
 import { EntityTags, TagFilterOptions } from 'Models';
 import { useCallback, useEffect, useMemo, useState } from 'react';
 import { useTranslation } from 'react-i18next';
@@ -177,12 +177,15 @@ const ModelTab = () => {
     field?: keyof Column
   ) => {
     const response = await updateDataModelColumn(columnFqn, column);
+    const cleanResponse = isEmpty(response.children)
+      ? omit(response, 'children')
+      : response;
 
     setPaginatedColumns((prev) =>
       prev.map((col) =>
         col.fullyQualifiedName === columnFqn
           ? // Have to omit the field which is being updated to avoid persisted old value
-            { ...omit(col, field ?? ''), ...response }
+            { ...omit(col, field ?? ''), ...cleanResponse }
           : col
       )
     );

openmetadata-ui/src/main/resources/ui/src/components/Database/SchemaTable/SchemaTable.component.tsx

@@ -314,12 +314,15 @@ const SchemaTable = () => {
     field?: keyof Column
   ) => {
     const response = await updateTableColumn(columnFqn, column);
+    const cleanResponse = isEmpty(response.children)
+      ? omit(response, 'children')
+      : response;
 
     setTableColumns((prev) =>
       prev.map((col) =>
         col.fullyQualifiedName === columnFqn
           ? // Have to omit the field which is being updated to avoid persisted old value
-            { ...omit(col, field ?? ''), ...response }
+            { ...omit(col, field ?? ''), ...cleanResponse }
           : col
       )
     );