From 5cbf8fc1ce86ebd6bf83d8f4a5bbdae2cc8da97a Mon Sep 17 00:00:00 2001 From: Yang Yang Date: Thu, 29 Aug 2024 09:15:21 -0700 Subject: [PATCH] remove tracking context from ingestAsset API (#408) Co-authored-by: Yang Yang --- .../metadata/restli/BaseAspectRoutingResource.java | 4 +++- .../metadata/restli/BaseEntityResource.java | 14 +++++++------- .../metadata/restli/BaseEntityResourceTest.java | 6 ++++-- 3 files changed, 14 insertions(+), 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 index 1af7ee660..3667f36db 100644 --- 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 ingestInternal(@Nonnull SNAPSHOT snapshot, @Nonnull @Override protected Task ingestInternalAsset(@Nonnull ASSET asset, - @Nonnull Set> aspectsToIgnore, @Nullable IngestionTrackingContext trackingContext, + @Nonnull Set> 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 index 05f2588cd..72fdc3956 100644 --- 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> batchGetWithErrors(@Nonnull Set 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 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 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 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 ingestInternal(@Nonnull SNAPSHOT snapshot, @Nonnull protected Task ingestInternalAsset(@Nonnull ASSET asset, - @Nonnull Set> aspectsToIgnore, @Nullable IngestionTrackingContext trackingContext, + @Nonnull Set> 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 index 7c2e67dd7..714a78f6a 100644 --- 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));