Fixes Test Suite Reference in Table Schema (#16129)

* Fixes Test Suite Reference in Table Schema

* fix: fix test suite to interact with entity reference

---------

Co-authored-by: Teddy Crepineau <teddy.crepineau@gmail.com>

bootstrap/sql/migrations/native/1.4.0/mysql/schemaChanges.sql

@@ -281,3 +281,5 @@ SET json = JSON_INSERT(
     '$.schemaDefinition',
     JSON_EXTRACT(json, '$.viewDefinition')
 );
+
+UPDATE table_entity SET json = JSON_REMOVE(json, '$.testSuite');

bootstrap/sql/migrations/native/1.4.0/postgres/schemaChanges.sql

@@ -278,3 +278,6 @@ SET json = jsonb_set(
         json->'viewDefinition'
     ) - 'viewDefinition'
 WHERE jsonb_exists(json::jsonb, 'viewDefinition') = true;
+
+UPDATE table_entity SET json = json - 'testSuite'; 
+

ingestion/src/metadata/data_quality/source/test_suite.py

@@ -139,7 +139,13 @@ class TestSuiteSource(Source):
                 )
             )
 
-        if table.testSuite and not table.testSuite.executable:
+        test_suite: Optional[TestSuite] = None
+        if table.testSuite:
+            test_suite = self.metadata.get_by_id(
+                entity=TestSuite, entity_id=table.testSuite.id.__root__
+            )
+
+        if test_suite and not test_suite.executable:
             yield Either(
                 left=StackTraceError(
                     name="Non-executable Test Suite",
@@ -149,7 +155,7 @@ class TestSuiteSource(Source):
             )
 
         else:
-            test_suite_cases = self._get_test_cases_from_test_suite(table.testSuite)
+            test_suite_cases = self._get_test_cases_from_test_suite(test_suite)
 
             yield Either(
                 right=TableAndTests(

ingestion/tests/integration/great_expectations/test_great_expectation_integration.py

@@ -30,6 +30,7 @@ from metadata.generated.schema.entity.services.databaseService import DatabaseSe
 from metadata.generated.schema.security.client.openMetadataJWTClientConfig import (
     OpenMetadataJWTClientConfig,
 )
+from metadata.generated.schema.tests.testSuite import TestSuite
 from metadata.ingestion.connections.session import create_and_bind_session
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.utils.time_utils import (
@@ -217,7 +218,10 @@ class TestGreatExpectationIntegration(TestCase):
         )
 
         assert table_entity.testSuite
-        assert len(table_entity.testSuite.tests) == 1
+        test_suite: TestSuite = self.metadata.get_by_id(
+            entity=TestSuite, entity_id=table_entity.testSuite.id, fields=["tests"]
+        )
+        assert len(test_suite.tests) == 1
 
         test_case_results = self.metadata.get_test_case_results(
             test_case_fqn=TEST_CASE_FQN,

openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/TableRepository.java

@@ -28,7 +28,7 @@ import static org.openmetadata.service.Entity.DATABASE_SCHEMA;
 import static org.openmetadata.service.Entity.FIELD_OWNER;
 import static org.openmetadata.service.Entity.FIELD_TAGS;
 import static org.openmetadata.service.Entity.TABLE;
-import static org.openmetadata.service.Entity.getEntity;
+import static org.openmetadata.service.Entity.TEST_SUITE;
 import static org.openmetadata.service.Entity.populateEntityFieldTags;
 import static org.openmetadata.service.util.EntityUtil.getLocalColumnName;
 import static org.openmetadata.service.util.FullyQualifiedName.getColumnName;
@@ -64,7 +64,6 @@ import org.openmetadata.schema.entity.data.DatabaseSchema;
 import org.openmetadata.schema.entity.data.Table;
 import org.openmetadata.schema.entity.feed.Suggestion;
 import org.openmetadata.schema.tests.CustomMetric;
-import org.openmetadata.schema.tests.TestSuite;
 import org.openmetadata.schema.type.Column;
 import org.openmetadata.schema.type.ColumnDataType;
 import org.openmetadata.schema.type.ColumnJoin;
@@ -324,22 +323,8 @@ public class TableRepository extends EntityRepository<Table> {
         TableProfilerConfig.class);
   }
 
-  public TestSuite getTestSuite(Table table) {
+  public EntityReference getTestSuite(Table table) {
-    List<CollectionDAO.EntityRelationshipRecord> entityRelationshipRecords =
+    return getToEntityRef(table.getId(), Relationship.CONTAINS, TEST_SUITE, false);
-        daoCollection
-            .relationshipDAO()
-            .findTo(table.getId(), TABLE, Relationship.CONTAINS.ordinal());
-    Optional<CollectionDAO.EntityRelationshipRecord> testSuiteRelationshipRecord =
-        entityRelationshipRecords.stream()
-            .filter(
-                entityRelationshipRecord ->
-                    entityRelationshipRecord.getType().equals(Entity.TEST_SUITE))
-            .findFirst();
-    return testSuiteRelationshipRecord
-        .<TestSuite>map(
-            entityRelationshipRecord ->
-                getEntity(Entity.TEST_SUITE, entityRelationshipRecord.getId(), "*", Include.ALL))
-        .orElse(null);
   }
 
   @Transaction

openmetadata-service/src/test/java/org/openmetadata/service/resources/dqtests/TestSuiteResourceTest.java

@@ -475,6 +475,7 @@ public class TestSuiteResourceTest extends EntityResourceTest<TestSuite, CreateT
   void get_execTestSuiteFromTable_200(TestInfo test) throws IOException {
     TableResourceTest tableResourceTest = new TableResourceTest();
     TestCaseResourceTest testCaseResourceTest = new TestCaseResourceTest();
+    TestSuiteResourceTest testSuiteResourceTest = new TestSuiteResourceTest();
     CreateTable tableReq =
         tableResourceTest
             .createRequest(test)
@@ -498,15 +499,25 @@ public class TestSuiteResourceTest extends EntityResourceTest<TestSuite, CreateT
       testCaseResourceTest.createAndCheckEntity(createTestCase, ADMIN_AUTH_HEADERS);
     }
 
+    Map<String, String> queryParams = new HashMap<>();
+    queryParams.put("include", Include.ALL.value());
+
     Table actualTable = tableResourceTest.getEntity(table.getId(), "testSuite", ADMIN_AUTH_HEADERS);
-    TestSuite tableTestSuite = actualTable.getTestSuite();
+    EntityReference tableTestSuiteRef = actualTable.getTestSuite();
-    assertEquals(testSuite.getId(), tableTestSuite.getId());
+    assertEquals(testSuite.getId(), tableTestSuiteRef.getId());
+    TestSuite tableTestSuite =
+        testSuiteResourceTest.getEntity(
+            tableTestSuiteRef.getId(), queryParams, "tests", ADMIN_AUTH_HEADERS);
     assertEquals(5, tableTestSuite.getTests().size());
 
     // Soft delete entity
     deleteExecutableTestSuite(tableTestSuite.getId(), true, false, ADMIN_AUTH_HEADERS);
-    actualTable = tableResourceTest.getEntity(actualTable.getId(), "testSuite", ADMIN_AUTH_HEADERS);
+    actualTable =
-    tableTestSuite = actualTable.getTestSuite();
+        tableResourceTest.getEntity(
+            actualTable.getId(), queryParams, "testSuite", ADMIN_AUTH_HEADERS);
+    tableTestSuiteRef = actualTable.getTestSuite();
+    tableTestSuite =
+        testSuiteResourceTest.getEntity(tableTestSuiteRef.getId(), "tests", ADMIN_AUTH_HEADERS);
     assertEquals(true, tableTestSuite.getDeleted());
 
     // Hard delete entity
@@ -518,6 +529,8 @@ public class TestSuiteResourceTest extends EntityResourceTest<TestSuite, CreateT
   @Test
   void get_execTestSuiteDeletedOnTableDeletion(TestInfo test) throws IOException {
     TableResourceTest tableResourceTest = new TableResourceTest();
+    TestSuiteResourceTest testSuiteResourceTest = new TestSuiteResourceTest();
+
     CreateTable tableReq =
         tableResourceTest
             .createRequest(test)
@@ -532,13 +545,18 @@ public class TestSuiteResourceTest extends EntityResourceTest<TestSuite, CreateT
     CreateTestSuite createTestSuite = createRequest(table.getFullyQualifiedName());
     TestSuite testSuite = createExecutableTestSuite(createTestSuite, ADMIN_AUTH_HEADERS);
 
+    HashMap<String, String> queryParams = new HashMap<>();
+    queryParams.put("include", Include.ALL.value());
+
     Table actualTable = tableResourceTest.getEntity(table.getId(), "testSuite", ADMIN_AUTH_HEADERS);
-    TestSuite actualTestSuite = actualTable.getTestSuite();
+    EntityReference actualTestSuiteRef = actualTable.getTestSuite();
+    TestSuite actualTestSuite =
+        testSuiteResourceTest.getEntity(
+            actualTestSuiteRef.getId(), queryParams, "tests", ADMIN_AUTH_HEADERS);
     assertEquals(actualTestSuite.getId(), testSuite.getId());
 
     tableResourceTest.deleteEntity(actualTable.getId(), true, false, ADMIN_AUTH_HEADERS);
-    HashMap<String, String> queryParams = new HashMap<>();
+
-    queryParams.put("include", Include.ALL.value());
     actualTestSuite =
         getEntityByName(testSuite.getFullyQualifiedName(), queryParams, "*", ADMIN_AUTH_HEADERS);
     assertEquals(true, actualTestSuite.getDeleted());

openmetadata-spec/src/main/resources/json/schema/entity/data/table.json

@@ -1039,7 +1039,7 @@
     },
     "testSuite": {
       "description": "Executable test suite associated with this table",
-      "$ref": "../../tests/testSuite.json"
+      "$ref": "../../type/entityReference.json"
     },
     "dataModel": {
       "description": "This captures information about how the table is modeled. Currently only DBT model is supported.",