forked from linkedin/datahub-gma
-
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.
Injecting PreIngestionLambda in BaseLocalDao (linkedin#410)
* Addressed comments * Added javadoc * Adding Registry to get lambda * Added unit tests * Update test name * Added one more unit test * Addressed all comments * Fixed unit typo * Resetting all changes * Adding unit tests again * Updated java doc * Indent * Fix java doc error * Java doc attempt 2 * Minor nitpicks --------- Co-authored-by: Rakhi Agrawal <rakagrawal@linkedin.com>
- Loading branch information
Showing
7 changed files
with
215 additions
and
4 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
40 changes: 40 additions & 0 deletions
40
.../main/java/com/linkedin/metadata/dao/ingestion/RestliCompliantPreUpdateRoutingClient.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,40 @@ | ||
package com.linkedin.metadata.dao.ingestion; | ||
|
||
import com.google.protobuf.Message; | ||
|
||
import com.linkedin.common.urn.Urn; | ||
import com.linkedin.data.template.RecordTemplate; | ||
|
||
/** | ||
* A restli client to route update request to the appropriate to custom APIs. | ||
* <p>This interface extends {@link PreUpdateRoutingClient} and provides additional methods for converting | ||
* * URNs and aspects between different representations (e.g., from Pegasus to Protobuf).</p> | ||
* * | ||
*/ | ||
|
||
public interface RestliCompliantPreUpdateRoutingClient<ASPECT extends Message> extends PreUpdateRoutingClient { | ||
|
||
/** | ||
* Converts a URN to a Protobuf message. | ||
* | ||
* @param pegasusUrn the URN to be converted | ||
* @return the converted Protobuf message | ||
*/ | ||
Message convertUrnToMessage(Urn pegasusUrn); | ||
|
||
/** | ||
* Converts a {@link RecordTemplate} aspect to a Protobuf message aspect. | ||
* | ||
* @param pegasusAspect the aspect to be converted | ||
* @return the converted Protobuf message aspect | ||
*/ | ||
ASPECT convertAspectToMessage(RecordTemplate pegasusAspect); | ||
|
||
/** | ||
* Converts a Protobuf message aspect to a {@link RecordTemplate} aspect. | ||
* | ||
* @param messageAspect the Protobuf message aspect to be converted | ||
* @return the converted {@link RecordTemplate} aspect | ||
*/ | ||
RecordTemplate convertAspectFromMessage(ASPECT messageAspect); | ||
} |
24 changes: 24 additions & 0 deletions
24
dao-api/src/main/java/com/linkedin/metadata/dao/ingestion/RestliPreUpdateAspectRegistry.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,24 @@ | ||
package com.linkedin.metadata.dao.ingestion; | ||
|
||
import com.linkedin.data.template.RecordTemplate; | ||
import javax.annotation.Nonnull; | ||
import javax.annotation.Nullable; | ||
|
||
|
||
/** | ||
* A registry which maintains mapping of aspects and their getPreUpdateRoutingClient. | ||
*/ | ||
public interface RestliPreUpdateAspectRegistry { | ||
|
||
/** | ||
* Get PreUpdateRoutingClient for an aspect. | ||
*/ | ||
@Nullable | ||
<ASPECT extends RecordTemplate> RestliCompliantPreUpdateRoutingClient getPreUpdateRoutingClient(@Nonnull final ASPECT aspect); | ||
|
||
/** | ||
* Check if PreUpdateRoutingClient is registered for an aspect. | ||
*/ | ||
<ASPECT extends RecordTemplate> boolean isRegistered(@Nonnull final Class<ASPECT> aspectClass); | ||
|
||
} |
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
28 changes: 28 additions & 0 deletions
28
.../src/test/java/com/linkedin/metadata/dao/ingestion/SamplePreUpdateAspectRegistryImpl.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,28 @@ | ||
package com.linkedin.metadata.dao.ingestion; | ||
|
||
import com.google.common.collect.ImmutableMap; | ||
import com.linkedin.data.template.RecordTemplate; | ||
import com.linkedin.testing.AspectFoo; | ||
import javax.annotation.Nonnull; | ||
import javax.annotation.Nullable; | ||
|
||
|
||
public class SamplePreUpdateAspectRegistryImpl implements RestliPreUpdateAspectRegistry { | ||
private final ImmutableMap<Class<? extends RecordTemplate>, RestliCompliantPreUpdateRoutingClient> registry; | ||
|
||
public SamplePreUpdateAspectRegistryImpl() { | ||
registry = new ImmutableMap.Builder<Class<? extends RecordTemplate>, RestliCompliantPreUpdateRoutingClient>() | ||
.put(AspectFoo.class, new SamplePreUpdateRoutingClient()) | ||
.build(); | ||
} | ||
@Nullable | ||
@Override | ||
public <ASPECT extends RecordTemplate> RestliCompliantPreUpdateRoutingClient getPreUpdateRoutingClient(@Nonnull ASPECT aspect) { | ||
return registry.get(aspect.getClass()); | ||
} | ||
|
||
@Override | ||
public <ASPECT extends RecordTemplate> boolean isRegistered(@Nonnull Class<ASPECT> aspectClass) { | ||
return registry.containsKey(aspectClass); | ||
} | ||
} |
40 changes: 40 additions & 0 deletions
40
dao-api/src/test/java/com/linkedin/metadata/dao/ingestion/SamplePreUpdateRoutingClient.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,40 @@ | ||
package com.linkedin.metadata.dao.ingestion; | ||
|
||
import com.google.protobuf.Any; | ||
import com.google.protobuf.Message; | ||
import com.google.protobuf.StringValue; | ||
import com.linkedin.common.urn.Urn; | ||
import com.linkedin.data.template.RecordTemplate; | ||
import com.linkedin.testing.AspectFoo; | ||
|
||
|
||
public class SamplePreUpdateRoutingClient implements RestliCompliantPreUpdateRoutingClient { | ||
@Override | ||
public Message routingLambda(Message urn, Message aspect) { | ||
// For testing, change the aspect value to "bar" | ||
return Any.pack(StringValue.of("bar")); | ||
} | ||
|
||
@Override | ||
public Message convertUrnToMessage(Urn urn) { | ||
// Directly wrap the URN string into a Protobuf message for testing | ||
return Any.pack(StringValue.of(urn.toString())); | ||
} | ||
|
||
@Override | ||
public Message convertAspectToMessage(RecordTemplate pegasusAspect) { | ||
// For testing, convert AspectFoo to a TestMessageProtos.AspectMessage | ||
// Assuming the aspect has a `value` field and its string representation can be used for now | ||
String aspectString = pegasusAspect.toString(); // Extracting the aspect as a string (e.g., {value=foo}) | ||
|
||
// Wrap the aspect string into a simple Protobuf message for testing | ||
return Any.pack(StringValue.of(aspectString)); | ||
} | ||
|
||
@Override | ||
public RecordTemplate convertAspectFromMessage(Message messageAspect) { | ||
// For testing, convert TestMessageProtos.AspectMessage back to AspectFoo | ||
// Create a new RecordTemplate (AspectFoo in this case) and set the value field | ||
return new AspectFoo().setValue("bar"); | ||
} | ||
} |