diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/DataContractRepository.java b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/DataContractRepository.java index 6ec8324a6ca..bae2933365d 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/DataContractRepository.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/DataContractRepository.java @@ -304,21 +304,27 @@ public class DataContractRepository extends EntityRepository { return fieldNames; } + private String getTestSuiteName(DataContract dataContract) { + return dataContract.getName() + " - Data Contract Expectations"; + } + private TestSuite createOrUpdateDataContractTestSuite(DataContract dataContract) { - if (nullOrEmpty(dataContract.getQualityExpectations())) { - return null; // No quality expectations, no test suite needed - } try { + // If we don't have quality expectations or a test suite, we don't need to create one + if (nullOrEmpty(dataContract.getQualityExpectations()) + && !contractHasTestSuite(dataContract)) { + return null; + } + + // If we had a test suite from older tests, but we removed them, we can delete the suite + if (nullOrEmpty(dataContract.getQualityExpectations())) { + deleteTestSuite(dataContract); + dataContract.setTestSuite(null); + return null; + } - TestCaseRepository testCaseRepository = - (TestCaseRepository) Entity.getEntityRepository(Entity.TEST_CASE); TestSuite testSuite = getOrCreateTestSuite(dataContract); - - // Collect test case references from quality expectations - List testCaseRefs = - dataContract.getQualityExpectations().stream().map(EntityReference::getId).toList(); - - testCaseRepository.addTestCasesToLogicalTestSuite(testSuite, testCaseRefs); + updateTestSuiteTests(dataContract, testSuite); // Add the test suite to the data contract dataContract.setTestSuite( @@ -335,8 +341,45 @@ public class DataContractRepository extends EntityRepository { } } + private void updateTestSuiteTests(DataContract dataContract, TestSuite testSuite) { + TestCaseRepository testCaseRepository = + (TestCaseRepository) Entity.getEntityRepository(Entity.TEST_CASE); + + // Collect test case references from quality expectations + List testCaseRefs = + dataContract.getQualityExpectations().stream().map(EntityReference::getId).toList(); + List currentTests = + testSuite.getTests() != null + ? testSuite.getTests().stream().map(EntityReference::getId).toList() + : Collections.emptyList(); + + // Add only new tests to the test suite + List newTestCases = + testCaseRefs.stream().filter(testCaseRef -> !currentTests.contains(testCaseRef)).toList(); + if (!nullOrEmpty(newTestCases)) { + testCaseRepository.addTestCasesToLogicalTestSuite(testSuite, newTestCases); + } + + // Then, remove any tests that are no longer in the quality expectations + List testsToRemove = + currentTests.stream().filter(testId -> !testCaseRefs.contains(testId)).toList(); + if (!nullOrEmpty(testsToRemove)) { + testsToRemove.forEach( + test -> { + testCaseRepository.deleteTestCaseFromLogicalTestSuite(testSuite.getId(), test); + }); + } + } + + private void deleteTestSuite(DataContract dataContract) { + TestSuiteRepository testSuiteRepository = + (TestSuiteRepository) Entity.getEntityRepository(Entity.TEST_SUITE); + TestSuite testSuite = getOrCreateTestSuite(dataContract); + testSuiteRepository.deleteLogicalTestSuite(ADMIN_USER_NAME, testSuite, true); + } + private TestSuite getOrCreateTestSuite(DataContract dataContract) { - String testSuiteName = dataContract.getName() + " - Data Contract Expectations"; + String testSuiteName = getTestSuiteName(dataContract); TestSuiteRepository testSuiteRepository = (TestSuiteRepository) Entity.getEntityRepository(Entity.TEST_SUITE); @@ -371,6 +414,20 @@ public class DataContractRepository extends EntityRepository { return maybeTestSuite.get(); } + private Boolean contractHasTestSuite(DataContract dataContract) { + TestSuiteRepository testSuiteRepository = + (TestSuiteRepository) Entity.getEntityRepository(Entity.TEST_SUITE); + + String testSuiteName = getTestSuiteName(dataContract); + + // Check if test suite already exists + Optional maybeTestSuite = + testSuiteRepository.getByNameOrNull( + null, testSuiteName, testSuiteRepository.getFields(""), Include.NON_DELETED, false); + + return maybeTestSuite.isPresent(); + } + // Prepare the Ingestion Pipeline from the test suite that will handle the execution private IngestionPipeline createIngestionPipeline(TestSuite testSuite) { IngestionPipelineRepository pipelineRepository = diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/TestSuiteRepository.java b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/TestSuiteRepository.java index 3b3d9c1984d..f110fc5440e 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/TestSuiteRepository.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/TestSuiteRepository.java @@ -527,10 +527,9 @@ public class TestSuiteRepository extends EntityRepository { } public RestUtil.DeleteResponse deleteLogicalTestSuite( - SecurityContext securityContext, TestSuite original, boolean hardDelete) { + String updatedBy, TestSuite original, boolean hardDelete) { // deleting a logical will delete the test suite and only remove the relationship to // test cases if hardDelete is true. Test Cases will not be deleted. - String updatedBy = securityContext.getUserPrincipal().getName(); preDelete(original, updatedBy); setFieldsInternal(original, putFields); deleteChildIngestionPipelines(original.getId(), hardDelete, updatedBy); @@ -581,7 +580,8 @@ public class TestSuiteRepository extends EntityRepository { () -> { try { RestUtil.DeleteResponse deleteResponse = - deleteLogicalTestSuite(securityContext, testSuite, hardDelete); + deleteLogicalTestSuite( + securityContext.getUserPrincipal().getName(), testSuite, hardDelete); deleteFromSearch(deleteResponse.entity(), hardDelete); WebsocketNotificationHandler.sendDeleteOperationCompleteNotification( diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/resources/dqtests/TestSuiteResource.java b/openmetadata-service/src/main/java/org/openmetadata/service/resources/dqtests/TestSuiteResource.java index 2da1d18033d..4cfd26bbe5f 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/resources/dqtests/TestSuiteResource.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/resources/dqtests/TestSuiteResource.java @@ -742,7 +742,8 @@ public class TestSuiteResource extends EntityResource response = - repository.deleteLogicalTestSuite(securityContext, testSuite, hardDelete); + repository.deleteLogicalTestSuite( + securityContext.getUserPrincipal().getName(), testSuite, hardDelete); repository.deleteFromSearch(response.entity(), hardDelete); addHref(uriInfo, response.entity()); return response.toResponse(); @@ -810,7 +811,8 @@ public class TestSuiteResource extends EntityResource response = - repository.deleteLogicalTestSuite(securityContext, testSuite, hardDelete); + repository.deleteLogicalTestSuite( + securityContext.getUserPrincipal().getName(), testSuite, hardDelete); addHref(uriInfo, response.entity()); return response.toResponse(); } diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/resources/data/DataContractResourceTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/resources/data/DataContractResourceTest.java index 1f6d9fcbf56..9ab64c0fdbc 100644 --- a/openmetadata-service/src/test/java/org/openmetadata/service/resources/data/DataContractResourceTest.java +++ b/openmetadata-service/src/test/java/org/openmetadata/service/resources/data/DataContractResourceTest.java @@ -15,6 +15,7 @@ package org.openmetadata.service.resources.data; import static jakarta.ws.rs.core.Response.Status.BAD_REQUEST; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertThrows; @@ -41,6 +42,7 @@ import jakarta.ws.rs.core.Response.Status; import java.io.IOException; import java.net.URISyntaxException; import java.util.ArrayList; +import java.util.Collections; import java.util.List; import java.util.Map; import java.util.UUID; @@ -112,6 +114,7 @@ public class DataContractResourceTest extends EntityResourceTest", table.getFullyQualifiedName()); + + CreateTestCase createTestCase1 = + testCaseResourceTest + .createRequest("test_case_completeness_" + test.getDisplayName()) + .withEntityLink(tableLink); + TestCase testCase1 = + testCaseResourceTest.createAndCheckEntity(createTestCase1, ADMIN_AUTH_HEADERS); + + CreateTestCase createTestCase2 = + testCaseResourceTest + .createRequest("test_case_validity_" + test.getDisplayName()) + .withEntityLink(tableLink); + TestCase testCase2 = + testCaseResourceTest.createAndCheckEntity(createTestCase2, ADMIN_AUTH_HEADERS); + + // Step 1: Add initial quality expectations with both test cases + List initialQualityExpectations = + List.of(testCase1.getEntityReference(), testCase2.getEntityReference()); + + create.withStatus(ContractStatus.Active).withQualityExpectations(initialQualityExpectations); + DataContract updated = updateDataContract(create); + + // Verify initial state - both test cases should be in TestSuite + assertEquals(2, updated.getQualityExpectations().size()); + assertNotNull(updated.getTestSuite()); + + // Get TestSuite and verify it contains both tests + TestSuite testSuite = + testSuiteResourceTest.getEntity( + updated.getTestSuite().getId(), "tests", ADMIN_AUTH_HEADERS); + + assertNotNull(testSuite.getTests()); + assertEquals(2, testSuite.getTests().size()); + assertTrue(testSuite.getTests().stream().anyMatch(t -> t.getId().equals(testCase1.getId()))); + assertTrue(testSuite.getTests().stream().anyMatch(t -> t.getId().equals(testCase2.getId()))); + + // Step 2: Create a third test case + CreateTestCase createTestCase3 = + testCaseResourceTest + .createRequest("test_case_accuracy_" + test.getDisplayName()) + .withEntityLink(tableLink); + TestCase testCase3 = + testCaseResourceTest.createAndCheckEntity(createTestCase3, ADMIN_AUTH_HEADERS); + + // Update to add new test case while keeping existing ones + List expandedQualityExpectations = + List.of( + testCase1.getEntityReference(), + testCase2.getEntityReference(), + testCase3.getEntityReference()); + + create.withQualityExpectations(expandedQualityExpectations); + DataContract updatedWithNewTest = updateDataContract(create); + + // Verify new test case was added + assertEquals(3, updatedWithNewTest.getQualityExpectations().size()); + + // Verify TestSuite now contains all three tests + TestSuite expandedTestSuite = + testSuiteResourceTest.getEntity( + updated.getTestSuite().getId(), "tests", ADMIN_AUTH_HEADERS); + + assertNotNull(expandedTestSuite.getTests()); + assertEquals(3, expandedTestSuite.getTests().size()); + assertTrue( + expandedTestSuite.getTests().stream().anyMatch(t -> t.getId().equals(testCase1.getId()))); + assertTrue( + expandedTestSuite.getTests().stream().anyMatch(t -> t.getId().equals(testCase2.getId()))); + assertTrue( + expandedTestSuite.getTests().stream().anyMatch(t -> t.getId().equals(testCase3.getId()))); + + // Step 3: Remove one test case (remove testCase2) + List reducedQualityExpectations = + List.of(testCase1.getEntityReference(), testCase3.getEntityReference()); + + create.withQualityExpectations(reducedQualityExpectations); + DataContract updatedWithRemovedTest = updateDataContract(create); + + // Verify test case was removed from quality expectations + assertEquals(2, updatedWithRemovedTest.getQualityExpectations().size()); + assertTrue( + updatedWithRemovedTest.getQualityExpectations().stream() + .anyMatch(ref -> ref.getId().equals(testCase1.getId()))); + assertTrue( + updatedWithRemovedTest.getQualityExpectations().stream() + .anyMatch(ref -> ref.getId().equals(testCase3.getId()))); + assertFalse( + updatedWithRemovedTest.getQualityExpectations().stream() + .anyMatch(ref -> ref.getId().equals(testCase2.getId()))); + + // Verify TestSuite also had the test removed + TestSuite reducedTestSuite = + testSuiteResourceTest.getEntity( + updated.getTestSuite().getId(), "tests", ADMIN_AUTH_HEADERS); + + assertNotNull(reducedTestSuite.getTests()); + assertEquals(2, reducedTestSuite.getTests().size()); + assertTrue( + reducedTestSuite.getTests().stream().anyMatch(t -> t.getId().equals(testCase1.getId()))); + assertFalse( + reducedTestSuite.getTests().stream().anyMatch(t -> t.getId().equals(testCase2.getId()))); + assertTrue( + reducedTestSuite.getTests().stream().anyMatch(t -> t.getId().equals(testCase3.getId()))); + + // Step 4: Remove all test cases + create.withQualityExpectations(Collections.emptyList()); + DataContract updatedWithNoTests = updateDataContract(create); + + // Verify all test cases were removed + assertTrue( + updatedWithNoTests.getQualityExpectations() == null + || updatedWithNoTests.getQualityExpectations().isEmpty()); + + // Verify TestSuite has been deleted + assertResponseContains( + () -> + testSuiteResourceTest.getEntity( + updated.getTestSuite().getId(), "tests", ADMIN_AUTH_HEADERS), + Status.NOT_FOUND, + "testSuite instance for " + updated.getTestSuite().getId().toString() + " not found"); + + // GET the data contract and verify persistence + DataContract retrieved = getDataContract(created.getId(), ""); + assertTrue( + retrieved.getQualityExpectations() == null || retrieved.getQualityExpectations().isEmpty()); + assertNull(retrieved.getTestSuite()); // We deleted the TestSuite when no tests remain + } + @Test @Execution(ExecutionMode.CONCURRENT) @Override diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/rules/RuleEngineTests.java b/openmetadata-service/src/test/java/org/openmetadata/service/rules/RuleEngineTests.java index 1edce81e38e..df9c36db312 100644 --- a/openmetadata-service/src/test/java/org/openmetadata/service/rules/RuleEngineTests.java +++ b/openmetadata-service/src/test/java/org/openmetadata/service/rules/RuleEngineTests.java @@ -3,7 +3,9 @@ package org.openmetadata.service.rules; import static org.junit.Assert.assertThrows; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.openmetadata.service.resources.EntityResourceTest.PII_SENSITIVE_TAG_LABEL; import static org.openmetadata.service.resources.EntityResourceTest.TEAM11_REF; +import static org.openmetadata.service.resources.EntityResourceTest.TIER1_TAG_LABEL; import static org.openmetadata.service.resources.EntityResourceTest.USER1_REF; import static org.openmetadata.service.resources.EntityResourceTest.USER2_REF; import static org.openmetadata.service.util.TestUtils.ADMIN_AUTH_HEADERS; @@ -433,4 +435,32 @@ public class RuleEngineTests extends OpenMetadataApplicationTest { // Ignored the table, should not pass Assertions.assertFalse(RuleEngine.getInstance().shouldApplyRule(table, ruleWithIgnoredTable)); } + + @Test + @Execution(ExecutionMode.CONCURRENT) + void testPIISensitiveTagRule(TestInfo test) { + Table table = getMockTable(test).withTags(List.of()); + + SemanticsRule piiRule = + new SemanticsRule() + .withName("Table has non-PII.Sensitive tags") + .withDescription("Table is not allowed to have PII.Sensitive tags") + .withRule( + "{\"!\":{\"some\":[{\"var\":\"tags\"},{\"==\":[{\"var\":\"tagFQN\"},\"PII.Sensitive\"]}]}}"); + + RuleEngine.getInstance().evaluate(table, List.of(piiRule), false, false); + + table.withTags(List.of(PII_SENSITIVE_TAG_LABEL)); + assertThrows( + RuleValidationException.class, + () -> RuleEngine.getInstance().evaluate(table, List.of(piiRule), false, false)); + + table.withTags(List.of(PII_SENSITIVE_TAG_LABEL, TIER1_TAG_LABEL)); + assertThrows( + RuleValidationException.class, + () -> RuleEngine.getInstance().evaluate(table, List.of(piiRule), false, false)); + + table.withTags(List.of(TIER1_TAG_LABEL)); + RuleEngine.getInstance().evaluate(table, List.of(piiRule), false, false); + } }