fix: ensure existence of urn when calling BaseEntityResource::get (#132)

* feat: add pre-update hooks to BaseLocalDAO Introduce infrastructure for pre-update hooks, motivated by the need for a pre-update hook to remove duplicate elements from certain aspects. * simplified adding pre- and post-update hooks to one helper method. * docs: add getting started section to contributing doc * fixed build style check failure * check if urn(s) exists when calling getInternalNonEmpty * cleaned up some code * cleaner approach * cleaned up code, added a todo for batch get * cleaned up some tests * added more tests to cover new expected behavior * added override tag for exists method * removed reference to internal ticket * added test cases for BaseEntityResourceTest Co-authored-by: Justin Donn <[email protected]>
linkedin · Nov 3, 2021 · 1d6f311 · 1d6f311
1 parent a9e6109
commit 1d6f311
Show file tree

Hide file tree

Showing 8 changed files with 100 additions and 24 deletions.
diff --git a/dao-api/src/main/java/com/linkedin/metadata/dao/BaseLocalDAO.java b/dao-api/src/main/java/com/linkedin/metadata/dao/BaseLocalDAO.java
@@ -594,6 +594,13 @@ protected abstract <ASPECT extends RecordTemplate> long getNextVersion(@Nonnull
   protected abstract void save(@Nonnull URN urn, @Nonnull RecordTemplate value, @Nonnull AuditStamp auditStamp,
       long version, boolean insert);
 
+  /**
+   * Returns a boolean representing if an Urn has any Aspects associated with it (i.e. if it exists in the DB).
+   * @param urn {@link Urn} for the entity
+   * @return boolean representing if entity associated with Urn exists
+   */
+  public abstract boolean exists(@Nonnull URN urn);
+
   /**
    * Applies version-based retention against a specific aspect type for an entity.
    *

diff --git a/dao-api/src/test/java/com/linkedin/metadata/dao/BaseLocalDAOTest.java b/dao-api/src/test/java/com/linkedin/metadata/dao/BaseLocalDAOTest.java
@@ -77,6 +77,11 @@ protected void save(FooUrn urn, RecordTemplate value, AuditStamp auditStamp, lon
 
     }
 
+    @Override
+    public boolean exists(FooUrn urn) {
+      return true;
+    }
+
     @Override
     protected <ASPECT extends RecordTemplate> void applyVersionBasedRetention(Class<ASPECT> aspectClass, FooUrn urn,
         VersionBasedRetention retention, long largestVersion) {

diff --git a/dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/EbeanLocalDAO.java b/dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/EbeanLocalDAO.java
@@ -560,6 +560,11 @@ protected <ASPECT extends RecordTemplate> void applyTimeBasedRetention(@Nonnull
     return result;
   }
 
+  @Override
+  public boolean exists(@Nonnull URN urn) {
+    return _server.find(EbeanMetadataAspect.class).where().eq(URN_COLUMN, urn.toString()).exists();
+  }
+
   public boolean existsInLocalIndex(@Nonnull URN urn) {
     return _server.find(EbeanMetadataIndex.class).where().eq(URN_COLUMN, urn.toString()).exists();
   }

diff --git a/dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/EbeanLocalDAOTest.java b/dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/EbeanLocalDAOTest.java
@@ -1458,6 +1458,21 @@ record = getRecordFromLocalIndex(recordId);
     assertEquals(record.getStringVal(), "valFoo");
   }
 
+  @Test
+  void testExists() {
+    // given
+    EbeanLocalDAO<EntityAspectUnion, FooUrn> dao = createDao(FooUrn.class);
+    FooUrn urn = makeFooUrn(1);
+
+    assertFalse(dao.exists(urn));
+
+    // add metadata
+    AspectFoo fooV1 = new AspectFoo().setValue("foo");
+    addMetadata(urn, AspectFoo.class.getCanonicalName(), 1, fooV1);
+
+    assertTrue(dao.exists(urn));
+  }
+
   @Test
   void testExistsInLocalIndex() {
     EbeanLocalDAO<EntityAspectUnion, BarUrn> dao = createDao(BarUrn.class);

diff --git a/restli-resources/src/main/java/com/linkedin/metadata/restli/BaseEntityResource.java b/restli-resources/src/main/java/com/linkedin/metadata/restli/BaseEntityResource.java
@@ -147,7 +147,10 @@ public Task<VALUE> get(@Nonnull KEY id,
 
     return RestliUtils.toTask(() -> {
       final URN urn = toUrn(id);
-      final VALUE value = getInternalNonEmpty(Collections.singleton(urn), parseAspectsParam(aspectNames)).get(urn);
+      if (!getLocalDAO().exists(urn)) {
+        throw RestliUtils.resourceNotFoundException();
+      }
+      final VALUE value = getInternal(Collections.singleton(urn), parseAspectsParam(aspectNames)).get(urn);
       if (value == null) {
         throw RestliUtils.resourceNotFoundException();
       }
@@ -163,6 +166,7 @@ public Task<VALUE> get(@Nonnull KEY id,
   public Task<Map<KEY, VALUE>> batchGet(
       @Nonnull Set<KEY> ids,
       @QueryParam(PARAM_ASPECTS) @Optional @Nullable String[] aspectNames) {
+    // TODO: discuss and sync with get()'s intended behavior (check if urn exists). https://github.com/linkedin/datahub-gma/issues/136
     return RestliUtils.toTask(() -> {
       final Map<URN, KEY> urnMap =
           ids.stream().collect(Collectors.toMap(id -> toUrn(id), Function.identity()));

diff --git a/restli-resources/src/test/java/com/linkedin/metadata/restli/BaseEntityResourceTest.java b/restli-resources/src/test/java/com/linkedin/metadata/restli/BaseEntityResourceTest.java
@@ -141,6 +141,7 @@ public void testGet() {
     AspectKey<FooUrn, AspectFoo> aspect1Key = new AspectKey<>(AspectFoo.class, urn, LATEST_VERSION);
     AspectKey<FooUrn, AspectBar> aspect2Key = new AspectKey<>(AspectBar.class, urn, LATEST_VERSION);
 
+    when(_mockLocalDAO.exists(urn)).thenReturn(true);
     when(_mockLocalDAO.get(new HashSet<>(Arrays.asList(aspect1Key, aspect2Key)))).thenReturn(
         Collections.singletonMap(aspect1Key, Optional.of(foo)));
 
@@ -151,22 +152,40 @@ public void testGet() {
   }
 
   @Test
-  public void testGetNotFound() {
+  public void testGetUrnNotFound() {
     FooUrn urn = makeFooUrn(1234);
 
     AspectKey<FooUrn, AspectFoo> aspect1Key = new AspectKey<>(AspectFoo.class, urn, LATEST_VERSION);
     AspectKey<FooUrn, AspectBar> aspect2Key = new AspectKey<>(AspectBar.class, urn, LATEST_VERSION);
 
+    when(_mockLocalDAO.exists(urn)).thenReturn(false);
     when(_mockLocalDAO.get(new HashSet<>(Arrays.asList(aspect1Key, aspect2Key)))).thenReturn(Collections.emptyMap());
 
     try {
       runAndWait(_resource.get(makeResourceKey(urn), new String[0]));
+      fail("An exception should've been thrown!");
     } catch (RestLiServiceException e) {
       assertEquals(e.getStatus(), HttpStatus.S_404_NOT_FOUND);
-      return;
     }
+  }
 
-    fail("No exception thrown");
+  @Test
+  public void testGetWithEmptyAspects() {
+    FooUrn urn = makeFooUrn(1234);
+
+    AspectKey<FooUrn, AspectFoo> aspect1Key = new AspectKey<>(AspectFoo.class, urn, LATEST_VERSION);
+    AspectKey<FooUrn, AspectBar> aspect2Key = new AspectKey<>(AspectBar.class, urn, LATEST_VERSION);
+
+    when(_mockLocalDAO.exists(urn)).thenReturn(true);
+    when(_mockLocalDAO.get(new HashSet<>(Arrays.asList(aspect1Key, aspect2Key)))).thenReturn(Collections.emptyMap());
+
+    try {
+      EntityValue value = runAndWait(_resource.get(makeResourceKey(urn), new String[0]));
+      assertFalse(value.hasFoo());
+      assertFalse(value.hasBar());
+    } catch (RestLiServiceException e) {
+      fail("No exception should be thrown!");
+    }
   }
 
   @Test
@@ -176,6 +195,7 @@ public void testGetSpecificAspect() {
     AspectKey<FooUrn, AspectFoo> aspect1Key = new AspectKey<>(AspectFoo.class, urn, LATEST_VERSION);
     String[] aspectNames = {AspectFoo.class.getCanonicalName()};
 
+    when(_mockLocalDAO.exists(urn)).thenReturn(true);
     when(_mockLocalDAO.get(new HashSet<>(Arrays.asList(aspect1Key)))).thenReturn(
         Collections.singletonMap(aspect1Key, Optional.of(foo)));
 
@@ -187,17 +207,17 @@ public void testGetSpecificAspect() {
   @Test
   public void testGetSpecificAspectNotFound() {
     FooUrn urn = makeFooUrn(1234);
-    AspectKey<FooUrn, AspectFoo> aspect1Key = new AspectKey<>(AspectFoo.class, urn, LATEST_VERSION);
     String[] aspectNames = {AspectFoo.class.getCanonicalName()};
+
+    when(_mockLocalDAO.exists(urn)).thenReturn(true);
+
     try {
-      runAndWait(_resource.get(makeResourceKey(urn), aspectNames));
+      EntityValue value = runAndWait(_resource.get(makeResourceKey(urn), aspectNames));
+      assertFalse(value.hasFoo());
+      assertFalse(value.hasBar());
     } catch (RestLiServiceException e) {
-      assertEquals(e.getStatus(), HttpStatus.S_404_NOT_FOUND);
-      verify(_mockLocalDAO, times(1)).get(Collections.singleton(aspect1Key));
-      verifyNoMoreInteractions(_mockLocalDAO);
-      return;
+      fail("No exception should be thrown!");
     }
-    fail("No exception thrown");
   }
 
   @Test

diff --git a/...resources/src/test/java/com/linkedin/metadata/restli/BaseEntitySimpleKeyResourceTest.java b/...resources/src/test/java/com/linkedin/metadata/restli/BaseEntitySimpleKeyResourceTest.java
@@ -56,6 +56,7 @@ public void testGet() {
     AspectKey<Urn, AspectFoo> aspect1Key = new AspectKey<>(AspectFoo.class, urn, LATEST_VERSION);
     AspectKey<Urn, AspectBar> aspect2Key = new AspectKey<>(AspectBar.class, urn, LATEST_VERSION);
 
+    when(_mockLocalDAO.exists(urn)).thenReturn(true);
     when(_mockLocalDAO.get(new HashSet<>(Arrays.asList(aspect1Key, aspect2Key))))
         .thenReturn(Collections.singletonMap(aspect1Key, Optional.of(foo)));
 
@@ -66,33 +67,50 @@ public void testGet() {
   }
 
   @Test
-  public void testGetNotFound() {
+  public void testGetUrnNotFound() {
     long id = 1234;
     Urn urn = makeUrn(id);
 
+    when(_mockLocalDAO.exists(urn)).thenReturn(false);
+
+    try {
+      runAndWait(_resource.get(id, new String[0]));
+      fail("An exception should've been thrown!");
+    } catch (RestLiServiceException e) {
+      assertEquals(e.getStatus(), HttpStatus.S_404_NOT_FOUND);
+    }
+  }
+
+  @Test
+  public void testGetWithEmptyAspects() {
+    long id = 1234;
+    Urn urn = makeUrn(id);
+    AspectFoo foo = new AspectFoo().setValue("foo");
     AspectKey<Urn, AspectFoo> aspect1Key = new AspectKey<>(AspectFoo.class, urn, LATEST_VERSION);
     AspectKey<Urn, AspectBar> aspect2Key = new AspectKey<>(AspectBar.class, urn, LATEST_VERSION);
 
-    when(_mockLocalDAO.get(new HashSet<>(Arrays.asList(aspect1Key, aspect2Key)))).thenReturn(
-        Collections.emptyMap());
+    when(_mockLocalDAO.exists(urn)).thenReturn(true);
+    when(_mockLocalDAO.get(new HashSet<>(Arrays.asList(aspect1Key, aspect2Key))))
+        .thenReturn(Collections.emptyMap());
 
     try {
-      runAndWait(_resource.get(id, new String[0]));
+      EntityValue value = runAndWait(_resource.get(id, new String[0]));
+      assertFalse(value.hasFoo());
+      assertFalse(value.hasBar());
     } catch (RestLiServiceException e) {
-      assertEquals(e.getStatus(), HttpStatus.S_404_NOT_FOUND);
-      return;
+      fail("No exception should be thrown!");
     }
   }
 
   @Test
   public void testGetSpecificAspect() {
     long id = 1234;
     Urn urn = makeUrn(id);
-
     AspectFoo foo = new AspectFoo().setValue("foo");
     AspectKey<Urn, AspectFoo> aspect1Key = new AspectKey<>(AspectFoo.class, urn, LATEST_VERSION);
     String[] aspectNames = {AspectFoo.class.getCanonicalName()};
 
+    when(_mockLocalDAO.exists(urn)).thenReturn(true);
     when(_mockLocalDAO.get(new HashSet<>(Collections.singletonList(aspect1Key))))
         .thenReturn(Collections.singletonMap(aspect1Key, Optional.of(foo)));
 
@@ -105,15 +123,16 @@ public void testGetSpecificAspect() {
   public void testGetSpecificAspectNotFound() {
     long id = 1234;
     Urn urn = makeUrn(id);
-
-    AspectKey<Urn, AspectFoo> aspect1Key = new AspectKey<>(AspectFoo.class, urn, LATEST_VERSION);
     String[] aspectNames = {AspectFoo.class.getCanonicalName()};
+
+    when(_mockLocalDAO.exists(urn)).thenReturn(true);
+
     try {
-      runAndWait(_resource.get(id, aspectNames));
+      EntityValue value = runAndWait(_resource.get(id, aspectNames));
+      assertFalse(value.hasFoo());
+      assertFalse(value.hasBar());
     } catch (RestLiServiceException e) {
-      assertEquals(e.getStatus(), HttpStatus.S_404_NOT_FOUND);
-      verify(_mockLocalDAO, times(1)).get(Collections.singleton(aspect1Key));
-      verifyNoMoreInteractions(_mockLocalDAO);
+      fail("No exception should be thrown!");
     }
   }
 

diff --git a/...c/test/java/com/linkedin/metadata/restli/BaseSingleAspectEntitySimpleKeyResourceTest.java b/...c/test/java/com/linkedin/metadata/restli/BaseSingleAspectEntitySimpleKeyResourceTest.java
@@ -61,6 +61,7 @@ public void testGet() throws URISyntaxException {
     SingleAspectEntityUrn urn = new SingleAspectEntityUrn(id1);
     AspectBar aspect = new AspectBar().setValue(field1);
     AspectKey<SingleAspectEntityUrn, AspectBar> aspectKey = new AspectKey<>(AspectBar.class, urn, LATEST_VERSION);
+    when(_mockLocalDao.exists(urn)).thenReturn(true);
     when(_mockLocalDao.get(Collections.singleton(aspectKey))).thenReturn(
         Collections.singletonMap(aspectKey, Optional.of(aspect)));