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

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,
   }) => {

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={
+            selectedSchema={contract?.schema ?? []}
-              contract?.schema?.map((column) => column.name) || []
-            }
             onChange={onFormChange}
             onNext={onNext}
             onPrev={onPrev}

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 [allColumnsData, setAllColumnData] = useState<Column[]>([]);
-  const [selectedKeys, setSelectedKeys] = useState<string[]>(selectedSchema);
+  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) =>
+        schema: selectedColumns,
-          selectedRowKeys.includes(column.name)
-        ),
       });
     },
-    [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>

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();
+    });
+  });
 });