fix contract schema pagination key selection not persisting (#23250)

* fix contract schema pagination key selection not persisting * fix the unit test * added playwrigt for it
2025-11-02 11:39:12 +00:00 · 2025-09-05 12:46:35 +05:30 · 2025-09-05 12:46:35 +05:30 · 13bf23939f
commit 13bf23939f
parent 999af800b3
4 changed files with 354 additions and 92 deletions
--- a/openmetadata-ui/src/main/resources/ui/playwright/e2e/Pages/DataContracts.spec.ts
+++ b/openmetadata-ui/src/main/resources/ui/playwright/e2e/Pages/DataContracts.spec.ts
@ -773,6 +773,266 @@ test.describe('Data Contracts', () => {
    );
  });

+  test('Pagination in Schema Tab with Selection Persistent', async ({
+    page,
+  }) => {
+    test.slow();
+
+    const entityFQN = 'sample_data.ecommerce_db.shopify.performance_test_table';
+
+    try {
+      await test.step('Redirect to Home Page and visit entity', async () => {
+        await redirectToHomePage(page);
+        await page.goto(`/table/${entityFQN}`);
+
+        await page.waitForLoadState('networkidle');
+        await page.waitForSelector('[data-testid="loader"]', {
+          state: 'detached',
+        });
+      });
+
+      await test.step(
+        'Open contract section and start adding contract',
+        async () => {
+          await page.click('[data-testid="contract"]');
+
+          await expect(page.getByTestId('no-data-placeholder')).toBeVisible();
+          await expect(page.getByTestId('add-contract-button')).toBeVisible();
+
+          await page.getByTestId('add-contract-button').click();
+
+          await expect(page.getByTestId('add-contract-card')).toBeVisible();
+        }
+      );
+
+      await test.step('Fill Contract Details form', async () => {
+        await page
+          .getByTestId('contract-name')
+          .fill(DATA_CONTRACT_DETAILS.name);
+      });
+
+      await test.step('Fill Contract Schema form', async () => {
+        const columnResponse = page.waitForResponse(
+          'api/v1/tables/name/sample_data.ecommerce_db.shopify.performance_test_table/columns?**'
+        );
+
+        await page.getByRole('button', { name: 'Schema' }).click();
+
+        await columnResponse;
+        await page.waitForSelector('[data-testid="loader"]', {
+          state: 'detached',
+        });
+
+        await page
+          .locator('input[type="checkbox"][aria-label="Select all"]')
+          .check();
+
+        await expect(
+          page.getByRole('checkbox', { name: 'Select all' })
+        ).toBeChecked();
+
+        // Move to 2nd Page and Select columns
+
+        const columnResponse2 = page.waitForResponse(
+          'api/v1/tables/name/sample_data.ecommerce_db.shopify.performance_test_table/columns?**'
+        );
+
+        await page.getByTestId('next').click();
+
+        await columnResponse2;
+        await page.waitForSelector('[data-testid="loader"]', {
+          state: 'detached',
+        });
+
+        await page
+          .locator('input[type="checkbox"][aria-label="Select all"]')
+          .check();
+
+        await expect(
+          page.getByRole('checkbox', { name: 'Select all' })
+        ).toBeChecked();
+
+        // Move to 3nd Page and Select columns
+
+        const columnResponse3 = page.waitForResponse(
+          'api/v1/tables/name/sample_data.ecommerce_db.shopify.performance_test_table/columns?**'
+        );
+
+        await page.getByTestId('next').click();
+
+        await columnResponse3;
+        await page.waitForSelector('[data-testid="loader"]', {
+          state: 'detached',
+        });
+
+        await page
+          .locator('input[type="checkbox"][aria-label="Select all"]')
+          .check();
+
+        await expect(
+          page.getByRole('checkbox', { name: 'Select all' })
+        ).toBeChecked();
+
+        // Now UnSelect the Selected Columns of 3rd Page
+
+        await page
+          .locator('input[type="checkbox"][aria-label="Select all"]')
+          .uncheck();
+
+        await expect(
+          page.getByRole('checkbox', { name: 'Select all' })
+        ).not.toBeChecked();
+      });
+
+      await test.step('Save contract and validate for schema', async () => {
+        const saveContractResponse = page.waitForResponse(
+          '/api/v1/dataContracts'
+        );
+        await page.getByTestId('save-contract-btn').click();
+
+        await saveContractResponse;
+
+        // Check all schema from 1 to 50
+        for (let i = 1; i <= 50; i++) {
+          if (i < 10) {
+            await expect(page.getByText(`test_col_000${i}`)).toBeVisible();
+          } else {
+            await expect(page.getByText(`test_col_00${i}`)).toBeVisible();
+          }
+        }
+
+        // Schema from 51 to 75 Should not be visible
+        for (let i = 51; i <= 75; i++) {
+          await expect(page.getByText(`test_col_00${i}`)).not.toBeVisible();
+        }
+      });
+
+      await test.step('Update the Schema and Validate', async () => {
+        await page.getByTestId('contract-edit-button').click();
+
+        const columnResponse = page.waitForResponse(
+          'api/v1/tables/name/sample_data.ecommerce_db.shopify.performance_test_table/columns?**'
+        );
+
+        await page.getByRole('button', { name: 'Schema' }).click();
+
+        await columnResponse;
+        await page.waitForSelector('[data-testid="loader"]', {
+          state: 'detached',
+        });
+
+        await page
+          .locator('input[type="checkbox"][aria-label="Select all"]')
+          .uncheck();
+
+        await expect(
+          page.getByRole('checkbox', { name: 'Select all' })
+        ).not.toBeChecked();
+
+        const saveContractResponse = page.waitForResponse(
+          '/api/v1/dataContracts'
+        );
+        await page.getByTestId('save-contract-btn').click();
+
+        await saveContractResponse;
+
+        await page.waitForLoadState('networkidle');
+        await page.waitForSelector('[data-testid="loader"]', {
+          state: 'detached',
+        });
+
+        // Check all schema from 26 to 50
+        for (let i = 26; i <= 50; i++) {
+          await expect(page.getByText(`test_col_00${i}`)).toBeVisible();
+        }
+      });
+
+      await test.step(
+        'Re-select some columns on page 1, save and validate',
+        async () => {
+          await page.getByTestId('contract-edit-button').click();
+
+          const columnResponse = page.waitForResponse(
+            'api/v1/tables/name/sample_data.ecommerce_db.shopify.performance_test_table/columns?**'
+          );
+
+          await page.getByRole('button', { name: 'Schema' }).click();
+
+          await columnResponse;
+          await page.waitForSelector('[data-testid="loader"]', {
+            state: 'detached',
+          });
+
+          for (let i = 1; i <= 5; i++) {
+            await page
+              .locator(
+                `[data-row-key="${entityFQN}.test_col_000${i}"] .ant-checkbox-input`
+              )
+              .click();
+          }
+
+          const saveContractResponse = page.waitForResponse(
+            '/api/v1/dataContracts'
+          );
+          await page.getByTestId('save-contract-btn').click();
+
+          await saveContractResponse;
+
+          await page.waitForLoadState('networkidle');
+          await page.waitForSelector('[data-testid="loader"]', {
+            state: 'detached',
+          });
+
+          // Check all schema from 1 to 5 and then, the one we didn't touch 26 to 50
+          for (let i = 1; i <= 5; i++) {
+            await expect(page.getByText(`test_col_000${i}`)).toBeVisible();
+          }
+
+          for (let i = 26; i <= 50; i++) {
+            await expect(page.getByText(`test_col_00${i}`)).toBeVisible();
+          }
+        }
+      );
+    } finally {
+      await test.step('Delete contract', async () => {
+        await redirectToHomePage(page);
+        await page.goto(`/table/${entityFQN}`);
+
+        await page.waitForLoadState('networkidle');
+        await page.waitForSelector('[data-testid="loader"]', {
+          state: 'detached',
+        });
+
+        await page.click('[data-testid="contract"]');
+
+        await page.waitForSelector('[data-testid="loader"]', {
+          state: 'detached',
+        });
+
+        const deleteContractResponse = page.waitForResponse(
+          'api/v1/dataContracts/*?hardDelete=true&recursive=true'
+        );
+
+        await page.getByTestId('delete-contract-button').click();
+
+        await expect(page.locator('.ant-modal-title')).toBeVisible();
+
+        await page.getByTestId('confirmation-text-input').click();
+        await page.getByTestId('confirmation-text-input').fill('DELETE');
+
+        await expect(page.getByTestId('confirm-button')).toBeEnabled();
+
+        await page.getByTestId('confirm-button').click();
+        await deleteContractResponse;
+
+        await toastNotification(page, '"Contract" deleted successfully!');
+
+        await expect(page.getByTestId('no-data-placeholder')).toBeVisible();
+        await expect(page.getByTestId('add-contract-button')).toBeVisible();
+      });
+    }
+  });
+
  test('should allow adding a semantic with multiple rules', async ({
    page,
  }) => {
--- a/openmetadata-ui/src/main/resources/ui/src/components/DataContract/AddDataContract/AddDataContract.tsx
+++ b/openmetadata-ui/src/main/resources/ui/src/components/DataContract/AddDataContract/AddDataContract.tsx
@ -151,9 +151,7 @@ const AddDataContract: React.FC<{
          <ContractSchemaFormTab
            nextLabel={t('label.semantic-plural')}
            prevLabel={t('label.contract-detail-plural')}
-            selectedSchema={
-              contract?.schema?.map((column) => column.name) || []
-            }
+            selectedSchema={contract?.schema ?? []}
            onChange={onFormChange}
            onNext={onNext}
            onPrev={onPrev}
--- a/openmetadata-ui/src/main/resources/ui/src/components/DataContract/ContractSchemaFormTab/ContractScehmaFormTab.tsx
+++ b/openmetadata-ui/src/main/resources/ui/src/components/DataContract/ContractSchemaFormTab/ContractScehmaFormTab.tsx
@ -13,7 +13,7 @@
 import Icon from '@ant-design/icons';
 import { Button, Card, Tag, Typography } from 'antd';
 import { ColumnsType } from 'antd/lib/table';
-import { isEmpty, pick } from 'lodash';
+import { isEmpty, pick, uniqBy } from 'lodash';
 import { Key, useCallback, useEffect, useMemo, useState } from 'react';
 import { useTranslation } from 'react-i18next';
 import { ReactComponent as LeftOutlined } from '../../../assets/svg/left-arrow.svg';
@ -44,7 +44,7 @@ import { TableCellRendered } from '../../Database/SchemaTable/SchemaTable.interf
 import TableTags from '../../Database/TableTags/TableTags.component';

 export const ContractSchemaFormTab: React.FC<{
-  selectedSchema: string[];
+  selectedSchema: Column[];
  onNext: () => void;
  onChange: (data: Partial<DataContract>) => void;
  onPrev: () => void;
@ -53,8 +53,9 @@ export const ContractSchemaFormTab: React.FC<{
 }> = ({ selectedSchema, onNext, onChange, onPrev, nextLabel, prevLabel }) => {
  const { t } = useTranslation();
  const { fqn } = useFqn();
-  const [allColumns, setAllColumns] = useState<Column[]>([]);
-  const [selectedKeys, setSelectedKeys] = useState<string[]>(selectedSchema);
+  const [allColumnsData, setAllColumnData] = useState<Column[]>([]);
+  const [columnsData, setColumnsData] = useState<Column[]>([]);
+  const [selectedKeys, setSelectedKeys] = useState<string[]>();
  const [isLoading, setIsLoading] = useState(false);

  const tableFqn = useMemo(
@ -79,13 +80,18 @@ export const ContractSchemaFormTab: React.FC<{
  const handleChangeTable = useCallback(
    (selectedRowKeys: Key[]) => {
      setSelectedKeys(selectedRowKeys as string[]);
+      const selectedColumns =
+        selectedRowKeys.length > 0
+          ? allColumnsData.filter((column) =>
+              selectedRowKeys.includes(column.fullyQualifiedName ?? '')
+            )
+          : [];
+
      onChange({
-        schema: allColumns.filter((column) =>
-          selectedRowKeys.includes(column.name)
-        ),
+        schema: selectedColumns,
      });
    },
-    [allColumns, onChange]
+    [allColumnsData, onChange]
  );

  const fetchTableColumns = useCallback(
@ -105,11 +111,17 @@ export const ContractSchemaFormTab: React.FC<{
        });

        const prunedColumns = pruneEmptyChildren(response.data);
-        setAllColumns(prunedColumns);
+        setColumnsData(prunedColumns);
+        setAllColumnData((prev) => {
+          const combined = [...prev, ...selectedSchema, ...prunedColumns];
+
+          return uniqBy(combined, 'fullyQualifiedName');
+        });
+
        handlePagingChange(response.paging);
      } catch {
        // Set empty state if API fails
-        setAllColumns([]);
+        setColumnsData([]);
        handlePagingChange({
          offset: 1,
          limit: pageSize,
@ -118,7 +130,7 @@ export const ContractSchemaFormTab: React.FC<{
      }
      setIsLoading(false);
    },
-    [tableFqn, pageSize]
+    [tableFqn, pageSize, selectedSchema, setAllColumnData]
  );

  const handleColumnsPageChange = useCallback(
@ -126,7 +138,7 @@ export const ContractSchemaFormTab: React.FC<{
      fetchTableColumns(currentPage);
      handlePageChange(currentPage);
    },
-    [fetchTableColumns]
+    [fetchTableColumns, handlePageChange]
  );

  const paginationProps = useMemo(
@ -273,6 +285,12 @@ export const ContractSchemaFormTab: React.FC<{
    [tableFqn]
  );

+  useEffect(() => {
+    setSelectedKeys(
+      selectedSchema.map((item) => (item as Column).fullyQualifiedName ?? '')
+    );
+  }, [selectedSchema]);
+
  useEffect(() => {
    fetchTableColumns();
  }, [fetchTableColumns]);
@ -291,13 +309,14 @@ export const ContractSchemaFormTab: React.FC<{
        <Table
          columns={columns}
          customPaginationProps={paginationProps}
-          dataSource={allColumns}
+          dataSource={columnsData}
          loading={isLoading}
          pagination={false}
-          rowKey="name"
+          rowKey="fullyQualifiedName"
          rowSelection={{
            selectedRowKeys: selectedKeys,
            onChange: handleChangeTable,
+            preserveSelectedRowKeys: true, // Preserve selections across page changes
          }}
        />
      </Card>
--- a/openmetadata-ui/src/main/resources/ui/src/components/DataContract/ContractSchemaFormTab/ContractSchemaFormTab.test.tsx
+++ b/openmetadata-ui/src/main/resources/ui/src/components/DataContract/ContractSchemaFormTab/ContractSchemaFormTab.test.tsx
@ -20,6 +20,7 @@ import {
  waitFor,
 } from '@testing-library/react';
 import { Column } from '../../../generated/entity/data/table';
+import { useFqn } from '../../../hooks/useFqn';
 import { getTableColumnsByFQN } from '../../../rest/tableAPI';
 import { ContractSchemaFormTab } from './ContractScehmaFormTab';

@ -70,7 +71,9 @@ jest.mock('../../common/Table/Table', () => {
            <span>{item.name}</span>
            <button
              data-testid={`select-row-${item.name}`}
-              onClick={() => rowSelection?.onChange?.([item.name])}>
+              onClick={() =>
+                rowSelection?.onChange?.([item.fullyQualifiedName])
+              }>
              Select {item.name}
            </button>
          </div>
@ -174,7 +177,7 @@ describe('ContractSchemaFormTab', () => {
    it('should render with selected schema columns', () => {
      render(
        <ContractSchemaFormTab
-          selectedSchema={['id', 'name']}
+          selectedSchema={[mockColumns[0], mockColumns[1]]}
          onChange={mockOnChange}
          onNext={mockOnNext}
          onPrev={mockOnPrev}
@ -201,59 +204,6 @@ describe('ContractSchemaFormTab', () => {
    });
  });

-  describe('Data Fetching', () => {
-    it('should fetch table columns on component mount', async () => {
-      render(
-        <ContractSchemaFormTab
-          selectedSchema={[]}
-          onChange={mockOnChange}
-          onNext={mockOnNext}
-          onPrev={mockOnPrev}
-        />
-      );
-
-      await waitFor(() => {
-        expect(getTableColumnsByFQN).toHaveBeenCalledWith(
-          'service.database.schema.table',
-          expect.objectContaining({ fields: 'tags', limit: 15, offset: 0 })
-        );
-      });
-    });
-
-    it('should show loading state during data fetch', async () => {
-      (getTableColumnsByFQN as jest.Mock).mockImplementationOnce(
-        () => new Promise((resolve) => setTimeout(resolve, 100))
-      );
-
-      render(
-        <ContractSchemaFormTab
-          selectedSchema={[]}
-          onChange={mockOnChange}
-          onNext={mockOnNext}
-          onPrev={mockOnPrev}
-        />
-      );
-
-      expect(screen.getByText('Loading: true')).toBeInTheDocument();
-    });
-
-    it('should not fetch data when table FQN is missing', () => {
-      const mockUseFqn = jest.requireMock('../../../hooks/useFqn').useFqn;
-      mockUseFqn.mockReturnValueOnce({ fqn: undefined });
-
-      render(
-        <ContractSchemaFormTab
-          selectedSchema={[]}
-          onChange={mockOnChange}
-          onNext={mockOnNext}
-          onPrev={mockOnPrev}
-        />
-      );
-
-      expect(getTableColumnsByFQN).not.toHaveBeenCalled();
-    });
-  });
-
  describe('Column Selection', () => {
    it('should handle column selection', async () => {
      render(
@ -285,7 +235,7 @@ describe('ContractSchemaFormTab', () => {
    it('should display selected rows in table', async () => {
      render(
        <ContractSchemaFormTab
-          selectedSchema={['id']}
+          selectedSchema={[mockColumns[0]]}
          onChange={mockOnChange}
          onNext={mockOnNext}
          onPrev={mockOnPrev}
@ -434,25 +384,6 @@ describe('ContractSchemaFormTab', () => {
    });
  });

-  describe('Data Processing', () => {
-    it('should process column data correctly', async () => {
-      render(
-        <ContractSchemaFormTab
-          selectedSchema={[]}
-          onChange={mockOnChange}
-          onNext={mockOnNext}
-          onPrev={mockOnPrev}
-        />
-      );
-
-      await waitFor(() => {
-        expect(screen.getByText('id')).toBeInTheDocument();
-        expect(screen.getByText('name')).toBeInTheDocument();
-        expect(screen.getByText('email')).toBeInTheDocument();
-      });
-    });
-  });
-
  describe('Error Handling', () => {
    it('should handle empty table FQN', () => {
      const mockUseFqn = jest.requireMock('../../../hooks/useFqn').useFqn;
@ -489,4 +420,58 @@ describe('ContractSchemaFormTab', () => {
      expect(prevButton).toBeInTheDocument();
    });
  });
+
+  describe('Data Fetching', () => {
+    it('should fetch table columns on component mount', async () => {
+      render(
+        <ContractSchemaFormTab
+          selectedSchema={[]}
+          onChange={mockOnChange}
+          onNext={mockOnNext}
+          onPrev={mockOnPrev}
+        />
+      );
+
+      await waitFor(() => {
+        expect(getTableColumnsByFQN).toHaveBeenCalledWith(
+          'service.database.schema.table',
+          expect.objectContaining({ fields: 'tags', limit: 15, offset: 0 })
+        );
+      });
+    });
+
+    it('should show loading state during data fetch', async () => {
+      (getTableColumnsByFQN as jest.Mock).mockImplementationOnce(
+        () => new Promise((resolve) => setTimeout(resolve, 100))
+      );
+
+      render(
+        <ContractSchemaFormTab
+          selectedSchema={[]}
+          onChange={mockOnChange}
+          onNext={mockOnNext}
+          onPrev={mockOnPrev}
+        />
+      );
+
+      expect(screen.getByText('Loading: true')).toBeInTheDocument();
+    });
+
+    it('should not fetch data when table FQN is missing', () => {
+      (useFqn as jest.Mock).mockImplementation(() => ({
+        fqn: undefined,
+      }));
+
+      render(
+        <ContractSchemaFormTab
+          selectedSchema={[]}
+          onChange={mockOnChange}
+          onNext={mockOnNext}
+          onPrev={mockOnPrev}
+        />
+      );
+
+      expect(getTableColumnsByFQN).not.toHaveBeenCalled();
+    });
+  });
 });