Prep v1.1.7 migrations to address test cases & suites (#13345)

* Prep v1.1.7 migrations to address test cases

* get or create executable suite

* Format

* Fix tests

* Add postgres

ingestion/setup.py

@@ -103,7 +103,7 @@ base_requirements = {
     VERSIONS["pymysql"],
     "python-dateutil>=2.8.1",
     "python-jose~=3.3",
-    "PyYAML",
+    "PyYAML~=6.0",
     "requests>=2.23",
     "requests-aws4auth~=1.1",  # Only depends on requests as external package. Leaving as base.
     "setuptools~=66.0.0",

openmetadata-service/src/main/java/org/openmetadata/service/Entity.java

@@ -344,7 +344,7 @@ public final class Entity {
   public static EntityRepository<? extends EntityInterface> getEntityRepository(@NonNull String entityType) {
     EntityRepository<? extends EntityInterface> entityRepository = ENTITY_REPOSITORY_MAP.get(entityType);
     if (entityRepository == null) {
-      throw EntityNotFoundException.byMessage(CatalogExceptionMessage.entityTypeNotFound(entityType));
+      throw EntityNotFoundException.byMessage(CatalogExceptionMessage.entityRepositoryNotFound(entityType));
     }
     return entityRepository;
   }

openmetadata-service/src/main/java/org/openmetadata/service/exception/CatalogExceptionMessage.java

@@ -95,6 +95,10 @@ public final class CatalogExceptionMessage {
     return String.format("Entity type %s not found", entityType);
   }
 
+  public static String entityRepositoryNotFound(String entityType) {
+    return String.format("Entity repository for %s not found. Is the ENTITY_TYPE_MAP initialized?", entityType);
+  }
+
   public static String entityRelationshipNotFound(
       String entityType, UUID id, String relationshipName, String toEntityType) {
     return String.format(

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

@@ -1357,6 +1357,10 @@ public abstract class EntityRepository<T extends EntityInterface> {
     return getFromEntityRef(toId, Relationship.CONTAINS, null, true);
   }
 
+  public EntityReference getContainer(UUID toId, String fromEntityType) {
+    return getFromEntityRef(toId, Relationship.CONTAINS, fromEntityType, true);
+  }
+
   public EntityReference getFromEntityRef(
       UUID toId, Relationship relationship, String fromEntityType, boolean mustHaveRelationship) {
     List<EntityRelationshipRecord> records = findFromRecords(toId, entityType, relationship, fromEntityType);

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

@@ -158,7 +158,7 @@ public class TableRepository extends EntityRepository<Table> {
   }
 
   private void setDefaultFields(Table table) {
-    EntityReference schemaRef = getContainer(table.getId());
+    EntityReference schemaRef = getContainer(table.getId(), DATABASE_SCHEMA);
     DatabaseSchema schema = Entity.getEntity(schemaRef, "", ALL);
     table.withDatabaseSchema(schemaRef).withDatabase(schema.getDatabase()).withService(schema.getService());
   }

openmetadata-service/src/main/java/org/openmetadata/service/migration/mysql/v117/Migration.java

@@ -0,0 +1,35 @@
+package org.openmetadata.service.migration.mysql.v117;
+
+import static org.openmetadata.service.migration.utils.V114.MigrationUtil.fixTestSuites;
+import static org.openmetadata.service.migration.utils.V117.MigrationUtil.fixTestCases;
+
+import lombok.SneakyThrows;
+import org.jdbi.v3.core.Handle;
+import org.openmetadata.service.jdbi3.CollectionDAO;
+import org.openmetadata.service.migration.api.MigrationProcessImpl;
+import org.openmetadata.service.migration.utils.MigrationFile;
+
+public class Migration extends MigrationProcessImpl {
+  private CollectionDAO collectionDAO;
+  private Handle handle;
+
+  public Migration(MigrationFile migrationFile) {
+    super(migrationFile);
+  }
+
+  @Override
+  public void initialize(Handle handle) {
+    super.initialize(handle);
+    this.handle = handle;
+    this.collectionDAO = handle.attach(CollectionDAO.class);
+  }
+
+  @Override
+  @SneakyThrows
+  public void runDataMigration() {
+    // tests cases coming from dbt for case-sensitive services are not properly linked to a table
+    fixTestCases(handle, collectionDAO);
+    // Try again the 1.1.6 test suite migration
+    fixTestSuites(collectionDAO);
+  }
+}

openmetadata-service/src/main/java/org/openmetadata/service/migration/utils/V114/MigrationUtil.java

@@ -1,22 +1,29 @@
 package org.openmetadata.service.migration.utils.V114;
 
 import static org.openmetadata.service.Entity.*;
+import static org.openmetadata.service.migration.utils.v110.MigrationUtil.getTestSuite;
 import static org.openmetadata.service.migration.utils.v110.MigrationUtil.groupTestCasesByTable;
 
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import lombok.extern.slf4j.Slf4j;
+import org.openmetadata.schema.api.tests.CreateTestSuite;
 import org.openmetadata.schema.tests.TestCase;
 import org.openmetadata.schema.tests.TestSuite;
 import org.openmetadata.schema.type.Include;
 import org.openmetadata.schema.type.Relationship;
+import org.openmetadata.service.Entity;
 import org.openmetadata.service.exception.EntityNotFoundException;
 import org.openmetadata.service.jdbi3.CollectionDAO;
 import org.openmetadata.service.jdbi3.ListFilter;
 import org.openmetadata.service.jdbi3.TestSuiteRepository;
+import org.openmetadata.service.resources.feeds.MessageParser;
 import org.openmetadata.service.util.EntityUtil;
+import org.openmetadata.service.util.FullyQualifiedName;
 
+@Slf4j
 public class MigrationUtil {
   private MigrationUtil() {
     /* Cannot create object  util class*/
@@ -54,57 +61,105 @@ public class MigrationUtil {
     // TestSuite
     Map<String, ArrayList<TestCase>> testCasesGroupByTable = groupTestCasesByTable(collectionDAO);
     for (String tableFQN : testCasesGroupByTable.keySet()) {
-      List<TestCase> testCases = testCasesGroupByTable.get(tableFQN);
+      try {
-      String executableTestSuiteFQN = tableFQN + ".testSuite";
+        List<TestCase> testCases = testCasesGroupByTable.get(tableFQN);
-      TestSuite executableTestSuite =
+        String executableTestSuiteFQN = tableFQN + ".testSuite";
-          testSuiteRepository.getDao().findEntityByName(executableTestSuiteFQN, "fqnHash", Include.ALL);
+        TestSuite executableTestSuite =
-      for (TestCase testCase : testCases) {
+            getOrCreateExecutableTestSuite(collectionDAO, testCases, testSuiteRepository, executableTestSuiteFQN);
-        // we are setting mustHaveRelationship to "false" to not throw any error.
+        for (TestCase testCase : testCases) {
-        List<CollectionDAO.EntityRelationshipRecord> existingRelations =
+          // we are setting mustHaveRelationship to "false" to not throw any error.
-            testSuiteRepository.findFromRecords(testCase.getId(), TEST_CASE, Relationship.CONTAINS, TEST_SUITE);
+          List<CollectionDAO.EntityRelationshipRecord> existingRelations =
-        boolean relationWithExecutableTestSuiteExists = false;
+              testSuiteRepository.findFromRecords(testCase.getId(), TEST_CASE, Relationship.CONTAINS, TEST_SUITE);
-        if (existingRelations != null) {
+          boolean relationWithExecutableTestSuiteExists = false;
-          for (CollectionDAO.EntityRelationshipRecord existingTestSuiteRel : existingRelations) {
+          if (existingRelations != null) {
-            try {
+            for (CollectionDAO.EntityRelationshipRecord existingTestSuiteRel : existingRelations) {
-              TestSuite existingTestSuite = testSuiteRepository.getDao().findEntityById(existingTestSuiteRel.getId());
+              try {
-              if (existingTestSuite.getExecutable()
+                TestSuite existingTestSuite = testSuiteRepository.getDao().findEntityById(existingTestSuiteRel.getId());
-                  && existingTestSuite.getFullyQualifiedName().equals(executableTestSuiteFQN)) {
+                if (Boolean.TRUE.equals(existingTestSuite.getExecutable())
-                // There is a native test suite associated with this testCase.
+                    && existingTestSuite.getFullyQualifiedName().equals(executableTestSuiteFQN)) {
-                relationWithExecutableTestSuiteExists = true;
+                  // There is a native test suite associated with this testCase.
+                  relationWithExecutableTestSuiteExists = true;
+                }
+              } catch (EntityNotFoundException ex) {
+                // if testsuite cannot be retrieved but the relation exists, then this is orphaned relation, we will
+                // delete the relation
+                testSuiteRepository.deleteRelationship(
+                    existingTestSuiteRel.getId(), TEST_SUITE, testCase.getId(), TEST_CASE, Relationship.CONTAINS);
               }
-            } catch (EntityNotFoundException ex) {
-              // if testsuite cannot be retrieved but the relation exists, then this is orphaned relation, we will
-              // delete the relation
-              testSuiteRepository.deleteRelationship(
-                  existingTestSuiteRel.getId(), TEST_SUITE, testCase.getId(), TEST_CASE, Relationship.CONTAINS);
             }
           }
+          // if we can't find any executable testSuite relationship add one
+          if (!relationWithExecutableTestSuiteExists) {
+            testSuiteRepository.addRelationship(
+                executableTestSuite.getId(), testCase.getId(), TEST_SUITE, TEST_CASE, Relationship.CONTAINS);
+          }
         }
-        // if we can't find any executable testSuite relationship add one
-        if (!relationWithExecutableTestSuiteExists) {
-          testSuiteRepository.addRelationship(
-              executableTestSuite.getId(), testCase.getId(), TEST_SUITE, TEST_CASE, Relationship.CONTAINS);
-        }
-      }
 
-      // check from table -> nativeTestSuite there should only one relation
+        // check from table -> nativeTestSuite there should only one relation
-      List<CollectionDAO.EntityRelationshipRecord> testSuiteRels =
+        List<CollectionDAO.EntityRelationshipRecord> testSuiteRels =
-          testSuiteRepository.findToRecords(
+            testSuiteRepository.findToRecords(
-              executableTestSuite.getExecutableEntityReference().getId(), TABLE, Relationship.CONTAINS, TEST_SUITE);
+                executableTestSuite.getExecutableEntityReference().getId(), TABLE, Relationship.CONTAINS, TEST_SUITE);
-      for (CollectionDAO.EntityRelationshipRecord testSuiteRel : testSuiteRels) {
+        for (CollectionDAO.EntityRelationshipRecord testSuiteRel : testSuiteRels) {
-        try {
+          try {
-          TestSuite existingTestSuite = testSuiteRepository.getDao().findEntityById(testSuiteRel.getId());
+            testSuiteRepository.getDao().findEntityById(testSuiteRel.getId());
-        } catch (EntityNotFoundException ex) {
+          } catch (EntityNotFoundException ex) {
-          // if testsuite cannot be retrieved but the relation exists, then this is orphaned relation, we will
+            // if testsuite cannot be retrieved but the relation exists, then this is orphaned relation, we will
-          // delete the relation
+            // delete the relation
-          testSuiteRepository.deleteRelationship(
+            testSuiteRepository.deleteRelationship(
-              executableTestSuite.getExecutableEntityReference().getId(),
+                executableTestSuite.getExecutableEntityReference().getId(),
-              TABLE,
+                TABLE,
-              testSuiteRel.getId(),
+                testSuiteRel.getId(),
-              TEST_SUITE,
+                TEST_SUITE,
-              Relationship.CONTAINS);
+                Relationship.CONTAINS);
+          }
         }
+      } catch (Exception exc) {
+        LOG.error(
+            String.format("Error trying to migrate tests from Table [%s] due to [%s]", tableFQN, exc.getMessage()));
       }
     }
   }
+
+  private static TestSuite getOrCreateExecutableTestSuite(
+      CollectionDAO collectionDAO,
+      List<TestCase> testCases,
+      TestSuiteRepository testSuiteRepository,
+      String executableTestSuiteFQN) {
+    try {
+      // Try to return the Executable Test Suite that should exist
+      return testSuiteRepository.getDao().findEntityByName(executableTestSuiteFQN, "fqnHash", Include.ALL);
+    } catch (EntityNotFoundException exc) {
+      // If it does not exist, create it and return it
+      MessageParser.EntityLink entityLink =
+          MessageParser.EntityLink.parse(testCases.stream().findFirst().get().getEntityLink());
+      TestSuite newExecutableTestSuite =
+          getTestSuite(
+                  collectionDAO,
+                  new CreateTestSuite()
+                      .withName(FullyQualifiedName.buildHash(executableTestSuiteFQN))
+                      .withDisplayName(executableTestSuiteFQN)
+                      .withExecutableEntityReference(entityLink.getEntityFQN()),
+                  "ingestion-bot")
+              .withExecutable(true)
+              .withFullyQualifiedName(executableTestSuiteFQN);
+      testSuiteRepository.prepareInternal(newExecutableTestSuite, false);
+      testSuiteRepository
+          .getDao()
+          .insert("fqnHash", newExecutableTestSuite, newExecutableTestSuite.getFullyQualifiedName());
+      // add relationship between executable TestSuite with Table
+      testSuiteRepository.addRelationship(
+          newExecutableTestSuite.getExecutableEntityReference().getId(),
+          newExecutableTestSuite.getId(),
+          Entity.TABLE,
+          TEST_SUITE,
+          Relationship.CONTAINS);
+
+      // add relationship between all the testCases that are created against a table with native test suite.
+      for (TestCase testCase : testCases) {
+        testSuiteRepository.addRelationship(
+            newExecutableTestSuite.getId(), testCase.getId(), TEST_SUITE, TEST_CASE, Relationship.CONTAINS);
+      }
+      return newExecutableTestSuite;
+    }
+  }
 }

openmetadata-service/src/main/java/org/openmetadata/service/migration/utils/V117/MigrationUtil.java

@@ -0,0 +1,65 @@
+package org.openmetadata.service.migration.utils.V117;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import org.jdbi.v3.core.Handle;
+import org.openmetadata.schema.entity.data.Table;
+import org.openmetadata.schema.tests.TestCase;
+import org.openmetadata.schema.type.Include;
+import org.openmetadata.service.jdbi3.CollectionDAO;
+import org.openmetadata.service.jdbi3.ListFilter;
+import org.openmetadata.service.jdbi3.TableRepository;
+import org.openmetadata.service.jdbi3.TestCaseRepository;
+import org.openmetadata.service.resources.databases.DatasourceConfig;
+import org.openmetadata.service.resources.feeds.MessageParser;
+import org.openmetadata.service.util.EntityUtil;
+import org.openmetadata.service.util.JsonUtils;
+
+public class MigrationUtil {
+  private MigrationUtil() {
+    /* Cannot create object  util class*/
+  }
+
+  private static final String MYSQL_LIST_TABLE_FQNS =
+      "SELECT JSON_UNQUOTE(JSON_EXTRACT(json, '$.fullyQualifiedName')) FROM table_entity";
+  private static final String POSTGRES_LIST_TABLE_FQNS = "SELECT json #>> '{fullyQualifiedName}' FROM table_entity";
+
+  public static void fixTestCases(Handle handle, CollectionDAO collectionDAO) {
+    TestCaseRepository testCaseRepository = new TestCaseRepository(collectionDAO);
+    TableRepository tableRepository = new TableRepository(collectionDAO);
+    List<TestCase> testCases =
+        testCaseRepository.listAll(new EntityUtil.Fields(Set.of("id")), new ListFilter(Include.ALL));
+
+    try {
+      List<String> fqnList;
+      if (Boolean.TRUE.equals(DatasourceConfig.getInstance().isMySQL())) {
+        fqnList = handle.createQuery(MYSQL_LIST_TABLE_FQNS).mapTo(String.class).list();
+      } else {
+        fqnList = handle.createQuery(POSTGRES_LIST_TABLE_FQNS).mapTo(String.class).list();
+      }
+      Map<String, String> tableMap = new HashMap<>();
+      for (String fqn : fqnList) {
+        tableMap.put(fqn.toLowerCase(), fqn);
+      }
+
+      for (TestCase testCase : testCases) {
+        // Create New Executable Test Suites
+        MessageParser.EntityLink entityLink = MessageParser.EntityLink.parse(testCase.getEntityLink());
+        String fqn = entityLink.getEntityFQN();
+        Table table = JsonUtils.readValue(tableRepository.getDao().findJsonByFqn(fqn, Include.ALL), Table.class);
+        if (table == null) {
+          String findTableFQN = tableMap.get(fqn.toLowerCase());
+          MessageParser.EntityLink newEntityLink =
+              new MessageParser.EntityLink(entityLink.getEntityType(), findTableFQN);
+          testCase.setEntityLink(newEntityLink.getLinkString());
+          testCase.setEntityFQN(findTableFQN);
+          collectionDAO.testCaseDAO().update(testCase);
+        }
+      }
+    } catch (Exception exc) {
+      throw exc;
+    }
+  }
+}

openmetadata-service/src/main/java/org/openmetadata/service/migration/utils/v120/MigrationUtil.java

@@ -58,7 +58,7 @@ public class MigrationUtil {
   public static void addQueryService(Handle handle, CollectionDAO collectionDAO) {
     QueryRepository queryRepository = new QueryRepository(collectionDAO);
 
-    try (handle) {
+    try {
       handle
           .createQuery(QUERY_LIST_SERVICE)
           .mapToMap()