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
This commit is contained in:
Pere Miquel Brull 2023-09-27 11:49:21 +02:00 committed by GitHub
parent d8ef497b9e
commit 2c3ff8dc08
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
13 changed files with 211 additions and 48 deletions

View File

@ -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",

View File

@ -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;
}

View File

@ -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(

View File

@ -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);

View File

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

View File

@ -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);
}
}

View File

@ -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);
String executableTestSuiteFQN = tableFQN + ".testSuite";
TestSuite executableTestSuite =
testSuiteRepository.getDao().findEntityByName(executableTestSuiteFQN, "fqnHash", Include.ALL);
for (TestCase testCase : testCases) {
// we are setting mustHaveRelationship to "false" to not throw any error.
List<CollectionDAO.EntityRelationshipRecord> existingRelations =
testSuiteRepository.findFromRecords(testCase.getId(), TEST_CASE, Relationship.CONTAINS, TEST_SUITE);
boolean relationWithExecutableTestSuiteExists = false;
if (existingRelations != null) {
for (CollectionDAO.EntityRelationshipRecord existingTestSuiteRel : existingRelations) {
try {
TestSuite existingTestSuite = testSuiteRepository.getDao().findEntityById(existingTestSuiteRel.getId());
if (existingTestSuite.getExecutable()
&& existingTestSuite.getFullyQualifiedName().equals(executableTestSuiteFQN)) {
// There is a native test suite associated with this testCase.
relationWithExecutableTestSuiteExists = true;
try {
List<TestCase> testCases = testCasesGroupByTable.get(tableFQN);
String executableTestSuiteFQN = tableFQN + ".testSuite";
TestSuite executableTestSuite =
getOrCreateExecutableTestSuite(collectionDAO, testCases, testSuiteRepository, executableTestSuiteFQN);
for (TestCase testCase : testCases) {
// we are setting mustHaveRelationship to "false" to not throw any error.
List<CollectionDAO.EntityRelationshipRecord> existingRelations =
testSuiteRepository.findFromRecords(testCase.getId(), TEST_CASE, Relationship.CONTAINS, TEST_SUITE);
boolean relationWithExecutableTestSuiteExists = false;
if (existingRelations != null) {
for (CollectionDAO.EntityRelationshipRecord existingTestSuiteRel : existingRelations) {
try {
TestSuite existingTestSuite = testSuiteRepository.getDao().findEntityById(existingTestSuiteRel.getId());
if (Boolean.TRUE.equals(existingTestSuite.getExecutable())
&& existingTestSuite.getFullyQualifiedName().equals(executableTestSuiteFQN)) {
// 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
List<CollectionDAO.EntityRelationshipRecord> testSuiteRels =
testSuiteRepository.findToRecords(
executableTestSuite.getExecutableEntityReference().getId(), TABLE, Relationship.CONTAINS, TEST_SUITE);
for (CollectionDAO.EntityRelationshipRecord testSuiteRel : testSuiteRels) {
try {
TestSuite existingTestSuite = testSuiteRepository.getDao().findEntityById(testSuiteRel.getId());
} catch (EntityNotFoundException ex) {
// if testsuite cannot be retrieved but the relation exists, then this is orphaned relation, we will
// delete the relation
testSuiteRepository.deleteRelationship(
executableTestSuite.getExecutableEntityReference().getId(),
TABLE,
testSuiteRel.getId(),
TEST_SUITE,
Relationship.CONTAINS);
// check from table -> nativeTestSuite there should only one relation
List<CollectionDAO.EntityRelationshipRecord> testSuiteRels =
testSuiteRepository.findToRecords(
executableTestSuite.getExecutableEntityReference().getId(), TABLE, Relationship.CONTAINS, TEST_SUITE);
for (CollectionDAO.EntityRelationshipRecord testSuiteRel : testSuiteRels) {
try {
testSuiteRepository.getDao().findEntityById(testSuiteRel.getId());
} catch (EntityNotFoundException ex) {
// if testsuite cannot be retrieved but the relation exists, then this is orphaned relation, we will
// delete the relation
testSuiteRepository.deleteRelationship(
executableTestSuite.getExecutableEntityReference().getId(),
TABLE,
testSuiteRel.getId(),
TEST_SUITE,
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;
}
}
}

View File

@ -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;
}
}
}

View File

@ -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()