diff --git a/.gitignore b/.gitignore index 9370b734506416..f47aeaee173c86 100644 --- a/.gitignore +++ b/.gitignore @@ -33,6 +33,8 @@ dmypy.json MANIFEST *.pyc .python-version +Pipfile +Pipfile.lock # Generated files **/bin diff --git a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/impl/DefaultUpgradeReport.java b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/impl/DefaultUpgradeReport.java index 913b0ff20e6ff9..fb76a4dea96fdd 100644 --- a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/impl/DefaultUpgradeReport.java +++ b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/impl/DefaultUpgradeReport.java @@ -1,5 +1,6 @@ package com.linkedin.datahub.upgrade.impl; +import com.google.common.base.Throwables; import com.linkedin.datahub.upgrade.UpgradeReport; import java.util.ArrayList; import java.util.List; @@ -20,6 +21,8 @@ public void addLine(String line) { public void addLine(String line, Exception e) { log.error(line, e); reportLines.add(line + String.format(": %s", e)); + reportLines.add( + String.format("Exception stack trace: %s", Throwables.getStackTraceAsString(e))); } @Override diff --git a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/restoreindices/SendMAEStep.java b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/restoreindices/SendMAEStep.java index 256b13b4283483..f095656df7bb41 100644 --- a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/restoreindices/SendMAEStep.java +++ b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/restoreindices/SendMAEStep.java @@ -1,19 +1,15 @@ package com.linkedin.datahub.upgrade.restoreindices; -import static com.linkedin.metadata.Constants.ASPECT_LATEST_VERSION; - import com.google.common.annotations.VisibleForTesting; import com.linkedin.datahub.upgrade.UpgradeContext; import com.linkedin.datahub.upgrade.UpgradeStep; import com.linkedin.datahub.upgrade.UpgradeStepResult; import com.linkedin.datahub.upgrade.impl.DefaultUpgradeStepResult; import com.linkedin.metadata.entity.EntityService; -import com.linkedin.metadata.entity.ebean.EbeanAspectV2; import com.linkedin.metadata.entity.restoreindices.RestoreIndicesArgs; import com.linkedin.metadata.entity.restoreindices.RestoreIndicesResult; import com.linkedin.upgrade.DataHubUpgradeState; import io.ebean.Database; -import io.ebean.ExpressionList; import java.util.ArrayList; import java.util.Arrays; import java.util.List; @@ -156,22 +152,8 @@ private RestoreIndicesArgs getArgs(UpgradeContext context) { } @VisibleForTesting - int getRowCount(RestoreIndicesArgs args) { - ExpressionList countExp = - _server - .find(EbeanAspectV2.class) - .where() - .eq(EbeanAspectV2.VERSION_COLUMN, ASPECT_LATEST_VERSION); - if (args.aspectName != null) { - countExp = countExp.eq(EbeanAspectV2.ASPECT_COLUMN, args.aspectName); - } - if (args.urn != null) { - countExp = countExp.eq(EbeanAspectV2.URN_COLUMN, args.urn); - } - if (args.urnLike != null) { - countExp = countExp.like(EbeanAspectV2.URN_COLUMN, args.urnLike); - } - return countExp.findCount(); + int getRowCount(UpgradeContext context, RestoreIndicesArgs args) { + return _entityService.countAspect(args, context.report()::addLine); } @Override @@ -184,7 +166,7 @@ public Function executable() { context.report().addLine("Sending MAE from local DB"); long startTime = System.currentTimeMillis(); - final int rowCount = getRowCount(args); + final int rowCount = getRowCount(context, args); context .report() .addLine( @@ -224,7 +206,7 @@ public Function executable() { context.report().addLine("End of data."); break; } else { - log.error("Failure processing restore indices batch.", e); + context.report().addLine("Exception while processing batch", e); return new DefaultUpgradeStepResult(id(), DataHubUpgradeState.FAILED); } } diff --git a/datahub-upgrade/src/test/java/com/linkedin/datahub/upgrade/UpgradeReportTest.java b/datahub-upgrade/src/test/java/com/linkedin/datahub/upgrade/UpgradeReportTest.java new file mode 100644 index 00000000000000..b8be334b6e957d --- /dev/null +++ b/datahub-upgrade/src/test/java/com/linkedin/datahub/upgrade/UpgradeReportTest.java @@ -0,0 +1,56 @@ +package com.linkedin.datahub.upgrade; + +import static org.testng.Assert.*; + +import com.linkedin.datahub.upgrade.impl.DefaultUpgradeReport; +import java.util.List; +import org.testng.annotations.BeforeMethod; +import org.testng.annotations.Test; + +public class UpgradeReportTest { + + private UpgradeReport upgradeReport; + + @BeforeMethod + public void setup() { + upgradeReport = new DefaultUpgradeReport(); + } + + @Test + public void testAddLine() { + // Given + String line1 = "Starting upgrade"; + String line2 = "Processing step 1"; + String line3 = "Upgrade completed"; + + // When + upgradeReport.addLine(line1); + upgradeReport.addLine(line2); + upgradeReport.addLine(line3); + + // Then + List lines = upgradeReport.lines(); + assertEquals(lines.size(), 3); + assertEquals(lines.get(0), line1); + assertEquals(lines.get(1), line2); + assertEquals(lines.get(2), line3); + } + + @Test + public void testAddLineWithException() { + // Given + String errorMessage = "Error occurred during upgrade"; + Exception testException = new RuntimeException("Test exception message"); + + // When + upgradeReport.addLine(errorMessage, testException); + + // Then + List lines = upgradeReport.lines(); + assertEquals(lines.size(), 2); + assertTrue(lines.get(0).contains(errorMessage)); + assertTrue(lines.get(0).contains("Test exception message")); + assertTrue(lines.get(1).contains("Exception stack trace:")); + assertTrue(lines.get(1).contains("RuntimeException")); + } +} diff --git a/datahub-upgrade/src/test/java/com/linkedin/datahub/upgrade/restoreindices/SendMAEStepTest.java b/datahub-upgrade/src/test/java/com/linkedin/datahub/upgrade/restoreindices/SendMAEStepTest.java index 0ef491de8d669d..614068408368ca 100644 --- a/datahub-upgrade/src/test/java/com/linkedin/datahub/upgrade/restoreindices/SendMAEStepTest.java +++ b/datahub-upgrade/src/test/java/com/linkedin/datahub/upgrade/restoreindices/SendMAEStepTest.java @@ -136,6 +136,9 @@ public void testExecutableWithDefaultArgs() { // Insert a few test rows insertTestRows(5, null); + // Mock countAspect to return 5 rows + when(mockEntityService.countAspect(any(RestoreIndicesArgs.class), any())).thenReturn(5); + // Execute UpgradeStepResult result = sendMAEStep.executable().apply(mockContext); @@ -175,6 +178,8 @@ public void testExecutableWithCustomArgs() { // Insert some test data insertTestRows(5, "testAspect"); + when(mockEntityService.countAspect(any(RestoreIndicesArgs.class), any())).thenReturn(5); + // Execute UpgradeStepResult result = sendMAEStep.executable().apply(mockContext); @@ -209,6 +214,8 @@ public void testExecutableWithUrnLike() { // Insert data that matches the URN pattern insertTestRows(3, null); + when(mockEntityService.countAspect(any(RestoreIndicesArgs.class), any())).thenReturn(3); + // Execute UpgradeStepResult result = sendMAEStep.executable().apply(mockContext); @@ -242,6 +249,8 @@ public void testExecutableWithPitEpochMs() { insertTestRow("urn:li:test:2", "testAspect", 0, oneHourAgo, "testUser"); // Edge of range insertTestRow("urn:li:test:3", "testAspect", 0, twoHoursAgo, "testUser"); // Outside range + when(mockEntityService.countAspect(any(RestoreIndicesArgs.class), any())).thenReturn(2); + // Execute UpgradeStepResult result = sendMAEStep.executable().apply(mockContext); @@ -267,14 +276,14 @@ public void testExecutableWithAspectNames() { insertTestRow("urn:li:test:2", "aspect2", 0, now, "testUser"); insertTestRow("urn:li:test:3", "aspect3", 0, now, "testUser"); + when(mockEntityService.countAspect(any(RestoreIndicesArgs.class), any())).thenReturn(3); + // Execute UpgradeStepResult result = sendMAEStep.executable().apply(mockContext); - // Verify aspect names parameter ArgumentCaptor argsCaptor = ArgumentCaptor.forClass(RestoreIndicesArgs.class); verify(mockEntityService).restoreIndices(eq(mockOpContext), argsCaptor.capture(), any()); - RestoreIndicesArgs capturedArgs = argsCaptor.getValue(); assertEquals(capturedArgs.aspectNames.size(), 3); assertTrue(capturedArgs.aspectNames.contains("aspect1")); @@ -291,6 +300,8 @@ public void testExecutableWithUrnBasedPagination() { // Insert enough rows to trigger pagination insertTestRows(5, null); + when(mockEntityService.countAspect(any(RestoreIndicesArgs.class), any())).thenReturn(5); + // Setup sequential results for pagination RestoreIndicesResult firstResult = new RestoreIndicesResult(); firstResult.rowsMigrated = 2; @@ -374,6 +385,8 @@ public void testExecutableWithError() { // Insert rows so the query returns data insertTestRows(10, null); + when(mockEntityService.countAspect(any(RestoreIndicesArgs.class), any())).thenReturn(10); + // Force the service to throw an exception when(mockEntityService.restoreIndices(eq(mockOpContext), any(RestoreIndicesArgs.class), any())) .thenThrow(new RuntimeException("Test exception")); @@ -387,6 +400,28 @@ public void testExecutableWithError() { assertEquals(result.stepId(), sendMAEStep.id()); } + @Test + public void testUrnBasedPaginationExecutableWithError() { + parsedArgs.put(RestoreIndices.URN_BASED_PAGINATION_ARG_NAME, Optional.of("true")); + + when(mockEntityService.countAspect(any(RestoreIndicesArgs.class), any())).thenReturn(10); + + // Force the service to throw an exception + when(mockEntityService.restoreIndices(eq(mockOpContext), any(RestoreIndicesArgs.class), any())) + .thenThrow(new RuntimeException("Test exception")); + + // Execute + UpgradeStepResult result = sendMAEStep.executable().apply(mockContext); + + // Verify failure + assertTrue(result instanceof DefaultUpgradeStepResult); + assertEquals(result.result(), DataHubUpgradeState.FAILED); + assertEquals(result.stepId(), sendMAEStep.id()); + + // verify exception reported + verify(mockReport, atLeastOnce()).addLine(anyString(), any(Exception.class)); + } + @Test public void testReportAddedLines() { // Insert some test data @@ -407,6 +442,8 @@ public void testExecutableWithCreateDefaultAspects() { // Insert test data insertTestRows(3, null); + when(mockEntityService.countAspect(any(RestoreIndicesArgs.class), any())).thenReturn(3); + // Execute UpgradeStepResult result = sendMAEStep.executable().apply(mockContext); diff --git a/metadata-io/src/main/java/com/linkedin/metadata/entity/AspectDao.java b/metadata-io/src/main/java/com/linkedin/metadata/entity/AspectDao.java index 0c72e919defaeb..f9d5d13f795079 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/entity/AspectDao.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/entity/AspectDao.java @@ -210,6 +210,9 @@ ListResult listUrns( @Nonnull Integer countAspect(@Nonnull final String aspectName, @Nullable String urnLike); + @Nonnull + Integer countAspect(final RestoreIndicesArgs args); + @Nonnull PartitionedStream streamAspectBatches(final RestoreIndicesArgs args); diff --git a/metadata-io/src/main/java/com/linkedin/metadata/entity/EntityServiceImpl.java b/metadata-io/src/main/java/com/linkedin/metadata/entity/EntityServiceImpl.java index c5547a94db63ec..f02f2909c7f752 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/entity/EntityServiceImpl.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/entity/EntityServiceImpl.java @@ -1781,6 +1781,12 @@ public Integer getCountAspect( return aspectDao.countAspect(aspectName, urnLike); } + @Override + public Integer countAspect(@Nonnull RestoreIndicesArgs args, @Nonnull Consumer logger) { + logger.accept(String.format("Args are %s", args)); + return aspectDao.countAspect(args); + } + @Nonnull @Override public List restoreIndices( diff --git a/metadata-io/src/main/java/com/linkedin/metadata/entity/cassandra/CassandraAspectDao.java b/metadata-io/src/main/java/com/linkedin/metadata/entity/cassandra/CassandraAspectDao.java index 3aa186f1416eae..5ed46252860e6e 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/entity/cassandra/CassandraAspectDao.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/entity/cassandra/CassandraAspectDao.java @@ -574,6 +574,13 @@ public Integer countAspect(@Nonnull String aspectName, @Nullable String urnLike) return -1; } + @Nonnull + @Override + public Integer countAspect(final RestoreIndicesArgs args) { + // Not implemented + return -1; + } + @Nonnull public PartitionedStream streamAspectBatches(final RestoreIndicesArgs args) { // Not implemented diff --git a/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/EbeanAspectDao.java b/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/EbeanAspectDao.java index 7182f478a27bdc..1f8e63af82302c 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/EbeanAspectDao.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/EbeanAspectDao.java @@ -552,6 +552,12 @@ public Integer countAspect(@Nonnull String aspectName, @Nullable String urnLike) return exp.findCount(); } + @Nonnull + @Override + public Integer countAspect(final RestoreIndicesArgs args) { + return buildExpressionList(args, true).findCount(); + } + /** * Warning this inner Streams must be closed * @@ -576,10 +582,47 @@ public PartitionedStream streamAspectBatches(final RestoreIndices */ public PartitionedStream streamAspectBatches( final RestoreIndicesArgs args, final TxIsolation isolationLevel) { + ExpressionList exp = buildExpressionList(args, false); + if (args.limit > 0) { + exp = exp.setMaxRows(args.limit); + } + + int start = args.urnBasedPagination ? 0 : args.start; + + // Execute with specific transaction isolation level + Stream stream; + if (isolationLevel == TxIsolation.READ_UNCOMMITTED) { + // Use explicit transaction scope for READ_UNCOMMITTED to override default + try (Transaction transaction = + server.beginTransaction(TxScope.requiresNew().setIsolation(isolationLevel))) { + stream = + exp.orderBy() + .asc(EbeanAspectV2.URN_COLUMN) + .orderBy() + .asc(EbeanAspectV2.ASPECT_COLUMN) + .setFirstRow(start) + .findStream(); // Transaction auto-closes when stream completes + } + } else { + // For READ_COMMITTED and other levels, use standard approach + stream = + exp.orderBy() + .asc(EbeanAspectV2.URN_COLUMN) + .orderBy() + .asc(EbeanAspectV2.ASPECT_COLUMN) + .setFirstRow(start) + .findStream(); + } + + return PartitionedStream.builder().delegateStream(stream).build(); + } + + private ExpressionList buildExpressionList( + RestoreIndicesArgs args, boolean forCount) { ExpressionList exp = server .find(EbeanAspectV2.class) - .select(EbeanAspectV2.ALL_COLUMNS) + .select(forCount ? EbeanAspectV2.KEY_ID : EbeanAspectV2.ALL_COLUMNS) .where() .eq(EbeanAspectV2.VERSION_COLUMN, ASPECT_LATEST_VERSION); if (args.aspectName != null) { @@ -604,9 +647,7 @@ public PartitionedStream streamAspectBatches( Timestamp.from(Instant.ofEpochMilli(args.lePitEpochMs))); } - int start = args.start; if (args.urnBasedPagination) { - start = 0; if (args.lastUrn != null && !args.lastUrn.isEmpty()) { exp = exp.where().ge(EbeanAspectV2.URN_COLUMN, args.lastUrn); @@ -622,37 +663,7 @@ public PartitionedStream streamAspectBatches( } } } - - if (args.limit > 0) { - exp = exp.setMaxRows(args.limit); - } - - // Execute with specific transaction isolation level - Stream stream; - if (isolationLevel == TxIsolation.READ_UNCOMMITTED) { - // Use explicit transaction scope for READ_UNCOMMITTED to override default - try (Transaction transaction = - server.beginTransaction(TxScope.requiresNew().setIsolation(isolationLevel))) { - stream = - exp.orderBy() - .asc(EbeanAspectV2.URN_COLUMN) - .orderBy() - .asc(EbeanAspectV2.ASPECT_COLUMN) - .setFirstRow(start) - .findStream(); // Transaction auto-closes when stream completes - } - } else { - // For READ_COMMITTED and other levels, use standard approach - stream = - exp.orderBy() - .asc(EbeanAspectV2.URN_COLUMN) - .orderBy() - .asc(EbeanAspectV2.ASPECT_COLUMN) - .setFirstRow(start) - .findStream(); - } - - return PartitionedStream.builder().delegateStream(stream).build(); + return exp; } /** diff --git a/metadata-io/src/test/java/com/linkedin/metadata/entity/AspectDaoTest.java b/metadata-io/src/test/java/com/linkedin/metadata/entity/AspectDaoTest.java index ff640d70364cb8..df6eefbf75a009 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/entity/AspectDaoTest.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/entity/AspectDaoTest.java @@ -455,6 +455,11 @@ public Integer countAspect(String aspectName, String urnLike) { return null; } + @Override + public Integer countAspect(RestoreIndicesArgs args) { + return null; + } + @Override public PartitionedStream streamAspectBatches(RestoreIndicesArgs args) { return null; diff --git a/metadata-io/src/test/java/com/linkedin/metadata/entity/EntityServiceTest.java b/metadata-io/src/test/java/com/linkedin/metadata/entity/EntityServiceTest.java index 238226d2586421..992a0576b01866 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/entity/EntityServiceTest.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/entity/EntityServiceTest.java @@ -3187,6 +3187,78 @@ public void testFailedAspectValidation() throws Exception { } } + @Test + public void testCountAspect() throws Exception { + if (!(this instanceof EbeanEntityServiceTest)) { + return; + } + + // Setup: Create test data with different URNs and aspects + Urn entityUrn1 = UrnUtils.getUrn("urn:li:corpuser:testCountAspect1"); + Urn entityUrn2 = UrnUtils.getUrn("urn:li:corpuser:testCountAspect2"); + Urn entityUrn3 = + UrnUtils.getUrn("urn:li:dataset:(urn:li:dataPlatform:test,testCountAspect3,PROD)"); + + CorpUserInfo userInfo1 = AspectGenerationUtils.createCorpUserInfo("test1@test.com"); + CorpUserInfo userInfo2 = AspectGenerationUtils.createCorpUserInfo("test2@test.com"); + DatasetProperties datasetProperties = new DatasetProperties(); + datasetProperties.setDescription("Test dataset"); + + SystemMetadata metadata = AspectGenerationUtils.createSystemMetadata(); + + // Ingest test aspects + _entityServiceImpl.ingestAspects( + opContext, + entityUrn1, + List.of(new Pair<>(AspectGenerationUtils.getAspectName(userInfo1), userInfo1)), + TEST_AUDIT_STAMP, + metadata); + + _entityServiceImpl.ingestAspects( + opContext, + entityUrn2, + List.of(new Pair<>(AspectGenerationUtils.getAspectName(userInfo2), userInfo2)), + TEST_AUDIT_STAMP, + metadata); + + _entityServiceImpl.ingestAspects( + opContext, + entityUrn3, + List.of( + new Pair<>(AspectGenerationUtils.getAspectName(datasetProperties), datasetProperties)), + TEST_AUDIT_STAMP, + metadata); + + List logMessages = new ArrayList<>(); + + // Test case 1: No filter - should return count of all aspects + RestoreIndicesArgs args1 = new RestoreIndicesArgs(); + int count1 = _entityServiceImpl.countAspect(args1, logMessages::add); + assertTrue(count1 >= 3, "Should have at least 3 aspects (corpUserInfo x2 + datasetProperties)"); + + // Test case 2: urnLike filter - should return count of aspects matching the URN pattern + RestoreIndicesArgs args2 = new RestoreIndicesArgs(); + args2.urnLike = "%corpuser:testCountAspect%"; + int count2 = _entityServiceImpl.countAspect(args2, logMessages::add); + assertTrue(count2 >= 2, "Should have at least 2 corpuser aspects"); + + // Test case 3: urnLike + aspectName filter - should return count of matching aspects + RestoreIndicesArgs args3 = new RestoreIndicesArgs(); + args3.urnLike = "%corpuser:testCountAspect%"; + args3.aspectName = "corpUserInfo"; + int count3 = _entityServiceImpl.countAspect(args3, logMessages::add); + assertEquals(count3, 2, "Should have exactly 2 corpUserInfo aspects for testCountAspect users"); + + // Test case 4: aspectName filter only + RestoreIndicesArgs args4 = new RestoreIndicesArgs(); + args4.aspectName = "datasetProperties"; + int count4 = _entityServiceImpl.countAspect(args4, logMessages::add); + assertTrue(count4 >= 1, "Should have at least 1 datasetProperties aspect"); + + // Verify logger was called + assertFalse(logMessages.isEmpty(), "Logger should have been called"); + } + @Nonnull protected com.linkedin.entity.Entity createCorpUserEntity(Urn entityUrn, String email) throws Exception { diff --git a/metadata-io/src/test/java/com/linkedin/metadata/entity/ebean/EbeanAspectDaoTest.java b/metadata-io/src/test/java/com/linkedin/metadata/entity/ebean/EbeanAspectDaoTest.java index 8ad11e5e240870..406de5dd87e7e1 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/entity/ebean/EbeanAspectDaoTest.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/entity/ebean/EbeanAspectDaoTest.java @@ -17,6 +17,7 @@ import io.datahubproject.metadata.context.OperationContext; import io.ebean.Database; import io.ebean.test.LoggedSql; +import java.sql.Timestamp; import java.util.List; import java.util.Map; import java.util.Set; @@ -143,4 +144,49 @@ public void testStreamAspectBatchesDefault() { var stream = testDao.streamAspectBatches(args); assertTrue(stream != null); } + + private void insertTestData() { + // Insert test data with different URN patterns and aspects + insertAspect("urn:li:test:test1", "testAspect1", 0, "test1"); + insertAspect("urn:li:test:test2", "testAspect1", 0, "test2"); + insertAspect("urn:li:test:test3", "testAspect2", 0, "test3"); + insertAspect("urn:li:other:test4", "testAspect1", 0, "test4"); + insertAspect("urn:li:other:test5", "testAspect2", 0, "test5"); + } + + private void insertAspect(String urn, String aspect, long version, String metadata) { + EbeanAspectV2 aspectRecord = new EbeanAspectV2(); + aspectRecord.setKey(new EbeanAspectV2.PrimaryKey(urn, aspect, version)); + aspectRecord.setMetadata(metadata); + aspectRecord.setCreatedBy("test"); + aspectRecord.setCreatedFor(null); + aspectRecord.setCreatedOn(new Timestamp(System.currentTimeMillis())); + aspectRecord.setSystemMetadata(null); + testDao.getServer().save(aspectRecord); + } + + @Test + public void testCountAspect() { + // Setup test data + insertTestData(); + + // Test case 1: No filter - should return count of all aspects + var args1 = new com.linkedin.metadata.entity.restoreindices.RestoreIndicesArgs(); + int count1 = testDao.countAspect(args1); + assertEquals(count1, 5, "Should return count of all aspects"); + + // Test case 2: urnLike filter - should return count of aspects matching the URN pattern + var args2 = new com.linkedin.metadata.entity.restoreindices.RestoreIndicesArgs(); + args2.urnLike = "%:test:%"; + int count2 = testDao.countAspect(args2); + assertEquals(count2, 3, "Should return count of aspects matching URN pattern '%:test:%'"); + + // Test case 3: urnLike + aspect filter - should return count of matching aspects + var args3 = new com.linkedin.metadata.entity.restoreindices.RestoreIndicesArgs(); + args3.urnLike = "%:test:%"; + args3.aspectName = "testAspect1"; + int count3 = testDao.countAspect(args3); + assertEquals( + count3, 2, "Should return count of aspects matching both URN pattern and aspect name"); + } } diff --git a/metadata-service/services/src/main/java/com/linkedin/metadata/entity/EntityService.java b/metadata-service/services/src/main/java/com/linkedin/metadata/entity/EntityService.java index bf319b6aa618dd..7bd30909bf41d0 100644 --- a/metadata-service/services/src/main/java/com/linkedin/metadata/entity/EntityService.java +++ b/metadata-service/services/src/main/java/com/linkedin/metadata/entity/EntityService.java @@ -423,6 +423,8 @@ String batchApplyRetention( Integer getCountAspect( @Nonnull OperationContext opContext, @Nonnull String aspectName, @Nullable String urnLike); + Integer countAspect(@Nonnull RestoreIndicesArgs args, @Nonnull Consumer logger); + // TODO: Extract this to a different service, doesn't need to be here List restoreIndices( @Nonnull OperationContext opContext,