remove tracking context from ingestAsset API (#408)

Co-authored-by: Yang Yang <[email protected]>
linkedin · Aug 29, 2024 · 5cbf8fc · 5cbf8fc
1 parent 62dda81
commit 5cbf8fc
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 10 deletions.
diff --git a/restli-resources/src/main/java/com/linkedin/metadata/restli/BaseAspectRoutingResource.java b/restli-resources/src/main/java/com/linkedin/metadata/restli/BaseAspectRoutingResource.java
@@ -345,12 +345,14 @@ protected Task<Void> ingestInternal(@Nonnull SNAPSHOT snapshot,
   @Nonnull
   @Override
   protected Task<Void> ingestInternalAsset(@Nonnull ASSET asset,
-      @Nonnull Set<Class<? extends RecordTemplate>> aspectsToIgnore, @Nullable IngestionTrackingContext trackingContext,
+      @Nonnull Set<Class<? extends RecordTemplate>> aspectsToIgnore,
       @Nullable IngestionParams ingestionParams) {
     // TODO: META-18950: add trackingContext to BaseAspectRoutingResource. currently the param is unused.
     return RestliUtils.toTask(() -> {
       final URN urn = (URN) ModelUtils.getUrnFromAsset(asset);
       final AuditStamp auditStamp = getAuditor().requestAuditStamp(getContext().getRawRequestContext());
+      IngestionTrackingContext trackingContext =
+          ingestionParams != null ? ingestionParams.getIngestionTrackingContext() : null;
       ModelUtils.getAspectsFromAsset(asset).forEach(aspect -> {
         if (!aspectsToIgnore.contains(aspect.getClass())) {
           if (getAspectRoutingGmsClientManager().hasRegistered(aspect.getClass())) {

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
@@ -309,7 +309,7 @@ private Task<BatchResult<KEY, VALUE>> batchGetWithErrors(@Nonnull Set<KEY> ids,
   }
 
   /**
-   * Deprecated to use {@link #ingestAsset(RecordTemplate, IngestionTrackingContext, IngestionParams)} instead.
+   * Deprecated to use {@link #ingestAsset(RecordTemplate, IngestionParams)} instead.
    * An action method for automated ingestion pipeline.
    */
   @Deprecated
@@ -320,7 +320,7 @@ public Task<Void> ingest(@ActionParam(PARAM_SNAPSHOT) @Nonnull SNAPSHOT snapshot
   }
 
   /**
-   * Deprecated to use {@link #ingestAsset(RecordTemplate, IngestionTrackingContext, IngestionParams)} instead.
+   * Deprecated to use {@link #ingestAsset(RecordTemplate, IngestionParams)} instead.
    * Same as {@link #ingest(RecordTemplate)} but with tracking context attached.
    * @param snapshot Snapshot of the metadata change to be ingested
    * @param trackingContext {@link IngestionTrackingContext} to 1) track DAO-level metrics and 2) to pass on to MAE emission
@@ -338,15 +338,13 @@ public Task<Void> ingestWithTracking(@ActionParam(PARAM_SNAPSHOT) @Nonnull SNAPS
   /**
    * An action method for automated ingestion pipeline.
    * @param asset Asset of the metadata change to be ingested
-   * @param trackingContext {@link IngestionTrackingContext} to 1) track DAO-level metrics and 2) to pass on to MAE emission
    * @return ingest task
    */
   @Action(name = ACTION_INGEST_ASSET)
   @Nonnull
   public Task<Void> ingestAsset(@ActionParam(PARAM_ASSET) @Nonnull ASSET asset,
-      @ActionParam(PARAM_TRACKING_CONTEXT) @Nonnull IngestionTrackingContext trackingContext,
       @Optional @ActionParam(PARAM_INGESTION_PARAMS) IngestionParams ingestionParams) {
-    return ingestInternalAsset(asset, Collections.emptySet(), trackingContext, ingestionParams);
+    return ingestInternalAsset(asset, Collections.emptySet(), ingestionParams);
   }
 
   @Nonnull
@@ -367,14 +365,16 @@ protected Task<Void> ingestInternal(@Nonnull SNAPSHOT snapshot,
 
   @Nonnull
   protected Task<Void> ingestInternalAsset(@Nonnull ASSET asset,
-      @Nonnull Set<Class<? extends RecordTemplate>> aspectsToIgnore, @Nullable IngestionTrackingContext trackingContext,
+      @Nonnull Set<Class<? extends RecordTemplate>> aspectsToIgnore,
       @Nullable IngestionParams ingestionParams) {
     return RestliUtils.toTask(() -> {
       final URN urn = (URN) ModelUtils.getUrnFromAsset(asset);
       final AuditStamp auditStamp = getAuditor().requestAuditStamp(getContext().getRawRequestContext());
+      IngestionTrackingContext ingestionTrackingContext =
+          ingestionParams != null ? ingestionParams.getIngestionTrackingContext() : null;
       ModelUtils.getAspectsFromAsset(asset).stream().forEach(aspect -> {
         if (!aspectsToIgnore.contains(aspect.getClass())) {
-          getLocalDAO().add(urn, aspect, auditStamp, trackingContext, ingestionParams);
+          getLocalDAO().add(urn, aspect, auditStamp, ingestionTrackingContext, ingestionParams);
         }
       });
       return null;

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
@@ -1459,13 +1459,15 @@ public void testIngestAsset() {
     IngestionTrackingContext trackingContext = new IngestionTrackingContext();
 
     IngestionParams ingestionParams1 = new IngestionParams().setTestMode(true);
-    runAndWait(_internalResource.ingestAsset(asset, trackingContext, ingestionParams1));
+    ingestionParams1.setIngestionTrackingContext(trackingContext);
+    runAndWait(_internalResource.ingestAsset(asset, ingestionParams1));
 
     verify(_mockLocalDAO, times(1)).add(eq(urn), eq(foo), any(), eq(trackingContext), eq(ingestionParams1));
     verify(_mockLocalDAO, times(1)).add(eq(urn), eq(bar), any(), eq(trackingContext), eq(ingestionParams1));
 
     IngestionParams ingestionParams2 = new IngestionParams().setIngestionMode(IngestionMode.LIVE);
-    runAndWait(_internalResource.ingestAsset(asset, trackingContext, ingestionParams2));
+    ingestionParams2.setIngestionTrackingContext(trackingContext);
+    runAndWait(_internalResource.ingestAsset(asset, ingestionParams2));
 
     verify(_mockLocalDAO, times(1)).add(eq(urn), eq(foo), any(), eq(trackingContext), eq(ingestionParams2));
     verify(_mockLocalDAO, times(1)).add(eq(urn), eq(bar), any(), eq(trackingContext), eq(ingestionParams2));