feat(ui): Add table2.keyColumns parameter for table diff validation (#23792)

* feat(ui): Add table2.keyColumns parameter for table diff validation

* address comments

* support key2 columns for edit test case

* fix failing test

openmetadata-ui/src/main/resources/ui/playwright/e2e/Pages/TestCases.spec.ts

@@ -75,6 +75,13 @@ test('Table difference test case', async ({ page }) => {
       );
       await page.getByTestId('tableDiff').click();
       await tableListSearchResponse;
+
+      const table2KeyColumnsInput = page.locator(
+        '#testCaseFormV1_params_table2\\.keyColumns_0_value'
+      );
+
+      await expect(table2KeyColumnsInput).toBeDisabled();
+
       await page.click('#testCaseFormV1_params_table2');
       await page.waitForSelector(`[data-id="tableDiff"]`, {
         state: 'visible',
@@ -107,6 +114,15 @@ test('Table difference test case', async ({ page }) => {
         table1.entity?.columns[0].name
       );
       await page.getByTitle(table1.entity?.columns[0].name).click();
+
+      await page.fill(
+        '#testCaseFormV1_params_table2\\.keyColumns_0_value',
+        table2.entity?.columns[0].name
+      );
+      await page.getByTitle(table2.entity?.columns[0].name).click();
+
+      await expect(table2KeyColumnsInput).not.toBeDisabled();
+
       await page.fill('#testCaseFormV1_params_threshold', testCase.threshold);
       await page.fill(
         '#testCaseFormV1_params_useColumns_0_value',
@@ -160,6 +176,44 @@ test('Table difference test case', async ({ page }) => {
         `Edit ${testCase.name}`
       );
 
+      // Wait for form to finish loading (isLoading becomes false)
+      await expect(page.getByTestId('edit-test-form')).toBeVisible();
+
+      // Verify Table 1's keyColumns is enabled and populated in edit mode
+      const table1KeyColumnsEditInput = page.locator(
+        '#tableTestForm_params_keyColumns_0_value'
+      );
+
+      // Wait for the input to be visible and enabled, and then check its value
+      await expect(table1KeyColumnsEditInput).toBeVisible();
+      await expect(table1KeyColumnsEditInput).not.toBeDisabled();
+
+      // Wait for the value to be populated
+      // Use data-testid to find the select component
+      const columnName = table1.entity?.columns[0].name;
+      const table1Select = page.getByTestId('keyColumns-select');
+
+      // Wait for the select to be visible and verify the selected value is displayed
+      await expect(table1Select).toBeVisible();
+      await expect(table1Select.getByText(columnName)).toBeVisible();
+
+      // Verify table2.keyColumns is enabled and populated in edit mode
+      const table2KeyColumnsEditInput = page.locator(
+        '#tableTestForm_params_table2\\.keyColumns_0_value'
+      );
+
+      // Wait for the input to be visible and enabled, and then check its value
+      await expect(table2KeyColumnsEditInput).toBeVisible();
+      await expect(table2KeyColumnsEditInput).not.toBeDisabled();
+
+      // Wait for the value to be populated
+      const table2ColumnName = table2.entity?.columns[0].name;
+      const table2Select = page.getByTestId('table2.keyColumns-select');
+
+      // Wait for the select to be visible and verify the selected value is displayed
+      await expect(table2Select).toBeVisible();
+      await expect(table2Select.getByText(table2ColumnName)).toBeVisible();
+
       await page
         .locator('label')
         .filter({ hasText: "Table 1's key columns" })

openmetadata-ui/src/main/resources/ui/src/components/DataQuality/AddDataQualityTest/components/ParameterForm.test.tsx

@@ -18,6 +18,7 @@ import { TestDefinition } from '../../../../generated/tests/testDefinition';
 import {
   MOCK_TABLE_COLUMN_NAME_TO_EXIST,
   MOCK_TABLE_CUSTOM_SQL_QUERY,
+  MOCK_TABLE_DIFF_DEFINITION,
   MOCK_TABLE_ROW_INSERTED_COUNT_TO_BE_BETWEEN,
   MOCK_TABLE_TEST_WITH_COLUMN,
   MOCK_TABLE_WITH_DATE_TIME_COLUMNS,
@@ -72,6 +73,12 @@ jest.mock('../../../../rest/searchAPI', () => {
   };
 });
 
+jest.mock('../../../../rest/tableAPI', () => {
+  return {
+    getTableDetailsByFQN: jest.fn(),
+  };
+});
+
 const renderWithForm = (component: React.ReactElement) => {
   return render(<Form>{component}</Form>);
 };
@@ -183,4 +190,58 @@ describe('ParameterForm component test', () => {
 
     expect(selectBox).toBeInTheDocument();
   });
+
+  describe('Table Diff functionality', () => {
+    beforeEach(() => {
+      jest.clearAllMocks();
+    });
+
+    it('Should render table2 parameter for table diff definition', async () => {
+      await act(async () => {
+        renderWithForm(
+          <ParameterForm
+            definition={MOCK_TABLE_DIFF_DEFINITION as TestDefinition}
+            table={MOCK_TABLE_WITH_DATE_TIME_COLUMNS}
+          />
+        );
+      });
+
+      const table2Select = await screen.findByTestId('table2');
+
+      expect(table2Select).toBeInTheDocument();
+    });
+
+    it('Should render all parameters when table has columns', async () => {
+      await act(async () => {
+        renderWithForm(
+          <ParameterForm
+            definition={MOCK_TABLE_DIFF_DEFINITION as TestDefinition}
+            table={MOCK_TABLE_WITH_DATE_TIME_COLUMNS}
+          />
+        );
+      });
+
+      const parameters = await screen.findAllByTestId('parameter');
+
+      expect(parameters.length).toBeGreaterThan(0);
+    });
+
+    it('Should have table2.keyColumns disabled when table2 is not selected', async () => {
+      await act(async () => {
+        renderWithForm(
+          <ParameterForm
+            definition={MOCK_TABLE_DIFF_DEFINITION as TestDefinition}
+            table={MOCK_TABLE_WITH_DATE_TIME_COLUMNS}
+          />
+        );
+      });
+
+      const keyColumnsInputs = screen.getAllByRole('combobox');
+      const table2KeyColumnsInput = keyColumnsInputs.find((input) =>
+        input.id.includes('table2.keyColumns')
+      );
+
+      expect(table2KeyColumnsInput).toBeDisabled();
+    });
+  });
 });

openmetadata-ui/src/main/resources/ui/src/components/DataQuality/AddDataQualityTest/components/ParameterForm.tsx

@@ -36,6 +36,7 @@ import { SUPPORTED_PARTITION_TYPE_FOR_DATE_TIME } from '../../../../constants/pr
 import { TABLE_DIFF } from '../../../../constants/TestSuite.constant';
 import { CSMode } from '../../../../enums/codemirror.enum';
 import { SearchIndex } from '../../../../enums/search.enum';
+import { Column } from '../../../../generated/entity/data/table';
 import {
   Rule,
   TestCaseParameterDefinition,
@@ -49,6 +50,7 @@ import { searchQuery } from '../../../../rest/searchAPI';
 import { getEntityName } from '../../../../utils/EntityUtils';
 import { getPopupContainer } from '../../../../utils/formUtils';
 import {
+  getSelectedColumnsSet,
   validateEquals,
   validateGreaterThanOrEquals,
   validateLessThanOrEquals,
@@ -320,24 +322,29 @@ const ParameterForm: React.FC<ParameterFormProps> = ({ definition, table }) => {
   };
 
   const TableDiffForm = () => {
+    const form = Form.useFormInstance();
     const [isOptionsLoading, setIsOptionsLoading] = useState(false);
     const [tableList, setTableList] = useState<
       SearchHitBody<
         SearchIndex.TABLE,
-        Pick<TableSearchSource, 'name' | 'displayName' | 'fullyQualifiedName'>
+        Pick<
+          TableSearchSource,
+          'name' | 'displayName' | 'fullyQualifiedName' | 'columns'
+        >
       >[]
     >([]);
+    const [table2Columns, setTable2Columns] = useState<Column[] | undefined>();
+
     const tableOptions = useMemo(
       () =>
-        tableList.map((hit) => {
-          return {
+        tableList.map((hit) => ({
           label: hit._source.fullyQualifiedName,
           value: hit._source.fullyQualifiedName,
-          };
-        }),
+        })),
       [tableList]
     );
-    const fetchTableData = async (search = WILD_CARD_CHAR) => {
+
+    const fetchTableData = useCallback(async (search = WILD_CARD_CHAR) => {
       setIsOptionsLoading(true);
       try {
         const response = await searchQuery({
@@ -346,25 +353,49 @@ const ParameterForm: React.FC<ParameterFormProps> = ({ definition, table }) => {
           pageSize: PAGE_SIZE_LARGE,
           searchIndex: SearchIndex.TABLE,
           fetchSource: true,
-          includeFields: ['name', 'fullyQualifiedName', 'displayName'],
+          includeFields: [
+            'name',
+            'fullyQualifiedName',
+            'displayName',
+            'columns',
+          ],
         });
-
         setTableList(response.hits.hits);
-      } catch (error) {
+      } catch {
         setTableList([]);
       } finally {
         setIsOptionsLoading(false);
       }
-    };
+    }, []);
 
-    const debounceFetchTableData = useCallback(debounce(fetchTableData, 1000), [
-      fetchTableData,
-    ]);
+    const debounceFetchTableData = useMemo(
+      () => debounce(fetchTableData, 1000),
+      [fetchTableData]
+    );
+
+    useEffect(() => {
+      fetchTableData();
+    }, [fetchTableData]);
+
+    useEffect(() => {
+      const table2Value = form.getFieldValue(['params', 'table2']);
+      if (table2Value && !table2Columns && tableList.length > 0) {
+        const selectedTable = tableList.find(
+          (hit) => hit._source.fullyQualifiedName === table2Value
+        );
+        if (selectedTable) {
+          setTable2Columns(selectedTable._source.columns);
+        }
+      }
+    }, [tableList, table2Columns, form]);
 
     const getFormData = (data: TestCaseParameterDefinition) => {
       switch (data.name) {
         case 'table2':
-          return prepareForm(
+          return (
+            <Form.Item noStyle shouldUpdate key={data.name}>
+              {({ setFieldsValue }) =>
+                prepareForm(
                   data,
                   <Select
                     allowClear
@@ -375,47 +406,76 @@ const ParameterForm: React.FC<ParameterFormProps> = ({ definition, table }) => {
                     options={tableOptions}
                     placeholder={t('label.table')}
                     popupClassName="no-wrap-option"
+                    onChange={(value) => {
+                      // Clear key columns when table2 changes
+                      setFieldsValue({
+                        params: { 'table2.keyColumns': [{ value: undefined }] },
+                      });
+
+                      // Update columns or clear them
+                      if (value) {
+                        const selectedTable = tableList.find(
+                          (hit) => hit._source.fullyQualifiedName === value
+                        );
+                        setTable2Columns(selectedTable?._source.columns);
+                      } else {
+                        setTable2Columns(undefined);
+                      }
+                    }}
                     onSearch={debounceFetchTableData}
                   />
+                )
+              }
+            </Form.Item>
           );
 
         case 'keyColumns':
+        case 'table2.keyColumns':
         case 'useColumns':
           return (
-            <Form.Item noStyle shouldUpdate>
+            <Form.Item noStyle shouldUpdate key={data.name}>
               {({ getFieldValue }) => {
-                // Convert selectedKeyColumn and selectedUseColumns to Sets for efficient lookup
-                const selectedKeyColumnSet = new Set(
-                  getFieldValue(['params', 'keyColumns'])?.map(
-                    (item: { value: string }) => item?.value
-                  )
-                );
-                const selectedUseColumnsSet = new Set(
-                  getFieldValue(['params', 'useColumns'])?.map(
-                    (item: { value: string }) => item?.value
-                  )
+                const isTable2KeyColumns = data.name === 'table2.keyColumns';
+                const table2Value = getFieldValue(['params', 'table2']);
+
+                let sourceColumns = table?.columns;
+                if (isTable2KeyColumns) {
+                  if (table2Value) {
+                    const selectedTable =
+                      tableList.find(
+                        (hit) => hit._source.fullyQualifiedName === table2Value
+                      ) ?? undefined;
+                    sourceColumns =
+                      selectedTable?._source.columns ?? table2Columns;
+                  } else {
+                    sourceColumns = undefined;
+                  }
+                }
+
+                // Disable when no table2 selected
+                const isDisabled = isTable2KeyColumns && !table2Value;
+
+                const selectedColumnsSet = getSelectedColumnsSet(
+                  data,
+                  getFieldValue
                 );
 
-                // Combine both Sets for a single lookup operation
-                const selectedColumnsSet = new Set([
-                  ...selectedKeyColumnSet,
-                  ...selectedUseColumnsSet,
-                ]);
-
-                const columns = table?.columns.map((column) => ({
+                const columnOptions =
+                  sourceColumns?.map((column) => ({
                     label: getEntityName(column),
                     value: column.name,
-                  // Check if column.name is in the combined Set to determine if it should be disabled
                     disabled: selectedColumnsSet.has(column.name),
-                }));
+                  })) ?? [];
 
                 return prepareForm(
                   data,
                   <Select
                     allowClear
                     showSearch
+                    data-testid={`${data.name}-select`}
+                    disabled={isDisabled}
                     getPopupContainer={getPopupContainer}
-                    options={columns}
+                    options={columnOptions}
                     placeholder={t('label.column')}
                   />
                 );
@@ -428,10 +488,6 @@ const ParameterForm: React.FC<ParameterFormProps> = ({ definition, table }) => {
       }
     };
 
-    useEffect(() => {
-      fetchTableData();
-    }, []);
-
     return <>{definition.parameterDefinition?.map(getFormData)}</>;
   };
 

openmetadata-ui/src/main/resources/ui/src/mocks/TestSuite.mock.ts

@@ -567,6 +567,52 @@ export const MOCK_TABLE_COLUMN_NAME_TO_EXIST = {
   deleted: false,
 };
 
+export const MOCK_TABLE_DIFF_DEFINITION = {
+  id: 'table-diff-test-id',
+  name: 'tableDiff',
+  displayName: 'Table Diff',
+  fullyQualifiedName: 'tableDiff',
+  description: 'Compare two tables',
+  entityType: 'TABLE',
+  testPlatforms: ['OpenMetadata'],
+  supportedDataTypes: [],
+  parameterDefinition: [
+    {
+      name: 'table2',
+      displayName: 'Table 2',
+      dataType: 'STRING',
+      description: 'Second table to compare',
+      required: true,
+    },
+    {
+      name: 'keyColumns',
+      displayName: 'Key Columns',
+      dataType: 'ARRAY',
+      description: 'Key columns for table 1',
+      required: false,
+    },
+    {
+      name: 'table2.keyColumns',
+      displayName: 'Table 2 Key Columns',
+      dataType: 'ARRAY',
+      description: 'Key columns for table 2',
+      required: false,
+    },
+    {
+      name: 'useColumns',
+      displayName: 'Use Columns',
+      dataType: 'ARRAY',
+      description: 'Additional columns to compare',
+      required: false,
+    },
+  ],
+  version: 0.1,
+  updatedAt: 1675211404184,
+  updatedBy: 'admin',
+  href: 'http://localhost:8585/api/v1/dataQuality/testDefinitions/table-diff-test-id',
+  deleted: false,
+};
+
 export const MOCK_TABLE_TEST_WITH_COLUMN = {
   id: 'id',
   name: 'tableColumnNameToExist',

openmetadata-ui/src/main/resources/ui/src/utils/ParameterForm/ParameterFormUtils.test.ts

@@ -10,7 +10,10 @@
  *  See the License for the specific language governing permissions and
  *  limitations under the License.
  */
+import { TestCaseParameterDefinition } from '../../generated/tests/testDefinition';
 import {
+  getColumnSet,
+  getSelectedColumnsSet,
   validateEquals,
   validateGreaterThanOrEquals,
   validateLessThanOrEquals,
@@ -113,4 +116,199 @@ describe('ParameterFormUtils', () => {
       );
     });
   });
+
+  describe('getColumnSet', () => {
+    it('should return empty Set when field value is undefined', () => {
+      const mockGetFieldValue = jest.fn().mockReturnValue(undefined);
+      const result = getColumnSet(mockGetFieldValue, 'keyColumns');
+
+      expect(result).toEqual(new Set());
+      expect(mockGetFieldValue).toHaveBeenCalledWith(['params', 'keyColumns']);
+    });
+
+    it('should return empty Set when field value is null', () => {
+      const mockGetFieldValue = jest.fn().mockReturnValue(null);
+      const result = getColumnSet(mockGetFieldValue, 'keyColumns');
+
+      expect(result).toEqual(new Set());
+    });
+
+    it('should return empty Set when field value is not an array', () => {
+      const mockGetFieldValue = jest.fn().mockReturnValue('not-an-array');
+      const result = getColumnSet(mockGetFieldValue, 'keyColumns');
+
+      expect(result).toEqual(new Set());
+    });
+
+    it('should return empty Set when field value is an empty array', () => {
+      const mockGetFieldValue = jest.fn().mockReturnValue([]);
+      const result = getColumnSet(mockGetFieldValue, 'keyColumns');
+
+      expect(result).toEqual(new Set());
+    });
+
+    it('should return Set with column values from array of objects', () => {
+      const mockGetFieldValue = jest
+        .fn()
+        .mockReturnValue([
+          { value: 'col1' },
+          { value: 'col2' },
+          { value: 'col3' },
+        ]);
+      const result = getColumnSet(mockGetFieldValue, 'keyColumns');
+
+      expect(result).toEqual(new Set(['col1', 'col2', 'col3']));
+      expect(result.size).toBe(3);
+    });
+
+    it('should handle objects with undefined value property', () => {
+      const mockGetFieldValue = jest
+        .fn()
+        .mockReturnValue([
+          { value: 'col1' },
+          { value: undefined },
+          { value: 'col2' },
+        ]);
+      const result = getColumnSet(mockGetFieldValue, 'keyColumns');
+
+      expect(result).toEqual(new Set(['col1', undefined, 'col2']));
+    });
+
+    it('should remove duplicates from column values', () => {
+      const mockGetFieldValue = jest
+        .fn()
+        .mockReturnValue([
+          { value: 'col1' },
+          { value: 'col2' },
+          { value: 'col1' },
+        ]);
+      const result = getColumnSet(mockGetFieldValue, 'keyColumns');
+
+      expect(result).toEqual(new Set(['col1', 'col2']));
+      expect(result.size).toBe(2);
+    });
+  });
+
+  describe('getSelectedColumnsSet', () => {
+    const mockGetFieldValue = jest.fn();
+
+    beforeEach(() => {
+      mockGetFieldValue.mockClear();
+    });
+
+    it('should return only table2.keyColumns for table2.keyColumns field', () => {
+      mockGetFieldValue.mockImplementation((path) => {
+        if (path[1] === 'table2.keyColumns') {
+          return [{ value: 't2col1' }, { value: 't2col2' }];
+        }
+
+        return [];
+      });
+
+      const data = { name: 'table2.keyColumns' } as TestCaseParameterDefinition;
+      const result = getSelectedColumnsSet(data, mockGetFieldValue);
+
+      expect(result).toEqual(new Set(['t2col1', 't2col2']));
+      expect(result.size).toBe(2);
+    });
+
+    it('should return merged keyColumns and useColumns for keyColumns field', () => {
+      mockGetFieldValue.mockImplementation((path) => {
+        if (path[1] === 'keyColumns') {
+          return [{ value: 'col1' }, { value: 'col2' }];
+        }
+        if (path[1] === 'useColumns') {
+          return [{ value: 'col3' }, { value: 'col4' }];
+        }
+
+        return [];
+      });
+
+      const data = { name: 'keyColumns' } as TestCaseParameterDefinition;
+      const result = getSelectedColumnsSet(data, mockGetFieldValue);
+
+      expect(result).toEqual(new Set(['col1', 'col2', 'col3', 'col4']));
+      expect(result.size).toBe(4);
+    });
+
+    it('should return merged keyColumns and useColumns for useColumns field', () => {
+      mockGetFieldValue.mockImplementation((path) => {
+        if (path[1] === 'keyColumns') {
+          return [{ value: 'col1' }];
+        }
+        if (path[1] === 'useColumns') {
+          return [{ value: 'col2' }];
+        }
+
+        return [];
+      });
+
+      const data = { name: 'useColumns' } as TestCaseParameterDefinition;
+      const result = getSelectedColumnsSet(data, mockGetFieldValue);
+
+      expect(result).toEqual(new Set(['col1', 'col2']));
+    });
+
+    it('should handle empty arrays for keyColumns field', () => {
+      mockGetFieldValue.mockReturnValue([]);
+
+      const data = { name: 'keyColumns' } as TestCaseParameterDefinition;
+      const result = getSelectedColumnsSet(data, mockGetFieldValue);
+
+      expect(result).toEqual(new Set());
+    });
+
+    it('should handle undefined values for table2.keyColumns field', () => {
+      mockGetFieldValue.mockReturnValue(undefined);
+
+      const data = { name: 'table2.keyColumns' } as TestCaseParameterDefinition;
+      const result = getSelectedColumnsSet(data, mockGetFieldValue);
+
+      expect(result).toEqual(new Set());
+    });
+
+    it('should remove duplicates when merging keyColumns and useColumns', () => {
+      mockGetFieldValue.mockImplementation((path) => {
+        if (path[1] === 'keyColumns') {
+          return [{ value: 'col1' }, { value: 'col2' }];
+        }
+        if (path[1] === 'useColumns') {
+          return [{ value: 'col2' }, { value: 'col3' }];
+        }
+
+        return [];
+      });
+
+      const data = { name: 'keyColumns' } as TestCaseParameterDefinition;
+      const result = getSelectedColumnsSet(data, mockGetFieldValue);
+
+      expect(result).toEqual(new Set(['col1', 'col2', 'col3']));
+      expect(result.size).toBe(3);
+    });
+
+    it('should handle table2.keyColumns being independent from other columns', () => {
+      mockGetFieldValue.mockImplementation((path) => {
+        if (path[1] === 'keyColumns') {
+          return [{ value: 'col1' }];
+        }
+        if (path[1] === 'table2.keyColumns') {
+          return [{ value: 't2col1' }];
+        }
+        if (path[1] === 'useColumns') {
+          return [{ value: 'col2' }];
+        }
+
+        return [];
+      });
+
+      const dataTable2 = {
+        name: 'table2.keyColumns',
+      } as TestCaseParameterDefinition;
+      const resultTable2 = getSelectedColumnsSet(dataTable2, mockGetFieldValue);
+
+      expect(resultTable2).toEqual(new Set(['t2col1']));
+      expect(resultTable2).not.toContain('col1');
+      expect(resultTable2).not.toContain('col2');
+    });
+  });
 });

openmetadata-ui/src/main/resources/ui/src/utils/ParameterForm/ParameterFormUtils.ts

@@ -10,6 +10,7 @@
  *  See the License for the specific language governing permissions and
  *  limitations under the License.
  */
+import { TestCaseParameterDefinition } from '../../generated/tests/testDefinition';
 import i18next from '../../utils/i18next/LocalUtil';
 
 export const validateGreaterThanOrEquals = (
@@ -56,3 +57,34 @@ export const validateNotEquals = (fieldValue: number, value: number) => {
 
   return Promise.resolve();
 };
+
+export const getColumnSet = (
+  getFieldValue: (path: (string | number)[]) => unknown,
+  fieldName: string
+): Set<string> => {
+  const columnValues = getFieldValue(['params', fieldName]);
+
+  if (!columnValues || !Array.isArray(columnValues)) {
+    return new Set();
+  }
+
+  return new Set(
+    (columnValues as { value: string }[]).map((item) => item?.value)
+  );
+};
+
+export const getSelectedColumnsSet = (
+  data: TestCaseParameterDefinition,
+  getFieldValue: (path: (string | number)[]) => unknown
+): Set<string> => {
+  const isTable2KeyColumns = data.name === 'table2.keyColumns';
+
+  if (isTable2KeyColumns) {
+    return getColumnSet(getFieldValue, 'table2.keyColumns');
+  }
+
+  const keyColumns = getColumnSet(getFieldValue, 'keyColumns');
+  const useColumns = getColumnSet(getFieldValue, 'useColumns');
+
+  return new Set<string>([...keyColumns, ...useColumns]);
+};