forked from datahub-project/datahub
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
fix(cassandra): fix Cassandra queries used by IngestDataPlatformInsta…
…ncesStep (datahub-project#5199)
- Loading branch information
1 parent
7cbaac9
commit 6751365
Showing
13 changed files
with
747 additions
and
334 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
69 changes: 69 additions & 0 deletions
69
metadata-io/src/test/java/com/linkedin/metadata/AspectGenerationUtils.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,69 @@ | ||
| package com.linkedin.metadata; | ||
|
|
||
| import com.linkedin.chart.ChartInfo; | ||
| import com.linkedin.common.AuditStamp; | ||
| import com.linkedin.common.ChangeAuditStamps; | ||
| import com.linkedin.common.urn.Urn; | ||
| import com.linkedin.common.urn.UrnUtils; | ||
| import com.linkedin.data.template.RecordTemplate; | ||
| import com.linkedin.identity.CorpUserInfo; | ||
| import com.linkedin.metadata.key.CorpUserKey; | ||
| import com.linkedin.metadata.utils.EntityKeyUtils; | ||
| import com.linkedin.metadata.utils.PegasusUtils; | ||
| import com.linkedin.mxe.SystemMetadata; | ||
| import javax.annotation.Nonnull; | ||
|
|
||
|
|
||
| public class AspectGenerationUtils { | ||
|
|
||
| private AspectGenerationUtils() { | ||
| } | ||
|
|
||
| @Nonnull | ||
| public static AuditStamp createAuditStamp() { | ||
| return new AuditStamp().setTime(123L).setActor(UrnUtils.getUrn("urn:li:corpuser:tester")); | ||
| } | ||
|
|
||
| @Nonnull | ||
| public static SystemMetadata createSystemMetadata() { | ||
| return createSystemMetadata(1625792689, "run-123"); | ||
| } | ||
|
|
||
| @Nonnull | ||
| public static SystemMetadata createSystemMetadata(long lastObserved, @Nonnull String runId) { | ||
| SystemMetadata metadata = new SystemMetadata(); | ||
| metadata.setLastObserved(lastObserved); | ||
| metadata.setRunId(runId); | ||
| return metadata; | ||
| } | ||
|
|
||
| @Nonnull | ||
| public static CorpUserKey createCorpUserKey(Urn urn) { | ||
| return (CorpUserKey) EntityKeyUtils.convertUrnToEntityKey(urn, new CorpUserKey().schema()); | ||
| } | ||
|
|
||
| @Nonnull | ||
| public static CorpUserInfo createCorpUserInfo(@Nonnull String email) { | ||
| CorpUserInfo corpUserInfo = new CorpUserInfo(); | ||
| corpUserInfo.setEmail(email); | ||
| corpUserInfo.setActive(true); | ||
| return corpUserInfo; | ||
| } | ||
|
|
||
| @Nonnull | ||
| public static ChartInfo createChartInfo(@Nonnull String title, @Nonnull String description) { | ||
| ChartInfo chartInfo = new ChartInfo(); | ||
| chartInfo.setTitle(title); | ||
| chartInfo.setDescription(description); | ||
| ChangeAuditStamps lastModified = new ChangeAuditStamps(); | ||
| lastModified.setCreated(createAuditStamp()); | ||
| lastModified.setLastModified(createAuditStamp()); | ||
| chartInfo.setLastModified(lastModified); | ||
| return chartInfo; | ||
| } | ||
|
|
||
| @Nonnull | ||
| public static String getAspectName(RecordTemplate record) { | ||
| return PegasusUtils.getAspectNameFromSchema(record.schema()); | ||
| } | ||
| } |
75 changes: 75 additions & 0 deletions
75
metadata-io/src/test/java/com/linkedin/metadata/AspectIngestionUtils.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,75 @@ | ||
| package com.linkedin.metadata; | ||
|
|
||
| import com.linkedin.chart.ChartInfo; | ||
| import com.linkedin.common.urn.Urn; | ||
| import com.linkedin.common.urn.UrnUtils; | ||
| import com.linkedin.identity.CorpUserInfo; | ||
| import com.linkedin.metadata.entity.EntityService; | ||
| import com.linkedin.metadata.key.CorpUserKey; | ||
| import java.util.HashMap; | ||
| import java.util.Map; | ||
| import javax.annotation.Nonnull; | ||
|
|
||
|
|
||
| public class AspectIngestionUtils { | ||
|
|
||
| private AspectIngestionUtils() { | ||
| } | ||
|
|
||
| @Nonnull | ||
| public static Map<Urn, CorpUserKey> ingestCorpUserKeyAspects(EntityService entityService, int aspectCount) { | ||
| return ingestCorpUserKeyAspects(entityService, aspectCount, 0); | ||
| } | ||
|
|
||
| @Nonnull | ||
| public static Map<Urn, CorpUserKey> ingestCorpUserKeyAspects(EntityService entityService, int aspectCount, int startIndex) { | ||
| String aspectName = AspectGenerationUtils.getAspectName(new CorpUserKey()); | ||
| Map<Urn, CorpUserKey> aspects = new HashMap<>(); | ||
| for (int i = startIndex; i < startIndex + aspectCount; i++) { | ||
| Urn urn = UrnUtils.getUrn(String.format("urn:li:corpuser:tester%d", i)); | ||
| CorpUserKey aspect = AspectGenerationUtils.createCorpUserKey(urn); | ||
| aspects.put(urn, aspect); | ||
| entityService.ingestAspect(urn, aspectName, aspect, AspectGenerationUtils.createAuditStamp(), AspectGenerationUtils.createSystemMetadata()); | ||
| } | ||
| return aspects; | ||
| } | ||
|
|
||
| @Nonnull | ||
| public static Map<Urn, CorpUserInfo> ingestCorpUserInfoAspects(@Nonnull final EntityService entityService, int aspectCount) { | ||
| return ingestCorpUserInfoAspects(entityService, aspectCount, 0); | ||
| } | ||
|
|
||
| @Nonnull | ||
| public static Map<Urn, CorpUserInfo> ingestCorpUserInfoAspects(@Nonnull final EntityService entityService, int aspectCount, int startIndex) { | ||
| String aspectName = AspectGenerationUtils.getAspectName(new CorpUserInfo()); | ||
| Map<Urn, CorpUserInfo> aspects = new HashMap<>(); | ||
| for (int i = startIndex; i < startIndex + aspectCount; i++) { | ||
| Urn urn = UrnUtils.getUrn(String.format("urn:li:corpuser:tester%d", i)); | ||
| String email = String.format("email%d@test.com", i); | ||
| CorpUserInfo aspect = AspectGenerationUtils.createCorpUserInfo(email); | ||
| aspects.put(urn, aspect); | ||
| entityService.ingestAspect(urn, aspectName, aspect, AspectGenerationUtils.createAuditStamp(), AspectGenerationUtils.createSystemMetadata()); | ||
| } | ||
| return aspects; | ||
| } | ||
|
|
||
| @Nonnull | ||
| public static Map<Urn, ChartInfo> ingestChartInfoAspects(@Nonnull final EntityService entityService, int aspectCount) { | ||
| return ingestChartInfoAspects(entityService, aspectCount, 0); | ||
| } | ||
|
|
||
| @Nonnull | ||
| public static Map<Urn, ChartInfo> ingestChartInfoAspects(@Nonnull final EntityService entityService, int aspectCount, int startIndex) { | ||
| String aspectName = AspectGenerationUtils.getAspectName(new ChartInfo()); | ||
| Map<Urn, ChartInfo> aspects = new HashMap<>(); | ||
| for (int i = startIndex; i < startIndex + aspectCount; i++) { | ||
| Urn urn = UrnUtils.getUrn(String.format("urn:li:chart:(looker,test%d)", i)); | ||
| String title = String.format("Test Title %d", i); | ||
| String description = String.format("Test description %d", i); | ||
| ChartInfo aspect = AspectGenerationUtils.createChartInfo(title, description); | ||
| aspects.put(urn, aspect); | ||
| entityService.ingestAspect(urn, aspectName, aspect, AspectGenerationUtils.createAuditStamp(), AspectGenerationUtils.createSystemMetadata()); | ||
| } | ||
| return aspects; | ||
| } | ||
| } |
98 changes: 98 additions & 0 deletions
98
metadata-io/src/test/java/com/linkedin/metadata/entity/AspectMigrationsDaoTest.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,98 @@ | ||
| package com.linkedin.metadata.entity; | ||
|
|
||
| import com.linkedin.common.urn.Urn; | ||
| import com.linkedin.common.urn.UrnUtils; | ||
| import com.linkedin.metadata.AspectIngestionUtils; | ||
| import com.linkedin.metadata.event.EventProducer; | ||
| import com.linkedin.metadata.key.CorpUserKey; | ||
| import com.linkedin.metadata.models.registry.ConfigEntityRegistry; | ||
| import com.linkedin.metadata.models.registry.EntityRegistry; | ||
| import com.linkedin.metadata.models.registry.EntityRegistryException; | ||
| import com.linkedin.metadata.models.registry.MergedEntityRegistry; | ||
| import com.linkedin.metadata.snapshot.Snapshot; | ||
| import java.util.ArrayList; | ||
| import java.util.List; | ||
| import java.util.Map; | ||
| import java.util.stream.Collectors; | ||
| import org.testcontainers.shaded.com.google.common.collect.ImmutableList; | ||
| import org.testng.annotations.Test; | ||
|
|
||
| import static com.linkedin.metadata.Constants.*; | ||
| import static org.testng.Assert.*; | ||
|
|
||
|
|
||
| abstract public class AspectMigrationsDaoTest<T extends AspectMigrationsDao> { | ||
|
|
||
| protected T _migrationsDao; | ||
|
|
||
| protected final EntityRegistry _snapshotEntityRegistry; | ||
| protected final EntityRegistry _configEntityRegistry; | ||
| protected final EntityRegistry _testEntityRegistry; | ||
| protected EventProducer _mockProducer; | ||
|
|
||
| protected EntityService _entityService; | ||
| protected RetentionService _retentionService; | ||
|
|
||
| protected AspectMigrationsDaoTest() throws EntityRegistryException { | ||
| _snapshotEntityRegistry = new TestEntityRegistry(); | ||
| _configEntityRegistry = new ConfigEntityRegistry(Snapshot.class.getClassLoader().getResourceAsStream("entity-registry.yml")); | ||
| _testEntityRegistry = new MergedEntityRegistry(_snapshotEntityRegistry).apply(_configEntityRegistry); | ||
| } | ||
|
|
||
| @Test | ||
| public void testListAllUrns() throws AssertionError { | ||
| final int totalAspects = 30; | ||
| final int pageSize = 25; | ||
| final int lastPageSize = 5; | ||
| Map<Urn, CorpUserKey> ingestedAspects = AspectIngestionUtils.ingestCorpUserKeyAspects(_entityService, totalAspects); | ||
| List<String> ingestedUrns = ingestedAspects.keySet().stream().map(Urn::toString).collect(Collectors.toList()); | ||
| List<String> seenUrns = new ArrayList<>(); | ||
|
|
||
| Iterable<String> page1 = _migrationsDao.listAllUrns(0, pageSize); | ||
| List<String> page1Urns = ImmutableList.copyOf(page1); | ||
|
|
||
| // validate first page | ||
| assertEquals(page1Urns.size(), pageSize); | ||
| for (String urn : page1Urns) { | ||
| assertNotNull(UrnUtils.getUrn(urn)); | ||
| seenUrns.add(urn); | ||
| } | ||
|
|
||
| Iterable<String> page2 = _migrationsDao.listAllUrns(pageSize, pageSize); | ||
| List<String> page2Urns = ImmutableList.copyOf(page2); | ||
|
|
||
| // validate last page | ||
| assertEquals(page2Urns.size(), lastPageSize); | ||
| for (String urn : page2Urns) { | ||
| assertNotNull(UrnUtils.getUrn(urn)); | ||
| seenUrns.add(urn); | ||
| } | ||
|
|
||
| // validate all ingested URNs were returned exactly once | ||
| for (String urn : ingestedUrns) { | ||
| assertEquals(seenUrns.stream().filter(u -> u.equals(urn)).count(), 1); | ||
| } | ||
| } | ||
|
|
||
| @Test | ||
| public void testCountEntities() throws AssertionError { | ||
| AspectIngestionUtils.ingestCorpUserInfoAspects(_entityService, 11); | ||
| AspectIngestionUtils.ingestChartInfoAspects(_entityService, 22); | ||
| final int expected = 33; | ||
|
|
||
| long actual = _migrationsDao.countEntities(); | ||
|
|
||
| assertEquals(actual, expected); | ||
| } | ||
|
|
||
| @Test | ||
| public void testCheckIfAspectExists() throws AssertionError { | ||
| boolean actual = _migrationsDao.checkIfAspectExists(CORP_USER_INFO_ASPECT_NAME); | ||
| assertFalse(actual); | ||
|
|
||
| AspectIngestionUtils.ingestCorpUserInfoAspects(_entityService, 1); | ||
|
|
||
| actual = _migrationsDao.checkIfAspectExists(CORP_USER_INFO_ASPECT_NAME); | ||
| assertTrue(actual); | ||
| } | ||
| } |
63 changes: 63 additions & 0 deletions
63
metadata-io/src/test/java/com/linkedin/metadata/entity/CassandraAspectMigrationsDaoTest.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,63 @@ | ||
| package com.linkedin.metadata.entity; | ||
|
|
||
| import com.datastax.oss.driver.api.core.CqlSession; | ||
| import com.linkedin.metadata.CassandraTestUtils; | ||
| import com.linkedin.metadata.entity.cassandra.CassandraAspectDao; | ||
| import com.linkedin.metadata.entity.cassandra.CassandraRetentionService; | ||
| import com.linkedin.metadata.event.EventProducer; | ||
| import com.linkedin.metadata.models.registry.EntityRegistryException; | ||
| import org.testcontainers.containers.CassandraContainer; | ||
| import org.testng.Assert; | ||
| import org.testng.annotations.AfterClass; | ||
| import org.testng.annotations.BeforeClass; | ||
| import org.testng.annotations.BeforeMethod; | ||
| import org.testng.annotations.Test; | ||
|
|
||
| import static org.mockito.Mockito.*; | ||
|
|
||
|
|
||
| public class CassandraAspectMigrationsDaoTest extends AspectMigrationsDaoTest<CassandraAspectDao> { | ||
|
|
||
| private CassandraContainer _cassandraContainer; | ||
|
|
||
| public CassandraAspectMigrationsDaoTest() throws EntityRegistryException { | ||
| } | ||
|
|
||
| @BeforeClass | ||
| public void setupContainer() { | ||
| _cassandraContainer = CassandraTestUtils.setupContainer(); | ||
| } | ||
|
|
||
| @AfterClass | ||
| public void tearDown() { | ||
| _cassandraContainer.stop(); | ||
| } | ||
|
|
||
| @BeforeMethod | ||
| public void setupTest() { | ||
| CassandraTestUtils.purgeData(_cassandraContainer); | ||
| configureComponents(); | ||
| } | ||
|
|
||
| private void configureComponents() { | ||
| CqlSession session = CassandraTestUtils.createTestSession(_cassandraContainer); | ||
| CassandraAspectDao dao = new CassandraAspectDao(session); | ||
| dao.setConnectionValidated(true); | ||
| _mockProducer = mock(EventProducer.class); | ||
| _entityService = new EntityService(dao, _mockProducer, _testEntityRegistry); | ||
| _retentionService = new CassandraRetentionService(_entityService, session, 1000); | ||
| _entityService.setRetentionService(_retentionService); | ||
|
|
||
| _migrationsDao = dao; | ||
| } | ||
|
|
||
| /** | ||
| * Ideally, all tests would be in the base class, so they're reused between all implementations. | ||
| * When that's the case - test runner will ignore this class (and its base!) so we keep this dummy test | ||
| * to make sure this class will always be discovered. | ||
| */ | ||
| @Test | ||
| public void obligatoryTest() throws AssertionError { | ||
| Assert.assertTrue(true); | ||
| } | ||
| } |
Oops, something went wrong.