diff --git a/README.md b/README.md
index 58764e2a789c9..22f57b6cbeb1d 100644
--- a/README.md
+++ b/README.md
@@ -67,7 +67,7 @@ HOSTED_DOCS_ONLY-->
 ## Introduction
 
 DataHub is an open-source metadata platform for the modern data stack. Read about the architectures of different metadata systems and why DataHub excels [here](https://engineering.linkedin.com/blog/2020/datahub-popular-metadata-architectures-explained). Also read our
-[LinkedIn Engineering blog post](https://engineering.linkedin.com/blog/2019/data-hub), check out our [Strata presentation](https://speakerdeck.com/shirshanka/the-evolution-of-metadata-linkedins-journey-strata-nyc-2019) and watch our [Crunch Conference Talk](https://www.youtube.com/watch?v=OB-O0Y6OYDE). You should also visit [DataHub Architecture](docs/architecture/architecture.md) to get a better understanding of how DataHub is implemented and [DataHub Onboarding Guide](docs/modeling/extending-the-metadata-model.md) to understand how to extend DataHub for your own use cases.
+[LinkedIn Engineering blog post](https://engineering.linkedin.com/blog/2019/data-hub), check out our [Strata presentation](https://speakerdeck.com/shirshanka/the-evolution-of-metadata-linkedins-journey-strata-nyc-2019) and watch our [Crunch Conference Talk](https://www.youtube.com/watch?v=OB-O0Y6OYDE). You should also visit [DataHub Architecture](docs/architecture/architecture.md) to get a better understanding of how DataHub is implemented.
 
 ## Quickstart
 
@@ -99,6 +99,10 @@ Check out DataHub's [Features](docs/features.md) & [Roadmap](https://feature-req
 
 We welcome contributions from the community. Please refer to our [Contributing Guidelines](docs/CONTRIBUTING.md) for more details. We also have a [contrib](contrib) directory for incubating experimental features.
 
+### Extending
+
+If you need to understand how to extend our model with custom types, please see [Extending the Metadata Model](docs/modeling/extending-the-metadata-model.md)
+
 ## Community
 
 Join our [slack workspace](https://slack.datahubproject.io) for discussions and important announcements. You can also find out more about our upcoming [town hall meetings](docs/townhalls.md) and view past recordings.
@@ -109,6 +113,7 @@ Here are the companies that have officially adopted DataHub. Please feel free to
 
 - [Adevinta](https://www.adevinta.com/)
 - [Banksalad](https://www.banksalad.com)
+- [Cabify](https://cabify.tech/)
 - [DefinedCrowd](http://www.definedcrowd.com)
 - [DFDS](https://www.dfds.com/)
 - [Expedia Group](http://expedia.com)
diff --git a/build.gradle b/build.gradle
index 80dcdbe7f72ab..e4b86ee03507d 100644
--- a/build.gradle
+++ b/build.gradle
@@ -77,7 +77,8 @@ project.ext.externalDependency = [
     'javaxInject' : 'javax.inject:javax.inject:1',
     'jerseyCore': 'org.glassfish.jersey.core:jersey-client:2.25.1',
     'jerseyGuava': 'org.glassfish.jersey.bundles.repackaged:jersey-guava:2.25.1',
-    'jettyJaas': 'org.eclipse.jetty:jetty-jaas:9.4.28.v20200408',
+    'jettyJaas': 'org.eclipse.jetty:jetty-jaas:9.4.32.v20200930',
+    'jgrapht': 'org.jgrapht:jgrapht-core:1.5.1',
     'jsonSimple': 'com.googlecode.json-simple:json-simple:1.1.1',
     'junitJupiterApi': "org.junit.jupiter:junit-jupiter-api:$junitJupiterVersion",
     'junitJupiterParams': "org.junit.jupiter:junit-jupiter-params:$junitJupiterVersion",
@@ -98,6 +99,7 @@ project.ext.externalDependency = [
     'neo4jJavaDriver': 'org.neo4j.driver:neo4j-java-driver:4.0.1',
     'opentelemetryApi': 'io.opentelemetry:opentelemetry-api:1.0.0',
     'opentelemetryAnnotations': 'io.opentelemetry:opentelemetry-extension-annotations:1.0.0',
+    'opentracingJdbc':'io.opentracing.contrib:opentracing-jdbc:0.2.15',
     'parseqTest': 'com.linkedin.parseq:parseq:3.0.7:test',
     'parquet': 'org.apache.parquet:parquet-avro:1.12.0',
     'picocli': 'info.picocli:picocli:4.5.0',
@@ -110,11 +112,12 @@ project.ext.externalDependency = [
     'pac4j': 'org.pac4j:pac4j-oidc:3.6.0',
     'playPac4j': 'org.pac4j:play-pac4j_2.11:7.0.1',
     'postgresql': 'org.postgresql:postgresql:42.3.3',
+    'protobuf': 'com.google.protobuf:protobuf-java:3.19.3',
     'reflections': 'org.reflections:reflections:0.9.9',
     'resilience4j': 'io.github.resilience4j:resilience4j-retry:1.7.1',
     'rythmEngine': 'org.rythmengine:rythm-engine:1.3.0',
     'servletApi': 'javax.servlet:javax.servlet-api:3.1.0',
-    'shiroCore': 'org.apache.shiro:shiro-core:1.7.1',
+    'shiroCore': 'org.apache.shiro:shiro-core:1.8.0',
     'sparkSql' : 'org.apache.spark:spark-sql_2.11:2.4.8',
     'sparkHive' : 'org.apache.spark:spark-hive_2.11:2.4.8',
     'springBeans': 'org.springframework:spring-beans:5.2.3.RELEASE',
@@ -140,6 +143,7 @@ project.ext.externalDependency = [
     'typesafeConfig':'com.typesafe:config:1.4.1',
     'wiremock':'com.github.tomakehurst:wiremock:2.10.0',
     'zookeeper': 'org.apache.zookeeper:zookeeper:3.4.14'
+    
 ]
 
 allprojects {
@@ -194,14 +198,27 @@ subprojects {
     }
   }
 
-  tasks.withType(JavaCompile).configureEach {
-    javaCompiler = javaToolchains.compilerFor {
-      languageVersion = JavaLanguageVersion.of(8)
+  if (project.name != 'datahub-protobuf') {
+    tasks.withType(JavaCompile).configureEach {
+      javaCompiler = javaToolchains.compilerFor {
+        languageVersion = JavaLanguageVersion.of(8)
+      }
     }
-  }
-  tasks.withType(Test).configureEach {
-    javaLauncher = javaToolchains.launcherFor {
-      languageVersion = JavaLanguageVersion.of(8)
+    tasks.withType(Test).configureEach {
+      javaLauncher = javaToolchains.launcherFor {
+        languageVersion = JavaLanguageVersion.of(8)
+      }
+    }
+  } else {
+    tasks.withType(JavaCompile).configureEach {
+      javaCompiler = javaToolchains.compilerFor {
+        languageVersion = JavaLanguageVersion.of(11)
+      }
+    }
+    tasks.withType(Test).configureEach {
+      javaLauncher = javaToolchains.launcherFor {
+        languageVersion = JavaLanguageVersion.of(11)
+      }
     }
   }
 
diff --git a/datahub-frontend/app/auth/sso/oidc/OidcCallbackLogic.java b/datahub-frontend/app/auth/sso/oidc/OidcCallbackLogic.java
index 22026681dde19..23776ff8f31e7 100644
--- a/datahub-frontend/app/auth/sso/oidc/OidcCallbackLogic.java
+++ b/datahub-frontend/app/auth/sso/oidc/OidcCallbackLogic.java
@@ -55,6 +55,7 @@
 import play.mvc.Result;
 import auth.sso.SsoManager;
 
+import static com.linkedin.metadata.Constants.*;
 import static play.mvc.Results.*;
 import static auth.AuthUtils.*;
 
@@ -125,16 +126,14 @@ private Result handleOidcCallback(
         if (oidcConfigs.isJitProvisioningEnabled()) {
           log.debug("Just-in-time provisioning is enabled. Beginning provisioning process...");
           CorpUserSnapshot extractedUser = extractUser(corpUserUrn, profile);
+          tryProvisionUser(extractedUser);
           if (oidcConfigs.isExtractGroupsEnabled()) {
             // Extract groups & provision them.
             List<CorpGroupSnapshot> extractedGroups = extractGroups(profile);
             tryProvisionGroups(extractedGroups);
-            if (extractedGroups.size() > 0) {
-              // Associate group with the user logging in.
-              extractedUser.getAspects().add(CorpUserAspect.create(createGroupMembership(extractedGroups)));
-            }
+            // Add users to groups on DataHub. Note that this clears existing group membership for a user if it already exists.
+            updateGroupMembership(corpUserUrn, createGroupMembership(extractedGroups));
           }
-          tryProvisionUser(extractedUser);
         } else if (oidcConfigs.isPreProvisioningRequired()) {
           // We should only allow logins for user accounts that have been pre-provisioned
           log.debug("Pre Provisioning is required. Beginning validation of extracted user...");
@@ -372,6 +371,21 @@ private void tryProvisionGroups(List<CorpGroupSnapshot> corpGroups) {
     }
   }
 
+  private void updateGroupMembership(Urn urn, GroupMembership groupMembership) {
+    log.debug(String.format("Updating group membership for user %s", urn));
+    final MetadataChangeProposal proposal = new MetadataChangeProposal();
+    proposal.setEntityUrn(urn);
+    proposal.setEntityType(CORP_USER_ENTITY_NAME);
+    proposal.setAspectName(GROUP_MEMBERSHIP_ASPECT_NAME);
+    proposal.setAspect(GenericAspectUtils.serializeAspect(groupMembership));
+    proposal.setChangeType(ChangeType.UPSERT);
+    try {
+      _entityClient.ingestProposal(proposal, _systemAuthentication);
+    } catch (RemoteInvocationException e) {
+      throw new RuntimeException(String.format("Failed to update group membership for user with urn %s", urn), e);
+    }
+  }
+
   private void verifyPreProvisionedUser(CorpuserUrn urn) {
     // Validate that the user exists in the system (there is more than just a key aspect for them, as of today).
     try {
diff --git a/datahub-frontend/app/auth/sso/oidc/OidcConfigs.java b/datahub-frontend/app/auth/sso/oidc/OidcConfigs.java
index e0b354e770ba8..1bf58572ad570 100644
--- a/datahub-frontend/app/auth/sso/oidc/OidcConfigs.java
+++ b/datahub-frontend/app/auth/sso/oidc/OidcConfigs.java
@@ -47,7 +47,7 @@ public class OidcConfigs extends SsoConfigs {
     private static final String DEFAULT_OIDC_CLIENT_AUTHENTICATION_METHOD = "client_secret_basic";
     private static final String DEFAULT_OIDC_JIT_PROVISIONING_ENABLED = "true";
     private static final String DEFAULT_OIDC_PRE_PROVISIONING_REQUIRED = "false";
-    private static final String DEFAULT_OIDC_EXTRACT_GROUPS_ENABLED = "true";
+    private static final String DEFAULT_OIDC_EXTRACT_GROUPS_ENABLED = "false"; // False since extraction of groups can overwrite existing group membership.
     private static final String DEFAULT_OIDC_GROUPS_CLAIM = "groups";
 
     private String clientId;
diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java
index 4e381201f9cb3..d33d9744a02c2 100644
--- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java
+++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java
@@ -12,6 +12,8 @@
 import com.linkedin.datahub.graphql.generated.AggregationMetadata;
 import com.linkedin.datahub.graphql.generated.Aspect;
 import com.linkedin.datahub.graphql.generated.Assertion;
+import com.linkedin.datahub.graphql.generated.AutoCompleteResultForEntity;
+import com.linkedin.datahub.graphql.generated.AutoCompleteResults;
 import com.linkedin.datahub.graphql.generated.BrowseResults;
 import com.linkedin.datahub.graphql.generated.Chart;
 import com.linkedin.datahub.graphql.generated.ChartInfo;
@@ -514,6 +516,10 @@ private void configureContainerResolvers(final RuntimeWiring.Builder builder) {
             .type("Container", typeWiring -> typeWiring
                 .dataFetcher("relationships", new EntityRelationshipsResultResolver(graphClient))
                 .dataFetcher("entities", new ContainerEntitiesResolver(entityClient))
+                .dataFetcher("domain", new LoadableTypeResolver<>(domainType, (env) -> {
+                    final Container container = env.getSource();
+                    return container.getDomain() != null ? container.getDomain().getUrn() : null;
+                }))
                 .dataFetcher("platform",
                     new LoadableTypeResolver<>(dataPlatformType,
                         (env) -> ((Container) env.getSource()).getPlatform().getUrn()))
@@ -730,7 +736,20 @@ private void configureGenericEntityResolvers(final RuntimeWiring.Builder builder
                         (env) -> ((ListDomainsResult) env.getSource()).getDomains().stream()
                             .map(Domain::getUrn)
                             .collect(Collectors.toList())))
+            )
+            .type("AutoCompleteResults", typeWiring -> typeWiring
+                .dataFetcher("entities",
+                    new EntityTypeBatchResolver(
+                        new ArrayList<>(entityTypes),
+                        (env) -> ((AutoCompleteResults) env.getSource()).getEntities()))
+            )
+            .type("AutoCompleteResultForEntity", typeWiring -> typeWiring
+                .dataFetcher("entities",
+                    new EntityTypeBatchResolver(
+                        new ArrayList<>(entityTypes),
+                        (env) -> ((AutoCompleteResultForEntity) env.getSource()).getEntities()))
             );
+        ;
     }
 
     /**
diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/analytics/resolver/GetChartsResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/analytics/resolver/GetChartsResolver.java
index 63592489b4959..c357af8bcfe69 100644
--- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/analytics/resolver/GetChartsResolver.java
+++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/analytics/resolver/GetChartsResolver.java
@@ -75,7 +75,7 @@ private List<AnalyticsChart> getProductAnalyticsCharts(Authentication authentica
 
     final List<NamedLine> wauTimeseries =
         _analyticsService.getTimeseriesChart(_analyticsService.getUsageIndexName(), twoMonthsDateRange, weeklyInterval,
-            Optional.empty(), ImmutableMap.of(), Optional.of("browserId"));
+            Optional.empty(), ImmutableMap.of(), Collections.emptyMap(), Optional.of("browserId"));
     charts.add(TimeSeriesChart.builder()
         .setTitle(wauTitle)
         .setDateRange(twoMonthsDateRange)
@@ -90,7 +90,8 @@ private List<AnalyticsChart> getProductAnalyticsCharts(Authentication authentica
 
     final List<NamedLine> searchesTimeseries =
         _analyticsService.getTimeseriesChart(_analyticsService.getUsageIndexName(), lastWeekDateRange, dailyInterval,
-            Optional.empty(), ImmutableMap.of("type", ImmutableList.of(searchEventType)), Optional.empty());
+            Optional.empty(), ImmutableMap.of("type", ImmutableList.of(searchEventType)), Collections.emptyMap(),
+            Optional.empty());
     charts.add(TimeSeriesChart.builder()
         .setTitle(searchesTitle)
         .setDateRange(lastWeekDateRange)
@@ -104,8 +105,8 @@ private List<AnalyticsChart> getProductAnalyticsCharts(Authentication authentica
 
     final List<Row> topSearchQueries =
         _analyticsService.getTopNTableChart(_analyticsService.getUsageIndexName(), Optional.of(lastWeekDateRange),
-            "query.keyword", ImmutableMap.of("type", ImmutableList.of(searchEventType)), Optional.empty(), 10,
-            AnalyticsUtil::buildCellWithSearchLandingPage);
+            "query.keyword", ImmutableMap.of("type", ImmutableList.of(searchEventType)), Collections.emptyMap(),
+            Optional.empty(), 10, AnalyticsUtil::buildCellWithSearchLandingPage);
     charts.add(TableChart.builder().setTitle(topSearchTitle).setColumns(columns).setRows(topSearchQueries).build());
 
     // Chart 4: Bar Graph Chart
@@ -113,7 +114,8 @@ private List<AnalyticsChart> getProductAnalyticsCharts(Authentication authentica
     final List<NamedBar> sectionViewsPerEntityType =
         _analyticsService.getBarChart(_analyticsService.getUsageIndexName(), Optional.of(lastWeekDateRange),
             ImmutableList.of("entityType.keyword", "section.keyword"),
-            ImmutableMap.of("type", ImmutableList.of("EntitySectionViewEvent")), Optional.empty(), true);
+            ImmutableMap.of("type", ImmutableList.of("EntitySectionViewEvent")), Collections.emptyMap(),
+            Optional.empty(), true);
     charts.add(BarChart.builder().setTitle(sectionViewsTitle).setBars(sectionViewsPerEntityType).build());
 
     // Chart 5: Bar Graph Chart
@@ -121,7 +123,8 @@ private List<AnalyticsChart> getProductAnalyticsCharts(Authentication authentica
     final List<NamedBar> eventsByEventType =
         _analyticsService.getBarChart(_analyticsService.getUsageIndexName(), Optional.of(lastWeekDateRange),
             ImmutableList.of("entityType.keyword", "actionType.keyword"),
-            ImmutableMap.of("type", ImmutableList.of("EntityActionEvent")), Optional.empty(), true);
+            ImmutableMap.of("type", ImmutableList.of("EntityActionEvent")), Collections.emptyMap(), Optional.empty(),
+            true);
     charts.add(BarChart.builder().setTitle(actionsByTypeTitle).setBars(eventsByEventType).build());
 
     // Chart 6: Table Chart
@@ -131,7 +134,7 @@ private List<AnalyticsChart> getProductAnalyticsCharts(Authentication authentica
     final List<Row> topViewedDatasets =
         _analyticsService.getTopNTableChart(_analyticsService.getUsageIndexName(), Optional.of(lastWeekDateRange),
             "entityUrn.keyword", ImmutableMap.of("type", ImmutableList.of("EntityViewEvent"), "entityType.keyword",
-                ImmutableList.of(EntityType.DATASET.name())), Optional.empty(), 10,
+                ImmutableList.of(EntityType.DATASET.name())), Collections.emptyMap(), Optional.empty(), 10,
             AnalyticsUtil::buildCellWithEntityLandingPage);
     AnalyticsUtil.hydrateDisplayNameForTable(_entityClient, topViewedDatasets, Constants.DATASET_ENTITY_NAME,
         ImmutableSet.of(Constants.DATASET_KEY_ASPECT_NAME), AnalyticsUtil::getDatasetName, authentication);
@@ -145,7 +148,8 @@ private List<AnalyticsChart> getGlobalMetadataAnalyticsCharts(Authentication aut
     // Chart 1: Entities per domain
     final List<NamedBar> entitiesPerDomain =
         _analyticsService.getBarChart(_analyticsService.getAllEntityIndexName(), Optional.empty(),
-            ImmutableList.of("domains.keyword", "platform.keyword"), Collections.emptyMap(), Optional.empty(), false);
+            ImmutableList.of("domains.keyword", "platform.keyword"), Collections.emptyMap(),
+            ImmutableMap.of("removed", ImmutableList.of("true")), Optional.empty(), false);
     AnalyticsUtil.hydrateDisplayNameForBars(_entityClient, entitiesPerDomain, Constants.DOMAIN_ENTITY_NAME,
         ImmutableSet.of(Constants.DOMAIN_PROPERTIES_ASPECT_NAME), AnalyticsUtil::getDomainName, authentication);
     AnalyticsUtil.hydrateDisplayNameForSegments(_entityClient, entitiesPerDomain, Constants.DATA_PLATFORM_ENTITY_NAME,
@@ -157,7 +161,8 @@ private List<AnalyticsChart> getGlobalMetadataAnalyticsCharts(Authentication aut
     // Chart 2: Entities per platform
     final List<NamedBar> entitiesPerPlatform =
         _analyticsService.getBarChart(_analyticsService.getAllEntityIndexName(), Optional.empty(),
-            ImmutableList.of("platform.keyword"), Collections.emptyMap(), Optional.empty(), false);
+            ImmutableList.of("platform.keyword"), Collections.emptyMap(),
+            ImmutableMap.of("removed", ImmutableList.of("true")), Optional.empty(), false);
     AnalyticsUtil.hydrateDisplayNameForBars(_entityClient, entitiesPerPlatform, Constants.DATA_PLATFORM_ENTITY_NAME,
         ImmutableSet.of(Constants.DATA_PLATFORM_INFO_ASPECT_NAME), AnalyticsUtil::getPlatformName, authentication);
     if (!entitiesPerPlatform.isEmpty()) {
@@ -167,7 +172,8 @@ private List<AnalyticsChart> getGlobalMetadataAnalyticsCharts(Authentication aut
     // Chart 3: Entities per term
     final List<NamedBar> entitiesPerTerm =
         _analyticsService.getBarChart(_analyticsService.getAllEntityIndexName(), Optional.empty(),
-            ImmutableList.of("glossaryTerms.keyword"), Collections.emptyMap(), Optional.empty(), false);
+            ImmutableList.of("glossaryTerms.keyword"), Collections.emptyMap(),
+            ImmutableMap.of("removed", ImmutableList.of("true")), Optional.empty(), false);
     AnalyticsUtil.hydrateDisplayNameForBars(_entityClient, entitiesPerTerm, Constants.GLOSSARY_TERM_ENTITY_NAME,
         ImmutableSet.of(Constants.GLOSSARY_TERM_KEY_ASPECT_NAME), AnalyticsUtil::getTermName, authentication);
     if (!entitiesPerTerm.isEmpty()) {
@@ -177,7 +183,8 @@ private List<AnalyticsChart> getGlobalMetadataAnalyticsCharts(Authentication aut
     // Chart 4: Entities per fabric type
     final List<NamedBar> entitiesPerEnv =
         _analyticsService.getBarChart(_analyticsService.getAllEntityIndexName(), Optional.empty(),
-            ImmutableList.of("origin.keyword"), Collections.emptyMap(), Optional.empty(), false);
+            ImmutableList.of("origin.keyword"), Collections.emptyMap(),
+            ImmutableMap.of("removed", ImmutableList.of("true")), Optional.empty(), false);
     if (entitiesPerEnv.size() > 1) {
       charts.add(BarChart.builder().setTitle("Entities per Environment").setBars(entitiesPerEnv).build());
     }
diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/analytics/service/AnalyticsService.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/analytics/service/AnalyticsService.java
index 17092234a8e91..48c017cdcaefd 100644
--- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/analytics/service/AnalyticsService.java
+++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/analytics/service/AnalyticsService.java
@@ -1,7 +1,6 @@
 package com.linkedin.datahub.graphql.analytics.service;
 
 import com.google.common.collect.ImmutableList;
-import com.google.common.collect.ImmutableMap;
 import com.linkedin.datahub.graphql.generated.BarSegment;
 import com.linkedin.datahub.graphql.generated.Cell;
 import com.linkedin.datahub.graphql.generated.DateInterval;
@@ -75,14 +74,14 @@ public String getUsageIndexName() {
 
   public List<NamedLine> getTimeseriesChart(String indexName, DateRange dateRange, DateInterval granularity,
       Optional<String> dimension, // Length 1 for now
-      Map<String, List<String>> filters, Optional<String> uniqueOn) {
+      Map<String, List<String>> filters, Map<String, List<String>> mustNotFilters, Optional<String> uniqueOn) {
 
     log.debug(
         String.format("Invoked getTimeseriesChart with indexName: %s, dateRange: %s, granularity: %s, dimension: %s,",
             indexName, dateRange, granularity, dimension) + String.format("filters: %s, uniqueOn: %s", filters,
             uniqueOn));
 
-    AggregationBuilder filteredAgg = getFilteredAggregation(filters, ImmutableMap.of(), Optional.of(dateRange));
+    AggregationBuilder filteredAgg = getFilteredAggregation(filters, mustNotFilters, Optional.of(dateRange));
 
     AggregationBuilder dateHistogram = AggregationBuilders.dateHistogram(DATE_HISTOGRAM)
         .field("timestamp")
@@ -128,13 +127,14 @@ private List<NumericDataPoint> extractPointsFromAggregations(Aggregations aggreg
 
   public List<NamedBar> getBarChart(String indexName, Optional<DateRange> dateRange, List<String> dimensions,
       // Length 1 or 2
-      Map<String, List<String>> filters, Optional<String> uniqueOn, boolean showMissing) {
+      Map<String, List<String>> filters, Map<String, List<String>> mustNotFilters, Optional<String> uniqueOn,
+      boolean showMissing) {
     log.debug(
         String.format("Invoked getBarChart with indexName: %s, dateRange: %s, dimensions: %s,", indexName, dateRange,
             dimensions) + String.format("filters: %s, uniqueOn: %s", filters, uniqueOn));
 
     assert (dimensions.size() == 1 || dimensions.size() == 2);
-    AggregationBuilder filteredAgg = getFilteredAggregation(filters, ImmutableMap.of(), dateRange);
+    AggregationBuilder filteredAgg = getFilteredAggregation(filters, mustNotFilters, dateRange);
 
     TermsAggregationBuilder termAgg = AggregationBuilders.terms(DIMENSION).field(dimensions.get(0));
     if (showMissing) {
@@ -142,8 +142,7 @@ public List<NamedBar> getBarChart(String indexName, Optional<DateRange> dateRang
     }
 
     if (dimensions.size() == 2) {
-      TermsAggregationBuilder secondTermAgg =
-          AggregationBuilders.terms(SECOND_DIMENSION).field(dimensions.get(1));
+      TermsAggregationBuilder secondTermAgg = AggregationBuilders.terms(SECOND_DIMENSION).field(dimensions.get(1));
       if (showMissing) {
         secondTermAgg.missing(NA);
       }
@@ -194,13 +193,13 @@ public Row buildRow(String groupByValue, Function<String, Cell> groupByValueToCe
   }
 
   public List<Row> getTopNTableChart(String indexName, Optional<DateRange> dateRange, String groupBy,
-      Map<String, List<String>> filters, Optional<String> uniqueOn, int maxRows,
-      Function<String, Cell> groupByValueToCell) {
+      Map<String, List<String>> filters, Map<String, List<String>> mustNotFilters, Optional<String> uniqueOn,
+      int maxRows, Function<String, Cell> groupByValueToCell) {
     log.debug(
         String.format("Invoked getTopNTableChart with indexName: %s, dateRange: %s, groupBy: %s", indexName, dateRange,
             groupBy) + String.format("filters: %s, uniqueOn: %s", filters, uniqueOn));
 
-    AggregationBuilder filteredAgg = getFilteredAggregation(filters, ImmutableMap.of(), dateRange);
+    AggregationBuilder filteredAgg = getFilteredAggregation(filters, mustNotFilters, dateRange);
 
     TermsAggregationBuilder termAgg = AggregationBuilders.terms(DIMENSION).field(groupBy).size(maxRows);
     if (uniqueOn.isPresent()) {
diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/AddOwnerResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/AddOwnerResolver.java
index 097edcfceb75c..2e5ebbd3c6f6c 100644
--- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/AddOwnerResolver.java
+++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/AddOwnerResolver.java
@@ -7,6 +7,7 @@
 import com.linkedin.datahub.graphql.exception.AuthorizationException;
 import com.linkedin.datahub.graphql.generated.AddOwnerInput;
 import com.linkedin.datahub.graphql.generated.OwnerEntityType;
+import com.linkedin.datahub.graphql.generated.OwnershipType;
 import com.linkedin.datahub.graphql.resolvers.mutate.util.OwnerUtils;
 import com.linkedin.metadata.entity.EntityService;
 import graphql.schema.DataFetcher;
@@ -30,6 +31,7 @@ public CompletableFuture<Boolean> get(DataFetchingEnvironment environment) throw
 
     Urn ownerUrn = Urn.createFromString(input.getOwnerUrn());
     OwnerEntityType ownerEntityType = input.getOwnerEntityType();
+    OwnershipType type = input.getType() == null ? OwnershipType.NONE : input.getType();
     Urn targetUrn = Urn.createFromString(input.getResourceUrn());
 
     if (!OwnerUtils.isAuthorizedToUpdateOwners(environment.getContext(), targetUrn)) {
@@ -50,6 +52,8 @@ public CompletableFuture<Boolean> get(DataFetchingEnvironment environment) throw
         Urn actor = CorpuserUrn.createFromString(((QueryContext) environment.getContext()).getActorUrn());
         OwnerUtils.addOwner(
             ownerUrn,
+            // Assumption Alert: Assumes that GraphQL ownership type === GMS ownership type
+            com.linkedin.common.OwnershipType.valueOf(type.name()),
             targetUrn,
             actor,
             _entityService
diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/OwnerUtils.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/OwnerUtils.java
index 134f36d920ea3..d7b9a3ec370e8 100644
--- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/OwnerUtils.java
+++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/OwnerUtils.java
@@ -18,6 +18,7 @@
 import com.linkedin.metadata.Constants;
 import com.linkedin.metadata.authorization.PoliciesConfig;
 import com.linkedin.metadata.entity.EntityService;
+import java.util.stream.Collectors;
 import javax.annotation.Nonnull;
 import lombok.extern.slf4j.Slf4j;
 
@@ -34,6 +35,7 @@ private OwnerUtils() { }
 
   public static void addOwner(
       Urn ownerUrn,
+      OwnershipType type,
       Urn resourceUrn,
       Urn actor,
       EntityService entityService
@@ -43,7 +45,7 @@ public static void addOwner(
         Constants.OWNERSHIP_ASPECT_NAME,
         entityService,
         new Ownership());
-    addOwner(ownershipAspect, ownerUrn);
+    addOwner(ownershipAspect, ownerUrn, type);
     persistAspect(resourceUrn, Constants.OWNERSHIP_ASPECT_NAME, ownershipAspect, actor, entityService);
   }
 
@@ -63,23 +65,22 @@ public static void removeOwner(
     persistAspect(resourceUrn, Constants.OWNERSHIP_ASPECT_NAME, ownershipAspect, actor, entityService);
   }
 
-  private static void addOwner(Ownership ownershipAspect, Urn ownerUrn) {
+  private static void addOwner(Ownership ownershipAspect, Urn ownerUrn, OwnershipType type) {
     if (!ownershipAspect.hasOwners()) {
       ownershipAspect.setOwners(new OwnerArray());
     }
 
-    OwnerArray ownerArray = ownershipAspect.getOwners();
-
-    // if owner exists, do not add it again
-    if (ownerArray.stream().anyMatch(association -> association.getOwner().equals(ownerUrn))) {
-      return;
-    }
+    final OwnerArray ownerArray = new OwnerArray(ownershipAspect.getOwners()
+        .stream()
+        .filter(owner ->  !owner.getOwner().equals(ownerUrn))
+        .collect(Collectors.toList()));
 
     Owner newOwner = new Owner();
-    newOwner.setType(OwnershipType.DATAOWNER);
+    newOwner.setType(type);
     newOwner.setSource(new OwnershipSource().setType(OwnershipSourceType.MANUAL));
     newOwner.setOwner(ownerUrn);
     ownerArray.add(newOwner);
+    ownershipAspect.setOwners(ownerArray);
   }
 
   private static void removeOwner(Ownership ownership, Urn ownerUrn) {
diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/search/AutoCompleteForMultipleResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/search/AutoCompleteForMultipleResolver.java
index 2c1292b03d635..911be3b9f2814 100644
--- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/search/AutoCompleteForMultipleResolver.java
+++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/search/AutoCompleteForMultipleResolver.java
@@ -18,6 +18,7 @@
 import org.slf4j.LoggerFactory;
 
 import static com.linkedin.datahub.graphql.resolvers.ResolverUtils.bindArgument;
+import static com.linkedin.datahub.graphql.resolvers.search.SearchUtils.*;
 import static org.apache.commons.lang3.StringUtils.isBlank;
 
 /**
@@ -27,11 +28,9 @@ public class AutoCompleteForMultipleResolver implements DataFetcher<CompletableF
 
     private static final Logger _logger = LoggerFactory.getLogger(AutoCompleteForMultipleResolver.class.getName());
 
-    private final List<SearchableEntityType<?>> _searchableEntities;
     private final Map<EntityType, SearchableEntityType<?>> _typeToEntity;
 
     public AutoCompleteForMultipleResolver(@Nonnull final List<SearchableEntityType<?>> searchableEntities) {
-        _searchableEntities = searchableEntities;
         _typeToEntity = searchableEntities.stream().collect(Collectors.toMap(
             SearchableEntityType::type,
             entity -> entity
@@ -51,10 +50,18 @@ public CompletableFuture<AutoCompleteMultipleResults> get(DataFetchingEnvironmen
 
         List<EntityType> types = input.getTypes();
         if (types != null && types.size() > 0) {
-            return AutocompleteUtils.batchGetAutocompleteResults(types.stream().map(type -> _typeToEntity.get(type)).collect(
-                Collectors.toList()), sanitizedQuery, input, environment);
+            return AutocompleteUtils.batchGetAutocompleteResults(
+                types.stream().map(_typeToEntity::get).collect(Collectors.toList()),
+                sanitizedQuery,
+                input,
+                environment);
         }
 
-        return AutocompleteUtils.batchGetAutocompleteResults(_searchableEntities, sanitizedQuery, input, environment);
+        // By default, autocomplete only against the set of Searchable Entity Types.
+        return AutocompleteUtils.batchGetAutocompleteResults(
+            AUTO_COMPLETE_ENTITY_TYPES.stream().map(_typeToEntity::get).collect(Collectors.toList()),
+            sanitizedQuery,
+            input,
+            environment);
     }
 }
diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/search/AutocompleteUtils.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/search/AutocompleteUtils.java
index 3fdeeadc2ed96..99e81cefd99b7 100644
--- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/search/AutocompleteUtils.java
+++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/search/AutocompleteUtils.java
@@ -7,7 +7,7 @@
 import com.linkedin.datahub.graphql.types.SearchableEntityType;
 import graphql.schema.DataFetchingEnvironment;
 import java.util.ArrayList;
-import java.util.Arrays;
+import java.util.Collections;
 import java.util.List;
 import java.util.concurrent.CompletableFuture;
 import java.util.stream.Collectors;
@@ -30,35 +30,35 @@ public static CompletableFuture<AutoCompleteMultipleResults> batchGetAutocomplet
   ) {
     final int limit = input.getLimit() != null ? input.getLimit() : DEFAULT_LIMIT;
 
-    final CompletableFuture<AutoCompleteResultForEntity>[] autoCompletesFuture = entities.stream().map(entity -> {
-      return CompletableFuture.supplyAsync(() -> {
-        try {
-          final AutoCompleteResults searchResult = entity.autoComplete(
-              sanitizedQuery,
-              input.getField(),
-              input.getFilters(),
-              limit,
-              environment.getContext()
-          );
-          final AutoCompleteResultForEntity autoCompleteResultForEntity =
-              new AutoCompleteResultForEntity(entity.type(), searchResult.getSuggestions());
-          return autoCompleteResultForEntity;
-        } catch (Exception e) {
-          _logger.error("Failed to execute autocomplete all: "
-              + String.format("field %s, query %s, filters: %s, limit: %s",
-              input.getField(),
-              input.getQuery(),
-              input.getFilters(),
-              input.getLimit()) + " "
-              + e.getMessage());
-          return new AutoCompleteResultForEntity(entity.type(), new ArrayList<>());
-        }
-      });
-    }).toArray(CompletableFuture[]::new);
-    return CompletableFuture.allOf(autoCompletesFuture)
+    final List<CompletableFuture<AutoCompleteResultForEntity>> autoCompletesFuture = entities.stream().map(entity -> CompletableFuture.supplyAsync(() -> {
+      try {
+        final AutoCompleteResults searchResult = entity.autoComplete(
+            sanitizedQuery,
+            input.getField(),
+            input.getFilters(),
+            limit,
+            environment.getContext()
+        );
+        return new AutoCompleteResultForEntity(
+            entity.type(),
+            searchResult.getSuggestions(),
+            searchResult.getEntities()
+        );
+      } catch (Exception e) {
+        _logger.error("Failed to execute autocomplete all: "
+            + String.format("field %s, query %s, filters: %s, limit: %s",
+            input.getField(),
+            input.getQuery(),
+            input.getFilters(),
+            input.getLimit()) + " "
+            + e.getMessage());
+        return new AutoCompleteResultForEntity(entity.type(), Collections.emptyList(), Collections.emptyList());
+      }
+    })).collect(Collectors.toList());
+    return CompletableFuture.allOf(autoCompletesFuture.toArray(new CompletableFuture[0]))
         .thenApplyAsync((res) -> {
           AutoCompleteMultipleResults result = new AutoCompleteMultipleResults(sanitizedQuery, new ArrayList<>());
-          result.setSuggestions(Arrays.stream(autoCompletesFuture)
+          result.setSuggestions(autoCompletesFuture.stream()
               .map(CompletableFuture::join)
               .filter(
                   autoCompleteResultForEntity ->
diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/search/SearchUtils.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/search/SearchUtils.java
index 398cea0c50f84..0153f5c39561c 100644
--- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/search/SearchUtils.java
+++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/search/SearchUtils.java
@@ -9,9 +9,41 @@ public class SearchUtils {
   private SearchUtils() {
   }
 
+  /**
+   * Entities that are searched by default in Search Across Entities
+   */
   public static final List<EntityType> SEARCHABLE_ENTITY_TYPES =
-      ImmutableList.of(EntityType.DATASET, EntityType.DASHBOARD, EntityType.CHART, EntityType.MLMODEL,
-          EntityType.MLMODEL_GROUP, EntityType.MLFEATURE_TABLE, EntityType.DATA_FLOW, EntityType.DATA_JOB,
-          EntityType.GLOSSARY_TERM, EntityType.TAG, EntityType.CORP_USER, EntityType.CORP_GROUP, EntityType.CONTAINER,
+      ImmutableList.of(
+          EntityType.DATASET,
+          EntityType.DASHBOARD,
+          EntityType.CHART,
+          EntityType.MLMODEL,
+          EntityType.MLMODEL_GROUP,
+          EntityType.MLFEATURE_TABLE,
+          EntityType.DATA_FLOW,
+          EntityType.DATA_JOB,
+          EntityType.GLOSSARY_TERM,
+          EntityType.TAG,
+          EntityType.CORP_USER,
+          EntityType.CORP_GROUP,
+          EntityType.CONTAINER,
           EntityType.DOMAIN);
+
+  /**
+   * Entities that are part of autocomplete by default in Auto Complete Across Entities
+   */
+  public static final List<EntityType> AUTO_COMPLETE_ENTITY_TYPES =
+      ImmutableList.of(
+          EntityType.DATASET,
+          EntityType.DASHBOARD,
+          EntityType.CHART,
+          EntityType.MLMODEL,
+          EntityType.MLMODEL_GROUP,
+          EntityType.MLFEATURE_TABLE,
+          EntityType.DATA_FLOW,
+          EntityType.DATA_JOB,
+          EntityType.GLOSSARY_TERM,
+          EntityType.TAG,
+          EntityType.CORP_USER,
+          EntityType.CORP_GROUP);
 }
diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/container/mappers/ContainerMapper.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/container/mappers/ContainerMapper.java
index 53beee6da93e1..1fc30a0234823 100644
--- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/container/mappers/ContainerMapper.java
+++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/container/mappers/ContainerMapper.java
@@ -27,6 +27,8 @@
 import com.linkedin.metadata.Constants;
 import javax.annotation.Nullable;
 
+import static com.linkedin.metadata.Constants.*;
+
 
 public class ContainerMapper {
 
@@ -43,8 +45,9 @@ public static Container map(final EntityResponse entityResponse) {
     if (envelopedPlatformInstance != null) {
       result.setPlatform(mapPlatform(new DataPlatformInstance(envelopedPlatformInstance.getValue().data())));
     } else {
-      // Containers must have DPI to be rendered.
-      return null;
+      final DataPlatform unknownPlatform = new DataPlatform();
+      unknownPlatform.setUrn(UNKNOWN_DATA_PLATFORM);
+      result.setPlatform(unknownPlatform);
     }
 
     final EnvelopedAspect envelopedContainerProperties = aspects.get(Constants.CONTAINER_PROPERTIES_ASPECT_NAME);
diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mappers/AutoCompleteResultsMapper.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mappers/AutoCompleteResultsMapper.java
index ea94173f73668..d575a81f4ae03 100644
--- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mappers/AutoCompleteResultsMapper.java
+++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mappers/AutoCompleteResultsMapper.java
@@ -1,8 +1,10 @@
 package com.linkedin.datahub.graphql.types.mappers;
 
 import com.linkedin.datahub.graphql.generated.AutoCompleteResults;
+import com.linkedin.datahub.graphql.types.common.mappers.UrnToEntityMapper;
 import com.linkedin.metadata.query.AutoCompleteResult;
 
+import java.util.stream.Collectors;
 import javax.annotation.Nonnull;
 
 
@@ -19,6 +21,8 @@ public AutoCompleteResults apply(@Nonnull final AutoCompleteResult input) {
         final AutoCompleteResults result = new AutoCompleteResults();
         result.setQuery(input.getQuery());
         result.setSuggestions(input.getSuggestions());
+        result.setEntities(input.getEntities().stream().map(entity -> UrnToEntityMapper.map(entity.getUrn())).collect(
+            Collectors.toList()));
         return result;
     }
 }
diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/tag/mappers/TagUpdateInputMapper.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/tag/mappers/TagUpdateInputMapper.java
index a09a2e4693018..264f4041f88cf 100644
--- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/tag/mappers/TagUpdateInputMapper.java
+++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/tag/mappers/TagUpdateInputMapper.java
@@ -47,7 +47,7 @@ public Collection<MetadataChangeProposal> apply(
     final Ownership ownership = new Ownership();
     final Owner owner = new Owner();
     owner.setOwner(actor);
-    owner.setType(OwnershipType.DATAOWNER);
+    owner.setType(OwnershipType.NONE);
     owner.setSource(new OwnershipSource().setType(OwnershipSourceType.SERVICE));
     ownership.setOwners(new OwnerArray(owner));
     ownership.setLastModified(auditStamp);
diff --git a/datahub-graphql-core/src/main/resources/entity.graphql b/datahub-graphql-core/src/main/resources/entity.graphql
index 7e38e71079c94..d2f43b5bf9256 100644
--- a/datahub-graphql-core/src/main/resources/entity.graphql
+++ b/datahub-graphql-core/src/main/resources/entity.graphql
@@ -2147,34 +2147,60 @@ Note that this field will soon become deprecated due to low usage
 """
 enum OwnershipType {
     """
-    A person or group that is in charge of developing the code
+    A person or group who is responsible for technical aspects of the asset.
+    """
+    TECHNICAL_OWNER
+
     """
-    DEVELOPER
+    A person or group who is responsible for logical, or business related, aspects of the asset.
+    """
+    BUSINESS_OWNER
 
     """
-    A person or group that is owning the data
+    A steward, expert, or delegate responsible for the asset.
     """
-    DATAOWNER
+    DATA_STEWARD
+
+    """
+    No specific type associated with the owner.
+    """
+    NONE
+
+    """
+    A person or group that owns the data.
+    Deprecated! This ownership type is no longer supported. Use TECHNICAL_OWNER instead.
+    """
+    DATAOWNER @deprecated
+
+    """
+    A person or group that is in charge of developing the code
+    Deprecated! This ownership type is no longer supported. Use TECHNICAL_OWNER instead.
+    """
+    DEVELOPER @deprecated
 
     """
     A person or a group that overseas the operation, eg a DBA or SRE
+    Deprecated! This ownership type is no longer supported. Use TECHNICAL_OWNER instead.
     """
-    DELEGATE
+    DELEGATE @deprecated
 
     """
     A person, group, or service that produces or generates the data
+    Deprecated! This ownership type is no longer supported. Use TECHNICAL_OWNER instead.
     """
-    PRODUCER
+    PRODUCER @deprecated
 
     """
-    A person, group, or service that consumes the data
+    A person or a group that has direct business interest
+    Deprecated! Use BUSINESS_OWNER instead.
     """
-    CONSUMER
+    STAKEHOLDER @deprecated
 
     """
-    A person or a group that has direct business interest
+    A person, group, or service that consumes the data
+    Deprecated! This ownership type is no longer supported.
     """
-    STAKEHOLDER
+    CONSUMER @deprecated
 }
 
 """
@@ -5105,6 +5131,11 @@ input AddOwnerInput {
     """
     ownerEntityType: OwnerEntityType!
 
+    """
+    The ownership type for the new owner. If none is provided, then a new NONE will be added.
+    """
+    type: OwnershipType
+
     """
     The urn of the resource or entity to attach or remove the owner from, for example a dataset urn
     """
diff --git a/datahub-graphql-core/src/main/resources/recommendation.graphql b/datahub-graphql-core/src/main/resources/recommendation.graphql
index 8b522a14f63c1..5f1340ba45254 100644
--- a/datahub-graphql-core/src/main/resources/recommendation.graphql
+++ b/datahub-graphql-core/src/main/resources/recommendation.graphql
@@ -42,7 +42,12 @@ enum ScenarioType {
   """
   Recommendations to show on an Entity Profile page
   """
-  ENTITY_PROFILE
+  ENTITY_PROFILE,
+
+  """
+  Recommendations to show on the search bar when clicked
+  """
+  SEARCH_BAR
 }
 
 """
diff --git a/datahub-graphql-core/src/main/resources/search.graphql b/datahub-graphql-core/src/main/resources/search.graphql
index 4d59332ae5748..b5deed8a35109 100644
--- a/datahub-graphql-core/src/main/resources/search.graphql
+++ b/datahub-graphql-core/src/main/resources/search.graphql
@@ -405,6 +405,11 @@ type AutoCompleteResultForEntity {
   The autocompletion results for specified entity type
   """
   suggestions: [String!]!
+
+  """
+  A list of entities to render in autocomplete
+  """
+  entities: [Entity!]!
 }
 
 """
@@ -435,6 +440,11 @@ type AutoCompleteResults {
 	The autocompletion results
 	"""
 	suggestions: [String!]!
+
+  """
+  A list of entities to render in autocomplete
+  """
+  entities: [Entity!]!
 }
 
 """
diff --git a/datahub-web-react/src/Mocks.tsx b/datahub-web-react/src/Mocks.tsx
index 2950928cc829b..6878271166d2d 100644
--- a/datahub-web-react/src/Mocks.tsx
+++ b/datahub-web-react/src/Mocks.tsx
@@ -282,6 +282,7 @@ const dataset2 = {
     downstream: null,
     health: null,
     assertions: null,
+    status: null,
 };
 
 export const dataset3 = {
@@ -484,6 +485,7 @@ export const dataset3 = {
     relationships: null,
     health: null,
     assertions: null,
+    status: null,
 } as Dataset;
 
 export const dataset4 = {
@@ -1082,6 +1084,7 @@ export const dataJob1 = {
         ],
     },
     domain: null,
+    status: null,
 } as DataJob;
 
 export const dataJob2 = {
@@ -1204,6 +1207,7 @@ export const dataJob3 = {
     domain: null,
     upstream: null,
     downstream: null,
+    status: null,
 } as DataJob;
 
 export const mlModel = {
@@ -1278,6 +1282,7 @@ export const mlModel = {
     outgoing: null,
     upstream: null,
     downstream: null,
+    status: null,
 } as MlModel;
 
 export const mlModelGroup = {
@@ -1340,6 +1345,7 @@ export const mlModelGroup = {
     outgoing: null,
     upstream: null,
     downstream: null,
+    status: null,
 } as MlModelGroup;
 
 export const recommendationModules = [
@@ -1594,7 +1600,8 @@ export const mocks = [
                     suggestions: [
                         {
                             type: EntityType.Dataset,
-                            suggestions: ['The Great Test Dataset', 'Some other test'],
+                            suggestions: ['The Great Test Dataset', 'Some Other Dataset'],
+                            entities: [dataset1, dataset2],
                         },
                     ],
                 },
@@ -1618,7 +1625,8 @@ export const mocks = [
                     suggestions: [
                         {
                             type: EntityType.Dataset,
-                            suggestions: ['The Great Test Dataset', 'Some other test'],
+                            suggestions: ['The Great Test Dataset', 'Some Other Dataset'],
+                            entities: [dataset1, dataset2],
                         },
                     ],
                 },
@@ -1640,6 +1648,7 @@ export const mocks = [
                 autoComplete: {
                     query: 'j',
                     suggestions: ['jjoyce'],
+                    entities: [user1],
                 },
             },
         },
@@ -2967,4 +2976,69 @@ export const mocks = [
             },
         },
     },
+    {
+        request: {
+            query: GetSearchResultsForMultipleDocument,
+            variables: {
+                input: {
+                    types: [],
+                    query: '*',
+                    start: 0,
+                    count: 20,
+                    filters: [],
+                },
+            },
+        },
+        result: {
+            data: {
+                __typename: 'Query',
+                searchAcrossEntities: {
+                    __typename: 'SearchResults',
+                    start: 0,
+                    count: 1,
+                    total: 1,
+                    searchResults: [
+                        {
+                            entity: {
+                                __typename: 'Dataset',
+                                ...dataset3,
+                            },
+                            matchedFields: [],
+                            insights: [],
+                        },
+                        {
+                            entity: {
+                                __typename: 'Dataset',
+                                ...dataset4,
+                            },
+                            matchedFields: [],
+                            insights: [],
+                        },
+                    ],
+                    facets: [
+                        {
+                            field: 'origin',
+                            displayName: 'origin',
+                            aggregations: [
+                                {
+                                    value: 'PROD',
+                                    count: 3,
+                                    entity: null,
+                                },
+                            ],
+                        },
+                        {
+                            field: 'platform',
+                            displayName: 'platform',
+                            aggregations: [
+                                { value: 'hdfs', count: 1, entity: null },
+                                { value: 'mysql', count: 1, entity: null },
+                                { value: 'kafka', count: 1, entity: null },
+                            ],
+                        },
+                    ],
+                },
+            } as GetSearchResultsForMultipleQuery,
+        },
+    },
 ];
diff --git a/datahub-web-react/src/app/browse/BrowseResultsPage.tsx b/datahub-web-react/src/app/browse/BrowseResultsPage.tsx
index 25e85f67b1825..485667b8f2634 100644
--- a/datahub-web-react/src/app/browse/BrowseResultsPage.tsx
+++ b/datahub-web-react/src/app/browse/BrowseResultsPage.tsx
@@ -61,7 +61,7 @@ export const BrowseResultsPage = () => {
                 <LegacyBrowsePath type={entityType} path={path} isBrowsable />
             </Affix>
             {loading && <Message type="loading" content="Loading..." style={{ marginTop: '10%' }} />}
-            {data && data.browse && (
+            {data && data.browse && !loading && (
                 <BrowseResults
                     type={entityType}
                     rootPath={rootPath}
diff --git a/datahub-web-react/src/app/entity/EntityRegistry.tsx b/datahub-web-react/src/app/entity/EntityRegistry.tsx
index 894265bbcdc87..331decf93c3ba 100644
--- a/datahub-web-react/src/app/entity/EntityRegistry.tsx
+++ b/datahub-web-react/src/app/entity/EntityRegistry.tsx
@@ -116,7 +116,13 @@ export default class EntityRegistry {
 
     getLineageVizConfig<T>(type: EntityType, data: T): FetchedEntity | undefined {
         const entity = validatedGet(type, this.entityTypeToEntity);
-        return entity.getLineageVizConfig?.(data) || undefined;
+        const genericEntityProperties = this.getGenericEntityProperties(type, data);
+        return (
+            ({
+                ...entity.getLineageVizConfig?.(data),
+                status: genericEntityProperties?.status,
+            } as FetchedEntity) || undefined
+        );
     }
 
     getDisplayName<T>(type: EntityType, data: T): string {
diff --git a/datahub-web-react/src/app/entity/dashboard/DashboardEntity.tsx b/datahub-web-react/src/app/entity/dashboard/DashboardEntity.tsx
index 9aadab5dd26e8..fb1135f6b3678 100644
--- a/datahub-web-react/src/app/entity/dashboard/DashboardEntity.tsx
+++ b/datahub-web-react/src/app/entity/dashboard/DashboardEntity.tsx
@@ -5,7 +5,7 @@ import {
     useGetDashboardQuery,
     useUpdateDashboardMutation,
 } from '../../../graphql/dashboard.generated';
-import { Dashboard, EntityType, PlatformType, SearchResult } from '../../../types.generated';
+import { Dashboard, EntityType, OwnershipType, PlatformType, SearchResult } from '../../../types.generated';
 import { EntityAndType } from '../../lineage/types';
 import { Entity, IconStyleType, PreviewType } from '../Entity';
 import { EntityProfile } from '../shared/containers/profile/EntityProfile';
@@ -104,6 +104,9 @@ export class DashboardEntity implements Entity<Dashboard> {
                 },
                 {
                     component: SidebarOwnerSection,
+                    properties: {
+                        defaultOwnerType: OwnershipType.TechnicalOwner,
+                    },
                 },
                 {
                     component: SidebarDomainSection,
diff --git a/datahub-web-react/src/app/entity/dataFlow/DataFlowEntity.tsx b/datahub-web-react/src/app/entity/dataFlow/DataFlowEntity.tsx
index e59856ea00abb..28832e25cf274 100644
--- a/datahub-web-react/src/app/entity/dataFlow/DataFlowEntity.tsx
+++ b/datahub-web-react/src/app/entity/dataFlow/DataFlowEntity.tsx
@@ -1,6 +1,6 @@
 import * as React from 'react';
 import { ShareAltOutlined } from '@ant-design/icons';
-import { DataFlow, EntityType, PlatformType, SearchResult } from '../../../types.generated';
+import { DataFlow, EntityType, OwnershipType, PlatformType, SearchResult } from '../../../types.generated';
 import { Preview } from './preview/Preview';
 import { Entity, IconStyleType, PreviewType } from '../Entity';
 import { EntityProfile } from '../shared/containers/profile/EntityProfile';
@@ -89,6 +89,9 @@ export class DataFlowEntity implements Entity<DataFlow> {
                 },
                 {
                     component: SidebarOwnerSection,
+                    properties: {
+                        defaultOwnerType: OwnershipType.TechnicalOwner,
+                    },
                 },
                 {
                     component: SidebarDomainSection,
diff --git a/datahub-web-react/src/app/entity/dataJob/DataJobEntity.tsx b/datahub-web-react/src/app/entity/dataJob/DataJobEntity.tsx
index 3b261446af898..2d66322c807cd 100644
--- a/datahub-web-react/src/app/entity/dataJob/DataJobEntity.tsx
+++ b/datahub-web-react/src/app/entity/dataJob/DataJobEntity.tsx
@@ -1,6 +1,6 @@
 import * as React from 'react';
 import { ConsoleSqlOutlined } from '@ant-design/icons';
-import { DataJob, EntityType, PlatformType, SearchResult } from '../../../types.generated';
+import { DataJob, EntityType, OwnershipType, PlatformType, SearchResult } from '../../../types.generated';
 import { Preview } from './preview/Preview';
 import { Entity, IconStyleType, PreviewType } from '../Entity';
 import { EntityProfile } from '../shared/containers/profile/EntityProfile';
@@ -101,6 +101,9 @@ export class DataJobEntity implements Entity<DataJob> {
                 },
                 {
                     component: SidebarOwnerSection,
+                    properties: {
+                        defaultOwnerType: OwnershipType.TechnicalOwner,
+                    },
                 },
                 {
                     component: SidebarDomainSection,
diff --git a/datahub-web-react/src/app/entity/dataset/DatasetEntity.tsx b/datahub-web-react/src/app/entity/dataset/DatasetEntity.tsx
index f952c15b891f7..b0b67f3ca6247 100644
--- a/datahub-web-react/src/app/entity/dataset/DatasetEntity.tsx
+++ b/datahub-web-react/src/app/entity/dataset/DatasetEntity.tsx
@@ -1,7 +1,7 @@
 import * as React from 'react';
 import { DatabaseFilled, DatabaseOutlined } from '@ant-design/icons';
 import { Typography } from 'antd';
-import { Dataset, EntityType, SearchResult } from '../../../types.generated';
+import { Dataset, EntityType, OwnershipType, SearchResult } from '../../../types.generated';
 import { Entity, IconStyleType, PreviewType } from '../Entity';
 import { Preview } from './preview/Preview';
 import { FIELDS_TO_HIGHLIGHT } from './search/highlights';
@@ -111,9 +111,13 @@ export class DatasetEntity implements Entity<Dataset> {
                     component: LineageTab,
                     display: {
                         visible: (_, _1) => true,
-                        enabled: (_, dataset: GetDatasetQuery) =>
-                            (dataset?.dataset?.upstream?.count || 0) > 0 ||
-                            (dataset?.dataset?.downstream?.count || 0) > 0,
+                        enabled: (_, dataset: GetDatasetQuery) => {
+                            console.log(dataset?.dataset?.upstream, dataset?.dataset?.downstream);
+                            return (
+                                (dataset?.dataset?.upstream?.total || 0) > 0 ||
+                                (dataset?.dataset?.downstream?.total || 0) > 0
+                            );
+                        },
                     },
                 },
                 {
@@ -175,6 +179,9 @@ export class DatasetEntity implements Entity<Dataset> {
                 },
                 {
                     component: SidebarOwnerSection,
+                    properties: {
+                        defaultOwnerType: OwnershipType.TechnicalOwner,
+                    },
                 },
                 {
                     component: SidebarDomainSection,
diff --git a/datahub-web-react/src/app/entity/domain/DomainEntity.tsx b/datahub-web-react/src/app/entity/domain/DomainEntity.tsx
index 92af08fb73cc7..edf6a4367a32d 100644
--- a/datahub-web-react/src/app/entity/domain/DomainEntity.tsx
+++ b/datahub-web-react/src/app/entity/domain/DomainEntity.tsx
@@ -79,6 +79,9 @@ export class DomainEntity implements Entity<Domain> {
                 },
                 {
                     component: SidebarOwnerSection,
+                    properties: {
+                        hideOwnerType: true,
+                    },
                 },
             ]}
         />
diff --git a/datahub-web-react/src/app/entity/group/GroupOwnerSideBarSection.tsx b/datahub-web-react/src/app/entity/group/GroupOwnerSideBarSection.tsx
index 6fad53f42fc81..d637bcc457f7f 100644
--- a/datahub-web-react/src/app/entity/group/GroupOwnerSideBarSection.tsx
+++ b/datahub-web-react/src/app/entity/group/GroupOwnerSideBarSection.tsx
@@ -52,6 +52,7 @@ export default function GroupOwnerSideBarSection({ urn, ownership, refetch }: Pr
             </SectionWrapper>
             <AddOwnerModal
                 urn={urn}
+                hideOwnerType
                 type={EntityType.CorpGroup}
                 visible={showAddModal}
                 refetch={refetch}
diff --git a/datahub-web-react/src/app/entity/shared/components/styled/ExpandedOwner.tsx b/datahub-web-react/src/app/entity/shared/components/styled/ExpandedOwner.tsx
index 446e734f4eac9..f79da080a449d 100644
--- a/datahub-web-react/src/app/entity/shared/components/styled/ExpandedOwner.tsx
+++ b/datahub-web-react/src/app/entity/shared/components/styled/ExpandedOwner.tsx
@@ -1,4 +1,4 @@
-import { message, Modal, Tag, Tooltip } from 'antd';
+import { message, Modal, Popover, Tag, Typography } from 'antd';
 import React from 'react';
 import { Link } from 'react-router-dom';
 import styled from 'styled-components';
@@ -9,10 +9,12 @@ import { CustomAvatar } from '../../../../shared/avatar';
 import { useEntityRegistry } from '../../../../useEntityRegistry';
 import analytics, { EventType, EntityActionType } from '../../../../analytics';
 import { useEntityData } from '../../EntityContext';
+import { getDescriptionFromType, getNameFromType } from '../../containers/profile/sidebar/Ownership/ownershipUtils';
 
 type Props = {
     entityUrn: string;
     owner: Owner;
+    hidePopOver?: boolean | undefined;
     refetch?: () => Promise<any>;
 };
 
@@ -24,7 +26,7 @@ const OwnerTag = styled(Tag)`
     align-items: center;
 `;
 
-export const ExpandedOwner = ({ entityUrn, owner, refetch }: Props) => {
+export const ExpandedOwner = ({ entityUrn, owner, hidePopOver, refetch }: Props) => {
     const entityRegistry = useEntityRegistry();
     const { entityType } = useEntityData();
     const [removeOwnerMutation] = useRemoveOwnerMutation();
@@ -85,9 +87,18 @@ export const ExpandedOwner = ({ entityUrn, owner, refetch }: Props) => {
         <OwnerTag onClose={onClose} closable>
             <Link to={`/${entityRegistry.getPathName(owner.owner.type)}/${owner.owner.urn}`}>
                 <CustomAvatar name={name} photoUrl={pictureLink} useDefaultAvatar={false} />
-                <Tooltip placement="top" title={owner.type}>
-                    {name}
-                </Tooltip>
+                {(hidePopOver && <>{name}</>) || (
+                    <Popover
+                        overlayStyle={{ maxWidth: 200 }}
+                        placement="left"
+                        title={<Typography.Text strong>{getNameFromType(owner.type)}</Typography.Text>}
+                        content={
+                            <Typography.Text type="secondary">{getDescriptionFromType(owner.type)}</Typography.Text>
+                        }
+                    >
+                        {name}
+                    </Popover>
+                )}
             </Link>
         </OwnerTag>
     );
diff --git a/datahub-web-react/src/app/entity/shared/containers/profile/EntityProfile.tsx b/datahub-web-react/src/app/entity/shared/containers/profile/EntityProfile.tsx
index 7fd0c888b8d52..91d1498b883da 100644
--- a/datahub-web-react/src/app/entity/shared/containers/profile/EntityProfile.tsx
+++ b/datahub-web-react/src/app/entity/shared/containers/profile/EntityProfile.tsx
@@ -231,6 +231,12 @@ export const EntityProfile = <T, U>({
         >
             <>
                 {showBrowseBar && <EntityProfileNavBar urn={urn} entityType={entityType} />}
+                {entityData?.status?.removed === true && (
+                    <Alert
+                        message="This entity has been soft deleted and is not discoverable via search or lineage graph"
+                        banner
+                    />
+                )}
                 {loading && <Message type="loading" content="Loading..." style={{ marginTop: '10%' }} />}
                 {!loading && error && (
                     <Alert type="error" message={error?.message || `Entity failed to load for urn ${urn}`} />
diff --git a/datahub-web-react/src/app/entity/shared/containers/profile/sidebar/Ownership/AddOwnerModal.tsx b/datahub-web-react/src/app/entity/shared/containers/profile/sidebar/Ownership/AddOwnerModal.tsx
index cae970b37b05a..b0003a642a6d5 100644
--- a/datahub-web-react/src/app/entity/shared/containers/profile/sidebar/Ownership/AddOwnerModal.tsx
+++ b/datahub-web-react/src/app/entity/shared/containers/profile/sidebar/Ownership/AddOwnerModal.tsx
@@ -4,11 +4,17 @@ import styled from 'styled-components';
 import { Link } from 'react-router-dom';
 import { useAddOwnerMutation } from '../../../../../../../graphql/mutations.generated';
 import { useGetSearchResultsLazyQuery } from '../../../../../../../graphql/search.generated';
-import { CorpUser, EntityType, OwnerEntityType, SearchResult } from '../../../../../../../types.generated';
+import {
+    CorpUser,
+    EntityType,
+    OwnerEntityType,
+    OwnershipType,
+    SearchResult,
+} from '../../../../../../../types.generated';
 import { useEntityRegistry } from '../../../../../../useEntityRegistry';
 import { CustomAvatar } from '../../../../../../shared/avatar';
 import analytics, { EventType, EntityActionType } from '../../../../../../analytics';
-import { useEnterKeyListener } from '../../../../../../shared/useEnterKeyListener';
+import { OWNERSHIP_DISPLAY_TYPES } from './ownershipUtils';
 
 const SearchResultContainer = styled.div`
     display: flex;
@@ -31,6 +37,8 @@ type Props = {
     urn: string;
     type: EntityType;
     visible: boolean;
+    defaultOwnerType?: OwnershipType;
+    hideOwnerType?: boolean | undefined;
     onClose: () => void;
     refetch?: () => Promise<any>;
 };
@@ -41,9 +49,10 @@ type SelectedActor = {
     urn: string;
 };
 
-export const AddOwnerModal = ({ urn, type, visible, onClose, refetch }: Props) => {
+export const AddOwnerModal = ({ urn, type, visible, hideOwnerType, defaultOwnerType, onClose, refetch }: Props) => {
     const entityRegistry = useEntityRegistry();
     const [selectedActor, setSelectedActor] = useState<SelectedActor | undefined>(undefined);
+    const [selectedOwnerType, setSelectedOwnerType] = useState<OwnershipType>(defaultOwnerType || OwnershipType.None);
     const [userSearch, { data: userSearchData }] = useGetSearchResultsLazyQuery();
     const [groupSearch, { data: groupSearchData }] = useGetSearchResultsLazyQuery();
     const [addOwnerMutation] = useAddOwnerMutation();
@@ -66,6 +75,7 @@ export const AddOwnerModal = ({ urn, type, visible, onClose, refetch }: Props) =
                 variables: {
                     input: {
                         ownerUrn: selectedActor.urn,
+                        type: selectedOwnerType,
                         resourceUrn: urn,
                         ownerEntityType,
                     },
@@ -112,6 +122,11 @@ export const AddOwnerModal = ({ urn, type, visible, onClose, refetch }: Props) =
         setSelectedActor(undefined);
     };
 
+    // When a user search result is selected, set the urn as the selected urn.
+    const onSelectOwnerType = (newType: OwnershipType) => {
+        setSelectedOwnerType(newType);
+    };
+
     // Invokes the search API as the user types
     const handleSearch = (entityType: EntityType, text: string, searchQuery: any) => {
         if (text.length > 2) {
@@ -172,16 +187,20 @@ export const AddOwnerModal = ({ urn, type, visible, onClose, refetch }: Props) =
     };
 
     const selectValue = (selectedActor && [selectedActor.displayName]) || [];
+    const ownershipTypes = OWNERSHIP_DISPLAY_TYPES;
 
     // Handle the Enter press
-    useEnterKeyListener({
-        querySelectorToExecuteClick: '#addOwnerButton',
-    });
+    // TODO: Allow user to be selected prior to executed the save.
+    // useEnterKeyListener({
+    //    querySelectorToExecuteClick: selectedActor && '#addOwnerButton',
+    // });
+
     return (
         <Modal
             title="Add Owner"
             visible={visible}
             onCancel={onClose}
+            keyboard
             footer={
                 <>
                     <Button onClick={onClose} type="text">
@@ -193,8 +212,9 @@ export const AddOwnerModal = ({ urn, type, visible, onClose, refetch }: Props) =
                 </>
             }
         >
-            <Form component={false}>
-                <Form.Item>
+            <Form layout="vertical" colon={false}>
+                <Form.Item name="owner" label={<Typography.Text strong>Owner</Typography.Text>}>
+                    <Typography.Paragraph>Find a user or group</Typography.Paragraph>
                     <Select
                         autoFocus
                         filterOption={false}
@@ -212,6 +232,23 @@ export const AddOwnerModal = ({ urn, type, visible, onClose, refetch }: Props) =
                         ))}
                     </Select>
                 </Form.Item>
+                {!hideOwnerType && (
+                    <Form.Item name="type" label={<Typography.Text strong>Type</Typography.Text>}>
+                        <Typography.Paragraph>Choose an owner type</Typography.Paragraph>
+                        <Select value={selectedOwnerType} onChange={onSelectOwnerType}>
+                            {ownershipTypes.map((ownerType) => (
+                                <Select.Option value={ownerType.type}>
+                                    <Typography.Text>{ownerType.name}</Typography.Text>
+                                    <div>
+                                        <Typography.Paragraph style={{ wordBreak: 'break-all' }} type="secondary">
+                                            {ownerType.description}
+                                        </Typography.Paragraph>
+                                    </div>
+                                </Select.Option>
+                            ))}
+                        </Select>
+                    </Form.Item>
+                )}
             </Form>
         </Modal>
     );
diff --git a/datahub-web-react/src/app/entity/shared/containers/profile/sidebar/Ownership/SidebarOwnerSection.tsx b/datahub-web-react/src/app/entity/shared/containers/profile/sidebar/Ownership/SidebarOwnerSection.tsx
index 20a45bd4de5ef..646c5017504c7 100644
--- a/datahub-web-react/src/app/entity/shared/containers/profile/sidebar/Ownership/SidebarOwnerSection.tsx
+++ b/datahub-web-react/src/app/entity/shared/containers/profile/sidebar/Ownership/SidebarOwnerSection.tsx
@@ -7,7 +7,7 @@ import { useEntityData, useRefetch } from '../../../../EntityContext';
 import { SidebarHeader } from '../SidebarHeader';
 import { AddOwnerModal } from './AddOwnerModal';
 
-export const SidebarOwnerSection = () => {
+export const SidebarOwnerSection = ({ properties }: { properties?: any }) => {
     const { urn, entityType, entityData } = useEntityData();
     const refetch = useRefetch();
     const [showAddModal, setShowAddModal] = useState(false);
@@ -32,6 +32,8 @@ export const SidebarOwnerSection = () => {
             </div>
             <AddOwnerModal
                 urn={urn}
+                defaultOwnerType={properties?.defaultOwnerType}
+                hideOwnerType={properties?.hideOwnerType || false}
                 type={entityType}
                 visible={showAddModal}
                 refetch={refetch}
diff --git a/datahub-web-react/src/app/entity/shared/containers/profile/sidebar/Ownership/ownershipUtils.ts b/datahub-web-react/src/app/entity/shared/containers/profile/sidebar/Ownership/ownershipUtils.ts
new file mode 100644
index 0000000000000..c5a250c760841
--- /dev/null
+++ b/datahub-web-react/src/app/entity/shared/containers/profile/sidebar/Ownership/ownershipUtils.ts
@@ -0,0 +1,41 @@
+import { OwnershipType } from '../../../../../../../types.generated';
+
+/**
+ * A mapping from OwnershipType to it's display name & description. In the future,
+ * we intend to make this configurable.
+ */
+export const OWNERSHIP_DISPLAY_TYPES = [
+    {
+        type: OwnershipType.TechnicalOwner,
+        name: 'Technical Owner',
+        description: 'Involved in the production, maintenance, or distribution of the asset.',
+    },
+    {
+        type: OwnershipType.BusinessOwner,
+        name: 'Business Owner',
+        description: 'Principle stakeholders or domain experts associated with the asset.',
+    },
+    {
+        type: OwnershipType.DataSteward,
+        name: 'Data Steward',
+        description: 'Involved in governance of the asset.',
+    },
+    {
+        type: OwnershipType.None,
+        name: 'None',
+        description: 'No ownership type specified.',
+    },
+];
+
+const ownershipTypeToDetails = new Map();
+OWNERSHIP_DISPLAY_TYPES.forEach((ownershipDetails) => {
+    ownershipTypeToDetails.set(ownershipDetails.type, ownershipDetails);
+});
+
+export const getNameFromType = (type: OwnershipType) => {
+    return ownershipTypeToDetails.get(type)?.name || type;
+};
+
+export const getDescriptionFromType = (type: OwnershipType) => {
+    return ownershipTypeToDetails.get(type)?.description || 'No description';
+};
diff --git a/datahub-web-react/src/app/entity/shared/types.ts b/datahub-web-react/src/app/entity/shared/types.ts
index 36e6909b45795..a51dfd3eeaaa3 100644
--- a/datahub-web-react/src/app/entity/shared/types.ts
+++ b/datahub-web-react/src/app/entity/shared/types.ts
@@ -23,6 +23,7 @@ import {
     SubTypes,
     Container,
     Health,
+    Status,
 } from '../../../types.generated';
 import { FetchedEntity } from '../../lineage/types';
 
@@ -71,6 +72,7 @@ export type GenericEntityProperties = {
     entityCount?: number;
     container?: Maybe<Container>;
     health?: Maybe<Health>;
+    status?: Maybe<Status>;
 };
 
 export type GenericEntityUpdate = {
diff --git a/datahub-web-react/src/app/entity/user/User.tsx b/datahub-web-react/src/app/entity/user/User.tsx
index ea488371e66f0..931e47c3c5ca2 100644
--- a/datahub-web-react/src/app/entity/user/User.tsx
+++ b/datahub-web-react/src/app/entity/user/User.tsx
@@ -41,9 +41,9 @@ export class UserEntity implements Entity<CorpUser> {
 
     getPathName: () => string = () => 'user';
 
-    getEntityName = () => 'User';
+    getEntityName = () => 'Person';
 
-    getCollectionName: () => string = () => 'Users';
+    getCollectionName: () => string = () => 'People';
 
     renderProfile: (urn: string) => JSX.Element = (_) => <UserProfile />;
 
diff --git a/datahub-web-react/src/app/home/HomePageHeader.tsx b/datahub-web-react/src/app/home/HomePageHeader.tsx
index ba54b92e3dfc5..7209facad1dcf 100644
--- a/datahub-web-react/src/app/home/HomePageHeader.tsx
+++ b/datahub-web-react/src/app/home/HomePageHeader.tsx
@@ -8,8 +8,10 @@ import { useGetAuthenticatedUser } from '../useGetAuthenticatedUser';
 import { useEntityRegistry } from '../useEntityRegistry';
 import { navigateToSearchUrl } from '../search/utils/navigateToSearchUrl';
 import { SearchBar } from '../search/SearchBar';
-import { GetSearchResultsQuery, useGetAutoCompleteMultipleResultsLazyQuery } from '../../graphql/search.generated';
-import { useGetAllEntitySearchResults } from '../../utils/customGraphQL/useGetAllEntitySearchResults';
+import {
+    useGetAutoCompleteMultipleResultsLazyQuery,
+    useGetSearchResultsForMultipleQuery,
+} from '../../graphql/search.generated';
 import { EntityType } from '../../types.generated';
 import analytics, { EventType } from '../analytics';
 import { AdminHeaderLinks } from '../shared/admin/AdminHeaderLinks';
@@ -96,28 +98,6 @@ const SearchBarContainer = styled.div`
     text-align: center;
 `;
 
-function getSuggestionFieldsFromResult(result: GetSearchResultsQuery | undefined): string[] {
-    return (
-        (result?.search?.searchResults
-            ?.map((searchResult) => searchResult.entity)
-            ?.map((entity) => {
-                switch (entity.__typename) {
-                    case 'Dataset':
-                        return entity.name.split('.').slice(-1)[0];
-                    case 'CorpUser':
-                        return entity.username;
-                    case 'Chart':
-                        return entity.properties?.name;
-                    case 'Dashboard':
-                        return entity.properties?.name;
-                    default:
-                        return undefined;
-                }
-            })
-            .filter(Boolean) as string[]) || []
-    );
-}
-
 function truncate(input, length) {
     if (input.length > length) {
         return `${input.substring(0, length)}...`;
@@ -166,38 +146,32 @@ export const HomePageHeader = () => {
         }
     };
 
-    // fetch some results from each entity to display search suggestions
-    const allSearchResultsByType = useGetAllEntitySearchResults({
-        query: '*',
-        start: 0,
-        count: 20,
-        filters: [],
+    // Fetch results
+    const { data: searchResultsData } = useGetSearchResultsForMultipleQuery({
+        variables: {
+            input: {
+                types: [],
+                query: '*',
+                start: 0,
+                count: 20,
+                filters: [],
+            },
+        },
     });
 
-    const suggestionsLoading = Object.keys(allSearchResultsByType).some((type) => {
-        return allSearchResultsByType[type].loading;
-    });
+    const searchResultsToShow = useMemo(() => {
+        let result: string[] | undefined = [];
+        if (searchResultsData) {
+            const entities = searchResultsData?.searchAcrossEntities?.searchResults.map((searchResult) => {
+                return searchResult?.entity;
+            });
 
-    const suggestionsToShow = useMemo(() => {
-        let result: string[] = [];
-        if (!suggestionsLoading) {
-            // TODO: Make this more dynamic.
-            // Add a ticket.
-            // Colored Tags: Feature Request...
-            // ...
-            [EntityType.Dashboard, EntityType.Chart, EntityType.Dataset].forEach((type) => {
-                const suggestionsToShowForEntity = getSuggestionFieldsFromResult(
-                    allSearchResultsByType[type]?.data,
-                ).sort(sortRandom);
-                const suggestionToAddToFront = suggestionsToShowForEntity?.pop();
-                result = [...result, ...suggestionsToShowForEntity];
-                if (suggestionToAddToFront) {
-                    result.splice(0, 0, suggestionToAddToFront);
-                }
+            result = entities?.map((entity) => {
+                return entityRegistry.getDisplayName(entity.type, entity);
             });
         }
-        return result;
-    }, [suggestionsLoading, allSearchResultsByType]);
+        return result?.sort(sortRandom);
+    }, [searchResultsData, entityRegistry]);
 
     return (
         <Background>
@@ -232,11 +206,11 @@ export const HomePageHeader = () => {
                         autoCompleteStyle={styles.searchBox}
                         entityRegistry={entityRegistry}
                     />
-                    {suggestionsToShow && suggestionsToShow.length > 0 && (
+                    {searchResultsToShow && searchResultsToShow.length > 0 && (
                         <SuggestionsContainer>
                             <SuggestedQueriesText strong>Try searching for</SuggestedQueriesText>
                             <SuggestionTagContainer>
-                                {suggestionsToShow.slice(0, 3).map((suggestion) => (
+                                {searchResultsToShow.slice(0, 3).map((suggestion) => (
                                     <SuggestionButton
                                         key={suggestion}
                                         type="link"
diff --git a/datahub-web-react/src/app/home/__tests__/HomePage.test.tsx b/datahub-web-react/src/app/home/__tests__/HomePage.test.tsx
index 58d457122e4f0..0d47c1d2cd083 100644
--- a/datahub-web-react/src/app/home/__tests__/HomePage.test.tsx
+++ b/datahub-web-react/src/app/home/__tests__/HomePage.test.tsx
@@ -61,8 +61,8 @@ describe('HomePage', () => {
         await waitFor(() => expect(searchInput).toBeInTheDocument());
         fireEvent.change(searchInput, { target: { value: 't' } });
 
-        await waitFor(() => expect(queryAllByText('The Great Test Dataset').length).toBeGreaterThanOrEqual(1));
-        expect(queryAllByText('Some other test').length).toBeGreaterThanOrEqual(1);
+        await waitFor(() => expect(queryAllByText('he Great Test Dataset').length).toBeGreaterThanOrEqual(1));
+        expect(queryAllByText('Some Other Dataset').length).toBeGreaterThanOrEqual(1);
     });
 
     it('renders search suggestions', async () => {
diff --git a/datahub-web-react/src/app/lineage/LineageEntityNode.tsx b/datahub-web-react/src/app/lineage/LineageEntityNode.tsx
index a715217e52341..ad6352ef8e17e 100644
--- a/datahub-web-react/src/app/lineage/LineageEntityNode.tsx
+++ b/datahub-web-react/src/app/lineage/LineageEntityNode.tsx
@@ -261,7 +261,9 @@ export default function LineageEntityNode({
                             |{' '}
                         </tspan>
                         <tspan dx=".25em" dy="-2px">
-                            {capitalizeFirstLetter(node.data.subtype || node.data.type)}
+                            {capitalizeFirstLetter(
+                                node.data.subtype || (node.data.type && entityRegistry.getEntityName(node.data.type)),
+                            )}
                         </tspan>
                     </UnselectableText>
                     {expandTitles ? (
diff --git a/datahub-web-react/src/app/lineage/__tests__/constructTree.test.ts b/datahub-web-react/src/app/lineage/__tests__/constructTree.test.ts
index 24c391c98936f..d8bb49ed82e06 100644
--- a/datahub-web-react/src/app/lineage/__tests__/constructTree.test.ts
+++ b/datahub-web-react/src/app/lineage/__tests__/constructTree.test.ts
@@ -76,6 +76,7 @@ describe('constructTree', () => {
                     children: [],
                     icon: undefined,
                     platform: 'Kafka',
+                    status: null,
                 },
             ],
         });
@@ -121,6 +122,7 @@ describe('constructTree', () => {
                     children: [],
                     icon: undefined,
                     platform: 'Kafka',
+                    status: null,
                 },
             ],
         });
@@ -166,6 +168,7 @@ describe('constructTree', () => {
                     countercurrentChildrenUrns: ['urn:li:dataset:3'],
                     icon: undefined,
                     platform: 'Kafka',
+                    status: null,
                     children: [
                         {
                             name: 'Sixth Test Dataset',
@@ -175,6 +178,7 @@ describe('constructTree', () => {
                             countercurrentChildrenUrns: ['urn:li:dataset:4'],
                             icon: undefined,
                             platform: 'Kafka',
+                            status: null,
                             children: [
                                 {
                                     name: 'Fifth Test Dataset',
@@ -189,6 +193,7 @@ describe('constructTree', () => {
                                     ],
                                     icon: undefined,
                                     platform: 'Kafka',
+                                    status: null,
                                 },
                             ],
                         },
@@ -201,6 +206,7 @@ describe('constructTree', () => {
                             countercurrentChildrenUrns: ['urn:li:dataset:7', 'urn:li:dataset:6', 'urn:li:dataset:4'],
                             icon: undefined,
                             platform: 'Kafka',
+                            status: null,
                         },
                     ],
                 },
@@ -276,6 +282,7 @@ describe('constructTree', () => {
                     countercurrentChildrenUrns: ['urn:li:dataset:3'],
                     icon: undefined,
                     platform: 'Kafka',
+                    status: null,
                 },
             ],
         });
diff --git a/datahub-web-react/src/app/lineage/types.ts b/datahub-web-react/src/app/lineage/types.ts
index eadee188818f7..a83e63a40fb59 100644
--- a/datahub-web-react/src/app/lineage/types.ts
+++ b/datahub-web-react/src/app/lineage/types.ts
@@ -9,6 +9,8 @@ import {
     MlFeature,
     MlModel,
     MlModelGroup,
+    Maybe,
+    Status,
 } from '../../types.generated';
 
 export type EntitySelectParams = {
@@ -33,6 +35,7 @@ export type FetchedEntity = {
     downstreamChildren?: Array<EntityAndType>;
     fullyFetched?: boolean;
     platform?: string;
+    status?: Maybe<Status>;
 };
 
 export type NodeData = {
@@ -47,6 +50,7 @@ export type NodeData = {
     // Currently our visualization does not support expanding in two directions
     countercurrentChildrenUrns?: string[];
     platform?: string;
+    status?: Maybe<Status>;
 };
 
 export type VizNode = {
diff --git a/datahub-web-react/src/app/lineage/utils/constructFetchedNode.ts b/datahub-web-react/src/app/lineage/utils/constructFetchedNode.ts
index 9e3ce2473c74a..4a766e081a964 100644
--- a/datahub-web-react/src/app/lineage/utils/constructFetchedNode.ts
+++ b/datahub-web-react/src/app/lineage/utils/constructFetchedNode.ts
@@ -35,6 +35,7 @@ export default function constructFetchedNode(
                 ) || [],
             children: [],
             platform: fetchedNode?.platform,
+            status: fetchedNode.status,
         };
 
         // eslint-disable-next-line no-param-reassign
diff --git a/datahub-web-react/src/app/lineage/utils/layoutTree.ts b/datahub-web-react/src/app/lineage/utils/layoutTree.ts
index 14888221f68d5..6c7d128236e66 100644
--- a/datahub-web-react/src/app/lineage/utils/layoutTree.ts
+++ b/datahub-web-react/src/app/lineage/utils/layoutTree.ts
@@ -44,10 +44,13 @@ export default function layoutTree(
             .filter((urn, pos) => urnsToAddInCurrentLayer.indexOf(urn) === pos)
             .filter((urn) => !nodesByUrn[urn || '']);
 
-        const layerSize = nodesToAddInCurrentLayer.length;
-
-        const layerHeight = nodesInCurrentLayer
+        const filteredNodesInCurrentLayer = nodesInCurrentLayer
             .filter(({ node }) => nodesToAddInCurrentLayer.indexOf(node.urn || '') > -1)
+            .filter(({ node }) => node.status?.removed !== true);
+
+        const layerSize = filteredNodesInCurrentLayer.length;
+
+        const layerHeight = filteredNodesInCurrentLayer
             .map(({ node }) => nodeHeightFromTitleLength(expandTitles ? node.name : undefined))
             .reduce((acc, height) => acc + height, 0);
 
@@ -58,10 +61,14 @@ export default function layoutTree(
             -((nodeHeightFromTitleLength(undefined) + VERTICAL_SPACE_BETWEEN_NODES) * (layerSize - 1)) / 2 +
             canvasHeight / 2 +
             HEADER_HEIGHT;
+
         // eslint-disable-next-line @typescript-eslint/no-loop-func
         nodesInCurrentLayer.forEach(({ node, parent }) => {
             if (!node.urn) return;
 
+            // don't show edges to soft deleted entities
+            if (node.status?.removed) return;
+
             let vizNodeForNode: VizNode;
 
             if (nodesByUrn[node.urn]) {
diff --git a/datahub-web-react/src/app/search/SearchBar.tsx b/datahub-web-react/src/app/search/SearchBar.tsx
index bcf5a1536b9dc..7a7bae024fd75 100644
--- a/datahub-web-react/src/app/search/SearchBar.tsx
+++ b/datahub-web-react/src/app/search/SearchBar.tsx
@@ -1,27 +1,35 @@
-import React, { useState } from 'react';
-import { Input, AutoComplete } from 'antd';
+import React, { useMemo, useState } from 'react';
+import { Input, AutoComplete, Image, Typography } from 'antd';
 import { SearchOutlined } from '@ant-design/icons';
 import styled from 'styled-components';
-
-import { AutoCompleteResultForEntity, EntityType } from '../../types.generated';
+import { useHistory } from 'react-router';
+import { AutoCompleteResultForEntity, CorpUser, Entity, EntityType, ScenarioType, Tag } from '../../types.generated';
 import { IconStyleType } from '../entity/Entity';
 import EntityRegistry from '../entity/EntityRegistry';
-import { SEARCH_FOR_ENTITY_PREFIX } from './utils/constants';
 import filterSearchQuery from './utils/filterSearchQuery';
 import { ANTD_GRAY } from '../entity/shared/constants';
+import { getEntityPath } from '../entity/shared/containers/profile/utils';
+import { EXACT_SEARCH_PREFIX } from './utils/constants';
+import { CustomAvatar } from '../shared/avatar';
+import { StyledTag } from '../entity/shared/components/styled/StyledTag';
+import { useListRecommendationsQuery } from '../../graphql/recommendations.generated';
+import { useGetAuthenticatedUserUrn } from '../useGetAuthenticatedUser';
 
 const SuggestionContainer = styled.div`
-    display: 'flex',
-    flex-direction: 'row',
-    align-items: 'center',
+    display: flex;
+    flex-direction: row;
+    justify-content: left;
+    align-items: center;
 `;
 
 const SuggestionText = styled.span`
-    margin-left: 10px;
+    margin-left: 8px;
+    margin-top: 2px;
+    margin-bottom: 2px;
+    color: ${ANTD_GRAY[9]};
 `;
 
 const ExploreForEntity = styled.span`
-    font-style: italic;
     font-weight: light;
     font-size: 11px;
 `;
@@ -48,16 +56,95 @@ const StyledSearchBar = styled(Input)`
     }
 `;
 
-const renderItem = (suggestion: string, icon: JSX.Element, type: string) => ({
-    value: suggestion,
-    label: (
-        <SuggestionContainer>
-            <span>{icon}</span>
-            <SuggestionText>{suggestion}</SuggestionText>
-        </SuggestionContainer>
-    ),
-    type,
-});
+const PreviewImage = styled(Image)`
+    max-height: 16px;
+    width: auto;
+    object-fit: contain;
+    background-color: transparent;
+`;
+
+const EXACT_AUTOCOMPLETE_OPTION_TYPE = 'exact_query';
+const RECOMMENDED_QUERY_OPTION_TYPE = 'recommendation';
+
+const renderTagSuggestion = (tag: Tag, registry: EntityRegistry) => {
+    return (
+        <>
+            <StyledTag $colorHash={tag?.urn} $color={tag?.properties?.colorHex}>
+                {registry.getDisplayName(EntityType.Tag, tag)}
+            </StyledTag>
+        </>
+    );
+};
+
+const renderUserSuggestion = (query: string, user: CorpUser, registry: EntityRegistry) => {
+    const displayName = registry.getDisplayName(EntityType.CorpUser, user);
+    const isPrefixMatch = displayName.toLowerCase().indexOf(query.toLowerCase()) > -1;
+    const matchedText = (isPrefixMatch && displayName.substring(0, query.length)) || '';
+    const unmatchedText = (isPrefixMatch && displayName.substring(query.length, displayName.length)) || displayName;
+    return (
+        <>
+            <CustomAvatar
+                size={20}
+                name={displayName}
+                photoUrl={user.editableProperties?.pictureLink || undefined}
+                useDefaultAvatar={false}
+                style={{
+                    marginRight: 0,
+                }}
+            />
+            <SuggestionText>
+                {matchedText}
+                <Typography.Text strong>{unmatchedText}</Typography.Text>
+            </SuggestionText>
+        </>
+    );
+};
+
+const renderEntitySuggestion = (query: string, entity: Entity, registry: EntityRegistry) => {
+    // Special rendering.
+    if (entity.type === EntityType.CorpUser) {
+        return renderUserSuggestion(query, entity as CorpUser, registry);
+    }
+    if (entity.type === EntityType.Tag) {
+        return renderTagSuggestion(entity as Tag, registry);
+    }
+    const genericEntityProps = registry.getGenericEntityProperties(entity.type, entity);
+    const platformName = genericEntityProps?.platform?.properties?.displayName || genericEntityProps?.platform?.name;
+    const platformLogoUrl = genericEntityProps?.platform?.properties?.logoUrl;
+    const displayName = registry.getDisplayName(entity.type, entity);
+    const icon =
+        (platformLogoUrl && <PreviewImage preview={false} src={platformLogoUrl} alt={platformName || ''} />) ||
+        registry.getIcon(entity.type, 12, IconStyleType.ACCENT);
+    const isPrefixMatch = displayName.toLowerCase().startsWith(query.toLowerCase());
+    const matchedText = isPrefixMatch ? displayName.substring(0, query.length) : '';
+    const unmatchedText = isPrefixMatch ? displayName.substring(query.length, displayName.length) : displayName;
+    return (
+        <>
+            {icon}
+            <SuggestionText>
+                <Typography.Text strong>{matchedText}</Typography.Text>
+                {unmatchedText}
+            </SuggestionText>
+        </>
+    );
+};
+
+const renderItem = (query: string, entity: Entity, registry: EntityRegistry) => {
+    return {
+        value: entity.urn,
+        label: <SuggestionContainer>{renderEntitySuggestion(query, entity, registry)}</SuggestionContainer>,
+        type: entity.type,
+        style: { paddingLeft: 16 },
+    };
+};
+
+const renderRecommendedQuery = (query: string) => {
+    return {
+        value: query,
+        label: query,
+        type: RECOMMENDED_QUERY_OPTION_TYPE,
+    };
+};
 
 interface Props {
     initialQuery?: string;
@@ -69,6 +156,7 @@ interface Props {
     inputStyle?: React.CSSProperties;
     autoCompleteStyle?: React.CSSProperties;
     entityRegistry: EntityRegistry;
+    fixAutoComplete?: boolean;
 }
 
 const defaultProps = {
@@ -88,52 +176,105 @@ export const SearchBar = ({
     style,
     inputStyle,
     autoCompleteStyle,
+    fixAutoComplete,
 }: Props) => {
+    const history = useHistory();
     const [searchQuery, setSearchQuery] = useState<string>();
     const [selected, setSelected] = useState<string>();
     const searchEntityTypes = entityRegistry.getSearchEntityTypes();
+    const userUrn = useGetAuthenticatedUserUrn();
+    const { data } = useListRecommendationsQuery({
+        variables: {
+            input: {
+                userUrn,
+                requestContext: {
+                    scenario: ScenarioType.SearchBar,
+                },
+                limit: 1,
+            },
+        },
+    });
+
+    const effectiveQuery = searchQuery !== undefined ? searchQuery : initialQuery || '';
+
+    const emptyQueryOptions = useMemo(() => {
+        // Map each module to a set of
+        return (
+            data?.listRecommendations?.modules.map((module) => ({
+                label: module.title,
+                options: [...module.content.map((content) => renderRecommendedQuery(content.value))],
+            })) || []
+        );
+    }, [data]);
+
+    const autoCompleteQueryOptions = useMemo(
+        () =>
+            (suggestions?.length > 0 &&
+                effectiveQuery.length > 0 && [
+                    {
+                        value: `${EXACT_SEARCH_PREFIX}${effectiveQuery}`,
+                        label: (
+                            <SuggestionContainer key={EXACT_AUTOCOMPLETE_OPTION_TYPE}>
+                                <ExploreForEntity>
+                                    View all results for <Typography.Text strong>{effectiveQuery}</Typography.Text>
+                                </ExploreForEntity>
+                            </SuggestionContainer>
+                        ),
+                        type: EXACT_AUTOCOMPLETE_OPTION_TYPE,
+                    },
+                ]) ||
+            [],
+        [suggestions, effectiveQuery],
+    );
+
+    const autoCompleteEntityOptions = useMemo(
+        () =>
+            suggestions.map((entity: AutoCompleteResultForEntity) => ({
+                label: entityRegistry.getCollectionName(entity.type),
+                options: [...entity.entities.map((e: Entity) => renderItem(effectiveQuery, e, entityRegistry))],
+            })),
+        [effectiveQuery, suggestions, entityRegistry],
+    );
 
-    const options = suggestions.map((entity: AutoCompleteResultForEntity) => ({
-        label: entityRegistry.getCollectionName(entity.type),
-        options: [
-            ...entity.suggestions.map((suggestion: string) =>
-                renderItem(suggestion, entityRegistry.getIcon(entity.type, 14, IconStyleType.TAB_VIEW), entity.type),
-            ),
-            ...(searchEntityTypes.indexOf(entity.type) >= 0
-                ? [
-                      {
-                          value: `${SEARCH_FOR_ENTITY_PREFIX}${searchQuery}__${entity.type}`,
-                          label: (
-                              <SuggestionContainer key={entity.type}>
-                                  <ExploreForEntity>
-                                      Explore all `{searchQuery}` in {entityRegistry.getCollectionName(entity.type)}
-                                  </ExploreForEntity>
-                              </SuggestionContainer>
-                          ),
-                          type: entity.type,
-                          key: `${searchQuery}-${entity.type}`,
-                      },
-                  ]
-                : []),
-        ],
-    }));
+    const options = useMemo(() => {
+        // Display recommendations when there is no search query, autocomplete suggestions otherwise.
+        if (autoCompleteEntityOptions.length > 0) {
+            return [...autoCompleteQueryOptions, ...autoCompleteEntityOptions];
+        }
+        return emptyQueryOptions;
+    }, [emptyQueryOptions, autoCompleteEntityOptions, autoCompleteQueryOptions]);
 
     return (
         <AutoCompleteContainer style={style}>
             <StyledAutoComplete
+                defaultActiveFirstOption={false}
                 style={autoCompleteStyle}
                 options={options}
-                onSelect={(value: string, option) =>
-                    onSearch(
-                        `"${filterSearchQuery(value)}"`,
-                        searchEntityTypes.indexOf(option.type) >= 0 ? option.type : undefined,
-                    )
-                }
+                filterOption={false}
+                onSelect={(value: string, option) => {
+                    // If the autocomplete option type is NOT an entity, then render as a normal search query.
+                    if (
+                        option.type === EXACT_AUTOCOMPLETE_OPTION_TYPE ||
+                        option.type === RECOMMENDED_QUERY_OPTION_TYPE
+                    ) {
+                        onSearch(
+                            `${filterSearchQuery(value)}`,
+                            searchEntityTypes.indexOf(option.type) >= 0 ? option.type : undefined,
+                        );
+                    } else {
+                        // Navigate directly to the entity profile.
+                        history.push(getEntityPath(option.type, value, entityRegistry, false));
+                    }
+                }}
                 onSearch={(value: string) => onQueryChange(value)}
                 defaultValue={initialQuery || undefined}
                 value={selected}
                 onChange={(v) => setSelected(filterSearchQuery(v))}
-                dropdownStyle={{ maxHeight: 1000, overflowY: 'visible', position: 'fixed' }}
+                dropdownStyle={{
+                    maxHeight: 1000,
+                    overflowY: 'visible',
+                    position: (fixAutoComplete && 'fixed') || 'relative',
+                }}
             >
                 <StyledSearchBar
                     placeholder={placeholderText}
diff --git a/datahub-web-react/src/app/search/SearchHeader.tsx b/datahub-web-react/src/app/search/SearchHeader.tsx
index 63e83415f020a..a19dbdfcd796b 100644
--- a/datahub-web-react/src/app/search/SearchHeader.tsx
+++ b/datahub-web-react/src/app/search/SearchHeader.tsx
@@ -88,6 +88,7 @@ export const SearchHeader = ({
                     onSearch={onSearch}
                     onQueryChange={onQueryChange}
                     entityRegistry={entityRegistry}
+                    fixAutoComplete
                 />
             </LogoSearchContainer>
             <NavGroup>
diff --git a/datahub-web-react/src/app/search/SearchablePage.tsx b/datahub-web-react/src/app/search/SearchablePage.tsx
index 3097caaa07d4d..1d747512d113f 100644
--- a/datahub-web-react/src/app/search/SearchablePage.tsx
+++ b/datahub-web-react/src/app/search/SearchablePage.tsx
@@ -1,4 +1,4 @@
-import React from 'react';
+import React, { useEffect } from 'react';
 import { useHistory } from 'react-router';
 import { useTheme } from 'styled-components';
 
@@ -72,6 +72,19 @@ export const SearchablePage = ({ initialQuery, onSearch, onAutoComplete, childre
         }
     };
 
+    // Load correct autocomplete results on initial page load.
+    useEffect(() => {
+        if (initialQuery && initialQuery.trim() !== '') {
+            getAutoCompleteResults({
+                variables: {
+                    input: {
+                        query: initialQuery,
+                    },
+                },
+            });
+        }
+    }, [initialQuery, getAutoCompleteResults]);
+
     return (
         <>
             <SearchHeader
diff --git a/datahub-web-react/src/app/search/utils/constants.ts b/datahub-web-react/src/app/search/utils/constants.ts
index b41575135ba76..bcf0e0e4bc95c 100644
--- a/datahub-web-react/src/app/search/utils/constants.ts
+++ b/datahub-web-react/src/app/search/utils/constants.ts
@@ -1,5 +1,6 @@
 export const FILTER_URL_PREFIX = 'filter_';
 export const SEARCH_FOR_ENTITY_PREFIX = 'SEARCH__';
+export const EXACT_SEARCH_PREFIX = 'EXACT__';
 
 export const ENTITY_FILTER_NAME = 'entity';
 export const TAG_FILTER_NAME = 'tags';
diff --git a/datahub-web-react/src/app/search/utils/filterSearchQuery.ts b/datahub-web-react/src/app/search/utils/filterSearchQuery.ts
index c3f6dab72c854..48df97dc77788 100644
--- a/datahub-web-react/src/app/search/utils/filterSearchQuery.ts
+++ b/datahub-web-react/src/app/search/utils/filterSearchQuery.ts
@@ -1,5 +1,5 @@
-import { SEARCH_FOR_ENTITY_PREFIX } from './constants';
+import { SEARCH_FOR_ENTITY_PREFIX, EXACT_SEARCH_PREFIX } from './constants';
 
 export default function filterSearchQuery(v: string) {
-    return v && v.startsWith(SEARCH_FOR_ENTITY_PREFIX) ? v.split('__')[1] : v;
+    return (v && v.startsWith(SEARCH_FOR_ENTITY_PREFIX)) || v.startsWith(EXACT_SEARCH_PREFIX) ? v.split('__')[1] : v;
 }
diff --git a/datahub-web-react/src/app/shared/TagStyleEntity.tsx b/datahub-web-react/src/app/shared/TagStyleEntity.tsx
index 1389a5a4408d5..6fda9643303a5 100644
--- a/datahub-web-react/src/app/shared/TagStyleEntity.tsx
+++ b/datahub-web-react/src/app/shared/TagStyleEntity.tsx
@@ -377,7 +377,7 @@ export default function TagStyleEntity({ urn, useGetSearchResults = useWrappedSe
                     <StatsLabel>Owners</StatsLabel>
                     <div>
                         {data?.tag?.ownership?.owners?.map((owner) => (
-                            <ExpandedOwner entityUrn={urn} owner={owner} refetch={refetch} />
+                            <ExpandedOwner entityUrn={urn} owner={owner} refetch={refetch} hidePopOver />
                         ))}
                         {ownersEmpty && (
                             <Typography.Paragraph type="secondary">
@@ -395,6 +395,7 @@ export default function TagStyleEntity({ urn, useGetSearchResults = useWrappedSe
                     </div>
                     <div>
                         <AddOwnerModal
+                            hideOwnerType
                             visible={showAddModal}
                             refetch={refetch}
                             onClose={() => {
diff --git a/datahub-web-react/src/app/useGetAuthenticatedUser.tsx b/datahub-web-react/src/app/useGetAuthenticatedUser.tsx
index 6d7b50ac1aa46..0b8e929b682df 100644
--- a/datahub-web-react/src/app/useGetAuthenticatedUser.tsx
+++ b/datahub-web-react/src/app/useGetAuthenticatedUser.tsx
@@ -16,3 +16,14 @@ export function useGetAuthenticatedUser() {
     }
     return data?.me;
 }
+
+/**
+ * Fetch an urn corresponding to the authenticated user.
+ */
+export function useGetAuthenticatedUserUrn() {
+    const userUrn = Cookies.get(CLIENT_AUTH_COOKIE);
+    if (!userUrn) {
+        throw new Error('Could not find logged in user.');
+    }
+    return userUrn;
+}
diff --git a/datahub-web-react/src/graphql-mock/resolver/getAutoCompleteAllResultsResolver.ts b/datahub-web-react/src/graphql-mock/resolver/getAutoCompleteAllResultsResolver.ts
index a5f33b7e9525c..52c4845302b6e 100644
--- a/datahub-web-react/src/graphql-mock/resolver/getAutoCompleteAllResultsResolver.ts
+++ b/datahub-web-react/src/graphql-mock/resolver/getAutoCompleteAllResultsResolver.ts
@@ -35,6 +35,9 @@ const findSuggestions = ({ query, type }: AutoCompleteInput): AutoCompleteResult
                 suggestions: results.map((r) => {
                     return r.name;
                 }),
+                entities: results.map((r) => {
+                    return r;
+                }),
                 __typename: 'AutoCompleteResultForEntity',
             },
         ];
@@ -50,6 +53,7 @@ const findSuggestions = ({ query, type }: AutoCompleteInput): AutoCompleteResult
                   suggestions: datasetResults.map((r) => {
                       return (r.entity as Dataset).name;
                   }),
+                  entities: [],
                   __typename: 'AutoCompleteResultForEntity',
               }
             : null;
@@ -63,6 +67,9 @@ const findSuggestions = ({ query, type }: AutoCompleteInput): AutoCompleteResult
                   suggestions: dashboardResults.map((r) => {
                       return (r.entity as Dashboard).info?.name || '';
                   }),
+                  entities: dashboardResults.map((r) => {
+                      return r.entity;
+                  }),
                   __typename: 'AutoCompleteResultForEntity',
               }
             : null;
@@ -76,6 +83,9 @@ const findSuggestions = ({ query, type }: AutoCompleteInput): AutoCompleteResult
                   suggestions: chartResults.map((r) => {
                       return (r.entity as Chart).info?.name || '';
                   }),
+                  entities: chartResults.map((r) => {
+                      return r.entity;
+                  }),
                   __typename: 'AutoCompleteResultForEntity',
               }
             : null;
@@ -89,6 +99,9 @@ const findSuggestions = ({ query, type }: AutoCompleteInput): AutoCompleteResult
                   suggestions: dataFlowResults.map((r) => {
                       return (r.entity as DataFlow).info?.name || '';
                   }),
+                  entities: dataFlowResults.map((r) => {
+                      return r.entity;
+                  }),
                   __typename: 'AutoCompleteResultForEntity',
               }
             : null;
@@ -102,6 +115,9 @@ const findSuggestions = ({ query, type }: AutoCompleteInput): AutoCompleteResult
                   suggestions: dataJobResults.map((r) => {
                       return (r.entity as DataJob).info?.name || '';
                   }),
+                  entities: dataJobResults.map((r) => {
+                      return r.entity;
+                  }),
                   __typename: 'AutoCompleteResultForEntity',
               }
             : null;
@@ -115,6 +131,9 @@ const findSuggestions = ({ query, type }: AutoCompleteInput): AutoCompleteResult
                   suggestions: userResults.map((r) => {
                       return (r.entity as CorpUser).info?.fullName || '';
                   }),
+                  entities: userResults.map((r) => {
+                      return r.entity;
+                  }),
                   __typename: 'AutoCompleteResultForEntity',
               }
             : null;
diff --git a/datahub-web-react/src/graphql-mock/resolver/getAutoCompleteResultsResolver.ts b/datahub-web-react/src/graphql-mock/resolver/getAutoCompleteResultsResolver.ts
index 4f952c2ffb538..68354c43983e1 100644
--- a/datahub-web-react/src/graphql-mock/resolver/getAutoCompleteResultsResolver.ts
+++ b/datahub-web-react/src/graphql-mock/resolver/getAutoCompleteResultsResolver.ts
@@ -86,6 +86,7 @@ export const getAutoCompleteResultsResolver = {
                     suggestions: suggestions.filter((s, i) => {
                         return suggestions.indexOf(s) === i;
                     }),
+                    entities: [],
                     __typename: 'AutoCompleteResults',
                 },
             },
diff --git a/datahub-web-react/src/graphql/chart.graphql b/datahub-web-react/src/graphql/chart.graphql
index 21ba0e65f8b1c..f1cc2e30d9fcb 100644
--- a/datahub-web-react/src/graphql/chart.graphql
+++ b/datahub-web-react/src/graphql/chart.graphql
@@ -62,6 +62,9 @@ query getChart($urn: String!) {
         downstream: lineage(input: { direction: DOWNSTREAM, start: 0, count: 100 }) {
             ...fullLineageResults
         }
+        status {
+            removed
+        }
     }
 }
 
diff --git a/datahub-web-react/src/graphql/dataFlow.graphql b/datahub-web-react/src/graphql/dataFlow.graphql
index 218acc0a64e9e..a78aeb0e3865e 100644
--- a/datahub-web-react/src/graphql/dataFlow.graphql
+++ b/datahub-web-react/src/graphql/dataFlow.graphql
@@ -35,6 +35,9 @@ fragment dataFlowFields on DataFlow {
     domain {
         ...entityDomain
     }
+    status {
+        removed
+    }
 }
 
 query getDataFlow($urn: String!) {
diff --git a/datahub-web-react/src/graphql/dataJob.graphql b/datahub-web-react/src/graphql/dataJob.graphql
index 820e976273e0f..123522b70fe89 100644
--- a/datahub-web-react/src/graphql/dataJob.graphql
+++ b/datahub-web-react/src/graphql/dataJob.graphql
@@ -25,6 +25,9 @@ query getDataJob($urn: String!) {
         downstream: lineage(input: { direction: DOWNSTREAM, start: 0, count: 100 }) {
             ...fullLineageResults
         }
+        status {
+            removed
+        }
     }
 }
 
diff --git a/datahub-web-react/src/graphql/dataset.graphql b/datahub-web-react/src/graphql/dataset.graphql
index 42847a3330681..e1c1f8d5bc748 100644
--- a/datahub-web-react/src/graphql/dataset.graphql
+++ b/datahub-web-react/src/graphql/dataset.graphql
@@ -137,6 +137,9 @@ query getDataset($urn: String!) {
                 key
             }
         }
+        status {
+            removed
+        }
     }
 }
 
diff --git a/datahub-web-react/src/graphql/fragments.graphql b/datahub-web-react/src/graphql/fragments.graphql
index cfc621383e375..c349c4e3f5beb 100644
--- a/datahub-web-react/src/graphql/fragments.graphql
+++ b/datahub-web-react/src/graphql/fragments.graphql
@@ -333,6 +333,9 @@ fragment dashboardFields on Dashboard {
     container {
         ...entityContainer
     }
+    status {
+        removed
+    }
 }
 
 fragment nonRecursiveMLFeature on MLFeature {
@@ -547,6 +550,9 @@ fragment nonRecursiveMLModel on MLModel {
     globalTags {
         ...globalTagsFields
     }
+    status {
+        removed
+    }
 }
 
 fragment nonRecursiveMLModelGroupFields on MLModelGroup {
@@ -561,6 +567,9 @@ fragment nonRecursiveMLModelGroupFields on MLModelGroup {
     ownership {
         ...ownershipFields
     }
+    status {
+        removed
+    }
 }
 
 fragment platformFields on DataPlatform {
diff --git a/datahub-web-react/src/graphql/lineage.graphql b/datahub-web-react/src/graphql/lineage.graphql
index 90b3f2e24aeba..d1fd26a182b94 100644
--- a/datahub-web-react/src/graphql/lineage.graphql
+++ b/datahub-web-react/src/graphql/lineage.graphql
@@ -6,6 +6,9 @@ fragment relationshipFields on EntityWithRelationships {
         editableProperties {
             description
         }
+        status {
+            removed
+        }
     }
     ... on DataFlow {
         orchestrator
@@ -34,6 +37,9 @@ fragment relationshipFields on EntityWithRelationships {
         domain {
             ...entityDomain
         }
+        status {
+            removed
+        }
     }
     ... on Dashboard {
         ...dashboardFields
@@ -43,6 +49,9 @@ fragment relationshipFields on EntityWithRelationships {
         platform {
             ...platformFields
         }
+        status {
+            removed
+        }
     }
     ... on Chart {
         tool
@@ -63,11 +72,14 @@ fragment relationshipFields on EntityWithRelationships {
         domain {
             ...entityDomain
         }
+        status {
+            removed
+        }
     }
     ... on Dataset {
         name
         properties {
-            name 
+            name
             description
         }
         editableProperties {
@@ -82,6 +94,9 @@ fragment relationshipFields on EntityWithRelationships {
         subTypes {
             typeNames
         }
+        status {
+            removed
+        }
     }
     ... on MLModelGroup {
         urn
@@ -95,6 +110,9 @@ fragment relationshipFields on EntityWithRelationships {
         ownership {
             ...ownershipFields
         }
+        status {
+            removed
+        }
     }
     ... on MLModel {
         urn
@@ -108,6 +126,9 @@ fragment relationshipFields on EntityWithRelationships {
         ownership {
             ...ownershipFields
         }
+        status {
+            removed
+        }
     }
     upstream: lineage(input: { direction: UPSTREAM, start: 0, count: 100 }) {
         ...leafLineageResults
diff --git a/datahub-web-react/src/graphql/search.graphql b/datahub-web-react/src/graphql/search.graphql
index 20a660a219efd..781a49bd8a9cd 100644
--- a/datahub-web-react/src/graphql/search.graphql
+++ b/datahub-web-react/src/graphql/search.graphql
@@ -2,6 +2,9 @@ query getAutoCompleteResults($input: AutoCompleteInput!) {
     autoComplete(input: $input) {
         query
         suggestions
+        entities {
+            ...searchResultFields
+        }
     }
 }
 
@@ -11,6 +14,9 @@ query getAutoCompleteMultipleResults($input: AutoCompleteMultipleInput!) {
         suggestions {
             type
             suggestions
+            entities {
+                ...searchResultFields
+            }
         }
     }
 }
diff --git a/datahub-web-react/src/images/datahublogo.png b/datahub-web-react/src/images/datahublogo.png
new file mode 100644
index 0000000000000..5e34e6425d23f
Binary files /dev/null and b/datahub-web-react/src/images/datahublogo.png differ
diff --git a/docker/datahub-frontend/README.md b/docker/datahub-frontend/README.md
index 3adbb5abe7534..bc1a153267522 100644
--- a/docker/datahub-frontend/README.md
+++ b/docker/datahub-frontend/README.md
@@ -23,8 +23,8 @@ You can sign in with `datahub` as username and password.
 
 ## Build instructions
 
-If you want to build the `datahub-frontend` Docker image yourself, you can run this command from the root directory of the DataHub repository you have locally:
+If you want to build the `datahub-frontend` Docker image yourself, you can run this command from the root directory of the DataHub repository you have locally (using Buildkit):
 
-`docker build -t your_datahub_frontend -f ./docker/datahub-frontend/Dockerfile .`
+`DOCKER_BUILDKIT=1 docker build -t your_datahub_frontend -f ./docker/datahub-frontend/Dockerfile .`
 
 Please note the final `.` and that the tag `your_datahub_frontend` is determined by you.
\ No newline at end of file
diff --git a/docker/datahub-frontend/env/docker.env b/docker/datahub-frontend/env/docker.env
index f513ce63f6e18..47024bdbb0bed 100644
--- a/docker/datahub-frontend/env/docker.env
+++ b/docker/datahub-frontend/env/docker.env
@@ -28,7 +28,7 @@ JAVA_OPTS=-Xms512m -Xmx512m -Dhttp.port=9002 -Dconfig.file=datahub-frontend/conf
 # Optional Provisioning Configs
 # AUTH_OIDC_JIT_PROVISIONING_ENABLED=true
 # AUTH_OIDC_PRE_PROVISIONING_REQUIRED=false
-# AUTH_OIDC_EXTRACT_GROUPS_ENABLED=true
+# AUTH_OIDC_EXTRACT_GROUPS_ENABLED=false
 # AUTH_OIDC_GROUPS_CLAIM=groups
 
 # Uncomment to disable JAAS username / password authentication (enabled by default)
diff --git a/docker/docker-compose-without-neo4j.yml b/docker/docker-compose-without-neo4j.yml
index cbc6a2d70bc05..0d0c3608c0595 100644
--- a/docker/docker-compose-without-neo4j.yml
+++ b/docker/docker-compose-without-neo4j.yml
@@ -57,6 +57,9 @@ services:
     hostname: elasticsearch
     ports:
       - "9200:9200"
+    environment:
+      - discovery.type=single-node
+      - xpack.security.enabled=false
     volumes:
       - esdata:/usr/share/elasticsearch/data
     healthcheck:
diff --git a/docker/docker-compose.yml b/docker/docker-compose.yml
index eb6d2f0a1eae0..0db24b0645971 100644
--- a/docker/docker-compose.yml
+++ b/docker/docker-compose.yml
@@ -59,6 +59,9 @@ services:
     hostname: elasticsearch
     ports:
       - "9200:9200"
+    environment:
+      - discovery.type=single-node
+      - xpack.security.enabled=false
     volumes:
       - esdata:/usr/share/elasticsearch/data
     healthcheck:
diff --git a/docker/elasticsearch/env/docker.env b/docker/elasticsearch/env/docker.env
index 511e627dd152b..038469425c25d 100644
--- a/docker/elasticsearch/env/docker.env
+++ b/docker/elasticsearch/env/docker.env
@@ -1,3 +1 @@
-discovery.type=single-node
-xpack.security.enabled=false
-ES_JAVA_OPTS=-Xms256m -Xmx256m -Dlog4j2.formatMsgNoLookups=true
+ES_JAVA_OPTS="-Xms256m -Xmx256m -Dlog4j2.formatMsgNoLookups=true"
diff --git a/docker/quickstart/generate_docker_quickstart.py b/docker/quickstart/generate_docker_quickstart.py
index 7576b1a22895e..3a54d8c21155d 100644
--- a/docker/quickstart/generate_docker_quickstart.py
+++ b/docker/quickstart/generate_docker_quickstart.py
@@ -4,6 +4,8 @@
 from collections.abc import Mapping
 from dotenv import dotenv_values
 from yaml import Loader
+from collections import OrderedDict
+
 
 # Generates a merged docker-compose file with env variables inlined.
 # Usage: python3 docker_compose_cli_gen.py ../docker-compose.yml ../docker-compose.override.yml ../docker-compose-gen.yml
@@ -46,23 +48,26 @@ def modify_docker_config(base_path, docker_yaml_config):
             # 3. Resolve the .env values
             env_vars = dotenv_values(env_file_path)
 
-            # 4. Add an "environment" block to YAML
-            service["environment"] = list(
-                f"{key}={value}" for key, value in env_vars.items()
-            )
+            # 4. Create an "environment" block if it does not exist
+            if "environment" not in service:
+                service["environment"] = list()
+
+            # 5. Append to an "environment" block to YAML
+            for key, value in env_vars.items():
+                service["environment"].append(f"{key}={value}")
 
-            # 5. Delete the "env_file" value
+            # 6. Delete the "env_file" value
             del service["env_file"]
 
-        # 6. Delete build instructions
+        # 7. Delete build instructions
         if "build" in service:
             del service["build"]
 
-        # 7. Set memory limits
+        # 8. Set memory limits
         if name in mem_limits:
             service["mem_limit"] = mem_limits[name]
 
-        # 8. Correct relative paths for volume mounts
+        # 9. Correct relative paths for volume mounts
         if "volumes" in service:
             volumes = service["volumes"]
             for i in range(len(volumes)):
@@ -72,7 +77,7 @@ def modify_docker_config(base_path, docker_yaml_config):
                 elif volumes[i].startswith("./"):
                     volumes[i] = "." + volumes[i]
 
-    # 8. Set docker compose version to 2.
+    # 9. Set docker compose version to 2.
     # We need at least this version, since we use features like start_period for
     # healthchecks and shell-like variable interpolation.
     docker_yaml_config["version"] = "2.3"
diff --git a/docs-website/sidebars.js b/docs-website/sidebars.js
index 4f7c4821238ef..1f32518bbfd91 100644
--- a/docs-website/sidebars.js
+++ b/docs-website/sidebars.js
@@ -82,6 +82,7 @@ module.exports = {
           "metadata-ingestion/as-a-library",
           "metadata-integration/java/as-a-library",
           "metadata-ingestion/integration_docs/great-expectations",
+          "metadata-integration/java/datahub-protobuf/README",
         ],
       },
       {
diff --git a/docs-website/src/components/Logos.js b/docs-website/src/components/Logos.js
index 2471296ec6968..04fc393425380 100644
--- a/docs-website/src/components/Logos.js
+++ b/docs-website/src/components/Logos.js
@@ -115,6 +115,11 @@ const companiesByIndustry = [
   {
     name: "And More",
     companies: [
+      {
+        name: "Cabify",
+        imageUrl: "/img/logos/companies/cabify.png",
+        size: "large",
+      },
       {
         name: "Viasat",
         imageUrl: "/img/logos/companies/viasat.png",
diff --git a/docs-website/static/img/logos/companies/cabify.png b/docs-website/static/img/logos/companies/cabify.png
new file mode 100755
index 0000000000000..fcae0616b2474
Binary files /dev/null and b/docs-website/static/img/logos/companies/cabify.png differ
diff --git a/docs/cli.md b/docs/cli.md
index 6bdc8fd0919c0..152b229701d02 100644
--- a/docs/cli.md
+++ b/docs/cli.md
@@ -138,7 +138,7 @@ Commands:
   check      Helper commands for checking various aspects of DataHub.
   delete     Delete metadata from datahub using a single urn or a combination of filters
   docker     Helper commands for setting up and interacting with a local DataHub instance using Docker.
-  get        Get metadata for an entity with an optional list of aspects to project
+  get        Get metadata for an entity with an optional list of aspects to project.
   ingest     Ingest metadata into DataHub.
   init       Configure which datahub instance to connect to
   put        Update a single aspect of an entity
@@ -199,7 +199,7 @@ datahub delete --urn "urn:li:dataset:(urn:li:dataPlatform:hive,SampleHiveDataset
 
 ### get
 
-The `get` command allows you to easily retrieve metadata from DataHub, by using the REST API.
+The `get` command allows you to easily retrieve metadata from DataHub, by using the REST API. This works for both versioned aspects and timeseries aspects. For timeseries aspects, it fetches the latest value.
 For example the following command gets the ownership aspect from the dataset `urn:li:dataset:(urn:li:dataPlatform:hive,SampleHiveDataset,PROD)`
 
 ```console
@@ -338,4 +338,4 @@ datahub timeline --urn "urn:li:dataset:(urn:li:dataPlatform:mysql,User.UserAccou
 	MODIFY TAG dataset:mysql:User.UserAccount : A change in aspect editableSchemaMetadata happened at time 2022-02-17 20:03:42.0
 2022-02-17 14:17:30 - 0.0.0-computed
 	MODIFY TAG dataset:mysql:User.UserAccount : A change in aspect editableSchemaMetadata happened at time 2022-02-17 20:17:30.118
-```
\ No newline at end of file
+```
diff --git a/docs/faq.md b/docs/faq.md
deleted file mode 100644
index 40902f7629383..0000000000000
--- a/docs/faq.md
+++ /dev/null
@@ -1,97 +0,0 @@
-# DataHub FAQs
-
-## Why should we use DataHub?
-DataHub is a self-service data portal which provides search and discovery capabilities (and more!) over the data assets of an organization. This tool can help improve productivity of data scientists, data analysts, engineers of organizations dealing with massive amounts of data. Also, the regulatory environment (GDPR, CCPA etc) requires a company to know what data it has, who is using it and how long it will be retained - DataHub provides a solution to these challenges by gathering metadata across a distributed data ecosystem and surfacing it as a data catalog thereby easing the burden of data privacy/compliance.
-
-## Would you recommend DataHub rather than available commercial solutions?
-Common problems with commercial solutions can be summarized as:
-- Lacks direct access to source code: Any feature gaps can only be closed by external parties, which can be both time consuming and expensive.
-- Dependency on larger proprietary systems or environments, e.g. AWS, Azure, Cloudera etc., making it infeasible to adopt if it doesn’t fit your environment.
-- Expensive to acquire, integrate and operate
-- Vendor Lock-in
-
-DataHub can be right for you if you want an open source unbundled solution (front-end application completely decoupled from a “battle-tested” metadata store), that you are free to modify, extend and integrate with your data ecosystem. In our experience at LinkedIn and talking to other companies in a similar situation, metadata always has a very company specific implementation and meaning. Commercial tools will typically drop-in and solve a few use-cases well out of the gate, but will need much more investment or will be impossible to extend for some specific kinds of metadata.
-
-## Who are the major contributors in the community?
-Currently LinkedIn engineers. However, we’re receiving more and more PRs from individuals working at various companies.
-
-## How big is the community?
-Check out our adoption [here](https://datahubproject.io/docs/#adoption)
-
-## Is there a contributor selection criteria?
-We welcome contributions from everyone in the community. Please read our [contributing guidelines](CONTRIBUTING.md). In general, we will review PRs with the same rigor as our internal code review process to maintain overall quality.
-
-## How does the project engage with the community?
-We organize public town hall meetings at a monthly cadence. We use [Slack](https://slack.datahubproject.io) as one of the main ways to support the community.
-
-## Should this be the platform we decide upon, we’d like to fully engage and work with LinkedIn and the community. What’s the best way and what level of engagement/involvement should we expect?
-The best way to engage is through the [Slack channel](https://slack.datahubproject.io). You’ll get to interact with the developers and the community. It is a vibrant community and most questions are answered within a few hours by the community. 
-
-For reproducible technical issues, bugs and code contributions, Github [issues](https://github.com/linkedin/datahub/issues) and [PRs](https://github.com/linkedin/datahub/pulls) are the preferred channel.
-
-## What’s the best way to ramp up the product knowledge to properly test and evaluate DataHub?
-[The docs](https://datahubproject.io/docs) are the best resource. We have documented the steps to install and test DataHub thoroughly there.
-
-The [DataHub Introduction](https://engineering.linkedin.com/blog/2019/data-hub) and [Open Sourcing Datahub](https://engineering.linkedin.com/blog/2020/open-sourcing-datahub--linkedins-metadata-search-and-discovery-p) blog posts are also useful resources for getting a high level understanding of the system.
-
-## Where can I learn about the roadmap?
-You can learn more about DataHub's [product roadmap](roadmap.md), which gets updated regularly.
-
-## Where can I learn about the current list of features/functionalities?
-You can learn more about the current [list of features](features.md).
-
-## Are the product strategy/vision/roadmap driven by the LinkedIn Engineering team, community, or a collaborative effort?
-Mixed of both LinkedIn DataHub team and the community. The roadmap will be a joint effort of both LinkedIn and the community. However, we’ll most likely prioritize tasks that align with the community's asks.
-
-## Does DataHub connect with Google Cloud Platform?
-LinkedIn is not using GCP so we cannot commit to building and testing that connectivity. However, we’ll be happy to accept community contributions for GCP integration. Also, our Slack channel and regularly scheduled town hall meetings are a good opportunity to meet with people from different companies who have similar requirements and might be interested in collaborating on these features.
-
-## How approachable would LinkedIn be to provide insights/support or collaborate on a functionality?
-Please take a look at our [roadmap](roadmap.md) & [features](features.md) to get a sense of what’s being open sourced in the near future. If there’s something missing from the list, we’re open to discussion. In fact, the town hall would be the perfect venue for such discussions.
-
-## How do LinkedIn Engineering team and the community ensure the quality of the community code for DataHub?
-All PRs are reviewed by the LinkedIn team. Any extension/contribution coming from the community which LinkedIn team doesn’t have any expertise on will be placed into a incuation directory first (`/contrib`). Once it’s blessed and adopted by the community, we’ll graduate it from incubation and move it into the main code base.
-
-In the beginning, LinkedIn will play a big role in terms of stewardship of the code base. We’ll re-evaluate this strategy based on the amount of engagement from the community. There is a large backlog of features that we have only for the internal DataHub. We are going through the effort of generalizing and open sourcing these features. This will lead to batches of large commits from LinkedIn in the near future until the two code bases get closely aligned. See our [blog post](https://engineering.linkedin.com/blog/2020/open-sourcing-datahub--linkedins-metadata-search-and-discovery-p) for more details.
-
-## How are the ingestion ETL processes scheduled at LinkedIn?
-It varies depending on the data platform. HDFS, MySQL, Oracle, Teradata, and LDAP are scheduled on a daily basis. We rely on real-time pushs to ingest from sveral data platforms such as Hive, Presto, Kafka, Pinot, [Espresso](https://engineering.linkedin.com/espresso/introducing-espresso-linkedins-hot-new-distributed-document-store), [Ambry](https://github.com/linkedin/ambry), [Galene](https://engineering.linkedin.com/search/did-you-mean-galene), [Venice](https://engineering.linkedin.com/blog/2017/02/building-venice-with-apache-helix), and more.
-
-## What are the options for the Kafka Key for MCE, MAE and FailedMCE Topic?
-URN is the only sensible option to ensure events for the same entity land in the same parition and get processed in the chronological order.
-
-## How is the Data Quality of the ingestion messages handled?
-In addition to leverage the Kafka schema validation to ensure the MXEs output from metadata producer, we also actively monitor the ingestion streaming pipeline on the snapshot level with status.
-
-## Can you give a high level overview about how Data Governance is handled? What privacy/governance use cases are supported in LinkedIn through DataHub?
-This talk ([slides](https://www.slideshare.net/ShirshankaDas/taming-the-everevolving-compliance-beast-lessons-learnt-at-linkedin-strata-nyc-2017), [video](https://www.youtube.com/watch?v=O1DI0fuY8PM)) describes the role of metadata (DataHub) in the data governance/privacy space at LinkedIn. Field-level, dataset-level classification, governed data movement, automated data deletion, data export etc. are the supported use cases. We have plans to open source some of the compliance capabilities, listed as part of our [roadmap](roadmap.md).
-
-## When using Kafka and Confluent Schema Registry, does DataHub support multiple schemas for the same topic?
-You can [configure](https://docs.confluent.io/current/schema-registry/develop/api.html#compatibility) compatibility level per topic at Confluent Schema Registry. The default being used is “Backward”. So, you’re only allowed to make backward compatible changes on the topic schema. You can also change this configuration and flex compatibility check. However, as a best practice, we would suggest not doing backward incompatible changes on the topic schema because this will fail your old metadata producers’ flows. Instead, you might consider creating a new Kafka topic (new version).
-
-## How do we better document and map transformations within an ETL process? How do we create institutional knowledge and processes to help create a paradigm for tribal knowledge?
-We plan to add “fine-grain lineage” in the near future, which should cover the transformation documentation. DataHub currently has a simple “Docs” feature that allows capturing of tribal knowledge. We also plan to expand it significantly going forward.
-
-## How do we advance the product from a Data Catalog Browser to a Data Collaboration environment like Alation?
-We are adding some “social features” and documentation captures to DataHub. However, we do welcome the community to contribute in this area.
-
-## Can you share how the catalog looks in LinkedIn production?
-It’s very similar to what you see on the community version. We have added screenshots of the internal version of the catalog in our [blog post](https://engineering.linkedin.com/blog/2019/data-hub).
-
-## Does the roadmap have provision for capturing the Data Quality Information of the Dataset?
-We’re working on a similar [feature](https://engineering.linkedin.com/blog/2020/data-sentinel-automating-data-validation) internally. Will evaluate and update the roadmap once we have a better idea of the timeline.
-
-## Is DataHub capturing/showing column level [constraints](https://www.w3schools.com/sql/sql_constraints.asp) set at table definition?
-The [SchemaField](https://github.com/linkedin/datahub/blob/master/metadata-models/src/main/pegasus/com/linkedin/schema/SchemaField.pdl) model currently does not capture any property/field corresponding to constraints defined in the table definition. However, it should be fairly easy to extend the model to support that if needed.
-
-## How does DataHub manage extracting metadata from stores residing in different security zones?
-MCE is the ideal way to push metadata from different security zones, assuming there is a common Kafka infrastructure that aggregates the events from various security zones.
-
-## What all data stores does DataHub backend support presently?
-Currently, DataHub supports all major database providers that are supported by Ebean as the document store i.e. Oracle, Postgres, MySQL, H2. We also support [Espresso](https://engineering.linkedin.com/espresso/introducing-espresso-linkedins-hot-new-distributed-document-store), which is LinkedIn's proprietary document store. Other than that, we support Elasticsearch and Neo4j for search and graph use cases, respectively. However, as data stores in the backend are all abstracted and accessed through DAOs, you should be able to easily support other data stores by plugging in your own DAO implementations. Please refer to [Metadata Serving](architecture/metadata-serving.md) for more details.
-
-## For which stores, you have discovery services?
-Supported data sources are listed [here](../metadata-ingestion/README.md). It's also fairly easy to add your own sources.
-
-## How is metadata ingested in DataHub? Is it real-time?
-You can call the [rest.li](https://github.com/linkedin/rest.li) API to ingest metadata in DataHub directly instead of using Kafka event. Metadata ingestion is real-time if you're updating via rest.li API. It's near real-time in the case of Kafka events due to the asynchronous nature of Kafka processing.
diff --git a/docs/how/auth/sso/configure-oidc-react-azure.md b/docs/how/auth/sso/configure-oidc-react-azure.md
index a3518f2baca81..fb7ee4c8b5ea6 100644
--- a/docs/how/auth/sso/configure-oidc-react-azure.md
+++ b/docs/how/auth/sso/configure-oidc-react-azure.md
@@ -85,7 +85,7 @@ AUTH_OIDC_CLIENT_ID=your-client-id
 AUTH_OIDC_CLIENT_SECRET=your-client-secret
 AUTH_OIDC_DISCOVERY_URI=https://login.microsoftonline.com/{tenant ID}/v2.0/.well-known/openid-configuration
 AUTH_OIDC_BASE_URL=your-datahub-url
-AUTH_OIDC_SCOPE="openid profile email groups"
+AUTH_OIDC_SCOPE="openid profile email"
 ```
 
 Replacing the placeholders above with the client id (step 5), client secret (step 3) and tenant ID (step 6) received from Microsoft Azure.
@@ -101,4 +101,4 @@ docker-compose -p datahub -f docker-compose.yml -f docker-compose.override.yml
 Navigate to your DataHub domain to see SSO in action.
 
 ## Resources
-- [OAuth 2.0 and OpenID Connect Overview](https://developer.okta.com/docs/concepts/oauth-openid/)
\ No newline at end of file
+- [Microsoft identity platform and OpenID Connect protocol](https://docs.microsoft.com/en-us/azure/active-directory/develop/v2-protocols-oidc/)
\ No newline at end of file
diff --git a/docs/how/auth/sso/configure-oidc-react.md b/docs/how/auth/sso/configure-oidc-react.md
index 8de70fb8c5e40..228f35c1d62f1 100644
--- a/docs/how/auth/sso/configure-oidc-react.md
+++ b/docs/how/auth/sso/configure-oidc-react.md
@@ -143,15 +143,13 @@ The default provisioning behavior can be customized using the following configs.
 # User and groups provisioning
 AUTH_OIDC_JIT_PROVISIONING_ENABLED=true
 AUTH_OIDC_PRE_PROVISIONING_REQUIRED=false
-AUTH_OIDC_EXTRACT_GROUPS_ENABLED=true
+AUTH_OIDC_EXTRACT_GROUPS_ENABLED=false
 AUTH_OIDC_GROUPS_CLAIM=<your-groups-claim-name>
 ```
 
 - `AUTH_OIDC_JIT_PROVISIONING_ENABLED`: Whether DataHub users & groups should be provisioned on login if they do not exist. Defaults to true.
 - `AUTH_OIDC_PRE_PROVISIONING_REQUIRED`: Whether the user should already exist in DataHub when they login, failing login if they are not. This is appropriate for situations in which users and groups are batch ingested and tightly controlled inside your environment. Defaults to false.
-  the userNameClaim field will contain an email address, and we want to omit the domain name suffix of the email, we can specify a custom
-  regex to do so. (e.g. `([^@]+)`)
-- `AUTH_OIDC_EXTRACT_GROUPS_ENABLED`: Only applies if `AUTH_OIDC_JIT_PROVISIONING_ENABLED` is set to true. This determines whether we should attempt to extract a list of group names from a particular claim in the OIDC attributes. Defaults to true. 
+- `AUTH_OIDC_EXTRACT_GROUPS_ENABLED`: Only applies if `AUTH_OIDC_JIT_PROVISIONING_ENABLED` is set to true. This determines whether we should attempt to extract a list of group names from a particular claim in the OIDC attributes. Note that if this is enabled, each login will re-sync group membership with the groups in your Identity Provider, clearing the group membership that has been assigned through the DataHub UI. Enable with care! Defaults to false.
 - `AUTH_OIDC_GROUPS_CLAIM`: Only applies if `AUTH_OIDC_EXTRACT_GROUPS_ENABLED` is set to true. This determines which OIDC claim will contain a list of string group names. Defaults to 'groups'
 
 
diff --git a/docs/how/delete-metadata.md b/docs/how/delete-metadata.md
index a72a49666e80f..1f9ec42268a06 100644
--- a/docs/how/delete-metadata.md
+++ b/docs/how/delete-metadata.md
@@ -99,3 +99,12 @@ datahub ingest rollback --run-id <run-id>
 ```
 
 to rollback all aspects added with this run and all entities created by this run.
+
+:::note
+
+Since datahub v0.8.29, the `rollback` endpoint will now perform a *soft delete* of the entities ingested by a given run `<run-id>`. 
+This was done to preserve potential changes that were made directly via DataHub's UI and not part of the ingestion run itself. Such that this information can be retrieved later on if a re-ingestion for the same deleted entity is done.
+
+If you wish to keep old behaviour (hard delete), please use the `--hard-delete` flag (short-hand: `-d`).
+
+:::
diff --git a/docs/modeling/extending-the-metadata-model.md b/docs/modeling/extending-the-metadata-model.md
index 06cea78ae30ad..872ca665c1982 100644
--- a/docs/modeling/extending-the-metadata-model.md
+++ b/docs/modeling/extending-the-metadata-model.md
@@ -20,13 +20,12 @@ We will refer to the two options as the **open-source fork** and **custom reposi
 ## This Guide
 
 This guide will outline what the experience of adding a new Entity should look like through a real example of adding the
-Dashboard Entity. If you want to extend an existing Entity, you can skip directly to [Step 4](#step_4).
+Dashboard Entity. If you want to extend an existing Entity, you can skip directly to [Step 3](#step_3).
 
 At a high level, an entity is made up of:
 
-1. A Key Aspect, which uniquely identifies an instance of an entity,
-2. A Snapshot, which combines the key aspect with a serialized key, or urn.
-3. Additional bundles of Aspects, groups of related attributes that are attached to an entity.
+1. A Key Aspect: Uniquely identifies an instance of an entity,
+2. A list of specified Aspects, groups of related attributes that are attached to an entity.
 
 ## Defining an Entity
 
@@ -79,74 +78,29 @@ Because they are aspects, keys need to be annotated with an @Aspect annotation,
 can be a part of.
 
 The key can also be annotated with the two index annotations: @Relationship and @Searchable. This instructs DataHub
-infra to use the fields in the key to create relationships and index fields for search. See [Step 4](#step_4) for more details on
+infra to use the fields in the key to create relationships and index fields for search. See [Step 3](#step_3) for more details on
 the annotation model.
 
 **Constraints**: Note that each field in a Key Aspect MUST be of String or Enum type.
 
 ### <a name="step_2"></a>Step 2: Create the new entity with its key aspect
 
-Create an Aspect union to define the key aspects an Entity is associated with. An aspect represents a related
-record of metadata about an entity. Any record appearing in the Union should be annotated with @Aspect. In this example below, we take the DashboardKey aspect and the BrowsePaths aspect and include them in this `DashboardAspect` union. For all other aspects, we will utilize a much more flexible way of attaching aspects to entities using a `yaml` file.
+Define the entity within an `entity-registry.yml` file. Depending on your approach, the location of this file may vary. More on that in steps [4](#step_4) and [5](#step_5).
 
-```
-namespace com.linkedin.metadata.aspect
-
-import com.linkedin.metadata.key.DashboardKey
-import com.linkedin.common.BrowsePaths
-
-/**
- * A union of all supported metadata aspects for a Dashboard
- */
-typeref DashboardAspect = union[
-    DashboardKey,
-    BrowsePaths
-]
-```
-
-The first aspect will be by convention the Entity’s key aspect.
-Previously, you were required to add all aspects for the entity into this Aspect union. You will see examples of this pattern throughout the code-base (e.g. `DatasetAspect`, `DashboardAspect` etc.). This is no longer required.
-
-### <a name="step_3"></a>Step 3: Define the Entity Snapshot
-
-The snapshot describes the format of how an entity is serialized (as a single snapshot record) for read and write operations to DataHub's metadata service (f.k.a. GMS). All snapshots have two fields:
-- `urn` of type `Urn`
-- `snapshot` of type `union[Aspect1, Aspect2, ...]`.
-
-The snapshot needs an `@Entity` annotation with the entity’s name. The name is used for specifying entity type when
-searching, using autocomplete, etc.
-
-```
-namespace com.linkedin.metadata.snapshot
-
-import com.linkedin.common.DashboardUrn
-import com.linkedin.metadata.aspect.DashboardAspect
-
-/**
- * A metadata snapshot for a specific Dashboard entity.
- */
-@Entity = {
-  "name": "dashboard"
-}
-record DashboardSnapshot {
-
-  /**
-   * URN for the entity the metadata snapshot is associated with.
-   */
-  urn: DashboardUrn
-
-  /**
-   * The list of metadata aspects associated with the dashboard.
-   */
-  aspects: array[DashboardAspect]
-}
+Example:
+```yaml
+  - name: dashboard
+    doc: A container of related data assets.
+    keyAspect: dashboardKey
 ```
 
-If you're extending an existing Entity, you can skip this step and go straight to step 4.
+- name: The entity name/type, this will be present as a part of the Urn.
+- doc: A brief description of the entity.
+- keyAspect: The name of the Key Aspect defined in step 1. This name must match the value in the PDL annotation.
 
 #
 
-### <a name="step_4"></a>Step 4: Define custom aspects or attach existing aspects to your entity
+### <a name="step_3"></a>Step 3: Define custom aspects or attach existing aspects to your entity
 
 Some aspects, like Ownership and GlobalTags, are reusable across entities. They can be included in an entity’s set of
 aspects freely. To include attributes that are not included in an existing Aspect, a new Aspect must be created.
@@ -238,14 +192,14 @@ The Aspect has four key components: its properties, the @Aspect annotation, the
   can be defined as PDL primitives, enums, records, or collections (
   see [pdl schema documentation](https://linkedin.github.io/rest.li/pdl_schema))
   references to other entities, of type Urn or optionally `<Entity>Urn`
-- **@Aspect annotation**: Declares record is an Aspect and includes it in an entity’s Snapshot. Unlike the following 
+- **@Aspect annotation**: Declares record is an Aspect and includes it when serializing an entity. Unlike the following 
   two annotations, @Aspect is applied to the entire record, rather than a specific field.  Note, you can mark an aspect 
   as a timeseries aspect. Check out this [doc](metadata-model.md#timeseries-aspects) for details.
 - **@Searchable annotation**: This annotation can be applied to any primitive field or a map field to indicate that it
   should be indexed in Elasticsearch and can be searched on. For a complete guide on using the search annotation, see
   the annotation docs further down in this document.
-- **@Relationship annotation**: These annotations create edges between the Snapshot’s Urn and the destination of the
-  annotated field when the snapshots are ingested. @Relationship annotations must be applied to fields of type Urn. In
+- **@Relationship annotation**: These annotations create edges between the Entity’s Urn and the destination of the
+  annotated field when the entities are ingested. @Relationship annotations must be applied to fields of type Urn. In
   the case of DashboardInfo, the `charts` field is an Array of Urns. The @Relationship annotation cannot be applied
   directly to an array of Urns. That’s why you see the use of an Annotation override (`”/*”:) to apply the @Relationship
   annotation to the Urn directly. Read more about overrides in the annotation docs further down on this page.
@@ -254,14 +208,14 @@ After you create your Aspect, you need to attach to all the entities that it app
 
 **Constraints**: Note that all aspects MUST be of type Record.
 
-### <a name="step_5"></a> Step 5: Choose a place to store your model extension
+### <a name="step_4"></a> Step 4: Choose a place to store your model extension
 
 At the beginning of this document, we walked you through a flow-chart that should help you decide whether you need to maintain a fork of the open source DataHub repo for your model extensions, or whether you can just use a model extension repository that can stay independent of the DataHub repo. Depending on what path you took, the place you store your aspect model files (the .pdl files) and the entity-registry files (the yaml file called `entity-registry.yaml` or `entity-registry.yml`) will vary.
 
-- Open source Fork: Aspect files go under [`metadata-models`](../../metadata-models) module in the main repo, entity registry goes into [`metadata-models/src/main/resources/entity-registry.yml`](../../metadata-models/src/main/resources/entity-registry.yml). Read on for more details in [Step 6](#step_6).
+- Open source Fork: Aspect files go under [`metadata-models`](../../metadata-models) module in the main repo, entity registry goes into [`metadata-models/src/main/resources/entity-registry.yml`](../../metadata-models/src/main/resources/entity-registry.yml). Read on for more details in [Step 5](#step_5).
 - Custom repository: Read the [metadata-models-custom](../../metadata-models-custom/README.md) documentation to learn how to store and version your aspect models and registry.
 
-### <a name="step_6"></a>Step 6: Attaching your non-key Aspect(s) to the Entity
+### <a name="step_5"></a>Step 5: Attaching your non-key Aspect(s) to the Entity
 
 Attaching non-key aspects to an entity can be done simply by adding them to the entity registry yaml file. The location of this file differs based on whether you are following the oss-fork path or the custom-repository path.
 
@@ -270,12 +224,15 @@ Here is an minimal example of adding our new `DashboardInfo` aspect to the `Dash
 ```yaml
 entities:
    - name: dashboard
+   - keyAspect: dashBoardKey
    aspects:
      # the name of the aspect must be the same as that on the @Aspect annotation on the class
      - dashboardInfo  
 ```
 
-### <a name="step_7"></a>Step 7 (Oss-Fork approach): Re-build DataHub to have access to your new or updated entity
+Previously, you were required to add all aspects for the entity into an Aspect union. You will see examples of this pattern throughout the code-base (e.g. `DatasetAspect`, `DashboardAspect` etc.). This is no longer required.
+
+### <a name="step_6"></a>Step 6 (Oss-Fork approach): Re-build DataHub to have access to your new or updated entity
 
 If you opted for the open-source fork approach, where you are editing models in the `metadata-models` repository of DataHub, you will need to re-build the DataHub metadata service using the steps below. If you are following the custom model repository approach, you just need to build your custom model repository and deploy it to a running metadata service instance to read and write metadata using your new model extensions.
 
@@ -288,16 +245,16 @@ before running `build`.
 Then, run `./gradlew build` from the repository root to rebuild Datahub with access to your new entity.
 
 Then, re-deploy metadata-service (gms), and mae-consumer and mce-consumer (optionally if you are running them unbundled). See [docker development](../../docker/README.md) for details on how
-to deploy during development. This will allow Datahub to read and write Snapshots of your new entity or extensions to existing entities, along with serving search and graph queries for that entity type.
+to deploy during development. This will allow Datahub to read and write your new entity or extensions to existing entities, along with serving search and graph queries for that entity type.
 
-To emit snapshots to ingest from the Datahub CLI tool, first install datahub cli
+To emit proposals to ingest from the Datahub CLI tool, first install datahub cli
 locally [following the instructions here](../../metadata-ingestion/developing.md). `./gradlew build` generated the avro
 schemas your local ingestion cli tool uses earlier. After following the developing guide, you should be able to emit
 your new event using the local datahub cli.
 
 Now you are ready to start ingesting metadata for your new entity!
 
-### <a name="step_8"></a>(Optional) Step 8: Extend the DataHub frontend to view your entity in GraphQL & React
+### <a name="step_7"></a>(Optional) Step 7: Extend the DataHub frontend to view your entity in GraphQL & React
 
 If you are extending an entity with additional aspects, and you can use the auto-render specifications to automatically render these aspects to your satisfaction, you do not need to write any custom code.
 
@@ -310,6 +267,7 @@ There are four core annotations that DataHub recognizes:
 
 #### @Entity
 
+**Legacy**
 This annotation is applied to each Entity Snapshot record, such as DashboardSnapshot.pdl. Each one that is included in
 the root Snapshot.pdl model must have this annotation.
 
@@ -329,7 +287,7 @@ It takes the following parameters:
 #### @Aspect
 
 This annotation is applied to each Aspect record, such as DashboardInfo.pdl. Each aspect that is included in an entity’s
-AspectUnion must have this annotation.
+set of aspects in the `entity-registry.yml` must have this annotation.
 
 It takes the following parameters:
 
@@ -468,7 +426,7 @@ record Owner {
 }
 ```
 
-This annotation says that when we ingest an Entity with an OwnershipAspect, DataHub will create an OwnedBy relationship
+This annotation says that when we ingest an Entity with an Ownership Aspect, DataHub will create an OwnedBy relationship
 between that entity and the CorpUser or CorpGroup who owns it. This will be queryable using the Relationships resource
 in both the forward and inverse directions.
 
diff --git a/entity-registry/build.gradle b/entity-registry/build.gradle
index 78590ee816736..986ffc7b42d7f 100644
--- a/entity-registry/build.gradle
+++ b/entity-registry/build.gradle
@@ -15,6 +15,8 @@ dependencies {
   compile externalDependency.mavenArtifact
 
   testCompile externalDependency.testng
+  testCompile externalDependency.mockito
+  testCompile externalDependency.mockitoInline
   testCompile project(':test-models')
 
 }
diff --git a/metadata-io/src/main/java/com/linkedin/metadata/graph/LineageRegistry.java b/entity-registry/src/main/java/com/linkedin/metadata/models/registry/LineageRegistry.java
similarity index 97%
rename from metadata-io/src/main/java/com/linkedin/metadata/graph/LineageRegistry.java
rename to entity-registry/src/main/java/com/linkedin/metadata/models/registry/LineageRegistry.java
index 16066e2480d3d..922ce6a81928b 100644
--- a/metadata-io/src/main/java/com/linkedin/metadata/graph/LineageRegistry.java
+++ b/entity-registry/src/main/java/com/linkedin/metadata/models/registry/LineageRegistry.java
@@ -1,7 +1,7 @@
-package com.linkedin.metadata.graph;
+package com.linkedin.metadata.models.registry;
 
+import com.linkedin.metadata.graph.LineageDirection;
 import com.linkedin.metadata.models.annotation.RelationshipAnnotation;
-import com.linkedin.metadata.models.registry.EntityRegistry;
 import com.linkedin.metadata.query.filter.RelationshipDirection;
 import java.util.ArrayList;
 import java.util.Collection;
@@ -91,7 +91,7 @@ public List<EdgeInfo> getLineageRelationships(String entityName, LineageDirectio
     if (spec == null) {
       return Collections.emptyList();
     }
-    
+
     if (direction == LineageDirection.UPSTREAM) {
       return spec.getUpstreamEdges();
     }
diff --git a/metadata-io/src/test/java/com/linkedin/metadata/graph/LineageRegistryTest.java b/entity-registry/src/test/java/com/linkedin/metadata/models/registry/LineageRegistryTest.java
similarity index 97%
rename from metadata-io/src/test/java/com/linkedin/metadata/graph/LineageRegistryTest.java
rename to entity-registry/src/test/java/com/linkedin/metadata/models/registry/LineageRegistryTest.java
index db3ca3be537d9..15f3df0fcd27d 100644
--- a/metadata-io/src/test/java/com/linkedin/metadata/graph/LineageRegistryTest.java
+++ b/entity-registry/src/test/java/com/linkedin/metadata/models/registry/LineageRegistryTest.java
@@ -1,10 +1,9 @@
-package com.linkedin.metadata.graph;
+package com.linkedin.metadata.models.registry;
 
 import com.google.common.collect.ImmutableList;
 import com.linkedin.metadata.models.EntitySpec;
 import com.linkedin.metadata.models.RelationshipFieldSpec;
 import com.linkedin.metadata.models.annotation.RelationshipAnnotation;
-import com.linkedin.metadata.models.registry.EntityRegistry;
 import com.linkedin.metadata.query.filter.RelationshipDirection;
 import java.util.Collections;
 import java.util.HashMap;
diff --git a/gradle.properties b/gradle.properties
index 66edd9e705008..e18585c793376 100644
--- a/gradle.properties
+++ b/gradle.properties
@@ -3,6 +3,9 @@ org.gradle.configureondemand=true
 org.gradle.parallel=true
 org.gradle.caching=false
 
+# Increase gradle JVM memory to 2GB to allow tests to run locally
+#org.gradle.jvmargs=-Xmx2000m
+
 # Increase retries to 5 (from default of 3) and increase interval from 125ms to 1s.
 org.gradle.internal.repository.max.retries=5
 org.gradle.internal.repository.max.tentatives=5
diff --git a/metadata-ingestion-examples/README.md b/metadata-ingestion-examples/README.md
index 0565c427bf657..d19438481f562 100644
--- a/metadata-ingestion-examples/README.md
+++ b/metadata-ingestion-examples/README.md
@@ -1,5 +1,8 @@
 # Metadata Ingestion
 
+**LEGACY**
+This is a legacy module. The examples here are not actively maintained and may not work as described. Please see the `metadata-ingestion` module for more up-to-date uses.
+
 This directory contains example apps for ingesting data into DataHub.
 
 You are more than welcome to use these examples directly, or use them as a reference for you own jobs.
diff --git a/metadata-ingestion/examples/mce_files/bootstrap_mce.json b/metadata-ingestion/examples/mce_files/bootstrap_mce.json
index 79b02447fc0d3..a73196d208c73 100644
--- a/metadata-ingestion/examples/mce_files/bootstrap_mce.json
+++ b/metadata-ingestion/examples/mce_files/bootstrap_mce.json
@@ -3139,5 +3139,101 @@
       "registryVersion": null,
       "properties": null
     }
+  },
+  {
+    "auditHeader":null,
+    "entityType":"container",
+    "entityUrn": "urn:li:container:DATABASE",
+    "changeType":"UPSERT",
+    "aspectName":"containerProperties",
+    "aspect":{
+      "value":"{\"name\": \"datahub_db\", \"description\": \"The Root DataHub DB\" }",
+      "contentType":"application/json"
+    },
+    "systemMetadata":null
+  },
+  {
+    "auditHeader":null,
+    "entityType":"container",
+    "entityUrn": "urn:li:container:DATABASE",
+    "changeType":"UPSERT",
+    "aspectName":"subTypes",
+    "aspect":{
+      "value":"{\"typeNames\": [ \"Database\" ] }",
+      "contentType":"application/json"
+    },
+    "systemMetadata":null
+  },
+  {
+    "auditHeader":null,
+    "entityType":"container",
+    "entityUrn": "urn:li:container:DATABASE",
+    "changeType":"UPSERT",
+    "aspectName":"dataPlatformInstance",
+    "aspect":{
+      "value":"{\"platform\": \"urn:li:dataPlatform:hive\" }",
+      "contentType":"application/json"
+    },
+    "systemMetadata":null
+  },
+  {
+    "auditHeader":null,
+    "entityType":"container",
+    "entityUrn": "urn:li:container:SCHEMA",
+    "changeType":"UPSERT",
+    "aspectName":"containerProperties",
+    "aspect":{
+      "value":"{\"name\": \"datahub_schema\", \"description\": \"The Root DataHub Schema\" }",
+      "contentType":"application/json"
+    },
+    "systemMetadata":null
+  },
+  {
+    "auditHeader":null,
+    "entityType":"container",
+    "entityUrn": "urn:li:container:SCHEMA",
+    "changeType":"UPSERT",
+    "aspectName":"dataPlatformInstance",
+    "aspect":{
+      "value":"{\"platform\": \"urn:li:dataPlatform:hive\" }",
+      "contentType":"application/json"
+    },
+    "systemMetadata":null
+  },
+  {
+    "auditHeader":null,
+    "entityType":"container",
+    "entityUrn": "urn:li:container:SCHEMA",
+    "changeType":"UPSERT",
+    "aspectName":"subTypes",
+    "aspect":{
+      "value":"{\"typeNames\": [ \"Schema\" ] }",
+      "contentType":"application/json"
+    },
+    "systemMetadata":null
+  },
+  {
+    "auditHeader":null,
+    "entityType":"dataset",
+    "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:hive,SampleHiveDataset,PROD)",
+    "changeType":"UPSERT",
+    "aspectName":"container",
+    "aspect":{
+      "value":"{\"container\": \"urn:li:container:SCHEMA\" }",
+      "contentType":"application/json"
+    },
+    "systemMetadata":null
+  },
+  {
+    "auditHeader":null,
+    "entityType":"container",
+    "entityUrn": "urn:li:container:SCHEMA",
+    "changeType":"UPSERT",
+    "aspectName":"container",
+    "aspect":{
+      "value":"{\"container\": \"urn:li:container:DATABASE\" }",
+      "contentType":"application/json"
+    },
+    "systemMetadata":null
   }
 ]
diff --git a/metadata-ingestion/examples/mce_files/data_platforms.json b/metadata-ingestion/examples/mce_files/data_platforms.json
index b131ac4c62bb9..0de966f7345d1 100644
--- a/metadata-ingestion/examples/mce_files/data_platforms.json
+++ b/metadata-ingestion/examples/mce_files/data_platforms.json
@@ -396,6 +396,46 @@
     },
     "proposedDelta": null
   },
+  {
+    "auditHeader": null,
+    "proposedSnapshot": {
+      "com.linkedin.pegasus2avro.metadata.snapshot.DataPlatformSnapshot": {
+        "urn": "urn:li:dataPlatform:superset",
+        "aspects": [
+          {
+            "com.linkedin.pegasus2avro.dataplatform.DataPlatformInfo": {
+              "datasetNameDelimiter": ".",
+              "name": "superset",
+              "displayName": "Superset",
+              "type": "RELATIONAL_DB",
+              "logoUrl": "/assets/platforms/supersetlogo.png"
+            }
+          }
+        ]
+      }
+    },
+    "proposedDelta": null
+  },
+  {
+    "auditHeader": null,
+    "proposedSnapshot": {
+      "com.linkedin.pegasus2avro.metadata.snapshot.DataPlatformSnapshot": {
+        "urn": "urn:li:dataPlatform:datahub",
+        "aspects": [
+          {
+            "com.linkedin.pegasus2avro.dataplatform.DataPlatformInfo": {
+              "datasetNameDelimiter": ".",
+              "name": "datahub",
+              "displayName": "DataHub",
+              "type": "RELATIONAL_DB",
+              "logoUrl": "/assets/platforms/datahublogo.png"
+            }
+          }
+        ]
+      }
+    },
+    "proposedDelta": null
+  },
   {
     "auditHeader": null,
     "proposedSnapshot": {
@@ -635,5 +675,24 @@
       }
     },
     "proposedDelta": null
+  },
+  {
+    "auditHeader": null,
+    "proposedSnapshot": {
+      "com.linkedin.pegasus2avro.metadata.snapshot.DataPlatformSnapshot": {
+        "urn": "urn:li:dataPlatform:unknown",
+        "aspects": [
+          {
+            "com.linkedin.pegasus2avro.dataplatform.DataPlatformInfo": {
+              "name": "Unknown Platform",
+              "displayName": "N/A",
+              "type": "OTHERS",
+              "datasetNameDelimiter": "."
+            }
+          }
+        ]
+      }
+    },
+    "proposedDelta": null
   }
-]
+]
\ No newline at end of file
diff --git a/metadata-ingestion/integration_docs/great-expectations.md b/metadata-ingestion/integration_docs/great-expectations.md
index b017d4bffac69..6d8b475cc838c 100644
--- a/metadata-ingestion/integration_docs/great-expectations.md
+++ b/metadata-ingestion/integration_docs/great-expectations.md
@@ -47,7 +47,8 @@ This integration does not support
     - `retry_status_codes` (optional): Retry HTTP request also on these status codes.
     - `retry_max_times` (optional): Maximum times to retry if HTTP request fails. The delay between retries is increased exponentially.
     - `extra_headers` (optional): Extra headers which will be added to the datahub request.
-   
+    - `parse_table_names_from_sql` (defaults to false): The integration can use an SQL parser to try to parse the datasets being asserted. This parsing is disabled by default, but can be enabled by setting `parse_table_names_from_sql: True`.  The parser is based on the [`sqllineage`](https://pypi.org/project/sqllineage/) package.
+    
 ## Learn more
 
 To see the Great Expectations in action, check out [this demo](https://www.loom.com/share/d781c9f0b270477fb5d6b0c26ef7f22d) from the Feb 2022 townhall. 
\ No newline at end of file
diff --git a/metadata-ingestion/setup.py b/metadata-ingestion/setup.py
index be0d0aa7f7a0b..3044bb5564720 100644
--- a/metadata-ingestion/setup.py
+++ b/metadata-ingestion/setup.py
@@ -51,7 +51,8 @@ def get_long_description():
     # Markupsafe breaking change broke Jinja and some other libs
     # Pinning it to a version which works even though we are not using explicitly
     # https://github.com/aws/aws-sam-cli/issues/3661
-    "markupsafe==2.0.1",
+    # Airflow compatibility: https://github.com/apache/airflow/blob/2.2.2/setup.cfg#L125
+    "markupsafe>=1.1.1,<=2.0.1",
     "Deprecated",
     "types-Deprecated",
 }
@@ -211,6 +212,8 @@ def get_long_description():
     "types-click==0.1.12",
     "boto3-stubs[s3,glue,sagemaker]",
     "types-tabulate",
+    # avrogen package requires this
+    "types-pytz",
 }
 
 base_dev_requirements = {
diff --git a/metadata-ingestion/source_docs/azure-ad.md b/metadata-ingestion/source_docs/azure-ad.md
index 84c3cae0aa5f8..3caeaf5374f7a 100644
--- a/metadata-ingestion/source_docs/azure-ad.md
+++ b/metadata-ingestion/source_docs/azure-ad.md
@@ -119,7 +119,7 @@ Note that a `.` is used to denote nested fields in the YAML configuration block.
 | `client_secret`                     | string   | ✅           |       | Client secret. Found in your app registration on Azure AD Portal       |                                  
 | `redirect`                     | string   | ✅          |       | Redirect URI.  Found in your app registration on Azure AD Portal       |                                  
 | `authority`                     | string   | ✅          |       | The [authority](https://docs.microsoft.com/en-us/azure/active-directory/develop/msal-client-application-configuration) is a URL that indicates a directory that MSAL can request tokens from. |
-| `token_url`                     | string   | ✅          |       | The token URL that acquires a token from Azure AD for authorizing requests |
+| `token_url`                     | string   | ✅          |       | The token URL that acquires a token from Azure AD for authorizing requests.  This source will only work with v1.0 endpoint. |
 | `graph_url`                     | string   | ✅          |       | [Microsoft Graph API endpoint](https://docs.microsoft.com/en-us/graph/use-the-api)
 | `ingest_users`                     | bool   |          | `True`      | Whether users should be ingested into DataHub.                                                                  |
 | `ingest_groups`                    | bool   |          | `True`      | Whether groups should be ingested into DataHub.                                                                 |
diff --git a/metadata-ingestion/source_docs/bigquery.md b/metadata-ingestion/source_docs/bigquery.md
index b4a93a321f57b..75a92b862fc77 100644
--- a/metadata-ingestion/source_docs/bigquery.md
+++ b/metadata-ingestion/source_docs/bigquery.md
@@ -1,14 +1,17 @@
 # BigQuery
 
+To get all metadata from BigQuery you need to use two plugins `bigquery` and `bigquery-usage`. Both of them are described in this page. These will require 2 separate recipes. We understand this is not ideal and we plan to make this easier in the future.
+
 For context on getting started with ingestion, check out our [metadata ingestion guide](../README.md).
 
-## Setup
+## `bigquery`
+### Setup
 
 To install this plugin, run `pip install 'acryl-datahub[bigquery]'`.
 
-## Prerequisites
-### Create a datahub profile in GCP:
-1. Create a custom role for datahub (https://cloud.google.com/iam/docs/creating-custom-roles#creating_a_custom_role)
+### Prerequisites
+#### Create a datahub profile in GCP
+1. Create a custom role for datahub as per [BigQuery docs](https://cloud.google.com/iam/docs/creating-custom-roles#creating_a_custom_role)
 2. Grant the following permissions to this role:
 ```   
    bigquery.datasets.get
@@ -27,9 +30,9 @@ To install this plugin, run `pip install 'acryl-datahub[bigquery]'`.
    logging.logEntries.list # Needs for lineage generation
    resourcemanager.projects.get
 ```
-### Create a service account:
+#### Create a service account
 
-1. Setup a ServiceAccount (https://cloud.google.com/iam/docs/creating-managing-service-accounts#iam-service-accounts-create-console)
+1. Setup a ServiceAccount as per [BigQuery docs](https://cloud.google.com/iam/docs/creating-managing-service-accounts#iam-service-accounts-create-console)
 and assign the previously created role to this service account.
 2. Download a service account JSON keyfile.
    Example credential file:
@@ -64,7 +67,7 @@ and assign the previously created role to this service account.
        client_id: "123456678890"
 ```
 
-## Capabilities
+### Capabilities
 
 This plugin extracts the following:
 
@@ -81,11 +84,11 @@ This plugin extracts the following:
 
 :::tip
 
-You can also get fine-grained usage statistics for BigQuery using the `bigquery-usage` source described below.
+You can also get fine-grained usage statistics for BigQuery using the `bigquery-usage` source described [below](#bigquery-usage-plugin).
 
 :::
 
-## Quickstart recipe
+### Quickstart recipe
 
 Check out the following recipe to get started with ingestion! See [below](#config-details) for full configuration options.
 
@@ -102,7 +105,7 @@ sink:
   # sink configs
 ```
 
-## Config details
+### Config details
 
 Note that a `.` is used to denote nested fields in the YAML recipe.
 
@@ -155,7 +158,7 @@ Note: the bigquery_audit_metadata_datasets parameter receives a list of datasets
 
 Note: Since bigquery source also supports dataset level lineage, the auth client will require additional permissions to be able to access the google audit logs. Refer the permissions section in bigquery-usage section below which also accesses the audit logs.
 
-## Profiling
+### Profiling
 Profiling can profile normal/partitioned and sharded tables as well but due to performance reasons, we only profile the latest partition for Partitioned tables and the latest shard for sharded tables.
 
 If limit/offset parameter is set or partitioning partitioned or sharded table Great Expectation (the profiling framework we use) needs to create temporary
@@ -175,11 +178,11 @@ Due to performance reasons, we only profile the latest partition for Partitioned
 You can set partition explicitly with `partition.partition_datetime` property if you want. (partition will be applied to all partitioned tables)
 :::
 
-# BigQuery Usage Stats
+## `bigquery-usage`
 
 For context on getting started with ingestion, check out our [metadata ingestion guide](../README.md).
 
-## Setup
+### Setup
 
 To install this plugin, run `pip install 'acryl-datahub[bigquery-usage]'`.
 
@@ -194,7 +197,7 @@ The Google Identity must have one of the following OAuth scopes granted to it:
 
 And should be authorized on all projects you'd like to ingest usage stats from. 
 
-## Capabilities
+### Capabilities
 
 This plugin extracts the following:
 
@@ -208,7 +211,7 @@ This plugin extracts the following:
 
 :::
 
-## Quickstart recipe
+### Quickstart recipe
 
 Check out the following recipe to get started with ingestion! See [below](#config-details) for full configuration options.
 
@@ -230,7 +233,7 @@ sink:
   # sink configs
 ```
 
-## Config details
+### Config details
 
 Note that a `.` is used to denote nested fields in the YAML recipe.
 
diff --git a/metadata-ingestion/source_docs/looker.md b/metadata-ingestion/source_docs/looker.md
index 6c6cbc9649ef0..fb9dafb63ca6d 100644
--- a/metadata-ingestion/source_docs/looker.md
+++ b/metadata-ingestion/source_docs/looker.md
@@ -96,6 +96,7 @@ Note that a `.` is used to denote nested fields in the YAML recipe.
 | `view_browse_pattern` |   | `/{env}/{platform}/{project}/views/{name}` | Pattern for providing browse paths to views. Allowed variables are `{project}`, `{model}`, `{name}`, `{platform}` and `{env}` | 
 | `explore_naming_pattern` |   | `{model}.explore.{name}` | Pattern for providing dataset names to explores. Allowed variables are `{project}`, `{model}`, `{name}` | 
 | `explore_browse_pattern` |   | `/{env}/{platform}/{project}/explores/{model}.{name}` | Pattern for providing browse paths to explores. Allowed variables are `{project}`, `{model}`, `{name}`, `{platform}` and `{env}` | 
+| `transport_options`                                |          |        |  Populates the [TransportOptions](https://github.com/looker-open-source/sdk-codegen/blob/94d6047a0d52912ac082eb91616c1e7c379ab262/python/looker_sdk/rtl/transport.py#L70) struct for looker client |
 | `max_threads`                                |          | `os.cpuCount or 40` |  Max parallelism for Looker API calls                   |
 
 
diff --git a/metadata-ingestion/source_docs/lookml.md b/metadata-ingestion/source_docs/lookml.md
index a6bb970009760..f1208cee20256 100644
--- a/metadata-ingestion/source_docs/lookml.md
+++ b/metadata-ingestion/source_docs/lookml.md
@@ -109,6 +109,7 @@ Note that a `.` is used to denote nested fields in the YAML recipe.
 | `github_info.branch`            |          | `main` | The default branch in your repo that you want urls to point to. Typically `main` or `master` |  
 | `github_info.base_url`          |          | `https://github.com` | The base url for your github coordinates | 
 | `sql_parser`                                   |          | `datahub.utilities.sql_parser.DefaultSQLParser`    | See note below.                                                         |
+| `transport_options`                                |          |        |  Populates the [TransportOptions](https://github.com/looker-open-source/sdk-codegen/blob/94d6047a0d52912ac082eb91616c1e7c379ab262/python/looker_sdk/rtl/transport.py#L70) struct for looker client |
 
 Note! The integration can use an SQL parser to try to parse the tables the views depends on. This parsing is disabled by default, 
 but can be enabled by setting `parse_table_names_from_sql: True`.  The default parser is based on the [`sqllineage`](https://pypi.org/project/sqllineage/) package. 
diff --git a/metadata-ingestion/source_docs/redshift.md b/metadata-ingestion/source_docs/redshift.md
index bce07dc5026ea..6cd9cd1fc5145 100644
--- a/metadata-ingestion/source_docs/redshift.md
+++ b/metadata-ingestion/source_docs/redshift.md
@@ -1,8 +1,12 @@
 # Redshift
 
+To get all metadata from BigQuery you need to use two plugins `redshift` and `redshift-usage`. Both of them are described in this page. These will require 2 separate recipes. We understand this is not ideal and we plan to make this easier in the future.
+
 For context on getting started with ingestion, check out our [metadata ingestion guide](../README.md).
 
-## Setup
+## `redshift`
+
+### Setup
 
 To install this plugin, run `pip install 'acryl-datahub[redshift]'`.
 
@@ -19,7 +23,7 @@ Giving a user unrestricted access to system tables gives the user visibility to
 
 :::
 
-## Capabilities
+### Capabilities
 
 This plugin extracts the following:
 
@@ -41,7 +45,7 @@ You can also get fine-grained usage statistics for Redshift using the `redshift-
 | Data Containers   | ✔️     |                                          |
 | Data Domains      | ✔️     | [link](../../docs/domains.md)            |
 
-## Quickstart recipe
+### Quickstart recipe
 
 Check out the following recipe to get started with ingestion! See [below](#config-details) for full configuration options.
 
@@ -93,7 +97,7 @@ sink:
 
 </details>
 
-## Config details
+### Config details
 
 Like all SQL-based sources, the Redshift integration supports:
 - Stale Metadata Deletion: See [here](./stateful_ingestion.md) for more details on configuration.
@@ -130,11 +134,11 @@ Note that a `.` is used to denote nested fields in the YAML recipe.
 | `domain.domain_key.deny`       |          |                    | List of regex patterns for tables/schemas to not assign domain_key. There can be multiple domain key specified.                                                                         |
 | `domain.domain_key.ignoreCase` |          | `True`             | Whether to ignore case sensitivity during pattern matching.There can be multiple domain key specified.                                                                                  |
 
-## Lineage
+### Lineage
 
 There are multiple lineage collector implementations as Redshift does not support table lineage out of the box.
 
-### stl_scan_based
+#### stl_scan_based
 The stl_scan based collector uses Redshift's [stl_insert](https://docs.aws.amazon.com/redshift/latest/dg/r_STL_INSERT.html) and [stl_scan](https://docs.aws.amazon.com/redshift/latest/dg/r_STL_SCAN.html) system tables to
 discover lineage between tables.
 Pros:
@@ -145,7 +149,7 @@ Cons:
 - Does not work with Spectrum/external tables because those scans do not show up in stl_scan table.
 - If a table is depending on a view then the view won't be listed as dependency. Instead the table will be connected with the view's dependencies.
 
-### sql_based
+#### sql_based
 The sql_based based collector uses Redshift's [stl_insert](https://docs.aws.amazon.com/redshift/latest/dg/r_STL_INSERT.html) to discover all the insert queries
 and uses sql parsing to discover the dependecies.
 
@@ -157,7 +161,7 @@ Cons:
 - Slow.
 - Less reliable as the query parser can fail on certain queries
 
-### mixed
+#### mixed
 Using both collector above and first applying the sql based and then the stl_scan based one.
 
 Pros:
@@ -169,10 +173,13 @@ Cons:
 - Slow
 - May be incorrect at times as the query parser can fail on certain queries
 
-# Note
-- The redshift stl redshift tables which are used for getting data lineage only retain approximately two to five days of log history. This means you cannot extract lineage from queries issued outside that window.
+:::note
+
+The redshift stl redshift tables which are used for getting data lineage only retain approximately two to five days of log history. This means you cannot extract lineage from queries issued outside that window.
+
+:::
 
-# Redshift Usage Stats
+## `redshift-usage`
 
 This plugin extracts usage statistics for datasets in Amazon Redshift. For context on getting started with ingestion, check out our [metadata ingestion guide](../README.md).
 
@@ -187,10 +194,10 @@ To grant access this plugin for all system tables, please alter your datahub Red
 ALTER USER datahub_user WITH SYSLOG ACCESS UNRESTRICTED;
 ```
 
-## Setup
+### Setup
 To install this plugin, run `pip install 'acryl-datahub[redshift-usage]'`.
 
-## Capabilities
+### Capabilities
 
 | Capability | Status | Details | 
 | -----------| ------ | ---- |
@@ -210,7 +217,7 @@ This source only does usage statistics. To get the tables, views, and schemas in
 
 :::
 
-## Quickstart recipe
+### Quickstart recipe
 
 Check out the following recipe to get started with ingestion! See [below](#config-details) for full configuration options.
 
@@ -233,7 +240,7 @@ sink:
 # sink configs
 ```
 
-## Config details
+### Config details
 Note that a `.` is used to denote nested fields in the YAML recipe.
 
 By default, we extract usage stats for the last day, with the recommendation that this source is executed every day.
diff --git a/metadata-ingestion/source_docs/snowflake.md b/metadata-ingestion/source_docs/snowflake.md
index 45c2eb2106131..6505d01f1704f 100644
--- a/metadata-ingestion/source_docs/snowflake.md
+++ b/metadata-ingestion/source_docs/snowflake.md
@@ -1,16 +1,22 @@
 # Snowflake
 
+To get all metadata from Snowflake you need to use two plugins `snowflake` and `snowflake-usage`. Both of them are described in this page. These will require 2 separate recipes. We understand this is not ideal and we plan to make this easier in the future.
+
 For context on getting started with ingestion, check out our [metadata ingestion guide](../README.md).
 
-## Setup
+## `snowflake`
+### Setup
 
 To install this plugin, run `pip install 'acryl-datahub[snowflake]'`.
 
 ### Prerequisites
 
 In order to execute this source, your Snowflake user will need to have specific privileges granted to it for reading metadata
-from your warehouse. You can create a DataHub-specific role, assign it the required privileges, and assign it to a new DataHub user 
-by executing the following Snowflake commands from a user with the `ACCOUNTADMIN` role: 
+from your warehouse. 
+
+You can use the `provision_role` block in the recipe to grant the requires roles. 
+
+If your system admins prefer running the commands themselves then they can follow this guide to create a DataHub-specific role, assign it the required privileges, and assign it to a new DataHub user by executing the following Snowflake commands from a user with the `ACCOUNTADMIN` role: 
 
 ```sql
 create or replace role datahub_role;
@@ -36,25 +42,21 @@ grant role datahub_role to user datahub_user;
 
 This represents the bare minimum privileges required to extract databases, schemas, views, tables from Snowflake. 
 
-If you plan to enable extraction of table lineage, via the `include_table_lineage` config flag, you'll also need to grant privileges
-to access the Snowflake Account Usage views. You can execute the following using the `ACCOUNTADMIN` role to do so:
+If you plan to enable extraction of table lineage, via the `include_table_lineage` config flag, you'll need to grant additional privileges. See [snowflake usage prerequisites](#prerequisites-1) as the same privilege is required for this purpose too.
 
-```sql
-grant imported privileges on database snowflake to role datahub_role;
-```
 
-## Capabilities
+### Capabilities
 
 This plugin extracts the following:
 
 - Metadata for databases, schemas, views and tables
 - Column types associated with each table
 - Table, row, and column statistics via optional [SQL profiling](./sql_profiles.md)
-- Table lineage.
+- Table lineage
 
 :::tip
 
-You can also get fine-grained usage statistics for Snowflake using the `snowflake-usage` source described below.
+You can also get fine-grained usage statistics for Snowflake using the `snowflake-usage` source described [below](#snowflake-usage-plugin).
 
 :::
 
@@ -64,7 +66,7 @@ You can also get fine-grained usage statistics for Snowflake using the `snowflak
 | Data Containers   | ✔️     |                                          |
 | Data Domains      | ✔️     | [link](../../docs/domains.md)            |
 
-## Quickstart recipe
+### Quickstart recipe
 
 Check out the following recipe to get started with ingestion! See [below](#config-details) for full configuration options.
 
@@ -74,20 +76,28 @@ For general pointers on writing and running a recipe, see our [main recipe guide
 source:
   type: snowflake
   config:
+
+    provision_role: # Optional
+      enabled: false
+      dry_run: true
+      run_ingestion: false
+      admin_username: "${SNOWFLAKE_ADMIN_USER}"
+      admin_password: "${SNOWFLAKE_ADMIN_PASS}"
+
     # Coordinates
     host_port: account_name
     warehouse: "COMPUTE_WH"
 
     # Credentials
-    username: user
-    password: pass
-    role: "accountadmin"
+    username: "${SNOWFLAKE_USER}"
+    password: "${SNOWFLAKE_PASS}"
+    role: "datahub_role"
 
 sink:
   # sink configs
 ```
 
-## Config details
+### Config details
 
 Like all SQL-based sources, the Snowflake integration supports:
 - Stale Metadata Deletion: See [here](./stateful_ingestion.md) for more details on configuration.
@@ -122,8 +132,8 @@ Note that a `.` is used to denote nested fields in the YAML recipe.
 | `view_pattern.ignoreCase`      |          | `True`                                                                     | Whether to ignore case sensitivity during pattern matching.                                                                                                                             |
 | `include_tables`               |          | `True`                                                                     | Whether tables should be ingested.                                                                                                                                                      |
 | `include_views`                |          | `True`                                                                     | Whether views should be ingested.                                                                                                                                                       |
-| `include_table_lineage`        |          | `True`                                                                     | If enabled, populates the snowflake table-to-table and s3-to-snowflake table lineage. Requires role to be `accountadmin`                                                                |
-| `include_view_lineage`         |          | `True`                                                                     | If enabled, populates the snowflake view->table and table->view lineages (no view->view lineage yet). Requires role to be `accountadmin`, and `include_table_lineage` to be `True`.     |
+| `include_table_lineage`        |          | `True`                                                                     | If enabled, populates the snowflake table-to-table and s3-to-snowflake table lineage. Requires appropriate grants given to the role.                                                                |
+| `include_view_lineage`         |          | `True`                                                                     | If enabled, populates the snowflake view->table and table->view lineages (no view->view lineage yet). Requires appropriate grants given to the role, and `include_table_lineage` to be `True`.     |
 | `bucket_duration`              |          | `"DAY"`                                                                    | Duration to bucket lineage data extraction by. Can be `"DAY"` or `"HOUR"`.                                                                                                              |
 | `start_time`                   |          | Start of last full day in UTC (or hour, depending on `bucket_duration`)    | Earliest time of lineage data to consider. For the bootstrap run, set it as far back in time as possible.                                                                               |
 | `end_time`                     |          | End of last full day in UTC (or hour, depending on `bucket_duration`)      | Latest time of lineage data to consider.                                                                                                                                                |
@@ -131,41 +141,46 @@ Note that a `.` is used to denote nested fields in the YAML recipe.
 | `domain.domain_key.allow`      |          |                                                                            | List of regex patterns for tables/schemas to set domain_key domain key (domain_key can be any string like `sales`. There can be multiple domain key specified.                          |
 | `domain.domain_key.deny`       |          |                                                                            | List of regex patterns for tables/schemas to not assign domain_key. There can be multiple domain key specified.                                                                         |
 | `domain.domain_key.ignoreCase` |          | `True`                                                                     | Whether to ignore case sensitivity during pattern matching.There can be multiple domain key specified.                                                                                  |
+| `provision_role.enabled`                |          | `False` | Whether provisioning of Snowflake role (used for ingestion) is enabled or not |
+| `provision_role.dry_run`                |          | `False` | If `provision_role` is enabled, whether to dry run the sql commands for system admins to see what sql grant commands would be run without actually running the grant commands |
+| `provision_role.drop_role_if_exists`    |          | `False` | Useful during testing to ensure you have a clean slate role. Not recommended for production use cases |
+| `provision_role.run_ingestion`         |          | `False`  | If system admins wish to skip actual ingestion of metadata during testing of the provisioning of `role` |
+| `provision_role.admin_role`             |          | `accountadmin` | The Snowflake role of admin user used for provisioning of the role specified by `role` config. System admins can audit the open source code and decide to use a different role |
+| `provision_role.admin_username`         |  ✅       |          | The username to be used for provisioning of role |
+| `provision_role.admin_password`         |  ✅       |          | The password to be used for provisioning of role |
 
 
-## Compatibility
-
-Table lineage requires Snowflake's [Access History](https://docs.snowflake.com/en/user-guide/access-history.html) feature.
-
-# Snowflake Usage Stats
+## `snowflake-usage`
 
 For context on getting started with ingestion, check out our [metadata ingestion guide](../README.md).
 
-## Setup
+### Setup
 
 To install this plugin, run `pip install 'acryl-datahub[snowflake-usage]'`.
 
 ### Prerequisites 
 
-In order to execute the snowflake-usage source, your Snowflake user will need to have specific privileges granted to it. Specifically,
-you'll need to grant access to the [Account Usage](https://docs.snowflake.com/en/sql-reference/account-usage.html) system tables, using which the DataHub source extracts information. Assuming
-you've followed the steps outlined above to create a DataHub-specific User & Role, you'll simply need to execute the following commands
-in Snowflake from a user with the `ACCOUNTADMIN` role: 
+:::note
+
+Table lineage requires Snowflake's [Access History](https://docs.snowflake.com/en/user-guide/access-history.html) feature. The "accountadmin" role has this by default.
+
+The underlying access history views that we use are only available in Snowflake's enterprise edition or higher.
+
+:::
+
+In order to execute the snowflake-usage source, your Snowflake user will need to have specific privileges granted to it. Specifically, you'll need to grant access to the [Account Usage](https://docs.snowflake.com/en/sql-reference/account-usage.html) system tables, using which the DataHub source extracts information. Assuming you've followed the steps outlined in `snowflake` plugin to create a DataHub-specific User & Role, you'll simply need to execute the following commands in Snowflake. This will require a user with the `ACCOUNTADMIN` role (or a role granted the IMPORT SHARES global privilege). Please see [Snowflake docs for more details](https://docs.snowflake.com/en/user-guide/data-share-consumers.html).
 
 ```sql
 grant imported privileges on database snowflake to role datahub_role;
 ```
 
-## Capabilities
+### Capabilities
 
 This plugin extracts the following:
 
 - Statistics on queries issued and tables and columns accessed (excludes views)
 - Aggregation of these statistics into buckets, by day or hour granularity
 
-Note: the user/role must have access to the account usage table. The "accountadmin" role has this by default, and other roles can be [granted this permission](https://docs.snowflake.com/en/sql-reference/account-usage.html#enabling-account-usage-for-other-roles).
-
-Note: the underlying access history views that we use are only available in Snowflake's enterprise edition or higher.
 
 :::note
 
@@ -173,7 +188,7 @@ This source only does usage statistics. To get the tables, views, and schemas in
 
 :::
 
-## Quickstart recipe
+### Quickstart recipe
 
 Check out the following recipe to get started with ingestion! See [below](#config-details) for full configuration options.
 
@@ -188,9 +203,9 @@ source:
     warehouse: "COMPUTE_WH"
 
     # Credentials
-    username: user
-    password: pass
-    role: "sysadmin"
+    username: "${SNOWFLAKE_USER}"
+    password: "${SNOWFLAKE_PASS}"
+    role: "datahub_role"
 
     # Options
     top_n_queries: 10
@@ -199,7 +214,7 @@ sink:
   # sink configs
 ```
 
-## Config details
+### Config details
 
 Snowflake integration also supports prevention of redundant reruns for the same data. See [here](./stateful_ingestion.md) for more details on configuration.
 
@@ -235,7 +250,7 @@ User's without email address will be ignored from usage if you don't set `email_
 
 
 
-# Compatibility
+## Compatibility
 
 Coming soon!
 
diff --git a/metadata-ingestion/src/datahub/cli/cli_utils.py b/metadata-ingestion/src/datahub/cli/cli_utils.py
index 00b05caeb89e4..1e1459e4f0c52 100644
--- a/metadata-ingestion/src/datahub/cli/cli_utils.py
+++ b/metadata-ingestion/src/datahub/cli/cli_utils.py
@@ -19,17 +19,22 @@
 from datahub.emitter.rest_emitter import _make_curl_command
 from datahub.emitter.serialization_helper import post_json_transform
 from datahub.metadata.schema_classes import (
+    AssertionRunEventClass,
     BrowsePathsClass,
     ChartInfoClass,
     ChartKeyClass,
+    DatahubIngestionCheckpointClass,
+    DatahubIngestionRunSummaryClass,
     DataJobInputOutputClass,
     DataJobKeyClass,
     DataPlatformInstanceClass,
     DataProcessInfoClass,
     DatasetDeprecationClass,
     DatasetKeyClass,
+    DatasetProfileClass,
     DatasetPropertiesClass,
     DatasetUpstreamLineageClass,
+    DatasetUsageStatisticsClass,
     EditableDatasetPropertiesClass,
     EditableSchemaMetadataClass,
     GlobalTagsClass,
@@ -39,6 +44,7 @@
     MLFeaturePropertiesClass,
     MLPrimaryKeyKeyClass,
     MLPrimaryKeyPropertiesClass,
+    OperationClass,
     OwnershipClass,
     SchemaMetadataClass,
     StatusClass,
@@ -470,30 +476,109 @@ def post_entity(
     ChartKeyClass: "chartKey",
 }
 
+timeseries_class_to_aspect_name_map: Dict[Type, str] = {
+    DatahubIngestionCheckpointClass: "datahubIngestionCheckpoint",
+    DatahubIngestionRunSummaryClass: "datahubIngestionRunSummary",
+    DatasetUsageStatisticsClass: "datasetUsageStatistics",
+    DatasetProfileClass: "datasetProfile",
+    AssertionRunEventClass: "assertionRunEvent",
+    OperationClass: "operation",
+}
+
 
 def _get_pydantic_class_from_aspect_name(aspect_name: str) -> Optional[Type[Aspect]]:
     candidates = [k for (k, v) in type_class_to_name_map.items() if v == aspect_name]
-    return candidates[0] or None
+    candidates.extend(
+        [
+            k
+            for (k, v) in timeseries_class_to_aspect_name_map.items()
+            if v == aspect_name
+        ]
+    )
+    return candidates[0] if candidates else None
 
 
 def _get_aspect_name_from_aspect_class(aspect_class: str) -> str:
     class_to_name_map = {
         k.RECORD_SCHEMA.fullname.replace("pegasus2avro.", ""): v  # type: ignore
-        for (k, v) in type_class_to_name_map.items()
+        for (k, v) in (
+            set(type_class_to_name_map.items())
+            | set(timeseries_class_to_aspect_name_map.items())
+        )
     }
     return class_to_name_map.get(aspect_class, "unknown")
 
 
+def get_latest_timeseries_aspect_values(
+    entity_urn: str,
+    timeseries_aspect_name: str,
+    cached_session_host: Optional[Tuple[Session, str]],
+) -> Dict:
+    if not cached_session_host:
+        session, gms_host = get_session_and_host()
+    else:
+        session, gms_host = cached_session_host
+    query_body = {
+        "urn": entity_urn,
+        "entity": guess_entity_type(entity_urn),
+        "aspect": timeseries_aspect_name,
+        "latestValue": True,
+    }
+    end_point = "/aspects?action=getTimeseriesAspectValues"
+    try:
+        response = session.post(url=gms_host + end_point, data=json.dumps(query_body))
+        response.raise_for_status()
+        return response.json()
+    except Exception:
+        # Ignore exceptions
+        return {}
+
+
 def get_aspects_for_entity(
     entity_urn: str,
     aspects: List[str],
     typed: bool = False,
     cached_session_host: Optional[Tuple[Session, str]] = None,
 ) -> Dict[str, Union[dict, DictWrapper]]:
-    entity_response = get_entity(entity_urn, aspects, cached_session_host)
+    # Process non-timeseries aspects
+    non_timeseries_aspects: List[str] = [
+        a for a in aspects if a not in timeseries_class_to_aspect_name_map.values()
+    ]
+    entity_response = get_entity(
+        entity_urn, non_timeseries_aspects, cached_session_host
+    )
     aspect_list: List[Dict[str, dict]] = list(entity_response["value"].values())[0][
         "aspects"
     ]
+
+    # Process timeseries aspects & append to aspect_list
+    timeseries_aspects: List[str] = [
+        a for a in aspects if a in timeseries_class_to_aspect_name_map.values()
+    ]
+    for timeseries_aspect in timeseries_aspects:
+        timeseries_response = get_latest_timeseries_aspect_values(
+            entity_urn, timeseries_aspect, cached_session_host
+        )
+        values: List[Dict] = timeseries_response.get("value", {}).get("values", [])
+        if values:
+            aspect_cls: Optional[Type] = _get_pydantic_class_from_aspect_name(
+                timeseries_aspect
+            )
+            if aspect_cls is not None:
+                aspect_value = values[0]
+                # Decode the json-encoded generic aspect value.
+                aspect_value["aspect"]["value"] = json.loads(
+                    aspect_value["aspect"]["value"]
+                )
+                aspect_list.append(
+                    # Follow the convention used for non-timeseries aspects.
+                    {
+                        aspect_cls.RECORD_SCHEMA.fullname.replace(
+                            "pegasus2avro.", ""
+                        ): aspect_value
+                    }
+                )
+
     aspect_map: Dict[str, Union[dict, DictWrapper]] = {}
     for a in aspect_list:
         aspect_class = list(a.keys())[0]
diff --git a/metadata-ingestion/src/datahub/cli/get_cli.py b/metadata-ingestion/src/datahub/cli/get_cli.py
index debcd64757c37..f20c451d7ab24 100644
--- a/metadata-ingestion/src/datahub/cli/get_cli.py
+++ b/metadata-ingestion/src/datahub/cli/get_cli.py
@@ -23,7 +23,10 @@
 @click.pass_context
 @telemetry.with_telemetry
 def get(ctx: Any, urn: Optional[str], aspect: List[str]) -> None:
-    """Get metadata for an entity with an optional list of aspects to project"""
+    """
+    Get metadata for an entity with an optional list of aspects to project.
+    This works for both versioned aspects and timeseries aspects. For timeseries aspects, it fetches the latest value.
+    """
 
     if urn is None:
         if not ctx.args:
diff --git a/metadata-ingestion/src/datahub/cli/ingest_cli.py b/metadata-ingestion/src/datahub/cli/ingest_cli.py
index c76434acdcc50..45f89099bdea5 100644
--- a/metadata-ingestion/src/datahub/cli/ingest_cli.py
+++ b/metadata-ingestion/src/datahub/cli/ingest_cli.py
@@ -122,8 +122,14 @@ def parse_restli_response(response):
 @ingest.command()
 @click.argument("page_offset", type=int, default=0)
 @click.argument("page_size", type=int, default=100)
+@click.option(
+    "--include-soft-deletes",
+    is_flag=True,
+    default=False,
+    help="If enabled, will list ingestion runs which have been soft deleted",
+)
 @telemetry.with_telemetry
-def list_runs(page_offset: int, page_size: int) -> None:
+def list_runs(page_offset: int, page_size: int, include_soft_deletes: bool) -> None:
     """List recent ingestion runs to datahub"""
 
     session, gms_host = get_session_and_host()
@@ -133,6 +139,7 @@ def list_runs(page_offset: int, page_size: int) -> None:
     payload_obj = {
         "pageOffset": page_offset,
         "pageSize": page_size,
+        "includeSoft": include_soft_deletes,
     }
 
     payload = json.dumps(payload_obj)
@@ -163,7 +170,7 @@ def list_runs(page_offset: int, page_size: int) -> None:
 def show(run_id: str) -> None:
     """Describe a provided ingestion run to datahub"""
 
-    payload_obj = {"runId": run_id, "dryRun": True}
+    payload_obj = {"runId": run_id, "dryRun": True, "hardDelete": True}
     structured_rows, entities_affected, aspects_affected = post_rollback_endpoint(
         payload_obj, "/runs?action=rollback"
     )
@@ -190,8 +197,9 @@ def show(run_id: str) -> None:
 @click.option("--run-id", required=True, type=str)
 @click.option("-f", "--force", required=False, is_flag=True)
 @click.option("--dry-run", "-n", required=False, is_flag=True, default=False)
+@click.option("--hard-delete", "-d", required=False, is_flag=True, default=False)
 @telemetry.with_telemetry
-def rollback(run_id: str, force: bool, dry_run: bool) -> None:
+def rollback(run_id: str, force: bool, dry_run: bool, hard_delete: bool) -> None:
     """Rollback a provided ingestion run to datahub"""
 
     cli_utils.test_connectivity_complain_exit("ingest")
@@ -202,7 +210,7 @@ def rollback(run_id: str, force: bool, dry_run: bool) -> None:
             abort=True,
         )
 
-    payload_obj = {"runId": run_id, "dryRun": dry_run}
+    payload_obj = {"runId": run_id, "dryRun": dry_run, "hardDelete": hard_delete}
     structured_rows, entities_affected, aspects_affected = post_rollback_endpoint(
         payload_obj, "/runs?action=rollback"
     )
diff --git a/metadata-ingestion/src/datahub/emitter/rest_emitter.py b/metadata-ingestion/src/datahub/emitter/rest_emitter.py
index e9beb06b9a9a2..5aadaf3765258 100644
--- a/metadata-ingestion/src/datahub/emitter/rest_emitter.py
+++ b/metadata-ingestion/src/datahub/emitter/rest_emitter.py
@@ -136,16 +136,12 @@ def __init__(
         self._session.mount("http://", adapter)
         self._session.mount("https://", adapter)
 
-    def test_connection(self) -> str:
+    def test_connection(self) -> dict:
         response = self._session.get(f"{self._gms_server}/config")
         if response.status_code == 200:
             config: dict = response.json()
             if config.get("noCode") == "true":
-                return (
-                    config.get("versions", {})
-                    .get("linkedin/datahub", {})
-                    .get("version", "")
-                )
+                return config
 
             else:
                 # Looks like we either connected to an old GMS or to some other service. Let's see if we can determine which before raising an error
diff --git a/metadata-ingestion/src/datahub/entrypoints.py b/metadata-ingestion/src/datahub/entrypoints.py
index f7b60622d3bdc..c897eb6bd79bb 100644
--- a/metadata-ingestion/src/datahub/entrypoints.py
+++ b/metadata-ingestion/src/datahub/entrypoints.py
@@ -1,5 +1,6 @@
 import logging
 import os
+import platform
 import sys
 
 import click
@@ -18,6 +19,7 @@
 from datahub.cli.timeline_cli import timeline
 from datahub.configuration import SensitiveError
 from datahub.telemetry import telemetry
+from datahub.utilities.server_config_util import get_gms_config
 
 logger = logging.getLogger(__name__)
 
@@ -156,4 +158,11 @@ def main(**kwargs):
                 **kwargs,
             )
         )
+        logger.info(
+            f"DataHub CLI version: {datahub_package.__version__} at {datahub_package.__file__}"
+        )
+        logger.info(
+            f"Python version: {sys.version} at {sys.executable} on {platform.platform()}"
+        )
+        logger.info(f"GMS config {get_gms_config()}")
         sys.exit(1)
diff --git a/metadata-ingestion/src/datahub/ingestion/api/committable.py b/metadata-ingestion/src/datahub/ingestion/api/committable.py
index a67bbe7d7b7ce..f1aada4477f1a 100644
--- a/metadata-ingestion/src/datahub/ingestion/api/committable.py
+++ b/metadata-ingestion/src/datahub/ingestion/api/committable.py
@@ -34,9 +34,9 @@ def commit(self) -> None:
 FilterType = TypeVar("FilterType")
 
 
-@dataclass
 class _StatefulCommittableConcrete(Generic[StateType]):
-    state_to_commit: StateType
+    def __init__(self, state_to_commit: StateType):
+        self.state_to_commit: StateType = state_to_commit
 
 
 class StatefulCommittable(
diff --git a/metadata-ingestion/src/datahub/ingestion/api/source.py b/metadata-ingestion/src/datahub/ingestion/api/source.py
index dbb1d9981ebfb..4b5c5068757be 100644
--- a/metadata-ingestion/src/datahub/ingestion/api/source.py
+++ b/metadata-ingestion/src/datahub/ingestion/api/source.py
@@ -1,7 +1,10 @@
+import platform
+import sys
 from abc import ABCMeta, abstractmethod
 from dataclasses import dataclass, field
 from typing import Dict, Generic, Iterable, List, TypeVar
 
+import datahub
 from datahub.ingestion.api.closeable import Closeable
 from datahub.ingestion.api.common import PipelineContext, RecordEnvelope, WorkUnit
 from datahub.ingestion.api.report import Report
@@ -14,7 +17,11 @@ class SourceReport(Report):
 
     warnings: Dict[str, List[str]] = field(default_factory=dict)
     failures: Dict[str, List[str]] = field(default_factory=dict)
-    cli_version: str = ""
+    cli_version: str = datahub.nice_version_name()
+    cli_entry_location: str = datahub.__file__
+    py_version: str = sys.version
+    py_exec_path: str = sys.executable
+    os_details: str = platform.platform()
 
     def report_workunit(self, wu: WorkUnit) -> None:
         self.workunits_produced += 1
diff --git a/metadata-ingestion/src/datahub/ingestion/run/pipeline.py b/metadata-ingestion/src/datahub/ingestion/run/pipeline.py
index 02b81e0e968e8..b2ebc3d984fac 100644
--- a/metadata-ingestion/src/datahub/ingestion/run/pipeline.py
+++ b/metadata-ingestion/src/datahub/ingestion/run/pipeline.py
@@ -8,7 +8,6 @@
 import click
 from pydantic import validator
 
-import datahub
 from datahub.configuration.common import (
     ConfigModel,
     DynamicTypedConfig,
@@ -117,6 +116,12 @@ def __init__(
             preview_mode=preview_mode,
         )
 
+        sink_type = self.config.sink.type
+        sink_class = sink_registry.get(sink_type)
+        sink_config = self.config.sink.dict().get("config", {})
+        self.sink: Sink = sink_class.create(sink_config, self.ctx)
+        logger.debug(f"Sink type:{self.config.sink.type},{sink_class} configured")
+
         source_type = self.config.source.type
         source_class = source_registry.get(source_type)
         self.source: Source = source_class.create(
@@ -124,12 +129,6 @@ def __init__(
         )
         logger.debug(f"Source type:{source_type},{source_class} configured")
 
-        sink_type = self.config.sink.type
-        sink_class = sink_registry.get(sink_type)
-        sink_config = self.config.sink.dict().get("config", {})
-        self.sink: Sink = sink_class.create(sink_config, self.ctx)
-        logger.debug(f"Sink type:{self.config.sink.type},{sink_class} configured")
-
         self.extractor_class = extractor_registry.get(self.config.source.extractor)
 
         self._configure_transforms()
@@ -179,7 +178,6 @@ def run(self) -> None:
 
         callback = LoggingCallback()
         extractor: Extractor = self.extractor_class()
-        self.source.get_report().cli_version = datahub.nice_version_name()
         for wu in itertools.islice(
             self.source.get_workunits(), 10 if self.preview_mode else None
         ):
diff --git a/metadata-ingestion/src/datahub/ingestion/sink/datahub_rest.py b/metadata-ingestion/src/datahub/ingestion/sink/datahub_rest.py
index 7e028b7d5eabc..2755751856c68 100644
--- a/metadata-ingestion/src/datahub/ingestion/sink/datahub_rest.py
+++ b/metadata-ingestion/src/datahub/ingestion/sink/datahub_rest.py
@@ -16,6 +16,7 @@
     MetadataChangeProposal,
 )
 from datahub.metadata.com.linkedin.pegasus2avro.usage import UsageAggregation
+from datahub.utilities.server_config_util import set_gms_config
 
 logger = logging.getLogger(__name__)
 
@@ -50,7 +51,14 @@ def __init__(self, ctx: PipelineContext, config: DatahubRestSinkConfig):
             extra_headers=self.config.extra_headers,
             ca_certificate_path=self.config.ca_certificate_path,
         )
-        self.report.gms_version = self.emitter.test_connection()
+        gms_config = self.emitter.test_connection()
+        self.report.gms_version = (
+            gms_config.get("versions", {})
+            .get("linkedin/datahub", {})
+            .get("version", "")
+        )
+        logger.info("Setting gms config")
+        set_gms_config(gms_config)
         self.executor = concurrent.futures.ThreadPoolExecutor(
             max_workers=self.config.max_threads
         )
diff --git a/metadata-ingestion/src/datahub/ingestion/source/feast_image/requirements.txt b/metadata-ingestion/src/datahub/ingestion/source/feast_image/requirements.txt
index 1c41862b0b9a7..6a2841e420e2f 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/feast_image/requirements.txt
+++ b/metadata-ingestion/src/datahub/ingestion/source/feast_image/requirements.txt
@@ -13,7 +13,7 @@ grpcio==1.38.0; python_full_version >= "3.7.0"
 idna==2.10; python_full_version >= "3.7.0"
 jinja2==3.0.1; python_version >= "3.6" and python_full_version >= "3.7.0"
 jsonschema==3.2.0; python_full_version >= "3.7.0"
-markupsafe==2.0.1; python_version >= "3.6" and python_full_version >= "3.7.0"
+markupsafe>==1.1.1,<=2.0.1; python_version >= "3.6" and python_full_version >= "3.7.0"
 mmh3==3.0.0; python_full_version >= "3.7.0"
 numpy==1.20.3; python_version >= "3.7" and python_full_version >= "3.7.1"
 packaging==20.9; python_full_version >= "3.7.0"
diff --git a/metadata-ingestion/src/datahub/ingestion/source/identity/azure_ad.py b/metadata-ingestion/src/datahub/ingestion/source/identity/azure_ad.py
index 56f41ec75b0e8..6413ed3497c9f 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/identity/azure_ad.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/identity/azure_ad.py
@@ -2,6 +2,7 @@
 import logging
 import re
 import urllib
+from collections import defaultdict
 from dataclasses import dataclass, field
 from typing import Any, Dict, Generator, Iterable, List
 
@@ -10,6 +11,7 @@
 
 from datahub.configuration import ConfigModel
 from datahub.configuration.common import AllowDenyPattern
+from datahub.emitter.mce_builder import make_group_urn, make_user_urn
 from datahub.ingestion.api.common import PipelineContext
 from datahub.ingestion.api.source import Source, SourceReport
 from datahub.ingestion.api.workunit import MetadataWorkUnit
@@ -58,13 +60,20 @@ class AzureADConfig(ConfigModel):
     users_pattern: AllowDenyPattern = AllowDenyPattern.allow_all()
     groups_pattern: AllowDenyPattern = AllowDenyPattern.allow_all()
 
+    # If enabled, report will contain names of filtered users and groups.
+    filtered_tracking: bool = True
+
 
 @dataclass
 class AzureADSourceReport(SourceReport):
     filtered: List[str] = field(default_factory=list)
+    filtered_tracking: bool = field(default=True, repr=False)
+    filtered_count: int = field(default=0)
 
     def report_filtered(self, name: str) -> None:
-        self.filtered.append(name)
+        self.filtered_count += 1
+        if self.filtered_tracking:
+            self.filtered.append(name)
 
 
 # Source that extracts Azure AD users, groups and group memberships using Microsoft Graph REST API
@@ -81,7 +90,9 @@ def create(cls, config_dict, ctx):
     def __init__(self, config: AzureADConfig, ctx: PipelineContext):
         super().__init__(ctx)
         self.config = config
-        self.report = AzureADSourceReport()
+        self.report = AzureADSourceReport(
+            filtered_tracking=self.config.filtered_tracking
+        )
         self.token_data = {
             "grant_type": "client_credentials",
             "client_id": self.config.client_id,
@@ -91,6 +102,8 @@ def __init__(self, config: AzureADConfig, ctx: PipelineContext):
             "scope": "https://graph.microsoft.com/.default",
         }
         self.token = self.get_token()
+        self.selected_azure_ad_groups: list = []
+        self.azure_ad_groups_users: list = []
 
     def get_token(self):
         token_response = requests.post(self.config.token_url, data=self.token_data)
@@ -107,9 +120,6 @@ def get_token(self):
             click.echo("Error: Token response invalid")
             exit()
 
-    selected_azure_ad_groups: list = []
-    azure_ad_groups_users: list = []
-
     def get_workunits(self) -> Iterable[MetadataWorkUnit]:
         # for future developers: The actual logic of this ingestion wants to be executed, in order:
         # 1) the groups
@@ -133,52 +143,26 @@ def get_workunits(self) -> Iterable[MetadataWorkUnit]:
                     yield wu
 
         # Populate GroupMembership Aspects for CorpUsers
-        datahub_corp_user_urn_to_group_membership: Dict[str, GroupMembershipClass] = {}
+        datahub_corp_user_urn_to_group_membership: Dict[
+            str, GroupMembershipClass
+        ] = defaultdict(lambda: GroupMembershipClass(groups=[]))
         if (
             self.config.ingest_group_membership
             and len(self.selected_azure_ad_groups) > 0
         ):
             # 2) the groups' membership
             for azure_ad_group in self.selected_azure_ad_groups:
+                # Azure supports nested groups, but not DataHub.  We need to explode the nested groups into a flat list.
                 datahub_corp_group_urn = self._map_azure_ad_group_to_urn(azure_ad_group)
                 if not datahub_corp_group_urn:
-                    error_str = "Failed to extract DataHub Group Name from Azure AD Group named {}. Skipping...".format(
-                        azure_ad_group.get("displayName")
-                    )
+                    error_str = f"Failed to extract DataHub Group Name from Azure AD Group named {azure_ad_group.get('displayName')}. Skipping..."
                     self.report.report_failure("azure_ad_group_mapping", error_str)
                     continue
-                # Extract and map users for each group
-                for azure_ad_group_users in self._get_azure_ad_group_users(
-                    azure_ad_group
-                ):
-                    # if group doesn't have any members, continue
-                    if not azure_ad_group_users:
-                        continue
-                    for azure_ad_user in azure_ad_group_users:
-                        datahub_corp_user_urn = self._map_azure_ad_user_to_urn(
-                            azure_ad_user
-                        )
-                        if not datahub_corp_user_urn:
-                            error_str = "Failed to extract DataHub Username from Azure ADUser {}. Skipping...".format(
-                                azure_ad_user.get("displayName")
-                            )
-                            self.report.report_failure(
-                                "azure_ad_user_mapping", error_str
-                            )
-                            continue
-                        self.azure_ad_groups_users.append(azure_ad_user)
-                        # update/create the GroupMembership aspect for this group member.
-                        if (
-                            datahub_corp_user_urn
-                            in datahub_corp_user_urn_to_group_membership
-                        ):
-                            datahub_corp_user_urn_to_group_membership[
-                                datahub_corp_user_urn
-                            ].groups.append(datahub_corp_group_urn)
-                        else:
-                            datahub_corp_user_urn_to_group_membership[
-                                datahub_corp_user_urn
-                            ] = GroupMembershipClass(groups=[datahub_corp_group_urn])
+                self._add_group_members_to_group_membership(
+                    datahub_corp_group_urn,
+                    azure_ad_group,
+                    datahub_corp_user_urn_to_group_membership,
+                )
 
         if (
             self.config.ingest_groups_users
@@ -205,6 +189,53 @@ def get_workunits(self) -> Iterable[MetadataWorkUnit]:
                     datahub_corp_user_urn_to_group_membership,
                 )
 
+    def _add_group_members_to_group_membership(
+        self,
+        parent_corp_group_urn: str,
+        azure_ad_group: dict,
+        user_urn_to_group_membership: Dict[str, GroupMembershipClass],
+    ) -> None:
+        # Extract and map members for each group
+        for azure_ad_group_members in self._get_azure_ad_group_members(azure_ad_group):
+            # if group doesn't have any members, continue
+            if not azure_ad_group_members:
+                continue
+            for azure_ad_member in azure_ad_group_members:
+                odata_type = azure_ad_member.get("@odata.type")
+                if odata_type == "#microsoft.graph.user":
+                    self._add_user_to_group_membership(
+                        parent_corp_group_urn,
+                        azure_ad_member,
+                        user_urn_to_group_membership,
+                    )
+                elif odata_type == "#microsoft.graph.group":
+                    # Since DataHub does not support nested group, we add the members to the parent group and not the nested one.
+                    self._add_group_members_to_group_membership(
+                        parent_corp_group_urn,
+                        azure_ad_member,
+                        user_urn_to_group_membership,
+                    )
+                else:
+                    raise ValueError(
+                        f"Unsupported @odata.type '{odata_type}' found in Azure group member"
+                    )
+
+    def _add_user_to_group_membership(
+        self,
+        group_urn: str,
+        azure_ad_user: dict,
+        user_urn_to_group_membership: Dict[str, GroupMembershipClass],
+    ) -> None:
+        user_urn = self._map_azure_ad_user_to_urn(azure_ad_user)
+        if not user_urn:
+            error_str = f"Failed to extract DataHub Username from Azure ADUser {azure_ad_user.get('displayName')}. Skipping..."
+            self.report.report_failure("azure_ad_user_mapping", error_str)
+        else:
+            self.azure_ad_groups_users.append(azure_ad_user)
+            # update/create the GroupMembership aspect for this group member.
+            if group_urn not in user_urn_to_group_membership[user_urn].groups:
+                user_urn_to_group_membership[user_urn].groups.append(group_urn)
+
     def ingest_ad_users(
         self,
         datahub_corp_user_snapshots: Generator[CorpUserSnapshot, Any, None],
@@ -240,7 +271,7 @@ def _get_azure_ad_groups(self) -> Iterable[List]:
     def _get_azure_ad_users(self) -> Iterable[List]:
         yield from self._get_azure_ad_data(kind="/users")
 
-    def _get_azure_ad_group_users(self, azure_ad_group: dict) -> Iterable[List]:
+    def _get_azure_ad_group_members(self, azure_ad_group: dict) -> Iterable[List]:
         group_id = azure_ad_group.get("id")
         kind = f"/groups/{group_id}/members"
         yield from self._get_azure_ad_data(kind=kind)
@@ -332,7 +363,7 @@ def _map_azure_ad_group_to_urn(self, azure_ad_group):
             return None
         # decode the group name to deal with URL encoding, and replace spaces with '_'
         url_encoded_group_name = urllib.parse.quote(group_name)
-        return self._make_corp_group_urn(url_encoded_group_name)
+        return make_group_urn(url_encoded_group_name)
 
     def _map_azure_ad_group_to_group_name(self, azure_ad_group):
         return self._extract_regex_match_from_dict_value(
@@ -371,7 +402,7 @@ def _map_azure_ad_user_to_urn(self, azure_ad_user):
         user_name = self._map_azure_ad_user_to_user_name(azure_ad_user)
         if not user_name:
             return None
-        return self._make_corp_user_urn(user_name)
+        return make_user_urn(user_name)
 
     def _map_azure_ad_user_to_corp_user(self, azure_ad_user):
         full_name = (
@@ -390,12 +421,6 @@ def _map_azure_ad_user_to_corp_user(self, azure_ad_user):
             countryCode=azure_ad_user.get("mobilePhone", None),
         )
 
-    def _make_corp_group_urn(self, groupname: str) -> str:
-        return f"urn:li:corpGroup:{groupname}"
-
-    def _make_corp_user_urn(self, username: str) -> str:
-        return f"urn:li:corpuser:{username}"
-
     def _extract_regex_match_from_dict_value(
         self, str_dict: Dict[str, str], key: str, pattern: str
     ) -> str:
diff --git a/metadata-ingestion/src/datahub/ingestion/source/kafka.py b/metadata-ingestion/src/datahub/ingestion/source/kafka.py
index d9d6b96dadce0..41158c8659524 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/kafka.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/kafka.py
@@ -20,7 +20,6 @@
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.emitter.mcp_builder import add_domain_to_entity_wu
 from datahub.ingestion.api.common import PipelineContext
-from datahub.ingestion.api.source import SourceReport
 from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.ingestion.extractor import schema_util
 from datahub.ingestion.source.state.checkpoint import Checkpoint
@@ -29,6 +28,7 @@
     JobId,
     StatefulIngestionConfig,
     StatefulIngestionConfigBase,
+    StatefulIngestionReport,
     StatefulIngestionSourceBase,
 )
 from datahub.metadata.com.linkedin.pegasus2avro.common import Status
@@ -71,7 +71,7 @@ class KafkaSourceConfig(StatefulIngestionConfigBase, DatasetSourceConfigBase):
 
 
 @dataclass
-class KafkaSourceReport(SourceReport):
+class KafkaSourceReport(StatefulIngestionReport):
     topics_scanned: int = 0
     filtered: List[str] = field(default_factory=list)
     soft_deleted_stale_entities: List[str] = field(default_factory=list)
diff --git a/metadata-ingestion/src/datahub/ingestion/source/looker.py b/metadata-ingestion/src/datahub/ingestion/source/looker.py
index 8dc260ca45bd7..47d94f620a294 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/looker.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/looker.py
@@ -21,6 +21,7 @@
 
 import looker_sdk
 from looker_sdk.error import SDKError
+from looker_sdk.rtl.transport import TransportOptions
 from looker_sdk.sdk.api31.methods import Looker31SDK
 from looker_sdk.sdk.api31.models import (
     Dashboard,
@@ -66,6 +67,7 @@ class LookerAPIConfig(ConfigModel):
     client_id: str
     client_secret: str
     base_url: str
+    transport_options: Optional[TransportOptions]
 
 
 class LookerAPI:
@@ -82,7 +84,7 @@ def __init__(self, config: LookerAPIConfig) -> None:
         # try authenticating current user to check connectivity
         # (since it's possible to initialize an invalid client without any complaints)
         try:
-            self.client.me()
+            self.client.me(transport_options=config.transport_options)
         except SDKError as e:
             raise ConfigurationError(
                 "Failed to initialize Looker client. Please check your configuration."
@@ -408,7 +410,7 @@ def _get_looker_dashboard_element(  # noqa: C901
                 fields = self._get_fields_from_query(element.look.query)
                 if element.look.query.view is not None:
                     explores = [element.look.query.view]
-                logger.info(
+                logger.debug(
                     "Element {}: Explores added: {}".format(element.title, explores)
                 )
                 for exp in explores:
@@ -576,8 +578,8 @@ def _make_explore_metadata_events(
                 events, explore_id, start_time, end_time = future.result()
                 explore_events.extend(events)
                 self.reporter.report_upstream_latency(start_time, end_time)
-                logger.info(
-                    f"Running time of fetch_one_explore for {explore_id}: {(end_time-start_time).total_seconds()}"
+                logger.debug(
+                    f"Running time of fetch_one_explore for {explore_id}: {(end_time - start_time).total_seconds()}"
                 )
 
         return explore_events
@@ -593,7 +595,11 @@ def fetch_one_explore(
         start_time = datetime.datetime.now()
         events: List[Union[MetadataChangeEvent, MetadataChangeProposalWrapper]] = []
         looker_explore = LookerExplore.from_api(
-            model, explore, self.client, self.reporter
+            model,
+            explore,
+            self.client,
+            self.reporter,
+            self.source_config.transport_options,
         )
         if looker_explore is not None:
             events = (
@@ -755,7 +761,9 @@ def process_dashboard(
                 "user_id",
             ]
             dashboard_object = self.client.dashboard(
-                dashboard_id=dashboard_id, fields=",".join(fields)
+                dashboard_id=dashboard_id,
+                fields=",".join(fields),
+                transport_options=self.source_config.transport_options,
             )
         except SDKError:
             # A looker dashboard could be deleted in between the list and the get
@@ -786,9 +794,15 @@ def process_dashboard(
         return workunits, dashboard_id, start_time, datetime.datetime.now()
 
     def get_workunits(self) -> Iterable[MetadataWorkUnit]:
-        dashboards = self.client.all_dashboards(fields="id")
+        dashboards = self.client.all_dashboards(
+            fields="id", transport_options=self.source_config.transport_options
+        )
         deleted_dashboards = (
-            self.client.search_dashboards(fields="id", deleted="true")
+            self.client.search_dashboards(
+                fields="id",
+                deleted="true",
+                transport_options=self.source_config.transport_options,
+            )
             if self.source_config.include_deleted
             else []
         )
@@ -809,8 +823,8 @@ def get_workunits(self) -> Iterable[MetadataWorkUnit]:
             ]
             for async_workunit in concurrent.futures.as_completed(async_workunits):
                 work_units, dashboard_id, start_time, end_time = async_workunit.result()
-                logger.info(
-                    f"Running time of process_dashboard for {dashboard_id} = {(end_time-start_time).total_seconds()}"
+                logger.debug(
+                    f"Running time of process_dashboard for {dashboard_id} = {(end_time - start_time).total_seconds()}"
                 )
                 self.reporter.report_upstream_latency(start_time, end_time)
                 for mwu in work_units:
diff --git a/metadata-ingestion/src/datahub/ingestion/source/looker_common.py b/metadata-ingestion/src/datahub/ingestion/source/looker_common.py
index e423666b9ef68..f366e061e78da 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/looker_common.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/looker_common.py
@@ -5,6 +5,7 @@
 from typing import Dict, Iterable, List, Optional, Tuple, Union
 
 from looker_sdk.error import SDKError
+from looker_sdk.rtl.transport import TransportOptions
 from looker_sdk.sdk.api31.methods import Looker31SDK
 from pydantic.class_validators import validator
 
@@ -309,10 +310,9 @@ def _get_field_type(
             # attempt Postgres modified type
             type_class = resolve_postgres_modified_type(native_type)
 
-        # if still not found, report a warning
+        # if still not found, log and continue
         if type_class is None:
-            reporter.report_warning(
-                native_type,
+            logger.info(
                 f"The type '{native_type}' is not recognized for field type, setting as NullTypeClass.",
             )
             type_class = NullTypeClass
@@ -500,9 +500,12 @@ def from_api(  # noqa: C901
         explore_name: str,
         client: Looker31SDK,
         reporter: SourceReport,
+        transport_options: Optional[TransportOptions],
     ) -> Optional["LookerExplore"]:  # noqa: C901
         try:
-            explore = client.lookml_model_explore(model, explore_name)
+            explore = client.lookml_model_explore(
+                model, explore_name, transport_options=transport_options
+            )
             views = set()
             if explore.joins is not None and explore.joins != []:
                 if explore.view_name is not None and explore.view_name != explore.name:
diff --git a/metadata-ingestion/src/datahub/ingestion/source/lookml.py b/metadata-ingestion/src/datahub/ingestion/source/lookml.py
index 0f2969ef1bee5..33ae1599b6c37 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/lookml.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/lookml.py
@@ -12,6 +12,7 @@
 
 import pydantic
 from looker_sdk.error import SDKError
+from looker_sdk.rtl.transport import TransportOptions
 from looker_sdk.sdk.api31.methods import Looker31SDK
 from looker_sdk.sdk.api31.models import DBConnection
 from pydantic import root_validator, validator
@@ -143,6 +144,7 @@ class LookMLSourceConfig(LookerCommonConfig):
     sql_parser: str = "datahub.utilities.sql_parser.DefaultSQLParser"
     api: Optional[LookerAPIConfig]
     project_name: Optional[str]
+    transport_options: Optional[TransportOptions]
 
     @validator("platform_instance")
     def platform_instance_not_supported(cls, v: str) -> str:
@@ -577,6 +579,7 @@ def from_looker_dict(
                 fields,
             )
             # also store the view logic and materialization
+            view_logic = looker_viewfile.raw_file_content
             if "sql" in derived_table:
                 view_logic = derived_table["sql"]
                 view_lang = "sql"
@@ -1003,7 +1006,11 @@ def get_project_name(self, model_name: str) -> str:
             self.looker_client is not None
         ), "Failed to find a configured Looker API client"
         try:
-            model = self.looker_client.lookml_model(model_name, "project_name")
+            model = self.looker_client.lookml_model(
+                model_name,
+                "project_name",
+                transport_options=self.source_config.transport_options,
+            )
             assert (
                 model.project_name is not None
             ), f"Failed to find a project name for model {model_name}"
diff --git a/metadata-ingestion/src/datahub/ingestion/source/openapi_parser.py b/metadata-ingestion/src/datahub/ingestion/source/openapi_parser.py
index c2f87b55d8871..830b6562755eb 100755
--- a/metadata-ingestion/src/datahub/ingestion/source/openapi_parser.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/openapi_parser.py
@@ -120,10 +120,14 @@ def get_endpoints(sw_dict: dict) -> dict:  # noqa: C901
         # will track only the "get" methods, which are the ones that give us data
         if "get" in p_o.keys():
 
-            try:
+            if "200" in p_o["get"]["responses"].keys():
                 base_res = p_o["get"]["responses"]["200"]
-            except KeyError:  # if you read a plain yml file the 200 will be an integer
+            elif 200 in p_o["get"]["responses"].keys():
+                # if you read a plain yml file the 200 will be an integer
                 base_res = p_o["get"]["responses"][200]
+            else:
+                # the endpoint does not have a 200 response
+                continue
 
             if "description" in p_o["get"].keys():
                 desc = p_o["get"]["description"]
diff --git a/metadata-ingestion/src/datahub/ingestion/source/sql/bigquery.py b/metadata-ingestion/src/datahub/ingestion/source/sql/bigquery.py
index 7cf432451b42d..cfea3ceaeebc6 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/sql/bigquery.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/sql/bigquery.py
@@ -30,6 +30,7 @@
 from datahub.ingestion.source.sql.sql_common import (
     SQLAlchemyConfig,
     SQLAlchemySource,
+    SQLSourceReport,
     SqlWorkUnit,
     make_sqlalchemy_type,
     register_custom_type,
@@ -284,13 +285,18 @@ class BigQueryDatasetKey(ProjectIdKey):
     dataset_id: str
 
 
-class BigQuerySource(SQLAlchemySource):
-    config: BigQueryConfig
-    maximum_shard_ids: Dict[str, str] = dict()
-    lineage_metadata: Optional[Dict[str, Set[str]]] = None
+@dataclass
+class BigQueryReport(SQLSourceReport):
+    pass
 
+
+class BigQuerySource(SQLAlchemySource):
     def __init__(self, config, ctx):
         super().__init__(config, ctx, "bigquery")
+        self.config: BigQueryConfig = config
+        self.report: BigQueryReport = BigQueryReport()
+        self.lineage_metadata: Optional[Dict[str, Set[str]]] = None
+        self.maximum_shard_ids: Dict[str, str] = dict()
 
     def get_db_name(self, inspector: Inspector = None) -> str:
         if self.config.project_id:
@@ -772,7 +778,6 @@ def prepare_profiler_args(
         partition: Optional[str],
         custom_sql: Optional[str] = None,
     ) -> dict:
-        self.config: BigQueryConfig
         return dict(
             schema=self.config.project_id,
             table=f"{schema}.{table}",
diff --git a/metadata-ingestion/src/datahub/ingestion/source/sql/redshift.py b/metadata-ingestion/src/datahub/ingestion/source/sql/redshift.py
index 1417b121cdd79..96934797f2d6e 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/sql/redshift.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/sql/redshift.py
@@ -1,12 +1,14 @@
+import logging
 from collections import defaultdict
 from dataclasses import dataclass, field
 from enum import Enum
-from typing import Dict, Iterable, List, Optional, Set, Tuple, Union
+from typing import Any, Dict, Iterable, List, Optional, Set, Tuple, Union
 from urllib.parse import urlparse
 
 # These imports verify that the dependencies are available.
 import psycopg2  # noqa: F401
 import pydantic  # noqa: F401
+import sqlalchemy
 import sqlalchemy_redshift  # noqa: F401
 from sqlalchemy import create_engine, inspect
 from sqlalchemy.engine import Connection, reflection
@@ -24,8 +26,8 @@
 from datahub.ingestion.source.sql.postgres import PostgresConfig
 from datahub.ingestion.source.sql.sql_common import (
     SQLAlchemySource,
+    SQLSourceReport,
     SqlWorkUnit,
-    logger,
 )
 
 # TRICKY: it's necessary to import the Postgres source because
@@ -41,6 +43,8 @@
     UpstreamClass,
 )
 
+logger: logging.Logger = logging.getLogger(__name__)
+
 
 class LineageMode(Enum):
     SQL_BASED = "sql_based"
@@ -336,18 +340,11 @@ def _get_schema_column_info(self, connection, schema=None, **kw):
 
 def _get_external_db_mapping(connection):
     # SQL query to get mapping of external schemas in redshift to its external database.
-    try:
-        result = connection.execute(
-            """
-            select * from svv_external_schemas
-            """
-        )
-        return result
-    except Exception as e:
-        logger.error(
-            "Error querying svv_external_schemas to get external database mapping.", e
-        )
-        return None
+    return connection.execute(
+        """
+        select * from svv_external_schemas
+        """
+    )
 
 
 # This monkey-patching enables us to batch fetch the table descriptions, rather than
@@ -360,15 +357,23 @@ def _get_external_db_mapping(connection):
 redshift_datetime_format = "%Y-%m-%d %H:%M:%S"
 
 
+@dataclass
+class RedshiftReport(SQLSourceReport):
+    # https://forums.aws.amazon.com/ann.jspa?annID=9105
+    saas_version: str = ""
+    upstream_lineage: Dict[str, List[str]] = field(default_factory=dict)
+
+
 class RedshiftSource(SQLAlchemySource):
-    config: RedshiftConfig
-    catalog_metadata: Dict = {}
     eskind_to_platform = {1: "glue", 2: "hive", 3: "postgres", 4: "redshift"}
 
     def __init__(self, config: RedshiftConfig, ctx: PipelineContext):
         super().__init__(config, ctx, "redshift")
+        self.catalog_metadata: Dict = {}
+        self.config: RedshiftConfig = config
         self._lineage_map: Optional[Dict[str, LineageItem]] = None
         self._all_tables_set: Optional[Set[str]] = None
+        self.report: RedshiftReport = RedshiftReport()
 
     @classmethod
     def create(cls, config_dict, ctx):
@@ -376,9 +381,12 @@ def create(cls, config_dict, ctx):
         return cls(config, ctx)
 
     def get_catalog_metadata(self, conn: Connection) -> None:
-        catalog_metadata = _get_external_db_mapping(conn)
-        if catalog_metadata is None:
+        try:
+            catalog_metadata = _get_external_db_mapping(conn)
+        except Exception as e:
+            self.error(logger, "external-svv_external_schemas", f"Error was {e}")
             return
+
         db_name = self.get_db_name()
 
         external_schema_mapping = {}
@@ -401,15 +409,30 @@ def get_catalog_metadata(self, conn: Connection) -> None:
     def get_inspectors(self) -> Iterable[Inspector]:
         # This method can be overridden in the case that you want to dynamically
         # run on multiple databases.
-        url = self.config.get_sql_alchemy_url()
-        logger.debug(f"sql_alchemy_url={url}")
-        engine = create_engine(url, **self.config.options)
+        engine = self.get_metadata_engine()
         with engine.connect() as conn:
             self.get_catalog_metadata(conn)
             inspector = inspect(conn)
             yield inspector
 
+    def get_metadata_engine(self) -> sqlalchemy.engine.Engine:
+        url = self.config.get_sql_alchemy_url()
+        logger.debug(f"sql_alchemy_url={url}")
+        return create_engine(url, **self.config.options)
+
+    def inspect_version(self) -> Any:
+        db_engine = self.get_metadata_engine()
+        logger.info("Checking current version")
+        for db_row in db_engine.execute("select version()"):
+            self.report.saas_version = db_row[0]
+
     def get_workunits(self) -> Iterable[Union[MetadataWorkUnit, SqlWorkUnit]]:
+        try:
+            self.inspect_version()
+        except Exception as e:
+            self.report.report_failure("version", f"Error: {e}")
+            return
+
         for wu in super().get_workunits():
             yield wu
             if (
@@ -471,7 +494,7 @@ def _get_all_tables(self) -> Set[str]:
             table_schema as schemaname,
             table_name as tablename
         from
-            information_schema.tables
+            pg_catalog.svv_tables
         where
             table_type = 'BASE TABLE'
             and table_schema not in ('information_schema', 'pg_catalog', 'pg_internal')
@@ -498,9 +521,7 @@ def _get_all_tables(self) -> Set[str]:
         db_name = self.get_db_name()
         all_tables_set = set()
 
-        url = self.config.get_sql_alchemy_url()
-        logger.debug(f"sql_alchemy_url={url}")
-        engine = create_engine(url, **self.config.options)
+        engine = self.get_metadata_engine()
         for db_row in engine.execute(all_tables_query):
             all_tables_set.add(
                 f'{db_name}.{db_row["schemaname"]}.{db_row["tablename"]}'
@@ -554,9 +575,7 @@ def _populate_lineage_map(
         if not self._all_tables_set:
             self._all_tables_set = self._get_all_tables()
 
-        url = self.config.get_sql_alchemy_url()
-        logger.debug(f"sql_alchemy_url={url}")
-        engine = create_engine(url, **self.config.options)
+        engine = self.get_metadata_engine()
 
         db_name = self.get_db_name()
 
@@ -593,17 +612,21 @@ def _populate_lineage_map(
                         )
                     except Exception as e:
                         target.query_parser_failed_sqls.append(db_row["ddl"])
-                        logger.warning(
+                        self.warn(
+                            logger,
+                            "parsing-query",
                             f'Error parsing query {db_row["ddl"]} for getting lineage .'
-                            f"\nError was {e}."
+                            f"\nError was {e}.",
                         )
                 else:
                     if lineage_type == lineage_type.COPY:
                         platform = LineageDatasetPlatform.S3
                         path = db_row["filename"].strip()
                         if urlparse(path).scheme != "s3":
-                            logger.warning(
-                                f"Only s3 source supported with copy. The source was: {path}.  ."
+                            self.warn(
+                                logger,
+                                "non-s3-lineage",
+                                f"Only s3 source supported with copy. The source was: {path}.",
                             )
                             continue
                     else:
@@ -624,7 +647,9 @@ def _populate_lineage_map(
                         source.platform == LineageDatasetPlatform.REDSHIFT
                         and source.path not in self._all_tables_set
                     ):
-                        logger.warning(f"{source.path} missing table")
+                        self.warn(
+                            logger, "missing-table", f"{source.path} missing table"
+                        )
                         continue
 
                     target.upstreams.add(source)
@@ -648,10 +673,7 @@ def _populate_lineage_map(
                 )
 
         except Exception as e:
-            logger.warning(
-                f"Extracting {lineage_type.name} lineage from Redshift failed."
-                f"Continuing...\nError was {e}."
-            )
+            self.warn(logger, f"extract-{lineage_type.name}", f"Error was {e}")
 
     def _populate_lineage(self) -> None:
 
@@ -957,7 +979,11 @@ def get_lineage_mcp(
         if custom_properties:
             properties = DatasetPropertiesClass(customProperties=custom_properties)
 
-        if not upstream_lineage:
+        if upstream_lineage:
+            self.report.upstream_lineage[dataset_urn] = [
+                u.dataset for u in upstream_lineage
+            ]
+        else:
             return None, properties
 
         mcp = MetadataChangeProposalWrapper(
diff --git a/metadata-ingestion/src/datahub/ingestion/source/sql/snowflake.py b/metadata-ingestion/src/datahub/ingestion/source/sql/snowflake.py
index 3e5299f06fd57..ab454ee20cd3e 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/sql/snowflake.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/sql/snowflake.py
@@ -1,7 +1,6 @@
 import json
 import logging
 from collections import defaultdict
-from dataclasses import dataclass, field
 from typing import Any, Dict, Iterable, List, Optional, Set, Tuple, Union
 
 import pydantic
@@ -9,35 +8,25 @@
 # This import verifies that the dependencies are available.
 import snowflake.sqlalchemy  # noqa: F401
 import sqlalchemy.engine
-from cryptography.hazmat.backends import default_backend
-from cryptography.hazmat.primitives import serialization
-from snowflake.connector.network import (
-    DEFAULT_AUTHENTICATOR,
-    EXTERNAL_BROWSER_AUTHENTICATOR,
-    KEY_PAIR_AUTHENTICATOR,
-)
 from snowflake.sqlalchemy import custom_types, snowdialect
 from sqlalchemy import create_engine, inspect
 from sqlalchemy.engine.reflection import Inspector
 from sqlalchemy.sql import sqltypes, text
 
 import datahub.emitter.mce_builder as builder
-from datahub.configuration.common import AllowDenyPattern
-from datahub.configuration.time_window_config import BaseTimeWindowConfig
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.ingestion.api.common import PipelineContext
 from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.ingestion.source.aws.s3_util import make_s3_urn
 from datahub.ingestion.source.sql.sql_common import (
     RecordTypeClass,
-    SQLAlchemyConfig,
     SQLAlchemySource,
-    SQLSourceReport,
     SqlWorkUnit,
     TimeTypeClass,
-    make_sqlalchemy_uri,
     register_custom_type,
 )
+from datahub.ingestion.source_config.sql.snowflake import SnowflakeConfig
+from datahub.ingestion.source_report.sql.snowflake import SnowflakeReport
 from datahub.metadata.com.linkedin.pegasus2avro.dataset import (
     DatasetLineageTypeClass,
     UpstreamClass,
@@ -54,146 +43,9 @@
 
 logger: logging.Logger = logging.getLogger(__name__)
 
-APPLICATION_NAME = "acryl_datahub"
-
 snowdialect.ischema_names["GEOGRAPHY"] = sqltypes.NullType
 
 
-@dataclass
-class SnowflakeReport(SQLSourceReport):
-    num_table_to_table_edges_scanned: int = 0
-    num_table_to_view_edges_scanned: int = 0
-    num_view_to_table_edges_scanned: int = 0
-    num_external_table_edges_scanned: int = 0
-    upstream_lineage: Dict[str, List[str]] = field(default_factory=dict)
-    # https://community.snowflake.com/s/topic/0TO0Z000000Unu5WAC/releases
-    saas_version: str = ""
-    role: str = ""
-    role_grants: List[str] = field(default_factory=list)
-
-
-class BaseSnowflakeConfig(BaseTimeWindowConfig):
-    # Note: this config model is also used by the snowflake-usage source.
-
-    scheme = "snowflake"
-
-    username: Optional[str] = None
-    password: Optional[pydantic.SecretStr] = pydantic.Field(default=None, exclude=True)
-    private_key_path: Optional[str]
-    private_key_password: Optional[pydantic.SecretStr] = pydantic.Field(
-        default=None, exclude=True
-    )
-    authentication_type: Optional[str] = "DEFAULT_AUTHENTICATOR"
-    host_port: str
-    warehouse: Optional[str]
-    role: Optional[str]
-    include_table_lineage: Optional[bool] = True
-    include_view_lineage: Optional[bool] = True
-
-    connect_args: Optional[dict]
-
-    @pydantic.validator("authentication_type", always=True)
-    def authenticator_type_is_valid(cls, v, values, **kwargs):
-        valid_auth_types = {
-            "DEFAULT_AUTHENTICATOR": DEFAULT_AUTHENTICATOR,
-            "EXTERNAL_BROWSER_AUTHENTICATOR": EXTERNAL_BROWSER_AUTHENTICATOR,
-            "KEY_PAIR_AUTHENTICATOR": KEY_PAIR_AUTHENTICATOR,
-        }
-        if v not in valid_auth_types.keys():
-            raise ValueError(
-                f"unsupported authenticator type '{v}' was provided,"
-                f" use one of {list(valid_auth_types.keys())}"
-            )
-        else:
-            if v == "KEY_PAIR_AUTHENTICATOR":
-                # If we are using key pair auth, we need the private key path and password to be set
-                if values.get("private_key_path") is None:
-                    raise ValueError(
-                        f"'private_key_path' was none "
-                        f"but should be set when using {v} authentication"
-                    )
-                if values.get("private_key_password") is None:
-                    raise ValueError(
-                        f"'private_key_password' was none "
-                        f"but should be set when using {v} authentication"
-                    )
-            logger.info(f"using authenticator type '{v}'")
-        return valid_auth_types.get(v)
-
-    @pydantic.validator("include_view_lineage")
-    def validate_include_view_lineage(cls, v, values):
-        if not values.get("include_table_lineage") and v:
-            raise ValueError(
-                "include_table_lineage must be True for include_view_lineage to be set."
-            )
-        return v
-
-    def get_sql_alchemy_url(self, database=None):
-        return make_sqlalchemy_uri(
-            self.scheme,
-            self.username,
-            self.password.get_secret_value() if self.password else None,
-            self.host_port,
-            f'"{database}"' if database is not None else database,
-            uri_opts={
-                # Drop the options if value is None.
-                key: value
-                for (key, value) in {
-                    "authenticator": self.authentication_type,
-                    "warehouse": self.warehouse,
-                    "role": self.role,
-                    "application": APPLICATION_NAME,
-                }.items()
-                if value
-            },
-        )
-
-    def get_sql_alchemy_connect_args(self) -> dict:
-        if self.authentication_type != KEY_PAIR_AUTHENTICATOR:
-            return {}
-        if self.connect_args is None:
-            if self.private_key_path is None:
-                raise ValueError("missing required private key path to read key from")
-            if self.private_key_password is None:
-                raise ValueError("missing required private key password")
-            with open(self.private_key_path, "rb") as key:
-                p_key = serialization.load_pem_private_key(
-                    key.read(),
-                    password=self.private_key_password.get_secret_value().encode(),
-                    backend=default_backend(),
-                )
-
-            pkb = p_key.private_bytes(
-                encoding=serialization.Encoding.DER,
-                format=serialization.PrivateFormat.PKCS8,
-                encryption_algorithm=serialization.NoEncryption(),
-            )
-            self.connect_args = {"private_key": pkb}
-        return self.connect_args
-
-
-class SnowflakeConfig(BaseSnowflakeConfig, SQLAlchemyConfig):
-    database_pattern: AllowDenyPattern = AllowDenyPattern(
-        deny=[r"^UTIL_DB$", r"^SNOWFLAKE$", r"^SNOWFLAKE_SAMPLE_DATA$"]
-    )
-
-    database: Optional[str]  # deprecated
-
-    @pydantic.validator("database")
-    def note_database_opt_deprecation(cls, v, values, **kwargs):
-        logger.warning(
-            "snowflake's `database` option has been deprecated; use database_pattern instead"
-        )
-        values["database_pattern"].allow = f"^{v}$"
-        return None
-
-    def get_sql_alchemy_url(self, database: str = None) -> str:
-        return super().get_sql_alchemy_url(database=database)
-
-    def get_sql_alchemy_connect_args(self) -> dict:
-        return super().get_sql_alchemy_connect_args()
-
-
 class SnowflakeSource(SQLAlchemySource):
     def __init__(self, config: SnowflakeConfig, ctx: PipelineContext):
         super().__init__(config, ctx, "snowflake")
@@ -201,6 +53,7 @@ def __init__(self, config: SnowflakeConfig, ctx: PipelineContext):
         self._external_lineage_map: Optional[Dict[str, Set[str]]] = None
         self.report: SnowflakeReport = SnowflakeReport()
         self.config: SnowflakeConfig = config
+        self.provision_role_in_progress: bool = False
 
     @classmethod
     def create(cls, config_dict, ctx):
@@ -210,7 +63,20 @@ def create(cls, config_dict, ctx):
     def get_metadata_engine(
         self, database: Optional[str] = None
     ) -> sqlalchemy.engine.Engine:
-        url = self.config.get_sql_alchemy_url(database=database)
+        if self.provision_role_in_progress and self.config.provision_role is not None:
+            username: Optional[str] = self.config.provision_role.admin_username
+            password: Optional[
+                pydantic.SecretStr
+            ] = self.config.provision_role.admin_password
+            role: Optional[str] = self.config.provision_role.admin_role
+        else:
+            username = self.config.username
+            password = self.config.password
+            role = self.config.role
+
+        url = self.config.get_sql_alchemy_url(
+            database=database, username=username, password=password, role=role
+        )
         logger.debug(f"sql_alchemy_url={url}")
         return create_engine(
             url,
@@ -643,8 +509,104 @@ def _get_upstream_lineage_info(
             return UpstreamLineage(upstreams=upstream_tables), column_lineage
         return None
 
+    def add_config_to_report(self):
+        self.report.cleaned_host_port = self.config.host_port
+
+        if self.config.provision_role is not None:
+            self.report.run_ingestion = self.config.provision_role.run_ingestion
+
+    def do_provision_role_internal(self):
+        provision_role_block = self.config.provision_role
+        if provision_role_block is None:
+            return
+        self.report.provision_role_done = not provision_role_block.dry_run
+
+        role = self.config.role
+        if role is None:
+            role = "datahub_role"
+            self.warn(
+                logger,
+                "role-grant",
+                f"role not specified during provision role using {role} as default",
+            )
+        self.report.role = role
+
+        warehouse = self.config.warehouse
+
+        logger.info("Creating connection for provision_role")
+        engine = self.get_metadata_engine(database=None)
+
+        sqls: List[str] = []
+        if provision_role_block.drop_role_if_exists:
+            sqls.append(f"DROP ROLE IF EXISTS {role}")
+
+        sqls.append(f"CREATE ROLE IF NOT EXISTS {role}")
+
+        if warehouse is None:
+            self.warn(
+                logger, "role-grant", "warehouse not specified during provision role"
+            )
+        else:
+            sqls.append(f"grant operate, usage on warehouse {warehouse} to role {role}")
+
+        for inspector in self.get_inspectors():
+            db_name = self.get_db_name(inspector)
+            sqls.extend(
+                [
+                    f"grant usage on DATABASE {db_name} to role {role}",
+                    f"grant usage on all schemas in database {db_name} to role {role}",
+                    f"grant select on all tables in database {db_name} to role {role}",
+                    f"grant select on all external tables in database {db_name} to role {role}",
+                    f"grant select on all views in database {db_name} to role {role}",
+                    f"grant usage on future schemas in database {db_name} to role {role}",
+                    f"grant select on future tables in database {db_name} to role {role}",
+                ]
+            )
+        if self.config.username is not None:
+            sqls.append(f"grant role {role} to user {self.config.username}")
+
+        sqls.append(f"grant imported privileges on database snowflake to role {role}")
+
+        dry_run_str = "[DRY RUN] " if provision_role_block.dry_run else ""
+        for sql in sqls:
+            logger.info(f"{dry_run_str} Attempting to run sql {sql}")
+            if provision_role_block.dry_run:
+                continue
+            try:
+                engine.execute(sql)
+            except Exception as e:
+                self.error(logger, "role-grant", f"Exception: {e}")
+
+        self.report.provision_role_success = not provision_role_block.dry_run
+
+    def do_provision_role(self):
+        if (
+            self.config.provision_role is None
+            or self.config.provision_role.enabled is False
+        ):
+            return
+
+        try:
+            self.provision_role_in_progress = True
+            self.do_provision_role_internal()
+        finally:
+            self.provision_role_in_progress = False
+
+    def should_run_ingestion(self) -> bool:
+        return (
+            self.config.provision_role is None
+            or self.config.provision_role.enabled is False
+            or self.config.provision_role.run_ingestion
+        )
+
     # Override the base class method.
     def get_workunits(self) -> Iterable[Union[MetadataWorkUnit, SqlWorkUnit]]:
+        self.add_config_to_report()
+
+        self.do_provision_role()
+        if not self.should_run_ingestion():
+            return
+
         try:
             self.inspect_version()
         except Exception as e:
diff --git a/metadata-ingestion/src/datahub/ingestion/source/sql/sql_common.py b/metadata-ingestion/src/datahub/ingestion/source/sql/sql_common.py
index 2ada238ae17bc..fbbdf93b7e0b1 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/sql/sql_common.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/sql/sql_common.py
@@ -42,7 +42,6 @@
     gen_containers,
 )
 from datahub.ingestion.api.common import PipelineContext
-from datahub.ingestion.api.source import SourceReport
 from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.ingestion.source.state.checkpoint import Checkpoint
 from datahub.ingestion.source.state.sql_common_state import (
@@ -52,6 +51,7 @@
     JobId,
     StatefulIngestionConfig,
     StatefulIngestionConfigBase,
+    StatefulIngestionReport,
     StatefulIngestionSourceBase,
 )
 from datahub.metadata.com.linkedin.pegasus2avro.common import StatusClass
@@ -161,7 +161,7 @@ class SqlContainerSubTypes(str, Enum):
 
 
 @dataclass
-class SQLSourceReport(SourceReport):
+class SQLSourceReport(StatefulIngestionReport):
     tables_scanned: int = 0
     views_scanned: int = 0
     entities_profiled: int = 0
@@ -413,7 +413,7 @@ def __init__(self, config: SQLAlchemyConfig, ctx: PipelineContext, platform: str
         super(SQLAlchemySource, self).__init__(config, ctx)
         self.config = config
         self.platform = platform
-        self.report = SQLSourceReport()
+        self.report: SQLSourceReport = SQLSourceReport()
 
         config_report = {
             config_option: config.dict().get(config_option)
@@ -447,7 +447,7 @@ def warn(self, log: logging.Logger, key: str, reason: str) -> Any:
 
     def error(self, log: logging.Logger, key: str, reason: str) -> Any:
         self.report.report_failure(key, reason)
-        log.error(reason)
+        log.error(f"{key} => {reason}")
 
     def get_inspectors(self) -> Iterable[Inspector]:
         # This method can be overridden in the case that you want to dynamically
diff --git a/metadata-ingestion/src/datahub/ingestion/source/state/stateful_ingestion_base.py b/metadata-ingestion/src/datahub/ingestion/source/state/stateful_ingestion_base.py
index a6968146313cf..22b6290f1f9eb 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/state/stateful_ingestion_base.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/state/stateful_ingestion_base.py
@@ -1,5 +1,6 @@
 import logging
 import platform
+from dataclasses import dataclass
 from datetime import datetime
 from typing import Any, Dict, Optional, Type, cast
 
@@ -20,7 +21,7 @@
 from datahub.ingestion.api.ingestion_job_reporting_provider_base import (
     IngestionReportingProviderBase,
 )
-from datahub.ingestion.api.source import Source
+from datahub.ingestion.api.source import Source, SourceReport
 from datahub.ingestion.source.state.checkpoint import Checkpoint, CheckpointStateBase
 from datahub.ingestion.source.state_provider.state_provider_registry import (
     ingestion_checkpoint_provider_registry,
@@ -65,6 +66,11 @@ class StatefulIngestionConfigBase(DatasetSourceConfigBase):
     stateful_ingestion: Optional[StatefulIngestionConfig] = None
 
 
+@dataclass
+class StatefulIngestionReport(SourceReport):
+    pass
+
+
 class StatefulIngestionSourceBase(Source):
     """
     Defines the base class for all stateful sources.
@@ -80,6 +86,15 @@ def __init__(
         self.cur_checkpoints: Dict[JobId, Optional[Checkpoint]] = {}
         self.run_summaries_to_report: Dict[JobId, DatahubIngestionRunSummaryClass] = {}
         self._initialize_checkpointing_state_provider()
+        self.report: StatefulIngestionReport = StatefulIngestionReport()
+
+    def warn(self, log: logging.Logger, key: str, reason: str) -> Any:
+        self.report.report_warning(key, reason)
+        log.warning(reason)
+
+    def error(self, log: logging.Logger, key: str, reason: str) -> Any:
+        self.report.report_failure(key, reason)
+        log.error(f"{key} => {reason}")
 
     #
     # Checkpointing specific support.
diff --git a/metadata-ingestion/src/datahub/ingestion/source/usage/bigquery_usage.py b/metadata-ingestion/src/datahub/ingestion/source/usage/bigquery_usage.py
index 2cec7afd7c1c7..1bda876415c9d 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/usage/bigquery_usage.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/usage/bigquery_usage.py
@@ -572,13 +572,10 @@ def report_dropped(self, key: str) -> None:
 
 
 class BigQueryUsageSource(Source):
-    config: BigQueryUsageConfig
-    report: BigQueryUsageSourceReport
-
     def __init__(self, config: BigQueryUsageConfig, ctx: PipelineContext):
         super().__init__(ctx)
-        self.config = config
-        self.report = BigQueryUsageSourceReport()
+        self.config: BigQueryUsageConfig = config
+        self.report: BigQueryUsageSourceReport = BigQueryUsageSourceReport()
 
     @classmethod
     def create(cls, config_dict: dict, ctx: PipelineContext) -> "BigQueryUsageSource":
diff --git a/metadata-ingestion/src/datahub/ingestion/source/usage/redshift_usage.py b/metadata-ingestion/src/datahub/ingestion/source/usage/redshift_usage.py
index 3c75665964be0..c75e00ff0fbd8 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/usage/redshift_usage.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/usage/redshift_usage.py
@@ -1,6 +1,7 @@
 import collections
 import dataclasses
 import logging
+from dataclasses import dataclass
 from datetime import datetime
 from typing import Any, Dict, Iterable, List, Optional, Union
 
@@ -13,6 +14,7 @@
 import datahub.emitter.mce_builder as builder
 from datahub.configuration.time_window_config import get_time_bucket
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
+from datahub.ingestion.api.common import PipelineContext
 from datahub.ingestion.api.source import Source, SourceReport
 from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.ingestion.source.sql.redshift import RedshiftConfig
@@ -87,15 +89,21 @@ def get_sql_alchemy_url(self):
         return super().get_sql_alchemy_url()
 
 
+@dataclass
+class RedshiftUsageReport(SourceReport):
+    pass
+
+
 @dataclasses.dataclass
 class RedshiftUsageSource(Source):
-    config: RedshiftUsageConfig
-    report: SourceReport = dataclasses.field(default_factory=SourceReport)
+    def __init__(self, config: RedshiftUsageConfig, ctx: PipelineContext):
+        self.config: RedshiftUsageConfig = config
+        self.report: RedshiftUsageReport = RedshiftUsageReport()
 
     @classmethod
     def create(cls, config_dict, ctx):
         config = RedshiftUsageConfig.parse_obj(config_dict)
-        return cls(ctx, config)
+        return cls(config, ctx)
 
     def get_workunits(self) -> Iterable[MetadataWorkUnit]:
         """Gets Redshift usage stats as work units"""
diff --git a/metadata-ingestion/src/datahub/ingestion/source/usage/snowflake_usage.py b/metadata-ingestion/src/datahub/ingestion/source/usage/snowflake_usage.py
index 8b20e92876832..a9c31517e7682 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/usage/snowflake_usage.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/usage/snowflake_usage.py
@@ -15,21 +15,16 @@
 from datahub.configuration.time_window_config import get_time_bucket
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.ingestion.api.common import PipelineContext
-from datahub.ingestion.api.source import SourceReport
 from datahub.ingestion.api.workunit import MetadataWorkUnit
-from datahub.ingestion.source.sql.snowflake import BaseSnowflakeConfig
 from datahub.ingestion.source.state.checkpoint import Checkpoint
 from datahub.ingestion.source.state.stateful_ingestion_base import (
     JobId,
-    StatefulIngestionConfig,
-    StatefulIngestionConfigBase,
     StatefulIngestionSourceBase,
 )
 from datahub.ingestion.source.state.usage_common_state import BaseUsageCheckpointState
-from datahub.ingestion.source.usage.usage_common import (
-    BaseUsageConfig,
-    GenericAggregatedDataset,
-)
+from datahub.ingestion.source.usage.usage_common import GenericAggregatedDataset
+from datahub.ingestion.source_config.usage.snowflake_usage import SnowflakeUsageConfig
+from datahub.ingestion.source_report.usage.snowflake_usage import SnowflakeUsageReport
 from datahub.metadata.schema_classes import (
     ChangeTypeClass,
     JobStatusClass,
@@ -37,6 +32,7 @@
     OperationTypeClass,
     TimeWindowSizeClass,
 )
+from datahub.utilities.perf_timer import PerfTimer
 
 logger = logging.getLogger(__name__)
 
@@ -126,50 +122,11 @@ class SnowflakeJoinedAccessEvent(PermissiveModel):
     role_name: str
 
 
-class SnowflakeStatefulIngestionConfig(StatefulIngestionConfig):
-    """
-    Specialization of basic StatefulIngestionConfig to adding custom config.
-    This will be used to override the stateful_ingestion config param of StatefulIngestionConfigBase
-    in the SnowflakeUsageConfig.
-    """
-
-    ignore_old_state = pydantic.Field(False, alias="force_rerun")
-
-
-class SnowflakeUsageConfig(
-    BaseSnowflakeConfig, BaseUsageConfig, StatefulIngestionConfigBase
-):
-    options: dict = {}
-    database_pattern: AllowDenyPattern = AllowDenyPattern(
-        deny=[r"^UTIL_DB$", r"^SNOWFLAKE$", r"^SNOWFLAKE_SAMPLE_DATA$"]
-    )
-    email_domain: Optional[str]
-    schema_pattern: AllowDenyPattern = AllowDenyPattern.allow_all()
-    table_pattern: AllowDenyPattern = AllowDenyPattern.allow_all()
-    view_pattern: AllowDenyPattern = AllowDenyPattern.allow_all()
-    apply_view_usage_to_tables: bool = False
-    stateful_ingestion: Optional[SnowflakeStatefulIngestionConfig] = None
-
-    @pydantic.validator("role", always=True)
-    def role_accountadmin(cls, v):
-        if not v or v.lower() != "accountadmin":
-            # This isn't an error, since the privileges can be delegated to other
-            # roles as well: https://docs.snowflake.com/en/sql-reference/account-usage.html#enabling-account-usage-for-other-roles
-            logger.info(
-                'snowflake usage tables are only accessible by role "accountadmin" by default; you set %s',
-                v,
-            )
-        return v
-
-    def get_sql_alchemy_url(self):
-        return super().get_sql_alchemy_url(database="snowflake")
-
-
 class SnowflakeUsageSource(StatefulIngestionSourceBase):
     def __init__(self, config: SnowflakeUsageConfig, ctx: PipelineContext):
         super(SnowflakeUsageSource, self).__init__(config, ctx)
         self.config: SnowflakeUsageConfig = config
-        self.report: SourceReport = SourceReport()
+        self.report: SnowflakeUsageReport = SnowflakeUsageReport()
         self.should_skip_this_run = self._should_skip_this_run()
 
     @classmethod
@@ -275,7 +232,23 @@ def update_default_job_summary(self) -> None:
             summary.config = self.config.json()
             summary.custom_summary = self.report.as_string()
 
+    def check_email_domain_missing(self) -> Any:
+        if self.config.email_domain is not None and self.config.email_domain != "":
+            return
+
+        self.warn(
+            logger,
+            "missing-email-domain",
+            "User's without email address will be ignored from usage if you don't set email_domain property",
+        )
+
+    def add_config_to_report(self):
+        self.report.start_time = self.config.start_time
+        self.report.end_time = self.config.end_time
+
     def get_workunits(self) -> Iterable[MetadataWorkUnit]:
+        self.add_config_to_report()
+        self.check_email_domain_missing()
         if not self.should_skip_this_run:
             # Initialize the checkpoints
             self._init_checkpoints()
@@ -316,107 +289,143 @@ def _make_sql_engine(self) -> Engine:
         )
         return engine
 
+    def _check_usage_date_ranges(self, engine: Engine) -> Any:
+
+        query = """
+            select
+                min(query_start_time) as min_time,
+                max(query_start_time) as max_time
+            from snowflake.account_usage.access_history
+            where ARRAY_SIZE(base_objects_accessed) > 0
+        """
+        with PerfTimer() as timer:
+            for db_row in engine.execute(query):
+                self.report.min_access_history_time = db_row[0].astimezone(
+                    tz=timezone.utc
+                )
+                self.report.max_access_history_time = db_row[1].astimezone(
+                    tz=timezone.utc
+                )
+                self.report.access_history_range_query_secs = round(
+                    timer.elapsed_seconds(), 2
+                )
+                break
+
+    def _is_unsupported_object_accessed(self, obj: Dict[str, Any]) -> bool:
+        unsupported_keys = ["locations"]
+        return any([obj.get(key) is not None for key in unsupported_keys])
+
+    def _is_object_valid(self, obj: Dict[str, Any]) -> bool:
+        if self._is_unsupported_object_accessed(
+            obj
+        ) or not self._is_dataset_pattern_allowed(
+            obj.get("objectName"), obj.get("objectDomain")
+        ):
+            return False
+        return True
+
+    def _is_dataset_pattern_allowed(
+        self, dataset_name: Optional[Any], dataset_type: Optional[Any]
+    ) -> bool:
+        # TODO: support table/view patterns for usage logs by pulling that information as well from the usage query
+        if not dataset_type or not dataset_name:
+            return True
+
+        table_or_view_pattern: Optional[AllowDenyPattern] = AllowDenyPattern.allow_all()
+        # Test domain type = external_table and then add it
+        table_or_view_pattern = (
+            self.config.table_pattern
+            if dataset_type.lower() in {"table"}
+            else (
+                self.config.view_pattern
+                if dataset_type.lower() in {"view", "materialized_view"}
+                else None
+            )
+        )
+        if table_or_view_pattern is None:
+            return True
+
+        dataset_params = dataset_name.split(".")
+        assert len(dataset_params) == 3
+        if (
+            not self.config.database_pattern.allowed(dataset_params[0])
+            or not self.config.schema_pattern.allowed(dataset_params[1])
+            or not table_or_view_pattern.allowed(dataset_params[2])
+        ):
+            return False
+        return True
+
+    def _process_snowflake_history_row(
+        self, row: Any
+    ) -> Iterable[SnowflakeJoinedAccessEvent]:
+        self.report.rows_processed += 1
+        # Make some minor type conversions.
+        if hasattr(row, "_asdict"):
+            # Compat with SQLAlchemy 1.3 and 1.4
+            # See https://docs.sqlalchemy.org/en/14/changelog/migration_14.html#rowproxy-is-no-longer-a-proxy-is-now-called-row-and-behaves-like-an-enhanced-named-tuple.
+            event_dict = row._asdict()
+        else:
+            event_dict = dict(row)
+
+        # no use processing events that don't have a query text
+        if not event_dict["query_text"]:
+            self.report.rows_missing_query_text += 1
+            return
+
+        event_dict["base_objects_accessed"] = [
+            obj
+            for obj in json.loads(event_dict["base_objects_accessed"])
+            if self._is_object_valid(obj)
+        ]
+        if len(event_dict["base_objects_accessed"]) == 0:
+            self.report.rows_zero_base_objects_accessed += 1
+
+        event_dict["direct_objects_accessed"] = [
+            obj
+            for obj in json.loads(event_dict["direct_objects_accessed"])
+            if self._is_object_valid(obj)
+        ]
+        if len(event_dict["direct_objects_accessed"]) == 0:
+            self.report.rows_zero_direct_objects_accessed += 1
+
+        event_dict["query_start_time"] = (event_dict["query_start_time"]).astimezone(
+            tz=timezone.utc
+        )
+
+        if not event_dict["email"] and self.config.email_domain:
+            if not event_dict["user_name"]:
+                self.report.report_warning("user-name-miss", f"Missing in {event_dict}")
+                logger.warning(
+                    f"The user_name is missing from {event_dict}. Skipping ...."
+                )
+                self.report.rows_missing_email += 1
+                return
+
+            event_dict[
+                "email"
+            ] = f'{event_dict["user_name"]}@{self.config.email_domain}'.lower()
+
+        try:  # big hammer try block to ensure we don't fail on parsing events
+            event = SnowflakeJoinedAccessEvent(**event_dict)
+            yield event
+        except Exception as e:
+            self.report.rows_parsing_error += 1
+            self.warn(logger, "usage", f"Failed to parse usage line {event_dict}, {e}")
+
     def _get_snowflake_history(self) -> Iterable[SnowflakeJoinedAccessEvent]:
-        query = self._make_usage_query()
         engine = self._make_sql_engine()
 
-        results = engine.execute(query)
+        logger.info("Checking usage date ranges")
+        self._check_usage_date_ranges(engine)
+
+        logger.info("Getting usage history")
+        with PerfTimer() as timer:
+            query = self._make_usage_query()
+            results = engine.execute(query)
+            self.report.access_history_query_secs = round(timer.elapsed_seconds(), 2)
 
         for row in results:
-            # Make some minor type conversions.
-            if hasattr(row, "_asdict"):
-                # Compat with SQLAlchemy 1.3 and 1.4
-                # See https://docs.sqlalchemy.org/en/14/changelog/migration_14.html#rowproxy-is-no-longer-a-proxy-is-now-called-row-and-behaves-like-an-enhanced-named-tuple.
-                event_dict = row._asdict()
-            else:
-                event_dict = dict(row)
-
-            # no use processing events that don't have a query text
-            if not event_dict["query_text"]:
-                continue
-
-            def is_unsupported_object_accessed(obj: Dict[str, Any]) -> bool:
-                unsupported_keys = ["locations"]
-                return any([obj.get(key) is not None for key in unsupported_keys])
-
-            def is_dataset_pattern_allowed(
-                dataset_name: Optional[Any], dataset_type: Optional[Any]
-            ) -> bool:
-                # TODO: support table/view patterns for usage logs by pulling that information as well from the usage query
-                if not dataset_type or not dataset_name:
-                    return True
-
-                table_or_view_pattern: Optional[
-                    AllowDenyPattern
-                ] = AllowDenyPattern.allow_all()
-                # Test domain type = external_table and then add it
-                table_or_view_pattern = (
-                    self.config.table_pattern
-                    if dataset_type.lower() in {"table"}
-                    else (
-                        self.config.view_pattern
-                        if dataset_type.lower() in {"view", "materialized_view"}
-                        else None
-                    )
-                )
-                if table_or_view_pattern is None:
-                    return True
-
-                dataset_params = dataset_name.split(".")
-                assert len(dataset_params) == 3
-                if (
-                    not self.config.database_pattern.allowed(dataset_params[0])
-                    or not self.config.schema_pattern.allowed(dataset_params[1])
-                    or not table_or_view_pattern.allowed(dataset_params[2])
-                ):
-                    return False
-                return True
-
-            def is_object_valid(obj: Dict[str, Any]) -> bool:
-                if is_unsupported_object_accessed(
-                    obj
-                ) or not is_dataset_pattern_allowed(
-                    obj.get("objectName"), obj.get("objectDomain")
-                ):
-                    return False
-                return True
-
-            event_dict["base_objects_accessed"] = [
-                obj
-                for obj in json.loads(event_dict["base_objects_accessed"])
-                if is_object_valid(obj)
-            ]
-            event_dict["direct_objects_accessed"] = [
-                obj
-                for obj in json.loads(event_dict["direct_objects_accessed"])
-                if is_object_valid(obj)
-            ]
-            event_dict["query_start_time"] = (
-                event_dict["query_start_time"]
-            ).astimezone(tz=timezone.utc)
-
-            if not event_dict["email"] and self.config.email_domain:
-                if not event_dict["user_name"]:
-                    self.report.report_warning(
-                        "user-name-miss", f"Missing in {event_dict}"
-                    )
-                    logger.warning(
-                        f"The user_name is missing from {event_dict}. Skipping ...."
-                    )
-                    continue
-
-                event_dict[
-                    "email"
-                ] = f'{event_dict["user_name"]}@{self.config.email_domain}'.lower()
-
-            try:  # big hammer try block to ensure we don't fail on parsing events
-                event = SnowflakeJoinedAccessEvent(**event_dict)
-                yield event
-            except Exception as e:
-                logger.warning(f"Failed to parse usage line {event_dict}", e)
-                self.report.report_warning(
-                    "usage", f"Failed to parse usage line {event_dict}"
-                )
+            yield from self._process_snowflake_history_row(row)
 
     def _get_operation_aspect_work_unit(
         self, event: SnowflakeJoinedAccessEvent
diff --git a/metadata-ingestion/src/datahub/ingestion/source_config/__init__.py b/metadata-ingestion/src/datahub/ingestion/source_config/__init__.py
new file mode 100644
index 0000000000000..e69de29bb2d1d
diff --git a/metadata-ingestion/src/datahub/ingestion/source_config/sql/__init__.py b/metadata-ingestion/src/datahub/ingestion/source_config/sql/__init__.py
new file mode 100644
index 0000000000000..e69de29bb2d1d
diff --git a/metadata-ingestion/src/datahub/ingestion/source_config/sql/snowflake.py b/metadata-ingestion/src/datahub/ingestion/source_config/sql/snowflake.py
new file mode 100644
index 0000000000000..87477cbbe8690
--- /dev/null
+++ b/metadata-ingestion/src/datahub/ingestion/source_config/sql/snowflake.py
@@ -0,0 +1,210 @@
+import logging
+from typing import Optional
+
+import pydantic
+from cryptography.hazmat.backends import default_backend
+from cryptography.hazmat.primitives import serialization
+from snowflake.connector.network import (
+    DEFAULT_AUTHENTICATOR,
+    EXTERNAL_BROWSER_AUTHENTICATOR,
+    KEY_PAIR_AUTHENTICATOR,
+)
+
+from datahub.configuration.common import AllowDenyPattern, ConfigModel
+from datahub.configuration.time_window_config import BaseTimeWindowConfig
+from datahub.ingestion.source.sql.sql_common import (
+    SQLAlchemyConfig,
+    make_sqlalchemy_uri,
+)
+from datahub.utilities.config_clean import (
+    remove_protocol,
+    remove_suffix,
+    remove_trailing_slashes,
+)
+
+APPLICATION_NAME = "acryl_datahub"
+
+logger: logging.Logger = logging.getLogger(__name__)
+
+
+class SnowflakeProvisionRoleConfig(ConfigModel):
+    enabled: bool = False
+
+    # Can be used by account admin to test what sql statements will be run
+    dry_run: bool = False
+
+    # Setting this to True is helpful in case you want a clean role without any extra privileges
+    # Not set to True by default because multiple parallel
+    #   snowflake ingestions can be dependent on single role
+    drop_role_if_exists: bool = False
+
+    # When Account admin is testing they might not want to actually do the ingestion
+    # Set this to False in case the account admin would want to
+    #   create role
+    #   grant role to user in main config
+    #   run ingestion as the user in main config
+    run_ingestion: bool = False
+
+    admin_role: Optional[str] = "accountadmin"
+
+    admin_username: str
+    admin_password: pydantic.SecretStr = pydantic.Field(default=None, exclude=True)
+
+    @pydantic.validator("admin_username", always=True)
+    def username_not_empty(cls, v, values, **kwargs):
+        v_str: str = str(v)
+        if v_str.strip() == "":
+            raise ValueError("username is empty")
+        return v
+
+
+class BaseSnowflakeConfig(BaseTimeWindowConfig):
+    # Note: this config model is also used by the snowflake-usage source.
+
+    scheme = "snowflake"
+
+    username: Optional[str] = None
+    password: Optional[pydantic.SecretStr] = pydantic.Field(default=None, exclude=True)
+    private_key_path: Optional[str]
+    private_key_password: Optional[pydantic.SecretStr] = pydantic.Field(
+        default=None, exclude=True
+    )
+    authentication_type: Optional[str] = "DEFAULT_AUTHENTICATOR"
+    host_port: str
+    warehouse: Optional[str]
+    role: Optional[str]
+    include_table_lineage: Optional[bool] = True
+    include_view_lineage: Optional[bool] = True
+
+    connect_args: Optional[dict]
+
+    @pydantic.validator("host_port", always=True)
+    def host_port_is_valid(cls, v, values, **kwargs):
+        v = remove_protocol(v)
+        v = remove_trailing_slashes(v)
+        v = remove_suffix(v, ".snowflakecomputing.com")
+        logger.info(f"Cleaned Host port is {v}")
+        return v
+
+    @pydantic.validator("authentication_type", always=True)
+    def authenticator_type_is_valid(cls, v, values, **kwargs):
+        valid_auth_types = {
+            "DEFAULT_AUTHENTICATOR": DEFAULT_AUTHENTICATOR,
+            "EXTERNAL_BROWSER_AUTHENTICATOR": EXTERNAL_BROWSER_AUTHENTICATOR,
+            "KEY_PAIR_AUTHENTICATOR": KEY_PAIR_AUTHENTICATOR,
+        }
+        if v not in valid_auth_types.keys():
+            raise ValueError(
+                f"unsupported authenticator type '{v}' was provided,"
+                f" use one of {list(valid_auth_types.keys())}"
+            )
+        else:
+            if v == "KEY_PAIR_AUTHENTICATOR":
+                # If we are using key pair auth, we need the private key path and password to be set
+                if values.get("private_key_path") is None:
+                    raise ValueError(
+                        f"'private_key_path' was none "
+                        f"but should be set when using {v} authentication"
+                    )
+                if values.get("private_key_password") is None:
+                    raise ValueError(
+                        f"'private_key_password' was none "
+                        f"but should be set when using {v} authentication"
+                    )
+            logger.info(f"using authenticator type '{v}'")
+        return valid_auth_types.get(v)
+
+    @pydantic.validator("include_view_lineage")
+    def validate_include_view_lineage(cls, v, values):
+        if not values.get("include_table_lineage") and v:
+            raise ValueError(
+                "include_table_lineage must be True for include_view_lineage to be set."
+            )
+        return v
+
+    def get_sql_alchemy_url(
+        self,
+        database: Optional[str] = None,
+        username: Optional[str] = None,
+        password: Optional[pydantic.SecretStr] = None,
+        role: Optional[str] = None,
+    ) -> str:
+        if username is None:
+            username = self.username
+        if password is None:
+            password = self.password
+        if role is None:
+            role = self.role
+        return make_sqlalchemy_uri(
+            self.scheme,
+            username,
+            password.get_secret_value() if password else None,
+            self.host_port,
+            f'"{database}"' if database is not None else database,
+            uri_opts={
+                # Drop the options if value is None.
+                key: value
+                for (key, value) in {
+                    "authenticator": self.authentication_type,
+                    "warehouse": self.warehouse,
+                    "role": role,
+                    "application": APPLICATION_NAME,
+                }.items()
+                if value
+            },
+        )
+
+    def get_sql_alchemy_connect_args(self) -> dict:
+        if self.authentication_type != KEY_PAIR_AUTHENTICATOR:
+            return {}
+        if self.connect_args is None:
+            if self.private_key_path is None:
+                raise ValueError("missing required private key path to read key from")
+            if self.private_key_password is None:
+                raise ValueError("missing required private key password")
+            with open(self.private_key_path, "rb") as key:
+                p_key = serialization.load_pem_private_key(
+                    key.read(),
+                    password=self.private_key_password.get_secret_value().encode(),
+                    backend=default_backend(),
+                )
+
+            pkb = p_key.private_bytes(
+                encoding=serialization.Encoding.DER,
+                format=serialization.PrivateFormat.PKCS8,
+                encryption_algorithm=serialization.NoEncryption(),
+            )
+            self.connect_args = {"private_key": pkb}
+        return self.connect_args
+
+
+class SnowflakeConfig(BaseSnowflakeConfig, SQLAlchemyConfig):
+    database_pattern: AllowDenyPattern = AllowDenyPattern(
+        deny=[r"^UTIL_DB$", r"^SNOWFLAKE$", r"^SNOWFLAKE_SAMPLE_DATA$"]
+    )
+
+    database: Optional[str]  # deprecated
+
+    provision_role: Optional[SnowflakeProvisionRoleConfig] = None
+
+    @pydantic.validator("database")
+    def note_database_opt_deprecation(cls, v, values, **kwargs):
+        logger.warning(
+            "snowflake's `database` option has been deprecated; use database_pattern instead"
+        )
+        values["database_pattern"].allow = f"^{v}$"
+        return None
+
+    def get_sql_alchemy_url(
+        self,
+        database: str = None,
+        username: Optional[str] = None,
+        password: Optional[pydantic.SecretStr] = None,
+        role: Optional[str] = None,
+    ) -> str:
+        return super().get_sql_alchemy_url(
+            database=database, username=username, password=password, role=role
+        )
+
+    def get_sql_alchemy_connect_args(self) -> dict:
+        return super().get_sql_alchemy_connect_args()
diff --git a/metadata-ingestion/src/datahub/ingestion/source_config/usage/__init__.py b/metadata-ingestion/src/datahub/ingestion/source_config/usage/__init__.py
new file mode 100644
index 0000000000000..e69de29bb2d1d
diff --git a/metadata-ingestion/src/datahub/ingestion/source_config/usage/snowflake_usage.py b/metadata-ingestion/src/datahub/ingestion/source_config/usage/snowflake_usage.py
new file mode 100644
index 0000000000000..415bfb163ab1e
--- /dev/null
+++ b/metadata-ingestion/src/datahub/ingestion/source_config/usage/snowflake_usage.py
@@ -0,0 +1,58 @@
+import logging
+from typing import Optional
+
+import pydantic
+
+from datahub.configuration.common import AllowDenyPattern
+from datahub.ingestion.source.state.stateful_ingestion_base import (
+    StatefulIngestionConfig,
+    StatefulIngestionConfigBase,
+)
+from datahub.ingestion.source.usage.usage_common import BaseUsageConfig
+from datahub.ingestion.source_config.sql.snowflake import BaseSnowflakeConfig
+
+logger = logging.getLogger(__name__)
+
+
+class SnowflakeStatefulIngestionConfig(StatefulIngestionConfig):
+    """
+    Specialization of basic StatefulIngestionConfig to adding custom config.
+    This will be used to override the stateful_ingestion config param of StatefulIngestionConfigBase
+    in the SnowflakeUsageConfig.
+    """
+
+    ignore_old_state = pydantic.Field(False, alias="force_rerun")
+
+
+class SnowflakeUsageConfig(
+    BaseSnowflakeConfig, BaseUsageConfig, StatefulIngestionConfigBase
+):
+    options: dict = {}
+    database_pattern: AllowDenyPattern = AllowDenyPattern(
+        deny=[r"^UTIL_DB$", r"^SNOWFLAKE$", r"^SNOWFLAKE_SAMPLE_DATA$"]
+    )
+    email_domain: Optional[str]
+    schema_pattern: AllowDenyPattern = AllowDenyPattern.allow_all()
+    table_pattern: AllowDenyPattern = AllowDenyPattern.allow_all()
+    view_pattern: AllowDenyPattern = AllowDenyPattern.allow_all()
+    apply_view_usage_to_tables: bool = False
+    stateful_ingestion: Optional[SnowflakeStatefulIngestionConfig] = None
+
+    @pydantic.validator("role", always=True)
+    def role_accountadmin(cls, v):
+        if not v or v.lower() != "accountadmin":
+            # This isn't an error, since the privileges can be delegated to other
+            # roles as well: https://docs.snowflake.com/en/sql-reference/account-usage.html#enabling-account-usage-for-other-roles
+            logger.info(
+                'snowflake usage tables are only accessible by role "accountadmin" by default; you set %s',
+                v,
+            )
+        return v
+
+    def get_sql_alchemy_url(self):
+        return super().get_sql_alchemy_url(
+            database="snowflake",
+            username=self.username,
+            password=self.password,
+            role=self.role,
+        )
diff --git a/metadata-ingestion/src/datahub/ingestion/source_report/__init__.py b/metadata-ingestion/src/datahub/ingestion/source_report/__init__.py
new file mode 100644
index 0000000000000..e69de29bb2d1d
diff --git a/metadata-ingestion/src/datahub/ingestion/source_report/sql/__init__.py b/metadata-ingestion/src/datahub/ingestion/source_report/sql/__init__.py
new file mode 100644
index 0000000000000..e69de29bb2d1d
diff --git a/metadata-ingestion/src/datahub/ingestion/source_report/sql/snowflake.py b/metadata-ingestion/src/datahub/ingestion/source_report/sql/snowflake.py
new file mode 100644
index 0000000000000..81b1a95d68270
--- /dev/null
+++ b/metadata-ingestion/src/datahub/ingestion/source_report/sql/snowflake.py
@@ -0,0 +1,29 @@
+from dataclasses import dataclass, field
+from typing import Dict, List
+
+from datahub.ingestion.source.sql.sql_common import SQLSourceReport
+from datahub.ingestion.source_report.time_window import BaseTimeWindowReport
+
+
+@dataclass
+class BaseSnowflakeReport(BaseTimeWindowReport):
+    pass
+
+
+@dataclass
+class SnowflakeReport(BaseSnowflakeReport, SQLSourceReport):
+    num_table_to_table_edges_scanned: int = 0
+    num_table_to_view_edges_scanned: int = 0
+    num_view_to_table_edges_scanned: int = 0
+    num_external_table_edges_scanned: int = 0
+    upstream_lineage: Dict[str, List[str]] = field(default_factory=dict)
+
+    cleaned_host_port: str = ""
+    run_ingestion: bool = False
+    provision_role_done: bool = False
+    provision_role_success: bool = False
+
+    # https://community.snowflake.com/s/topic/0TO0Z000000Unu5WAC/releases
+    saas_version: str = ""
+    role: str = ""
+    role_grants: List[str] = field(default_factory=list)
diff --git a/metadata-ingestion/src/datahub/ingestion/source_report/time_window.py b/metadata-ingestion/src/datahub/ingestion/source_report/time_window.py
new file mode 100644
index 0000000000000..db8cbbca5ad0b
--- /dev/null
+++ b/metadata-ingestion/src/datahub/ingestion/source_report/time_window.py
@@ -0,0 +1,9 @@
+from dataclasses import dataclass
+from datetime import datetime
+from typing import Optional
+
+
+@dataclass
+class BaseTimeWindowReport:
+    end_time: Optional[datetime] = None
+    start_time: Optional[datetime] = None
diff --git a/metadata-ingestion/src/datahub/ingestion/source_report/usage/__init__.py b/metadata-ingestion/src/datahub/ingestion/source_report/usage/__init__.py
new file mode 100644
index 0000000000000..e69de29bb2d1d
diff --git a/metadata-ingestion/src/datahub/ingestion/source_report/usage/snowflake_usage.py b/metadata-ingestion/src/datahub/ingestion/source_report/usage/snowflake_usage.py
new file mode 100644
index 0000000000000..5f7962fc36710
--- /dev/null
+++ b/metadata-ingestion/src/datahub/ingestion/source_report/usage/snowflake_usage.py
@@ -0,0 +1,23 @@
+from dataclasses import dataclass
+from datetime import datetime
+from typing import Optional
+
+from datahub.ingestion.source.state.stateful_ingestion_base import (
+    StatefulIngestionReport,
+)
+from datahub.ingestion.source_report.sql.snowflake import BaseSnowflakeReport
+
+
+@dataclass
+class SnowflakeUsageReport(BaseSnowflakeReport, StatefulIngestionReport):
+    min_access_history_time: Optional[datetime] = None
+    max_access_history_time: Optional[datetime] = None
+    access_history_range_query_secs: float = -1
+    access_history_query_secs: float = -1
+
+    rows_processed: int = 0
+    rows_missing_query_text: int = 0
+    rows_zero_base_objects_accessed: int = 0
+    rows_zero_direct_objects_accessed: int = 0
+    rows_missing_email: int = 0
+    rows_parsing_error: int = 0
diff --git a/metadata-ingestion/src/datahub/integrations/great_expectations/action.py b/metadata-ingestion/src/datahub/integrations/great_expectations/action.py
index 9c97fa1c5b9c1..a19b7d7b22e78 100644
--- a/metadata-ingestion/src/datahub/integrations/great_expectations/action.py
+++ b/metadata-ingestion/src/datahub/integrations/great_expectations/action.py
@@ -51,7 +51,7 @@
 from datahub.metadata.com.linkedin.pegasus2avro.common import DataPlatformInstance
 from datahub.metadata.com.linkedin.pegasus2avro.events.metadata import ChangeType
 from datahub.metadata.schema_classes import PartitionSpecClass, PartitionTypeClass
-from datahub.utilities.sql_parser import MetadataSQLSQLParser
+from datahub.utilities.sql_parser import DefaultSQLParser
 
 logger = logging.getLogger(__name__)
 
@@ -71,6 +71,7 @@ def __init__(
         retry_status_codes: Optional[List[int]] = None,
         retry_max_times: Optional[int] = None,
         extra_headers: Optional[Dict[str, str]] = None,
+        parse_table_names_from_sql: bool = False,
     ):
         super().__init__(data_context)
         self.server_url = server_url
@@ -82,6 +83,7 @@ def __init__(
         self.retry_status_codes = retry_status_codes
         self.retry_max_times = retry_max_times
         self.extra_headers = extra_headers
+        self.parse_table_names_from_sql = parse_table_names_from_sql
 
     def _run(
         self,
@@ -598,6 +600,12 @@ def get_dataset_partitions(self, batch_identifier, data_asset):
                     }
                 )
             elif isinstance(ge_batch_spec, RuntimeQueryBatchSpec):
+                if not self.parse_table_names_from_sql:
+                    warn(
+                        "Enable parse_table_names_from_sql in DatahubValidationAction config\
+                            to try to parse the tables being asserted from SQL query"
+                    )
+                    return []
                 query = data_asset.batches[
                     batch_identifier
                 ].batch_request.runtime_parameters["query"]
@@ -610,11 +618,12 @@ def get_dataset_partitions(self, batch_identifier, data_asset):
                     query=query,
                     customProperties=batchSpecProperties,
                 )
-                tables = MetadataSQLSQLParser(query).get_tables()
+                tables = DefaultSQLParser(query).get_tables()
                 if len(set(tables)) != 1:
                     warn(
                         "DataHubValidationAction does not support cross dataset assertions."
                     )
+                    return []
                 for table in tables:
                     dataset_urn = make_dataset_urn_from_sqlalchemy_uri(
                         sqlalchemy_uri,
diff --git a/metadata-ingestion/src/datahub/telemetry/telemetry.py b/metadata-ingestion/src/datahub/telemetry/telemetry.py
index 746cf9c20638d..bd98f576d7dd2 100644
--- a/metadata-ingestion/src/datahub/telemetry/telemetry.py
+++ b/metadata-ingestion/src/datahub/telemetry/telemetry.py
@@ -29,6 +29,7 @@ class Telemetry:
 
     client_id: str
     enabled: bool = True
+    tracking_init: bool = False
 
     def __init__(self):
 
@@ -47,14 +48,6 @@ def __init__(self):
                 self.mp = Mixpanel(
                     MIXPANEL_TOKEN, consumer=Consumer(request_timeout=int(TIMEOUT))
                 )
-                self.mp.people_set(
-                    self.client_id,
-                    {
-                        "datahub_version": datahub_package.nice_version_name(),
-                        "os": platform.system(),
-                        "python_version": platform.python_version(),
-                    },
-                )
             except Exception as e:
                 logger.debug(f"Error connecting to mixpanel: {e}")
 
@@ -62,6 +55,7 @@ def update_config(self) -> None:
         """
         Update the config file with the current client ID and enabled status.
         """
+        logger.info("Updating telemetry config")
 
         if not DATAHUB_FOLDER.exists():
             os.makedirs(DATAHUB_FOLDER)
@@ -124,6 +118,24 @@ def load_config(self):
                     f"{CONFIG_FILE} had an IOError, please inspect this file for issues."
                 )
 
+    def init_tracking(self) -> None:
+        if not self.enabled or self.mp is None or self.tracking_init is True:
+            return
+
+        logger.info("Sending init Telemetry")
+        try:
+            self.mp.people_set(
+                self.client_id,
+                {
+                    "datahub_version": datahub_package.nice_version_name(),
+                    "os": platform.system(),
+                    "python_version": platform.python_version(),
+                },
+            )
+        except Exception as e:
+            logger.debug(f"Error reporting telemetry: {e}")
+        self.init_track = True
+
     def ping(
         self,
         action: str,
@@ -144,6 +156,7 @@ def ping(
 
         # send event
         try:
+            logger.info("Sending Telemetry")
             self.mp.track(self.client_id, action, properties)
 
         except Exception as e:
@@ -155,6 +168,13 @@ def ping(
 T = TypeVar("T")
 
 
+def set_telemetry_enable(enable: bool) -> Any:
+    telemetry_instance.enabled = enable
+    if not enable:
+        logger.info("Disabling Telemetry locally due to server config")
+    telemetry_instance.update_config()
+
+
 def get_full_class_name(obj):
     module = obj.__class__.__module__
     if module is None or module == str.__class__.__module__:
@@ -168,6 +188,7 @@ def wrapper(*args: Any, **kwargs: Any) -> Any:
 
         action = f"function:{func.__module__}.{func.__name__}"
 
+        telemetry_instance.init_tracking()
         telemetry_instance.ping(action)
         try:
             res = func(*args, **kwargs)
diff --git a/metadata-ingestion/src/datahub/utilities/config_clean.py b/metadata-ingestion/src/datahub/utilities/config_clean.py
index 8208772b6021f..64df846ae5bde 100644
--- a/metadata-ingestion/src/datahub/utilities/config_clean.py
+++ b/metadata-ingestion/src/datahub/utilities/config_clean.py
@@ -1,10 +1,16 @@
 import re
 
 
+def remove_suffix(inp: str, suffix: str, remove_all: bool = False) -> str:
+    while suffix and inp.endswith(suffix):
+        inp = inp[: -len(suffix)]
+        if not remove_all:
+            break
+    return inp
+
+
 def remove_trailing_slashes(url: str) -> str:
-    while url.endswith("/"):
-        url = url[:-1]
-    return url
+    return remove_suffix(url, "/", remove_all=True)
 
 
 def remove_protocol(url: str) -> str:
diff --git a/metadata-ingestion/src/datahub/utilities/server_config_util.py b/metadata-ingestion/src/datahub/utilities/server_config_util.py
new file mode 100644
index 0000000000000..c919a1356f264
--- /dev/null
+++ b/metadata-ingestion/src/datahub/utilities/server_config_util.py
@@ -0,0 +1,22 @@
+from typing import Any, Dict, Optional
+
+from datahub.telemetry.telemetry import set_telemetry_enable
+
+# Only to be written to for logging server related information
+global_debug: Dict[str, Any] = dict()
+
+
+def set_gms_config(config: Dict) -> Any:
+    global_debug["gms_config"] = config
+
+    cli_telemtry_enabled = is_cli_telemetry_enabled()
+    if cli_telemtry_enabled is not None:
+        set_telemetry_enable(cli_telemtry_enabled)
+
+
+def get_gms_config() -> Dict:
+    return global_debug.get("gms_config", {})
+
+
+def is_cli_telemetry_enabled() -> Optional[bool]:
+    return get_gms_config().get("telemetry", {}).get("enabledCli", None)
diff --git a/metadata-ingestion/src/datahub/utilities/sql_lineage_parser_impl.py b/metadata-ingestion/src/datahub/utilities/sql_lineage_parser_impl.py
new file mode 100644
index 0000000000000..49f77a08f7fdd
--- /dev/null
+++ b/metadata-ingestion/src/datahub/utilities/sql_lineage_parser_impl.py
@@ -0,0 +1,134 @@
+import logging
+import re
+import unittest
+import unittest.mock
+from typing import List, Set
+
+from sqllineage.core.holders import Column, SQLLineageHolder
+
+try:
+    import sqlparse
+    from networkx import DiGraph
+    from sqllineage.core import LineageAnalyzer
+
+    import datahub.utilities.sqllineage_patch
+except ImportError:
+    pass
+
+logger = logging.getLogger(__name__)
+
+
+class SqlLineageSQLParserImpl:
+    _DATE_SWAP_TOKEN = "__d_a_t_e"
+    _TIMESTAMP_SWAP_TOKEN = "__t_i_m_e_s_t_a_m_p"
+    _MYVIEW_SQL_TABLE_NAME_TOKEN = "__my_view__.__sql_table_name__"
+    _MYVIEW_LOOKER_TOKEN = "my_view.SQL_TABLE_NAME"
+
+    def __init__(self, sql_query: str) -> None:
+        original_sql_query = sql_query
+
+        # SqlLineageParser makes mistakes on lateral flatten queries, use the prefix
+        if "lateral flatten" in sql_query:
+            sql_query = sql_query[: sql_query.find("lateral flatten")]
+
+        # SqlLineageParser makes mistakes on columns called "date", rename them
+        sql_query = re.sub(
+            r"(\bdate\b)", rf"{self._DATE_SWAP_TOKEN}", sql_query, flags=re.IGNORECASE
+        )
+
+        # SqlLineageParser lowercarese tablenames and we need to replace Looker specific token which should be uppercased
+        sql_query = re.sub(
+            rf"(\${{{self._MYVIEW_LOOKER_TOKEN}}})",
+            rf"{self._MYVIEW_SQL_TABLE_NAME_TOKEN}",
+            sql_query,
+        )
+
+        # SqlLineageParser makes mistakes on columns called "timestamp", rename them
+        sql_query = re.sub(
+            r"(\btimestamp\b)",
+            rf"{self._TIMESTAMP_SWAP_TOKEN}",
+            sql_query,
+            flags=re.IGNORECASE,
+        )
+
+        # SqlLineageParser does not handle "encode" directives well. Remove them
+        sql_query = re.sub(r"\sencode [a-zA-Z]*", "", sql_query, flags=re.IGNORECASE)
+
+        # Replace lookml templates with the variable otherwise sqlparse can't parse ${
+        sql_query = re.sub(r"(\${)(.+)(})", r"\2", sql_query)
+        if sql_query != original_sql_query:
+            logger.debug(f"rewrote original query {original_sql_query} as {sql_query}")
+
+        self._sql = sql_query
+
+        self._stmt = [
+            s
+            for s in sqlparse.parse(
+                # first apply sqlparser formatting just to get rid of comments, which cause
+                # inconsistencies in parsing output
+                sqlparse.format(
+                    self._sql.strip(),
+                    strip_comments=True,
+                    use_space_around_operators=True,
+                ),
+            )
+            if s.token_first(skip_cm=True)
+        ]
+
+        with unittest.mock.patch(
+            "sqllineage.core.handlers.source.SourceHandler.end_of_query_cleanup",
+            datahub.utilities.sqllineage_patch.end_of_query_cleanup_patch,
+        ):
+            with unittest.mock.patch(
+                "sqllineage.core.holders.SubQueryLineageHolder.add_column_lineage",
+                datahub.utilities.sqllineage_patch.add_column_lineage_patch,
+            ):
+                self._stmt_holders = [
+                    LineageAnalyzer().analyze(stmt) for stmt in self._stmt
+                ]
+                self._sql_holder = SQLLineageHolder.of(*self._stmt_holders)
+
+    def get_tables(self) -> List[str]:
+        result: List[str] = list()
+        for table in self._sql_holder.source_tables:
+            table_normalized = re.sub(r"^<default>.", "", str(table))
+            result.append(str(table_normalized))
+
+        # We need to revert TOKEN replacements
+        result = ["date" if c == self._DATE_SWAP_TOKEN else c for c in result]
+        result = [
+            "timestamp" if c == self._TIMESTAMP_SWAP_TOKEN else c for c in list(result)
+        ]
+        result = [
+            self._MYVIEW_LOOKER_TOKEN if c == self._MYVIEW_SQL_TABLE_NAME_TOKEN else c
+            for c in result
+        ]
+
+        # Sort tables to make the list deterministic
+        result.sort()
+
+        return result
+
+    def get_columns(self) -> List[str]:
+        graph: DiGraph = self._sql_holder.graph  # For mypy attribute checking
+        column_nodes = [n for n in graph.nodes if isinstance(n, Column)]
+        column_graph = graph.subgraph(column_nodes)
+
+        target_columns = {column for column, deg in column_graph.out_degree if deg == 0}
+
+        result: Set[str] = set()
+        for column in target_columns:
+            # Let's drop all the count(*) and similard columns which are expression actually if it does not have an alias
+            if not any(ele in column.raw_name for ele in ["*", "(", ")"]):
+                result.add(str(column.raw_name))
+
+        # Reverting back all the previously renamed words which confuses the parser
+        result = set(["date" if c == self._DATE_SWAP_TOKEN else c for c in result])
+        result = set(
+            [
+                "timestamp" if c == self._TIMESTAMP_SWAP_TOKEN else c
+                for c in list(result)
+            ]
+        )
+        # swap back renamed date column
+        return list(result)
diff --git a/metadata-ingestion/src/datahub/utilities/sql_parser.py b/metadata-ingestion/src/datahub/utilities/sql_parser.py
index 8492bd16b91c1..eb0bc0ec8262f 100644
--- a/metadata-ingestion/src/datahub/utilities/sql_parser.py
+++ b/metadata-ingestion/src/datahub/utilities/sql_parser.py
@@ -1,19 +1,16 @@
 import logging
+import multiprocessing
 import re
-import unittest
-import unittest.mock
+import sys
+import traceback
 from abc import ABCMeta, abstractmethod
-from typing import List, Set
+from multiprocessing import Process, Queue
+from typing import List, Optional, Tuple, Type
 
-from sqllineage.core.holders import Column, SQLLineageHolder
+from datahub.utilities.sql_lineage_parser_impl import SqlLineageSQLParserImpl
 
 try:
-    import sqlparse
-    from networkx import DiGraph
     from sql_metadata import Parser as MetadataSQLParser
-    from sqllineage.core import LineageAnalyzer
-
-    import datahub.utilities.sqllineage_patch
 except ImportError:
     pass
 
@@ -85,122 +82,67 @@ def get_columns(self) -> List[str]:
         return ["date" if c == self._DATE_SWAP_TOKEN else c for c in filtered_cols]
 
 
-class SqlLineageSQLParser(SQLParser):
-    _DATE_SWAP_TOKEN = "__d_a_t_e"
-    _TIMESTAMP_SWAP_TOKEN = "__t_i_m_e_s_t_a_m_p"
-    _MYVIEW_SQL_TABLE_NAME_TOKEN = "__my_view__.__sql_table_name__"
-    _MYVIEW_LOOKER_TOKEN = "my_view.SQL_TABLE_NAME"
+def sql_lineage_parser_impl_func_wrapper(
+    queue: multiprocessing.Queue,
+    sql_query: str,
+) -> None:
+    """
+    The wrapper function that computes the tables and columns using the SqlLineageSQLParserImpl
+    and puts the results on the shared IPC queue. This is used to isolate SqlLineageSQLParserImpl
+    functionality in a separate process, and hence protect our sources from memory leaks originating in
+    the sqllineage module.
+    :param queue: The shared IPC queue on to which the results will be put.
+    :param sql_query: The SQL query to extract the tables & columns from.
+    :return: None.
+    """
+    exception_details: Optional[Tuple[Optional[Type[BaseException]], str]] = None
+    tables: List[str] = []
+    columns: List[str] = []
+    try:
+        parser = SqlLineageSQLParserImpl(sql_query)
+        tables = parser.get_tables()
+        columns = parser.get_columns()
+    except BaseException:
+        exc_info = sys.exc_info()
+        exc_msg: str = str(exc_info[1]) + "".join(traceback.format_tb(exc_info[2]))
+        exception_details = (exc_info[0], exc_msg)
+        logger.error(exc_msg)
+    finally:
+        queue.put((tables, columns, exception_details))
+
 
+class SqlLineageSQLParser(SQLParser):
     def __init__(self, sql_query: str) -> None:
         super().__init__(sql_query)
-
-        original_sql_query = sql_query
-
-        # SqlLineageParser makes mistakes on lateral flatten queries, use the prefix
-        if "lateral flatten" in sql_query:
-            sql_query = sql_query[: sql_query.find("lateral flatten")]
-
-        # SqlLineageParser makes mistakes on columns called "date", rename them
-        sql_query = re.sub(
-            r"(\bdate\b)", rf"{self._DATE_SWAP_TOKEN}", sql_query, flags=re.IGNORECASE
-        )
-
-        # SqlLineageParser lowercarese tablenames and we need to replace Looker specific token which should be uppercased
-        sql_query = re.sub(
-            rf"(\${{{self._MYVIEW_LOOKER_TOKEN}}})",
-            rf"{self._MYVIEW_SQL_TABLE_NAME_TOKEN}",
-            sql_query,
-        )
-
-        # SqlLineageParser makes mistakes on columns called "timestamp", rename them
-        sql_query = re.sub(
-            r"(\btimestamp\b)",
-            rf"{self._TIMESTAMP_SWAP_TOKEN}",
-            sql_query,
-            flags=re.IGNORECASE,
+        self.tables, self.columns = self._get_tables_columns_process_wrapped(sql_query)
+
+    @staticmethod
+    def _get_tables_columns_process_wrapped(
+        sql_query: str,
+    ) -> Tuple[List[str], List[str]]:
+        # Invoke sql_lineage_parser_impl_func_wrapper in a separate process to avoid
+        # memory leaks from sqllineage module used by SqlLineageSQLParserImpl. This will help
+        # shield our sources like lookml & redash, that need to parse a large number of SQL statements,
+        # from causing significant memory leaks in the datahub cli during ingestion.
+        queue: multiprocessing.Queue = Queue()
+        process: multiprocessing.Process = Process(
+            target=sql_lineage_parser_impl_func_wrapper,
+            args=(
+                queue,
+                sql_query,
+            ),
         )
-
-        # SqlLineageParser does not handle "encode" directives well. Remove them
-        sql_query = re.sub(r"\sencode [a-zA-Z]*", "", sql_query, flags=re.IGNORECASE)
-
-        # Replace lookml templates with the variable otherwise sqlparse can't parse ${
-        sql_query = re.sub(r"(\${)(.+)(})", r"\2", sql_query)
-        if sql_query != original_sql_query:
-            logger.debug(f"rewrote original query {original_sql_query} as {sql_query}")
-
-        self._sql = sql_query
-
-        self._stmt = [
-            s
-            for s in sqlparse.parse(
-                # first apply sqlparser formatting just to get rid of comments, which cause
-                # inconsistencies in parsing output
-                sqlparse.format(
-                    self._sql.strip(),
-                    strip_comments=True,
-                    use_space_around_operators=True,
-                ),
-            )
-            if s.token_first(skip_cm=True)
-        ]
-
-        with unittest.mock.patch(
-            "sqllineage.core.handlers.source.SourceHandler.end_of_query_cleanup",
-            datahub.utilities.sqllineage_patch.end_of_query_cleanup_patch,
-        ):
-            with unittest.mock.patch(
-                "sqllineage.core.holders.SubQueryLineageHolder.add_column_lineage",
-                datahub.utilities.sqllineage_patch.add_column_lineage_patch,
-            ):
-                self._stmt_holders = [
-                    LineageAnalyzer().analyze(stmt) for stmt in self._stmt
-                ]
-                self._sql_holder = SQLLineageHolder.of(*self._stmt_holders)
+        process.start()
+        tables, columns, exception_details = queue.get(block=True)
+        if exception_details is not None:
+            raise exception_details[0](f"Sub-process exception: {exception_details[1]}")
+        return tables, columns
 
     def get_tables(self) -> List[str]:
-        result: List[str] = list()
-        for table in self._sql_holder.source_tables:
-            table_normalized = re.sub(r"^<default>.", "", str(table))
-            result.append(str(table_normalized))
-
-        # We need to revert TOKEN replacements
-        result = ["date" if c == self._DATE_SWAP_TOKEN else c for c in result]
-        result = [
-            "timestamp" if c == self._TIMESTAMP_SWAP_TOKEN else c for c in list(result)
-        ]
-        result = [
-            self._MYVIEW_LOOKER_TOKEN if c == self._MYVIEW_SQL_TABLE_NAME_TOKEN else c
-            for c in result
-        ]
-
-        # Sort tables to make the list deterministic
-        result.sort()
-
-        return result
+        return self.tables
 
     def get_columns(self) -> List[str]:
-        graph: DiGraph = self._sql_holder.graph  # For mypy attribute checking
-        column_nodes = [n for n in graph.nodes if isinstance(n, Column)]
-        column_graph = graph.subgraph(column_nodes)
-
-        target_columns = {column for column, deg in column_graph.out_degree if deg == 0}
-
-        result: Set[str] = set()
-        for column in target_columns:
-            # Let's drop all the count(*) and similard columns which are expression actually if it does not have an alias
-            if not any(ele in column.raw_name for ele in ["*", "(", ")"]):
-                result.add(str(column.raw_name))
-
-        # Reverting back all the previously renamed words which confuses the parser
-        result = set(["date" if c == self._DATE_SWAP_TOKEN else c for c in result])
-        result = set(
-            [
-                "timestamp" if c == self._TIMESTAMP_SWAP_TOKEN else c
-                for c in list(result)
-            ]
-        )
-        # swap back renamed date column
-        return list(result)
+        return self.columns
 
 
 class DefaultSQLParser(SQLParser):
diff --git a/metadata-ingestion/src/datahub/utilities/sqlalchemy_query_combiner.py b/metadata-ingestion/src/datahub/utilities/sqlalchemy_query_combiner.py
index 7105d9f0f4393..0474f4ec7d3d6 100644
--- a/metadata-ingestion/src/datahub/utilities/sqlalchemy_query_combiner.py
+++ b/metadata-ingestion/src/datahub/utilities/sqlalchemy_query_combiner.py
@@ -213,7 +213,13 @@ def _handle_execute(
 
         # Figure out how many columns this query returns.
         # This also implicitly ensures that the typing is generally correct.
-        assert len(get_query_columns(query)) > 0
+        try:
+            assert len(get_query_columns(query)) > 0
+        except AttributeError as e:
+            logger.debug(
+                f"Query of type: '{type(query)}' does not contain attributes required by 'get_query_columns()'. AttributeError: {e}"
+            )
+            return False, None
 
         # Add query to the queue.
         queue = self._get_queue(main_greenlet)
diff --git a/metadata-ingestion/tests/integration/azure_ad/azure_ad_groups.json b/metadata-ingestion/tests/integration/azure_ad/azure_ad_groups.json
index 761a3c05a7237..e4b9ecfa7571b 100644
--- a/metadata-ingestion/tests/integration/azure_ad/azure_ad_groups.json
+++ b/metadata-ingestion/tests/integration/azure_ad/azure_ad_groups.json
@@ -1,6 +1,5 @@
 [
     {
-        "@odata.id": "https://graph.microsoft.com/v2/00000000-0000-0000-0000-000000000000/directoryObjects/005203a5-d73b-4b21-8077-d96ee309b454/Microsoft.DirectoryServices.Group",
         "id": "00000000-0000-0000-0000-000000000000",
         "deletedDateTime": null,
         "classification": null,
@@ -35,7 +34,6 @@
         "onPremisesProvisioningErrors": []
     },
     {
-        "@odata.id": "https://graph.microsoft.com/v2/00000000-0000-0000-0000-000000000001/directoryObjects/005203a5-d73b-4b21-8077-d96ee309b454/Microsoft.DirectoryServices.Group",
         "id": "00000000-0000-0000-0000-0000000000001",
         "deletedDateTime": null,
         "classification": null,
diff --git a/metadata-ingestion/tests/integration/azure_ad/azure_ad_mces_golden_nested_groups.json b/metadata-ingestion/tests/integration/azure_ad/azure_ad_mces_golden_nested_groups.json
new file mode 100644
index 0000000000000..0f9d738b1ca93
--- /dev/null
+++ b/metadata-ingestion/tests/integration/azure_ad/azure_ad_mces_golden_nested_groups.json
@@ -0,0 +1,150 @@
+[
+    {
+        "auditHeader": null,
+        "proposedSnapshot": {
+            "com.linkedin.pegasus2avro.metadata.snapshot.CorpGroupSnapshot": {
+                "urn": "urn:li:corpGroup:nestedgroupDisplayName1",
+                "aspects": [
+                    {
+                        "com.linkedin.pegasus2avro.identity.CorpGroupInfo": {
+                            "displayName": "nestedgroupDisplayName1",
+                            "email": null,
+                            "admins": [],
+                            "members": [],
+                            "groups": [],
+                            "description": "Placeholder group with members from nested groups."
+                        }
+                    }
+                ]
+            }
+        },
+        "proposedDelta": null,
+        "systemMetadata": {
+            "lastObserved": 1629795600000,
+            "runId": "test-azure-ad",
+            "registryName": null,
+            "registryVersion": null,
+            "properties": null
+        }
+    },
+    {
+        "auditHeader": null,
+        "proposedSnapshot": {
+            "com.linkedin.pegasus2avro.metadata.snapshot.CorpUserSnapshot": {
+                "urn": "urn:li:corpuser:foobar@acryl.io",
+                "aspects": [
+                    {
+                        "com.linkedin.pegasus2avro.identity.CorpUserInfo": {
+                            "active": true,
+                            "displayName": "Foo Bar",
+                            "email": "foobar@acryl.io",
+                            "title": null,
+                            "managerUrn": null,
+                            "departmentId": null,
+                            "departmentName": null,
+                            "firstName": "Foo",
+                            "lastName": "Bar",
+                            "fullName": "Foo Bar",
+                            "countryCode": null
+                        }
+                    },
+                    {
+                        "com.linkedin.pegasus2avro.identity.GroupMembership": {
+                            "groups": [
+                                "urn:li:corpGroup:nestedgroupDisplayName1"
+                            ]
+                        }
+                    }
+                ]
+            }
+        },
+        "proposedDelta": null,
+        "systemMetadata": {
+            "lastObserved": 1629795600000,
+            "runId": "test-azure-ad",
+            "registryName": null,
+            "registryVersion": null,
+            "properties": null
+        }
+    },
+    {
+        "auditHeader": null,
+        "proposedSnapshot": {
+            "com.linkedin.pegasus2avro.metadata.snapshot.CorpUserSnapshot": {
+                "urn": "urn:li:corpuser:johngreen@acryl.io",
+                "aspects": [
+                    {
+                        "com.linkedin.pegasus2avro.identity.CorpUserInfo": {
+                            "active": true,
+                            "displayName": "John Green",
+                            "email": "johngreen@acryl.io",
+                            "title": null,
+                            "managerUrn": null,
+                            "departmentId": null,
+                            "departmentName": null,
+                            "firstName": "John",
+                            "lastName": "Green",
+                            "fullName": "John Green",
+                            "countryCode": null
+                        }
+                    },
+                    {
+                        "com.linkedin.pegasus2avro.identity.GroupMembership": {
+                            "groups": [
+                                "urn:li:corpGroup:nestedgroupDisplayName1"
+                            ]
+                        }
+                    }
+                ]
+            }
+        },
+        "proposedDelta": null,
+        "systemMetadata": {
+            "lastObserved": 1629795600000,
+            "runId": "test-azure-ad",
+            "registryName": null,
+            "registryVersion": null,
+            "properties": null
+        }
+    },
+    {
+        "auditHeader": null,
+        "proposedSnapshot": {
+            "com.linkedin.pegasus2avro.metadata.snapshot.CorpUserSnapshot": {
+                "urn": "urn:li:corpuser:adamhall@acryl.io",
+                "aspects": [
+                    {
+                        "com.linkedin.pegasus2avro.identity.CorpUserInfo": {
+                            "active": true,
+                            "displayName": "Adam Hall",
+                            "email": "adamhall@acryl.io",
+                            "title": null,
+                            "managerUrn": null,
+                            "departmentId": null,
+                            "departmentName": null,
+                            "firstName": "Adam",
+                            "lastName": "Hall",
+                            "fullName": "Adam Hall",
+                            "countryCode": null
+                        }
+                    },
+                    {
+                        "com.linkedin.pegasus2avro.identity.GroupMembership": {
+                            "groups": [
+                                "urn:li:corpGroup:nestedgroupDisplayName1"
+                            ]
+                        }
+                    }
+                ]
+            }
+        },
+        "proposedDelta": null,
+        "systemMetadata": {
+            "lastObserved": 1629795600000,
+            "runId": "test-azure-ad",
+            "registryName": null,
+            "registryVersion": null,
+            "properties": null
+        }
+    }
+]
\ No newline at end of file
diff --git a/metadata-ingestion/tests/integration/azure_ad/azure_ad_nested_group.json b/metadata-ingestion/tests/integration/azure_ad/azure_ad_nested_group.json
new file mode 100644
index 0000000000000..f598f42afe912
--- /dev/null
+++ b/metadata-ingestion/tests/integration/azure_ad/azure_ad_nested_group.json
@@ -0,0 +1,36 @@
+[
+    {
+        "id": "99999999-9999-9999-9999-999999999999",
+        "deletedDateTime": null,
+        "classification": null,
+        "createdDateTime": "2021-08-20 11: 00: 00",
+        "creationOptions": [],
+        "description": "Placeholder group with members from nested groups.",
+        "displayName": "nestedgroupDisplayName1",
+        "expirationDateTime": null,
+        "groupTypes": [],
+        "isAssignableToRole": null,
+        "mail": null,
+        "mailEnabled": false,
+        "mailNickname": "nestedgroupDisplayName1",
+        "membershipRule": null,
+        "membershipRuleProcessingState": null,
+        "onPremisesDomainName": null,
+        "onPremisesLastSyncDateTime": null,
+        "onPremisesNetBiosName": null,
+        "onPremisesSamAccountName": null,
+        "onPremisesSecurityIdentifier": null,
+        "onPremisesSyncEnabled": null,
+        "preferredDataLocation": null,
+        "preferredLanguage": null,
+        "proxyAddresses": [],
+        "renewedDateTime": "2021-08-20 11:00:00",
+        "resourceBehaviorOptions": [],
+        "resourceProvisioningOptions": [],
+        "securityEnabled": true,
+        "securityIdentifier": "xxxxx",
+        "theme": null,
+        "visibility": null,
+        "onPremisesProvisioningErrors": []
+    }
+]
\ No newline at end of file
diff --git a/metadata-ingestion/tests/integration/azure_ad/azure_ad_nested_groups_users.json b/metadata-ingestion/tests/integration/azure_ad/azure_ad_nested_groups_users.json
new file mode 100644
index 0000000000000..37f558b6c1d82
--- /dev/null
+++ b/metadata-ingestion/tests/integration/azure_ad/azure_ad_nested_groups_users.json
@@ -0,0 +1,85 @@
+[
+    {
+        "@odata.type": "#microsoft.graph.user",
+        "displayName": "Foo Bar",
+        "surname": "Bar",
+        "givenName": "Foo",
+        "id": "fa379b0d5435ef75",
+        "userPrincipalName": "foobar@acryl.io",
+        "businessPhones": [],
+        "jobTitle": null,
+        "mail": "foobar@acryl.io",
+        "mobilePhone": null,
+        "officeLocation": null,
+        "preferredLanguage": null
+    },
+    {
+        "@odata.type": "#microsoft.graph.group",
+        "id": "00000000-0000-0000-0000-000000000000",
+        "deletedDateTime": null,
+        "classification": null,
+        "createdDateTime": "2021-08-20 11: 00: 00",
+        "creationOptions": [],
+        "description": null,
+        "displayName": "groupDisplayName1",
+        "expirationDateTime": null,
+        "groupTypes": [],
+        "isAssignableToRole": null,
+        "mail": null,
+        "mailEnabled": false,
+        "mailNickname": "groupDisplayName1",
+        "membershipRule": null,
+        "membershipRuleProcessingState": null,
+        "onPremisesDomainName": null,
+        "onPremisesLastSyncDateTime": null,
+        "onPremisesNetBiosName": null,
+        "onPremisesSamAccountName": null,
+        "onPremisesSecurityIdentifier": null,
+        "onPremisesSyncEnabled": null,
+        "preferredDataLocation": null,
+        "preferredLanguage": null,
+        "proxyAddresses": [],
+        "renewedDateTime": "2021-08-20 11:00:00",
+        "resourceBehaviorOptions": [],
+        "resourceProvisioningOptions": [],
+        "securityEnabled": true,
+        "securityIdentifier": "xxxxx",
+        "theme": null,
+        "visibility": null,
+        "onPremisesProvisioningErrors": []
+    },
+    {
+        "@odata.type": "#microsoft.graph.group",
+        "id": "00000000-0000-0000-0000-0000000000001",
+        "deletedDateTime": null,
+        "classification": null,
+        "createdDateTime": "2021-08-20 11: 00: 00",
+        "creationOptions": [],
+        "description": "This is an interesting description",
+        "displayName": "groupDisplayName2",
+        "expirationDateTime": null,
+        "groupTypes": [],
+        "isAssignableToRole": null,
+        "mail": "groupDisplayName2@onmicrosoft.com",
+        "mailEnabled": false,
+        "mailNickname": "groupDisplayName2",
+        "membershipRule": null,
+        "membershipRuleProcessingState": null,
+        "onPremisesDomainName": null,
+        "onPremisesLastSyncDateTime": null,
+        "onPremisesNetBiosName": null,
+        "onPremisesSamAccountName": null,
+        "onPremisesSecurityIdentifier": null,
+        "onPremisesSyncEnabled": null,
+        "preferredDataLocation": null,
+        "preferredLanguage": null,
+        "proxyAddresses": [],
+        "renewedDateTime": "2021-08-20 11:00:00",
+        "resourceBehaviorOptions": [],
+        "resourceProvisioningOptions": [],
+        "securityEnabled": true,
+        "securityIdentifier": "xxxxx",
+        "theme": null,
+        "visibility": null,
+        "onPremisesProvisioningErrors": []
+    }]
\ No newline at end of file
diff --git a/metadata-ingestion/tests/integration/azure_ad/azure_ad_users.json b/metadata-ingestion/tests/integration/azure_ad/azure_ad_users.json
index 09cdfa117b606..8a6252cf81015 100644
--- a/metadata-ingestion/tests/integration/azure_ad/azure_ad_users.json
+++ b/metadata-ingestion/tests/integration/azure_ad/azure_ad_users.json
@@ -1,5 +1,6 @@
 [
     {
+        "@odata.type": "#microsoft.graph.user",
         "displayName": "John Green",
         "surname": "Green",
         "givenName": "John",
@@ -13,6 +14,7 @@
         "preferredLanguage": null
     },
     {
+        "@odata.type": "#microsoft.graph.user",
         "displayName": "Adam Hall",
         "surname": "Hall",
         "givenName": "Adam",
diff --git a/metadata-ingestion/tests/integration/azure_ad/test_azure_ad.py b/metadata-ingestion/tests/integration/azure_ad/test_azure_ad.py
index 8c4492bddef5f..21e92b46bf6df 100644
--- a/metadata-ingestion/tests/integration/azure_ad/test_azure_ad.py
+++ b/metadata-ingestion/tests/integration/azure_ad/test_azure_ad.py
@@ -1,5 +1,6 @@
 import json
 import pathlib
+from typing import List
 from unittest.mock import patch
 
 from freezegun import freeze_time
@@ -65,7 +66,7 @@ def test_azure_ad_source_default_configs(pytestconfig, tmp_path):
     ) as mock_users, patch(
         "datahub.ingestion.source.identity.azure_ad.AzureADSource._get_azure_ad_groups"
     ) as mock_groups, patch(
-        "datahub.ingestion.source.identity.azure_ad.AzureADSource._get_azure_ad_group_users"
+        "datahub.ingestion.source.identity.azure_ad.AzureADSource._get_azure_ad_group_members"
     ) as mock_group_users:
         mocked_functions(
             test_resources_dir, mock_token, mock_users, mock_groups, mock_group_users
@@ -107,6 +108,67 @@ def test_azure_ad_source_default_configs(pytestconfig, tmp_path):
     )
 
 
+@freeze_time(FROZEN_TIME)
+def test_azure_ad_source_nested_groups(pytestconfig, tmp_path):
+
+    test_resources_dir: pathlib.Path = (
+        pytestconfig.rootpath / "tests/integration/azure_ad"
+    )
+
+    with patch(
+        "datahub.ingestion.source.identity.azure_ad.AzureADSource.get_token"
+    ) as mock_token, patch(
+        "datahub.ingestion.source.identity.azure_ad.AzureADSource._get_azure_ad_users"
+    ) as mock_users, patch(
+        "datahub.ingestion.source.identity.azure_ad.AzureADSource._get_azure_ad_groups"
+    ) as mock_groups, patch(
+        "datahub.ingestion.source.identity.azure_ad.AzureADSource._get_azure_ad_group_members"
+    ) as mock_group_users:
+        mocked_functions(
+            test_resources_dir,
+            mock_token,
+            mock_users,
+            mock_groups,
+            mock_group_users,
+            True,
+        )
+        # Run an azure usage ingestion run.
+        pipeline = Pipeline.create(
+            {
+                "run_id": "test-azure-ad",
+                "source": {
+                    "type": "azure-ad",
+                    "config": {
+                        "client_id": "00000000-0000-0000-0000-000000000000",
+                        "tenant_id": "00000000-0000-0000-0000-000000000000",
+                        "client_secret": "client_secret",
+                        "redirect": "https://login.microsoftonline.com/common/oauth2/nativeclient",
+                        "authority": "https://login.microsoftonline.com/00000000-0000-0000-0000-000000000000",
+                        "token_url": "https://login.microsoftonline.com/00000000-0000-0000-0000-000000000000/oauth2/token",
+                        "graph_url": "https://graph.microsoft.com/v1.0",
+                        "ingest_group_membership": True,
+                        "ingest_groups": True,
+                        "ingest_users": False,
+                    },
+                },
+                "sink": {
+                    "type": "file",
+                    "config": {
+                        "filename": f"{tmp_path}/azure_ad_mces_nested_groups.json",
+                    },
+                },
+            }
+        )
+        pipeline.run()
+        pipeline.raise_from_status()
+
+    mce_helpers.check_golden_file(
+        pytestconfig,
+        output_path=tmp_path / "azure_ad_mces_nested_groups.json",
+        golden_path=test_resources_dir / "azure_ad_mces_golden_nested_groups.json",
+    )
+
+
 @freeze_time(FROZEN_TIME)
 def test_azure_source_ingestion_disabled(pytestconfig, tmp_path):
 
@@ -121,7 +183,7 @@ def test_azure_source_ingestion_disabled(pytestconfig, tmp_path):
     ) as mock_users, patch(
         "datahub.ingestion.source.identity.azure_ad.AzureADSource._get_azure_ad_groups"
     ) as mock_groups, patch(
-        "datahub.ingestion.source.identity.azure_ad.AzureADSource._get_azure_ad_group_users"
+        "datahub.ingestion.source.identity.azure_ad.AzureADSource._get_azure_ad_group_members"
     ) as mock_group_users:
         mocked_functions(
             test_resources_dir, mock_token, mock_users, mock_groups, mock_group_users
@@ -167,6 +229,10 @@ def test_azure_source_ingestion_disabled(pytestconfig, tmp_path):
 def load_test_resources(test_resources_dir):
     azure_ad_users_json_file = test_resources_dir / "azure_ad_users.json"
     azure_ad_groups_json_file = test_resources_dir / "azure_ad_groups.json"
+    azure_ad_nested_group_json_file = test_resources_dir / "azure_ad_nested_group.json"
+    azure_ad_nested_groups_users_json_file = (
+        test_resources_dir / "azure_ad_nested_groups_users.json"
+    )
 
     with azure_ad_users_json_file.open() as azure_ad_users_json:
         reference_users = json.loads(azure_ad_users_json.read())
@@ -174,24 +240,58 @@ def load_test_resources(test_resources_dir):
     with azure_ad_groups_json_file.open() as azure_ad_groups_json:
         reference_groups = json.loads(azure_ad_groups_json.read())
 
-    return reference_users, reference_groups
+    with azure_ad_nested_group_json_file.open() as azure_ad_nested_group_json:
+        reference_nested_group = json.loads(azure_ad_nested_group_json.read())
+
+    with azure_ad_nested_groups_users_json_file.open() as azure_ad_nested_groups_users_json:
+        reference_nested_groups_users = json.loads(
+            azure_ad_nested_groups_users_json.read()
+        )
+
+    return (
+        reference_users,
+        reference_groups,
+        reference_nested_group,
+        reference_nested_groups_users,
+    )
 
 
+#
+# Azure offers a tool called 'graph-explorer' that you can use to generate fake data,
+# or use your tenant's Azure AD to generate real data:
+# https://developer.microsoft.com/graph/graph-explorer
+#
 def mocked_functions(
-    test_resources_dir, mock_token, mock_users, mock_groups, mock_groups_users
+    test_resources_dir,
+    mock_token,
+    mock_users,
+    mock_groups,
+    mock_groups_users,
+    return_nested_group=False,
 ):
     # mock token response
     mock_token.return_value = "xxxxxxxx"
 
     # mock users and groups response
-    users, groups = load_test_resources(test_resources_dir)
+    users, groups, nested_group, nested_group_members = load_test_resources(
+        test_resources_dir
+    )
     mock_users.return_value = iter(list([users]))
-    mock_groups.return_value = iter(list([groups]))
+    mock_groups.return_value = (
+        iter(list([nested_group])) if return_nested_group else iter(list([groups]))
+    )
 
     # For simplicity, each user is placed in ALL groups.
     # Create a separate response mock for each group in our sample data.
-    mock_groups_users.return_value = [users]
-    # r = []
-    # for _ in groups:
-    #     r.append(users)
-    # mock_groups_users.return_value = iter(r)
+    # mock_groups_users.return_value = [users]
+    def mocked_group_members(azure_ad_group: dict) -> List:
+        group_id = azure_ad_group.get("id")
+        if group_id == "00000000-0000-0000-0000-000000000000":
+            return [users]
+        if group_id == "00000000-0000-0000-0000-0000000000001":
+            return [users]
+        if group_id == "99999999-9999-9999-9999-999999999999":
+            return [nested_group_members]
+        raise ValueError(f"Unexpected Azure AD group ID {group_id}")
+
+    mock_groups_users.side_effect = mocked_group_members
diff --git a/metadata-ingestion/tests/integration/great-expectations/setup/great_expectations/checkpoints/test_checkpoint.yml b/metadata-ingestion/tests/integration/great-expectations/setup/great_expectations/checkpoints/test_checkpoint.yml
index 1c5a8ebcca22a..7b12c5f2cf987 100644
--- a/metadata-ingestion/tests/integration/great-expectations/setup/great_expectations/checkpoints/test_checkpoint.yml
+++ b/metadata-ingestion/tests/integration/great-expectations/setup/great_expectations/checkpoints/test_checkpoint.yml
@@ -23,6 +23,7 @@ action_list:
       class_name: DataHubValidationAction
       server_url: http://localhost:8080
       graceful_exceptions: False
+      parse_table_names_from_sql: True
       platform_instance_map:
         my_postgresql_datasource: postgres1
         runtime_postgres_datasource: postgres1
diff --git a/metadata-ingestion/tests/unit/test_config_clean.py b/metadata-ingestion/tests/unit/test_config_clean.py
index 218ddfc90b3d3..178030c773e7f 100644
--- a/metadata-ingestion/tests/unit/test_config_clean.py
+++ b/metadata-ingestion/tests/unit/test_config_clean.py
@@ -1,6 +1,20 @@
 from datahub.utilities import config_clean
 
 
+def test_remove_suffix():
+    assert (
+        config_clean.remove_suffix(
+            "xaaabcdef.snowflakecomputing.com", ".snowflakecomputing.com"
+        )
+        == "xaaabcdef"
+    )
+
+    assert (
+        config_clean.remove_suffix("xaaabcdef", ".snowflakecomputing.com")
+        == "xaaabcdef"
+    )
+
+
 def test_url_without_slash_suffix():
     assert (
         config_clean.remove_trailing_slashes("http://example.com")
diff --git a/metadata-ingestion/tests/unit/test_openapi.py b/metadata-ingestion/tests/unit/test_openapi.py
index b4b61e685718b..b6c14426dbb38 100644
--- a/metadata-ingestion/tests/unit/test_openapi.py
+++ b/metadata-ingestion/tests/unit/test_openapi.py
@@ -261,6 +261,13 @@ class TestGetEndpoints(unittest.TestCase):
                         }
                     ]
                    }
+  /redirect:
+    get:
+      operationId: redirectSomewhere
+      summary: Redirect to a different endpoint
+      responses:
+        '302':
+          description: 302 response
   /v2:
     get:
       operationId: getVersionDetailsv2
diff --git a/metadata-ingestion/tests/unit/test_snowflake_source.py b/metadata-ingestion/tests/unit/test_snowflake_source.py
index 49b0a379b62ed..071eb267263df 100644
--- a/metadata-ingestion/tests/unit/test_snowflake_source.py
+++ b/metadata-ingestion/tests/unit/test_snowflake_source.py
@@ -1,7 +1,3 @@
-import pytest
-
-
-@pytest.mark.integration
 def test_snowflake_uri_default_authentication():
     from datahub.ingestion.source.sql.snowflake import SnowflakeConfig
 
@@ -23,7 +19,6 @@ def test_snowflake_uri_default_authentication():
     )
 
 
-@pytest.mark.integration
 def test_snowflake_uri_external_browser_authentication():
     from datahub.ingestion.source.sql.snowflake import SnowflakeConfig
 
@@ -45,7 +40,6 @@ def test_snowflake_uri_external_browser_authentication():
     )
 
 
-@pytest.mark.integration
 def test_snowflake_uri_key_pair_authentication():
     from datahub.ingestion.source.sql.snowflake import SnowflakeConfig
 
diff --git a/metadata-integration/java/datahub-protobuf/README.md b/metadata-integration/java/datahub-protobuf/README.md
new file mode 100644
index 0000000000000..3a805091d36d5
--- /dev/null
+++ b/metadata-integration/java/datahub-protobuf/README.md
@@ -0,0 +1,221 @@
+# Protobuf Integration
+
+This module is designed to be used with the Java Emitter, the input is a compiled protobuf binary `*.protoc` files
+and optionally the corresponding `*.proto` source code. In addition, you can supply the root message in cases
+where a single protobuf source file includes multiple non-nested messages.
+
+## Supported Features
+
+The following protobuf features are supported and are translated into descriptions, tags, properties and terms on a
+dataset.
+
+    * C++/C style code comments on Messages and Fields
+    * Nested Types
+    * Scalar Values
+    * Well Known Type Wrappers (i.e. DoubleValue, FloatValue, StringValue)
+    * Enumerations
+    * Oneof
+    * Maps
+    * Extensions
+    * Web links
+    * Parsing of GitHub team names and slack channel references
+
+## Usage
+
+### Protobuf Compile Options
+
+In order to support parsing comments the following option flags should be used during `protoc` compilation.
+
+    protoc --include_imports --include_source_info --descriptor_set_out=MyProto.protoc MyProto.proto
+
+### Code Example
+
+Given an input stream of the `protoc` binary and the emitter the minimal code is shown below.
+
+```java
+import com.linkedin.common.FabricType;
+import com.linkedin.common.AuditStamp;
+import com.linkedin.common.urn.CorpuserUrn;
+import datahub.client.rest.RestEmitter;
+import datahub.protobuf.ProtobufDataset;
+
+RestEmitter emitter;
+InputStream protocInputStream;
+
+AuditStamp auditStamp = new AuditStamp()
+    .setTime(System.currentTimeMillis())
+    .setActor(new CorpuserUrn("datahub"));
+
+ProtobufDataset dataset = ProtobufDataset.builder()
+    .setDataPlatformUrn(new DataPlatformUrn("kafka"))
+    .setProtocIn(protocInputStream)
+    .setAuditStamp(auditStamp)
+    .setFabricType(FabricType.DEV)
+    .build();
+
+dataset.getAllMetadataChangeProposals().forEach(mcpw -> emitter.emit(mcpw, null).get());
+```
+
+Additionally, the raw protobuf source can be included as well as information to allow parsing of additional
+references to GitHub and Slack in the source code comments.
+
+```java
+ProtobufDataset dataset = ProtobufDataset.builder()
+    .setDataPlatformUrn(new DataPlatformUrn("kafka"))
+    .setSchema(" my raw protobuf schema ")
+    .setProtocIn(protocInputStream)
+    .setAuditStamp(auditStamp)
+    .setFabricType(FabricType.DEV)
+    .setGithubOrganization("myOrg")
+    .setSlackTeamId("SLACK123")
+    .build();
+```
+
+### Protobuf Extensions
+
+In order to extract even more metadata from the protobuf schema we can extend the FieldOptions and MessageOptions
+to be able to annotate Messages and Fields with arbitrary information. This information can then be emitted
+as DataHub primary key information, tags, glossary terms or properties on the dataset.
+
+An annotated protobuf schema would look like the following, except for the `is_primary_key` all
+annotations are configurable for individual needs.
+
+*Note*: Extending FieldOptions and MessageOptions does not change the messages themselves. The metadata is not included
+in messages being sent over the wire.
+
+```protobuf
+syntax = "proto3";
+import "meta.proto";
+
+message Department {
+    int32 id = 1 [(meta.fld.is_primary_key) = true];
+    string name = 2;
+}
+
+message Person {
+    option(meta.msg.type) = ENTITY;
+    option(meta.msg.classification_enum) = HighlyConfidential;
+    option(meta.msg.team) = "TeamB";
+    option(meta.msg.bool_feature) = true;
+    option(meta.msg.alert_channel) = "#alerts";
+
+    string name = 1 [(meta.fld.classification) = "Classification.HighlyConfidential"];
+
+    int32 id = 2
+    [(meta.fld.is_primary_key) = true];
+
+    string email = 3
+    [(meta.fld.classification_enum) = Confidential];
+    
+    Department dept = 4;
+    
+    string test_coverage = 5
+    [(meta.fld.product_type_bool) = true, (meta.fld.product_type) = "my type", (meta.fld.product_type_enum) = EVENT];
+}
+```
+
+#### meta.proto 
+
+In order to use the annotations above, create a proto file called `meta.proto`. Feel free to customize the kinds of
+metadata and how it is emitted to DataHub for your use cases.
+
+```protobuf
+syntax = "proto3";
+package meta;
+
+import "google/protobuf/descriptor.proto";
+
+/*
+   This is assigned to metadata fields. It describes how the metadata field should be represented
+   in DataHub. This enum must be used in the `meta` package. Multiple can be used for the same
+   metadata annotation. This allows a single piece of information to be captured in DataHub
+   as a property, tag and/or term.
+
+   Tags can be strings, enums, or booleans
+   Terms can be strings or enums
+   Properties should be strings
+
+*/
+enum DataHubMetadataType {
+  PROPERTY = 0; // Datahub Custom Property
+  TAG      = 1; // Datahub Tag
+  TERM     = 2; // Datahub Term
+}
+
+/*
+   Example below: The following is not required for annotation processing. This is an example
+   of creating an annotation using an enum.
+ */
+
+enum MetaEnumExample {
+  UNKNOWN = 0;
+  ENTITY = 1;
+  EVENT = 2;
+}
+
+// Assuming Glossary Term defined from bootstrap example
+enum Classification {
+  HighlyConfidential = 0;
+  Confidential = 1;
+  Sensitive = 2;
+}
+```
+
+#### FieldOptions
+
+Define possible annotations on fields and how they are exported to DataHub.
+
+```protobuf
+
+message fld {
+  extend google.protobuf.FieldOptions {
+    // Required: Mark option field with how to export to DataHub in one or more places.
+    repeated meta.DataHubMetadataType type = 6000;
+
+    /*
+       Examples below: The following is not required for annotation processing.
+     */
+
+    // Set true if the field is a primary key. This works for any boolean with `primary_key` in it.
+    bool is_primary_key = 6010;
+
+    // Extract classification field option as a Term, either works
+    string classification = 6001 [(meta.fld.type) = TERM];
+    meta.Classification classification_enum = 6002 [(meta.fld.type) = TERM];
+
+    // Expose this option as a tag on the field.
+    string product_type = 70004 [(meta.fld.type) = TAG];
+    bool product_type_bool = 70005 [(meta.fld.type) = TAG];
+    meta.MetaEnumExample product_type_enum = 70006 [(meta.fld.type) = TAG];
+  }
+}
+```
+
+#### MessageOptions
+
+Define possible annotations on messages and how they are exported to DataHub.
+
+```protobuf
+
+message msg {
+  extend google.protobuf.MessageOptions {
+    /*
+       Examples below: The following is not required for annotation processing.
+     */
+
+    // Place the classification term at the Message/Dataset level, either string or enum is supported
+    string classification = 4000 [(meta.fld.type) = TERM, (meta.fld.type) = PROPERTY];
+    meta.Classification classification_enum = 4001 [(meta.fld.type) = TERM, (meta.fld.type) = PROPERTY];
+
+    // Attach these Message/Dataset options as a tag and property.
+    string product = 5001 [(meta.fld.type) = TAG, (meta.fld.type) = PROPERTY];
+    string project = 5002 [(meta.fld.type) = TAG, (meta.fld.type) = PROPERTY];
+    string team = 5003 [(meta.fld.type) = TAG, (meta.fld.type) = PROPERTY];
+
+    string domain = 60003 [(meta.fld.type) = TAG, (meta.fld.type) = PROPERTY];
+    meta.MetaEnumExample type = 60004 [(meta.fld.type) = TAG, (meta.fld.type) = PROPERTY];
+    bool bool_feature = 60005 [(meta.fld.type) = TAG];
+    string alert_channel = 60007 [(meta.fld.type) = PROPERTY];
+  }
+}
+```
diff --git a/metadata-integration/java/datahub-protobuf/build.gradle b/metadata-integration/java/datahub-protobuf/build.gradle
new file mode 100644
index 0000000000000..8a0c09d533e8b
--- /dev/null
+++ b/metadata-integration/java/datahub-protobuf/build.gradle
@@ -0,0 +1,53 @@
+plugins {
+    id("com.palantir.git-version") apply false
+}
+apply plugin: 'java'
+apply plugin: 'jacoco'
+
+afterEvaluate {
+    if (project.plugins.hasPlugin('java')) {
+        sourceCompatibility = 11
+        targetCompatibility = 11
+    }
+}
+
+dependencies {
+    implementation project(':metadata-models')
+    implementation project(path: ':metadata-integration:java:datahub-client', configuration: 'shadow')
+
+    implementation externalDependency.protobuf
+    implementation externalDependency.jgrapht
+    compileOnly externalDependency.lombok
+    annotationProcessor externalDependency.lombok
+    testImplementation externalDependency.junitJupiterApi
+    testRuntimeOnly externalDependency.junitJupiterEngine
+}
+
+import java.nio.file.Path
+task compileProtobuf {
+    doLast {
+        def basePath = Path.of("${projectDir}/src/test/resources")
+        [
+                fileTree("${projectDir}/src/test/resources/protobuf") { include "*.proto" },
+                fileTree("${projectDir}/src/test/resources/extended_protobuf") { include "*.proto" }
+        ].collectMany { it.collect() }.each { f ->
+            def input = basePath.relativize(Path.of(f.getAbsolutePath()))
+            println(input.toString() + " => " + input.toString().replace(".proto", ".protoc"))
+            exec {
+                workingDir "${projectDir}/src/test/resources"
+                commandLine 'protoc', '--proto_path=.', '--include_imports', '--include_source_info',
+                        "--descriptor_set_out=${input.toString().replace(".proto", ".protoc")}",
+                        input
+            }
+        }
+    }
+}
+
+jacocoTestReport {
+    dependsOn test // tests are required to run before generating the report
+}
+
+test {
+    useJUnit()
+    finalizedBy jacocoTestReport
+}
\ No newline at end of file
diff --git a/metadata-integration/java/datahub-protobuf/src/main/java/com/google/protobuf/ExtensionRegistry.java b/metadata-integration/java/datahub-protobuf/src/main/java/com/google/protobuf/ExtensionRegistry.java
new file mode 100644
index 0000000000000..e6f93eb1a4f0c
--- /dev/null
+++ b/metadata-integration/java/datahub-protobuf/src/main/java/com/google/protobuf/ExtensionRegistry.java
@@ -0,0 +1,388 @@
+// Protocol Buffers - Google's data interchange format
+// Copyright 2008 Google Inc.  All rights reserved.
+// https://developers.google.com/protocol-buffers/
+//
+// Redistribution and use in source and binary forms, with or without
+// modification, are permitted provided that the following conditions are
+// met:
+//
+//     * Redistributions of source code must retain the above copyright
+// notice, this list of conditions and the following disclaimer.
+//     * Redistributions in binary form must reproduce the above
+// copyright notice, this list of conditions and the following disclaimer
+// in the documentation and/or other materials provided with the
+// distribution.
+//     * Neither the name of Google Inc. nor the names of its
+// contributors may be used to endorse or promote products derived from
+// this software without specific prior written permission.
+//
+// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+// "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+// LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+// A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+// OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+// LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+// DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+// THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+package com.google.protobuf;
+
+import com.google.protobuf.Descriptors.Descriptor;
+import com.google.protobuf.Descriptors.FieldDescriptor;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Set;
+
+/**
+ * A table of known extensions, searchable by name or field number. When parsing a protocol message
+ * that might have extensions, you must provide an {@code ExtensionRegistry} in which you have
+ * registered any extensions that you want to be able to parse. Otherwise, those extensions will
+ * just be treated like unknown fields.
+ *
+ * <p>For example, if you had the {@code .proto} file:
+ *
+ * <pre>
+ * option java_class = "MyProto";
+ *
+ * message Foo {
+ *   extensions 1000 to max;
+ * }
+ *
+ * extend Foo {
+ *   optional int32 bar;
+ * }
+ * </pre>
+ *
+ * Then you might write code like:
+ *
+ * <pre>
+ * ExtensionRegistry registry = ExtensionRegistry.newInstance();
+ * registry.add(MyProto.bar);
+ * MyProto.Foo message = MyProto.Foo.parseFrom(input, registry);
+ * </pre>
+ *
+ * <p>Background:
+ *
+ * <p>You might wonder why this is necessary. Two alternatives might come to mind. First, you might
+ * imagine a system where generated extensions are automatically registered when their containing
+ * classes are loaded. This is a popular technique, but is bad design; among other things, it
+ * creates a situation where behavior can change depending on what classes happen to be loaded. It
+ * also introduces a security vulnerability, because an unprivileged class could cause its code to
+ * be called unexpectedly from a privileged class by registering itself as an extension of the right
+ * type.
+ *
+ * <p>Another option you might consider is lazy parsing: do not parse an extension until it is first
+ * requested, at which point the caller must provide a type to use. This introduces a different set
+ * of problems. First, it would require a mutex lock any time an extension was accessed, which would
+ * be slow. Second, corrupt data would not be detected until first access, at which point it would
+ * be much harder to deal with it. Third, it could violate the expectation that message objects are
+ * immutable, since the type provided could be any arbitrary message class. An unprivileged user
+ * could take advantage of this to inject a mutable object into a message belonging to privileged
+ * code and create mischief.
+ *
+ * @author kenton@google.com Kenton Varda
+ */
+public class ExtensionRegistry extends ExtensionRegistryLite {
+    /** Construct a new, empty instance. */
+    public static ExtensionRegistry newInstance() {
+        return new ExtensionRegistry();
+    }
+
+    /** Get the unmodifiable singleton empty instance. */
+    public static ExtensionRegistry getEmptyRegistry() {
+        return EMPTY_REGISTRY;
+    }
+
+
+    /** Returns an unmodifiable view of the registry. */
+    @Override
+    public ExtensionRegistry getUnmodifiable() {
+        return new ExtensionRegistry(this);
+    }
+
+    /** A (Descriptor, Message) pair, returned by lookup methods. */
+    public static final class ExtensionInfo {
+        /** The extension's descriptor. */
+        public final FieldDescriptor descriptor;
+
+        /**
+         * A default instance of the extension's type, if it has a message type. Otherwise, {@code
+         * null}.
+         */
+        public final Message defaultInstance;
+
+        private ExtensionInfo(final FieldDescriptor descriptor) {
+            this.descriptor = descriptor;
+            defaultInstance = null;
+        }
+
+        private ExtensionInfo(final FieldDescriptor descriptor, final Message defaultInstance) {
+            this.descriptor = descriptor;
+            this.defaultInstance = defaultInstance;
+        }
+    }
+
+    /** Deprecated. Use {@link #findImmutableExtensionByName(String)} instead. */
+    @Deprecated
+    public ExtensionInfo findExtensionByName(final String fullName) {
+        return findImmutableExtensionByName(fullName);
+    }
+
+    /**
+     * Find an extension for immutable APIs by fully-qualified field name, in the proto namespace.
+     * i.e. {@code result.descriptor.fullName()} will match {@code fullName} if a match is found.
+     *
+     * @return Information about the extension if found, or {@code null} otherwise.
+     */
+    public ExtensionInfo findImmutableExtensionByName(final String fullName) {
+        return immutableExtensionsByName.get(fullName);
+    }
+
+    /**
+     * Find an extension for mutable APIs by fully-qualified field name, in the proto namespace. i.e.
+     * {@code result.descriptor.fullName()} will match {@code fullName} if a match is found.
+     *
+     * @return Information about the extension if found, or {@code null} otherwise.
+     */
+    public ExtensionInfo findMutableExtensionByName(final String fullName) {
+        return mutableExtensionsByName.get(fullName);
+    }
+
+    /** Deprecated. Use {@link #findImmutableExtensionByNumber( Descriptors.Descriptor, int)} */
+    @Deprecated
+    public ExtensionInfo findExtensionByNumber(
+            final Descriptor containingType, final int fieldNumber) {
+        return findImmutableExtensionByNumber(containingType, fieldNumber);
+    }
+
+    /**
+     * Find an extension by containing type and field number for immutable APIs.
+     *
+     * @return Information about the extension if found, or {@code null} otherwise.
+     */
+    public ExtensionInfo findImmutableExtensionByNumber(
+            final Descriptor containingType, final int fieldNumber) {
+        return immutableExtensionsByNumber.get(new DescriptorIntPair(containingType, fieldNumber));
+    }
+
+    /**
+     * Find an extension by containing type and field number for mutable APIs.
+     *
+     * @return Information about the extension if found, or {@code null} otherwise.
+     */
+    public ExtensionInfo findMutableExtensionByNumber(
+            final Descriptor containingType, final int fieldNumber) {
+        return mutableExtensionsByNumber.get(new DescriptorIntPair(containingType, fieldNumber));
+    }
+
+    /**
+     * Find all extensions for mutable APIs by fully-qualified name of extended class. Note that this
+     * method is more computationally expensive than getting a single extension by name or number.
+     *
+     * @return Information about the extensions found, or {@code null} if there are none.
+     */
+    public Set<ExtensionInfo> getAllMutableExtensionsByExtendedType(final String fullName) {
+        HashSet<ExtensionInfo> extensions = new HashSet<ExtensionInfo>();
+        for (DescriptorIntPair pair : mutableExtensionsByNumber.keySet()) {
+            if (pair.descriptor.getFullName().equals(fullName)) {
+                extensions.add(mutableExtensionsByNumber.get(pair));
+            }
+        }
+        return extensions;
+    }
+
+    /**
+     * Find all extensions for immutable APIs by fully-qualified name of extended class. Note that
+     * this method is more computationally expensive than getting a single extension by name or
+     * number.
+     *
+     * @return Information about the extensions found, or {@code null} if there are none.
+     */
+    public Set<ExtensionInfo> getAllImmutableExtensionsByExtendedType(final String fullName) {
+        HashSet<ExtensionInfo> extensions = new HashSet<ExtensionInfo>();
+        for (DescriptorIntPair pair : immutableExtensionsByNumber.keySet()) {
+            if (pair.descriptor.getFullName().equals(fullName)) {
+                extensions.add(immutableExtensionsByNumber.get(pair));
+            }
+        }
+        return extensions;
+    }
+
+    /** Add an extension from a generated file to the registry. */
+    public void add(final Extension<?, ?> extension) {
+        if (extension.getExtensionType() != Extension.ExtensionType.IMMUTABLE
+                && extension.getExtensionType() != Extension.ExtensionType.MUTABLE) {
+            // do not support other extension types. ignore
+            return;
+        }
+        add(newExtensionInfo(extension), extension.getExtensionType());
+    }
+
+    /** Add an extension from a generated file to the registry. */
+    public void add(final GeneratedMessage.GeneratedExtension<?, ?> extension) {
+        add((Extension<?, ?>) extension);
+    }
+
+    static ExtensionInfo newExtensionInfo(final Extension<?, ?> extension) {
+        if (extension.getDescriptor().getJavaType() == FieldDescriptor.JavaType.MESSAGE) {
+            if (extension.getMessageDefaultInstance() == null) {
+                throw new IllegalStateException(
+                        "Registered message-type extension had null default instance: "
+                                + extension.getDescriptor().getFullName());
+            }
+            return new ExtensionInfo(
+                    extension.getDescriptor(), (Message) extension.getMessageDefaultInstance());
+        } else {
+            return new ExtensionInfo(extension.getDescriptor(), null);
+        }
+    }
+
+    /** Add a non-message-type extension to the registry by descriptor. */
+    public void add(final FieldDescriptor type) {
+        if (type.getJavaType() == FieldDescriptor.JavaType.MESSAGE) {
+            throw new IllegalArgumentException(
+                    "ExtensionRegistry.add() must be provided a default instance when "
+                            + "adding an embedded message extension.");
+        }
+        ExtensionInfo info = new ExtensionInfo(type, null);
+        add(info, Extension.ExtensionType.IMMUTABLE);
+        add(info, Extension.ExtensionType.MUTABLE);
+    }
+
+    /** Add a message-type extension to the registry by descriptor. */
+    public void add(final FieldDescriptor type, final Message defaultInstance) {
+        if (type.getJavaType() != FieldDescriptor.JavaType.MESSAGE) {
+            throw new IllegalArgumentException(
+                    "ExtensionRegistry.add() provided a default instance for a non-message extension.");
+        }
+        add(new ExtensionInfo(type, defaultInstance), Extension.ExtensionType.IMMUTABLE);
+    }
+
+    // =================================================================
+    // Private stuff.
+
+    private ExtensionRegistry() {
+        this.immutableExtensionsByName = new HashMap<String, ExtensionInfo>();
+        this.mutableExtensionsByName = new HashMap<String, ExtensionInfo>();
+        this.immutableExtensionsByNumber = new HashMap<DescriptorIntPair, ExtensionInfo>();
+        this.mutableExtensionsByNumber = new HashMap<DescriptorIntPair, ExtensionInfo>();
+    }
+
+    private ExtensionRegistry(ExtensionRegistry other) {
+        super(other);
+        this.immutableExtensionsByName = Collections.unmodifiableMap(other.immutableExtensionsByName);
+        this.mutableExtensionsByName = Collections.unmodifiableMap(other.mutableExtensionsByName);
+        this.immutableExtensionsByNumber =
+                Collections.unmodifiableMap(other.immutableExtensionsByNumber);
+        this.mutableExtensionsByNumber = Collections.unmodifiableMap(other.mutableExtensionsByNumber);
+    }
+
+    private final Map<String, ExtensionInfo> immutableExtensionsByName;
+    private final Map<String, ExtensionInfo> mutableExtensionsByName;
+    private final Map<DescriptorIntPair, ExtensionInfo> immutableExtensionsByNumber;
+    private final Map<DescriptorIntPair, ExtensionInfo> mutableExtensionsByNumber;
+
+    ExtensionRegistry(boolean empty) {
+        super(EMPTY_REGISTRY_LITE);
+        this.immutableExtensionsByName = Collections.<String, ExtensionInfo>emptyMap();
+        this.mutableExtensionsByName = Collections.<String, ExtensionInfo>emptyMap();
+        this.immutableExtensionsByNumber = Collections.<DescriptorIntPair, ExtensionInfo>emptyMap();
+        this.mutableExtensionsByNumber = Collections.<DescriptorIntPair, ExtensionInfo>emptyMap();
+    }
+
+    static final ExtensionRegistry EMPTY_REGISTRY = new ExtensionRegistry(true);
+
+    private void add(final ExtensionInfo extension, final Extension.ExtensionType extensionType) {
+        if (!extension.descriptor.isExtension()) {
+            throw new IllegalArgumentException(
+                    "ExtensionRegistry.add() was given a FieldDescriptor for a regular "
+                            + "(non-extension) field.");
+        }
+
+        Map<String, ExtensionInfo> extensionsByName;
+        Map<DescriptorIntPair, ExtensionInfo> extensionsByNumber;
+        switch (extensionType) {
+            case IMMUTABLE:
+                extensionsByName = immutableExtensionsByName;
+                extensionsByNumber = immutableExtensionsByNumber;
+                break;
+            case MUTABLE:
+                extensionsByName = mutableExtensionsByName;
+                extensionsByNumber = mutableExtensionsByNumber;
+                break;
+            default:
+                // Ignore the unknown supported type.
+                return;
+        }
+
+        extensionsByName.put(extension.descriptor.getFullName(), extension);
+        extensionsByNumber.put(
+                new DescriptorIntPair(
+                        extension.descriptor.getContainingType(), extension.descriptor.getNumber()),
+                extension);
+
+        final FieldDescriptor field = extension.descriptor;
+        if (field.getContainingType().getOptions().getMessageSetWireFormat()
+                && field.getType() == FieldDescriptor.Type.MESSAGE
+                && field.isOptional()
+                && field.getExtensionScope() == field.getMessageType()) {
+            // This is an extension of a MessageSet type defined within the extension
+            // type's own scope.  For backwards-compatibility, allow it to be looked
+            // up by type name.
+            extensionsByName.put(field.getMessageType().getFullName(), extension);
+        }
+    }
+
+    /**
+     *
+     * DataHub modification of hashcode/equals based on full name. The upstream
+     * project uses the descriptor and in our use of the registry results
+     * in objects that are practically identical except for the `jsonName` field.
+     * This is a difference generated by internal components and is not under
+     * our control.
+     *
+     * A (GenericDescriptor, int) pair, used as a map key.
+     *
+     * */
+    private static final class DescriptorIntPair {
+        private final String fullName;
+        private final Descriptor descriptor;
+        private final int number;
+
+        DescriptorIntPair(final Descriptor descriptor, final int number) {
+            this.descriptor = descriptor;
+            this.fullName = descriptor.getFullName();
+            this.number = number;
+        }
+
+        @Override
+        public boolean equals(Object o) {
+            if (this == o) {
+                return true;
+            }
+            if (o == null || getClass() != o.getClass()) {
+                return false;
+            }
+
+            DescriptorIntPair that = (DescriptorIntPair) o;
+
+            if (number != that.number) {
+                return false;
+            }
+            return fullName.equals(that.fullName);
+        }
+
+        @Override
+        public int hashCode() {
+            int result = fullName.hashCode();
+            result = 31 * result + number;
+            return result;
+        }
+    }
+}
diff --git a/metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/ProtobufDataset.java b/metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/ProtobufDataset.java
new file mode 100644
index 0000000000000..7ae3925c4c407
--- /dev/null
+++ b/metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/ProtobufDataset.java
@@ -0,0 +1,255 @@
+package datahub.protobuf;
+
+import com.google.protobuf.DescriptorProtos.FileDescriptorSet;
+import com.linkedin.common.AuditStamp;
+import com.linkedin.common.FabricType;
+import com.linkedin.common.Status;
+import com.linkedin.common.urn.DataPlatformUrn;
+import com.linkedin.common.urn.DatasetUrn;
+import com.linkedin.data.template.RecordTemplate;
+import com.linkedin.events.metadata.ChangeType;
+import com.linkedin.schema.KafkaSchema;
+import com.linkedin.schema.SchemaField;
+import com.linkedin.schema.SchemaFieldArray;
+import com.linkedin.schema.SchemaMetadata;
+import com.linkedin.util.Pair;
+import datahub.protobuf.model.ProtobufGraph;
+import datahub.protobuf.visitors.ProtobufModelVisitor;
+import datahub.protobuf.visitors.VisitContext;
+import datahub.protobuf.visitors.dataset.DatasetVisitor;
+import datahub.protobuf.visitors.dataset.InstitutionalMemoryVisitor;
+import datahub.protobuf.visitors.dataset.KafkaTopicPropertyVisitor;
+import datahub.protobuf.visitors.dataset.ProtobufExtensionPropertyVisitor;
+import datahub.protobuf.visitors.dataset.ProtobufExtensionTagAssocVisitor;
+import datahub.protobuf.visitors.dataset.ProtobufExtensionTermAssocVisitor;
+import datahub.protobuf.visitors.field.SchemaFieldVisitor;
+import datahub.event.MetadataChangeProposalWrapper;
+import datahub.protobuf.visitors.field.ProtobufExtensionFieldVisitor;
+import datahub.protobuf.visitors.tags.ProtobufExtensionTagVisitor;
+
+import javax.annotation.Nullable;
+import java.io.IOException;
+import java.io.InputStream;
+import java.util.Base64;
+import java.util.Collection;
+import java.util.Comparator;
+import java.util.Optional;
+import java.util.List;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+
+
+public class ProtobufDataset {
+
+    public static ProtobufDataset.Builder builder() {
+        return new Builder();
+    }
+
+    public static class Builder {
+        private DataPlatformUrn dataPlatformUrn;
+        private DatasetUrn datasetUrn;
+        private FabricType fabricType;
+        private AuditStamp auditStamp;
+        private byte[] protocBytes;
+        private String messageName;
+        private String schema;
+        private String githubOrganization;
+        private String slackTeamId;
+
+        public Builder setGithubOrganization(@Nullable String githubOrganization) {
+            this.githubOrganization = githubOrganization;
+            return this;
+        }
+
+        public Builder setSlackTeamId(@Nullable String slackTeamId) {
+            this.slackTeamId = slackTeamId;
+            return this;
+        }
+
+        public Builder setProtocIn(InputStream protocIn) throws IOException {
+            return setProtocBytes(protocIn.readAllBytes());
+        }
+
+        public Builder setDataPlatformUrn(@Nullable DataPlatformUrn dataPlatformUrn) {
+            this.dataPlatformUrn = dataPlatformUrn;
+            return this;
+        }
+
+        public Builder setDatasetUrn(@Nullable DatasetUrn datasetUrn) {
+            this.datasetUrn = datasetUrn;
+            return this;
+        }
+
+        public Builder setProtocBytes(byte[] protocBytes) {
+            this.protocBytes = protocBytes;
+            return this;
+        }
+
+        public Builder setFabricType(FabricType fabricType) {
+            this.fabricType = fabricType;
+            return this;
+        }
+
+        public Builder setAuditStamp(AuditStamp auditStamp) {
+            this.auditStamp = auditStamp;
+            return this;
+        }
+
+        public Builder setMessageName(@Nullable String messageName) {
+            this.messageName = messageName;
+            return this;
+        }
+
+        public Builder setSchema(@Nullable String schema) {
+            this.schema = schema;
+            return this;
+        }
+
+        public ProtobufDataset build() throws IOException {
+            FileDescriptorSet fileSet = FileDescriptorSet.parseFrom(protocBytes);
+
+            return new ProtobufDataset(
+                    this,
+                    Optional.ofNullable(dataPlatformUrn).orElse(new DataPlatformUrn("kafka")),
+                    datasetUrn,
+                    new ProtobufGraph(fileSet, messageName), schema,  auditStamp, fabricType)
+                    .setMetadataChangeProposalVisitors(
+                            List.of(
+                                    new ProtobufExtensionTagVisitor()
+                            )
+                    )
+                    .setFieldVisitor(new ProtobufExtensionFieldVisitor())
+                    .setDatasetVisitor(DatasetVisitor.builder()
+                            .protocBase64(Base64.getEncoder().encodeToString(protocBytes))
+                            .datasetPropertyVisitors(
+                                    List.of(
+                                            new KafkaTopicPropertyVisitor(),
+                                            new ProtobufExtensionPropertyVisitor()
+                                    )
+                            )
+                            .institutionalMemoryMetadataVisitors(
+                                    List.of(
+                                            new InstitutionalMemoryVisitor(slackTeamId, githubOrganization)
+                                    )
+                            )
+                            .tagAssociationVisitors(
+                                    List.of(
+                                            new ProtobufExtensionTagAssocVisitor()
+                                    )
+                            )
+                            .termAssociationVisitors(
+                                    List.of(
+                                            new ProtobufExtensionTermAssocVisitor()
+                                    )
+                            )
+                            .build()
+                    );
+        }
+    }
+
+    private final DatasetUrn datasetUrn;
+    private final Optional<String> schemaSource;
+    private final ProtobufGraph graph;
+    private final AuditStamp auditStamp;
+    private final VisitContext.VisitContextBuilder contextBuilder;
+    private final ProtobufDataset.Builder builder;
+
+    private DatasetVisitor datasetVisitor;
+    private ProtobufModelVisitor<Pair<SchemaField, Double>> fieldVisitor;
+    private List<ProtobufModelVisitor<MetadataChangeProposalWrapper<? extends RecordTemplate>>> mcpwVisitors;
+
+    public ProtobufDataset(DataPlatformUrn dataPlatformUrn, DatasetUrn datasetUrn, ProtobufGraph graph, String schema,
+                           AuditStamp auditStamp, FabricType fabricType) {
+        this(null, dataPlatformUrn, datasetUrn, graph, schema, auditStamp, fabricType);
+    }
+
+    public ProtobufDataset(ProtobufDataset.Builder builder, DataPlatformUrn dataPlatformUrn, DatasetUrn datasetUrn, ProtobufGraph graph,
+                           String schema, AuditStamp auditStamp, FabricType fabricType) {
+        this.builder = builder;
+        this.schemaSource = Optional.ofNullable(schema);
+        this.auditStamp = auditStamp;
+        this.graph = graph;
+
+        // Default - non-protobuf extension
+        fieldVisitor = new SchemaFieldVisitor();
+        mcpwVisitors = List.of();
+
+        this.datasetUrn = datasetUrn != null ? datasetUrn : new DatasetUrn(dataPlatformUrn, this.graph.getFullName(), fabricType);
+        this.contextBuilder = VisitContext.builder().datasetUrn(this.datasetUrn).auditStamp(this.auditStamp);
+    }
+
+    public ProtobufDataset setMetadataChangeProposalVisitors(List<ProtobufModelVisitor<MetadataChangeProposalWrapper<? extends RecordTemplate>>> visitors) {
+        this.mcpwVisitors = visitors;
+        return this;
+    }
+
+    public ProtobufDataset setDatasetVisitor(DatasetVisitor datasetVisitor) {
+        this.datasetVisitor = datasetVisitor;
+        return this;
+    }
+
+    public ProtobufDataset setFieldVisitor(ProtobufModelVisitor<Pair<SchemaField, Double>> visitor) {
+        this.fieldVisitor = visitor;
+        return this;
+    }
+
+    public ProtobufDataset.Builder toBuilder() {
+        return builder;
+    }
+
+    public ProtobufGraph getGraph() {
+        return graph;
+    }
+
+    public AuditStamp getAuditStamp() {
+        return auditStamp;
+    }
+
+    public DatasetUrn getDatasetUrn() {
+        return datasetUrn;
+    }
+
+    public Stream<Collection<MetadataChangeProposalWrapper<? extends RecordTemplate>>> getAllMetadataChangeProposals() {
+        return Stream.of(getVisitorMCPs(), getDatasetMCPs());
+    }
+
+    public List<MetadataChangeProposalWrapper<? extends RecordTemplate>> getVisitorMCPs() {
+        return graph.accept(contextBuilder, mcpwVisitors).collect(Collectors.toList());
+    }
+
+    public List<MetadataChangeProposalWrapper<? extends RecordTemplate>> getDatasetMCPs() {
+        return Stream.concat(
+                this.graph.accept(contextBuilder, List.of(datasetVisitor)),
+                Stream.of(
+                        new MetadataChangeProposalWrapper<>(DatasetUrn.ENTITY_TYPE, datasetUrn.toString(), ChangeType.UPSERT,
+                                getSchemaMetadata(), "schemaMetadata"),
+                        new MetadataChangeProposalWrapper<>(DatasetUrn.ENTITY_TYPE, datasetUrn.toString(), ChangeType.UPSERT,
+                                new Status().setRemoved(false), "status")
+                )
+        ).collect(Collectors.toList());
+    }
+
+    public SchemaMetadata getSchemaMetadata() {
+        SchemaMetadata.PlatformSchema platformSchema = new SchemaMetadata.PlatformSchema();
+        schemaSource.ifPresent(schemaStr -> platformSchema.setKafkaSchema(new KafkaSchema().setDocumentSchema(schemaStr)));
+
+        List<SchemaField> schemaFields = graph.accept(contextBuilder, List.of(fieldVisitor))
+                .sorted(COMPARE_BY_ROOT_MESSAGE_FIELD_WEIGHT.thenComparing(COMPARE_BY_FIELD_PATH))
+                .map(Pair::getFirst)
+                .collect(Collectors.toList());
+
+        return new SchemaMetadata()
+                .setSchemaName(graph.getFullName())
+                .setPlatform(datasetUrn.getPlatformEntity())
+                .setCreated(auditStamp)
+                .setLastModified(auditStamp)
+                .setVersion(graph.getMajorVersion())
+                .setHash(graph.getHash())
+                .setPlatformSchema(platformSchema)
+                .setFields(new SchemaFieldArray(schemaFields));
+    }
+
+    public static final Comparator<Pair<SchemaField, Double>> COMPARE_BY_ROOT_MESSAGE_FIELD_WEIGHT = Comparator.comparing(Pair::getSecond);
+    public static final Comparator<Pair<SchemaField, Double>> COMPARE_BY_FIELD_PATH = Comparator
+            .comparing(p -> p.getFirst().getFieldPath());
+}
diff --git a/metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/ProtobufUtils.java b/metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/ProtobufUtils.java
new file mode 100644
index 0000000000000..347bf593dc068
--- /dev/null
+++ b/metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/ProtobufUtils.java
@@ -0,0 +1,101 @@
+package datahub.protobuf;
+
+import com.google.common.collect.ImmutableList;
+import com.google.protobuf.DescriptorProtos;
+import com.google.protobuf.Descriptors;
+import com.google.protobuf.ExtensionRegistry;
+
+import java.nio.charset.StandardCharsets;
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Objects;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+
+public class ProtobufUtils {
+    private ProtobufUtils() { }
+
+    public static String collapseLocationComments(DescriptorProtos.SourceCodeInfo.Location location) {
+       String orig = Stream.concat(location.getLeadingDetachedCommentsList().stream(),
+                       Stream.of(location.getLeadingComments(), location.getTrailingComments()))
+               .filter(Objects::nonNull)
+               .flatMap(line -> Arrays.stream(line.split("\n")))
+               .map(line -> line.replaceFirst("^[*/ ]+", ""))
+               .collect(Collectors.joining("\n"))
+               .trim();
+
+        /*
+         * Sometimes DataHub doesn't like these strings. Not sure if its DataHub
+         * or protobuf issue: https://github.com/protocolbuffers/protobuf/issues/4691
+         *
+         * We essentially smash utf8 chars to ascii here
+         */
+        return new String(orig.getBytes(StandardCharsets.ISO_8859_1));
+    }
+
+    public static ExtensionRegistry buildRegistry(DescriptorProtos.FileDescriptorSet fileSet) {
+        ExtensionRegistry registry = ExtensionRegistry.newInstance();
+        Map<String, DescriptorProtos.FileDescriptorProto> descriptorProtoMap = fileSet.getFileList().stream()
+                .collect(Collectors.toMap(DescriptorProtos.FileDescriptorProto::getName, Function.identity()));
+        Map<String, Descriptors.FileDescriptor> descriptorCache = new HashMap<>();
+
+        fileSet.getFileList().forEach(fdp -> {
+            try {
+                Descriptors.FileDescriptor file = descriptorFromProto(fdp, descriptorProtoMap, descriptorCache);
+                Stream.concat(file.getExtensions().stream(), file.getMessageTypes().stream().flatMap(msg -> msg.getExtensions().stream()))
+                        .forEach(ext -> addToRegistry(fdp, ext, registry));
+            } catch (Descriptors.DescriptorValidationException e) {
+                e.printStackTrace();
+            }
+        });
+        return registry;
+    }
+
+    private static void addToRegistry(DescriptorProtos.FileDescriptorProto fileDescriptorProto,
+                                      Descriptors.FieldDescriptor fieldDescriptor, ExtensionRegistry registry) {
+        if (fieldDescriptor.getJavaType() != Descriptors.FieldDescriptor.JavaType.MESSAGE) {
+            registry.add(fieldDescriptor);
+        } else {
+            fileDescriptorProto.getMessageTypeList().stream()
+                    .filter(typ -> typ.getName().equals(fieldDescriptor.getMessageType().getName()))
+                    .findFirst().ifPresent(messageType -> registry.add(fieldDescriptor, messageType.getDefaultInstanceForType()));
+            fieldDescriptor.getMessageType().getFields()
+                    .stream().filter(Descriptors.FieldDescriptor::isExtension)
+                    .forEach(f -> addToRegistry(fileDescriptorProto, f, registry));
+        }
+    }
+
+    /**
+     * Recursively constructs file descriptors for all dependencies of the supplied proto and returns
+     * a {@link Descriptors.FileDescriptor} for the supplied proto itself. For maximal efficiency, reuse the
+     * descriptorCache argument across calls.
+     */
+    private static Descriptors.FileDescriptor descriptorFromProto(
+            DescriptorProtos.FileDescriptorProto descriptorProto,
+            Map<String, DescriptorProtos.FileDescriptorProto> descriptorProtoIndex,
+            Map<String, Descriptors.FileDescriptor> descriptorCache) throws Descriptors.DescriptorValidationException {
+        // First, check the cache.
+        String descriptorName = descriptorProto.getName();
+        if (descriptorCache.containsKey(descriptorName)) {
+            return descriptorCache.get(descriptorName);
+        }
+
+        // Then, fetch all the required dependencies recursively.
+        ImmutableList.Builder<Descriptors.FileDescriptor> dependencies = ImmutableList.builder();
+        for (String dependencyName : descriptorProto.getDependencyList()) {
+            if (!descriptorProtoIndex.containsKey(dependencyName)) {
+                throw new IllegalArgumentException("Could not find dependency: " + dependencyName);
+            }
+            DescriptorProtos.FileDescriptorProto dependencyProto = descriptorProtoIndex.get(dependencyName);
+            dependencies.add(descriptorFromProto(dependencyProto, descriptorProtoIndex, descriptorCache));
+        }
+
+        // Finally, construct the actual descriptor.
+        Descriptors.FileDescriptor[] empty = new Descriptors.FileDescriptor[0];
+        return Descriptors.FileDescriptor.buildFrom(descriptorProto, dependencies.build().toArray(empty), false);
+    }
+
+}
+
diff --git a/metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/model/FieldTypeEdge.java b/metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/model/FieldTypeEdge.java
new file mode 100644
index 0000000000000..7926ba0702762
--- /dev/null
+++ b/metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/model/FieldTypeEdge.java
@@ -0,0 +1,55 @@
+package datahub.protobuf.model;
+
+import lombok.Builder;
+import lombok.Getter;
+import org.jgrapht.graph.DefaultDirectedGraph;
+import org.jgrapht.graph.DefaultEdge;
+
+
+@Builder
+@Getter
+public class FieldTypeEdge extends DefaultEdge {
+    @Builder.Default
+    private final String type = "";
+    @Builder.Default
+    private final boolean isMessageType = false;
+    private final transient ProtobufElement edgeSource;
+    private final transient ProtobufElement edgeTarget;
+
+    public FieldTypeEdge inGraph(DefaultDirectedGraph<ProtobufElement, FieldTypeEdge> g) {
+        g.addEdge(edgeSource, edgeTarget, this);
+        return this;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) {
+            return true;
+        }
+        if (o == null || getClass() != o.getClass()) {
+            return false;
+        }
+
+        FieldTypeEdge that = (FieldTypeEdge) o;
+
+        if (isMessageType() != that.isMessageType()) {
+            return false;
+        }
+        if (!getType().equals(that.getType())) {
+            return false;
+        }
+        if (!getEdgeSource().equals(that.getEdgeSource())) {
+            return false;
+        }
+        return getEdgeTarget().equals(that.getEdgeTarget());
+    }
+
+    @Override
+    public int hashCode() {
+        int result = getType().hashCode();
+        result = 31 * result + (isMessageType() ? 1 : 0);
+        result = 31 * result + getEdgeSource().hashCode();
+        result = 31 * result + getEdgeTarget().hashCode();
+        return result;
+    }
+}
diff --git a/metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/model/ProtobufElement.java b/metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/model/ProtobufElement.java
new file mode 100644
index 0000000000000..91c76fe16b73f
--- /dev/null
+++ b/metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/model/ProtobufElement.java
@@ -0,0 +1,32 @@
+package datahub.protobuf.model;
+
+import com.google.protobuf.DescriptorProtos.DescriptorProto;
+import com.google.protobuf.DescriptorProtos.FileDescriptorProto;
+import com.google.protobuf.DescriptorProtos.SourceCodeInfo;
+import datahub.protobuf.visitors.ProtobufModelVisitor;
+import datahub.protobuf.visitors.VisitContext;
+
+import java.util.List;
+import java.util.stream.Stream;
+
+
+public interface ProtobufElement {
+    String name();
+    String fullName();
+    String nativeType();
+    String comment();
+    String fieldPathType();
+
+    FileDescriptorProto fileProto();
+    DescriptorProto messageProto();
+
+    default Stream<SourceCodeInfo.Location> messageLocations() {
+        List<SourceCodeInfo.Location> fileLocations = fileProto().getSourceCodeInfo().getLocationList();
+        return fileLocations.stream()
+                .filter(loc -> loc.getPathCount() > 1
+                        && loc.getPath(0) == FileDescriptorProto.MESSAGE_TYPE_FIELD_NUMBER
+                        && messageProto() == fileProto().getMessageType(loc.getPath(1)));
+    }
+
+    <T> Stream<T> accept(ProtobufModelVisitor<T> v, VisitContext context);
+}
diff --git a/metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/model/ProtobufEnum.java b/metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/model/ProtobufEnum.java
new file mode 100644
index 0000000000000..ff894112d0d51
--- /dev/null
+++ b/metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/model/ProtobufEnum.java
@@ -0,0 +1,87 @@
+package datahub.protobuf.model;
+
+import com.google.protobuf.DescriptorProtos.DescriptorProto;
+import com.google.protobuf.DescriptorProtos.EnumDescriptorProto;
+import com.google.protobuf.DescriptorProtos.FileDescriptorProto;
+import com.linkedin.schema.EnumType;
+import com.linkedin.schema.SchemaFieldDataType;
+import datahub.protobuf.ProtobufUtils;
+import lombok.Builder;
+import lombok.Getter;
+
+import java.util.stream.Collectors;
+
+
+@Getter
+public class ProtobufEnum extends ProtobufMessage {
+
+    private final EnumDescriptorProto enumProto;
+
+    @Builder(builderMethodName = "enumBuilder")
+    public ProtobufEnum(FileDescriptorProto fileProto,
+                        DescriptorProto messageProto,
+                        EnumDescriptorProto enumProto) {
+        super(messageProto, null, fileProto);
+        this.enumProto = enumProto;
+    }
+
+    @Override
+    public String name() {
+        return enumProto.getName();
+    }
+
+    @Override
+    public String fieldPathType() {
+        return "[type=enum]";
+    }
+
+    @Override
+    public String nativeType() {
+        return "enum";
+    }
+
+    @Override
+    public SchemaFieldDataType schemaFieldDataType() throws IllegalStateException {
+        return new SchemaFieldDataType().setType(SchemaFieldDataType.Type.create(new EnumType()));
+    }
+
+    @Override
+    public String comment() {
+        return messageLocations()
+                .filter(loc -> loc.getPathCount() > 3
+                        && loc.getPath(2) == DescriptorProto.ENUM_TYPE_FIELD_NUMBER
+                        && enumProto == messageProto().getEnumType(loc.getPath(3)))
+                .map(ProtobufUtils::collapseLocationComments)
+                .collect(Collectors.joining("\n"))
+                .trim();
+    }
+
+    @Override
+    public String toString() {
+        return String.format("ProtobufEnum[%s]", fullName());
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) {
+            return true;
+        }
+        if (o == null || getClass() != o.getClass()) {
+            return false;
+        }
+        if (!super.equals(o)) {
+            return false;
+        }
+
+        ProtobufEnum that = (ProtobufEnum) o;
+
+        return getEnumProto().equals(that.getEnumProto());
+    }
+
+    @Override
+    public int hashCode() {
+        int result = super.hashCode();
+        result = 31 * result + getEnumProto().hashCode();
+        return result;
+    }
+}
diff --git a/metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/model/ProtobufField.java b/metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/model/ProtobufField.java
new file mode 100644
index 0000000000000..61b0572290606
--- /dev/null
+++ b/metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/model/ProtobufField.java
@@ -0,0 +1,239 @@
+package datahub.protobuf.model;
+
+import com.google.protobuf.DescriptorProtos.DescriptorProto;
+import com.google.protobuf.DescriptorProtos.FieldDescriptorProto;
+import com.google.protobuf.DescriptorProtos.FileDescriptorProto;
+import com.google.protobuf.DescriptorProtos.OneofDescriptorProto;
+import com.linkedin.data.template.StringArray;
+import com.linkedin.schema.ArrayType;
+import com.linkedin.schema.BooleanType;
+import com.linkedin.schema.BytesType;
+import com.linkedin.schema.EnumType;
+import com.linkedin.schema.FixedType;
+import com.linkedin.schema.NumberType;
+import com.linkedin.schema.RecordType;
+import com.linkedin.schema.SchemaFieldDataType;
+import com.linkedin.schema.StringType;
+import datahub.protobuf.ProtobufUtils;
+import datahub.protobuf.visitors.ProtobufModelVisitor;
+import datahub.protobuf.visitors.VisitContext;
+import lombok.AllArgsConstructor;
+import lombok.Builder;
+import lombok.Getter;
+
+import java.util.Optional;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+
+
+@Builder(toBuilder = true)
+@Getter
+@AllArgsConstructor
+public class ProtobufField implements ProtobufElement {
+
+    private final ProtobufMessage protobufMessage;
+    private final FieldDescriptorProto fieldProto;
+    private final String nativeType;
+    private final String fieldPathType;
+    private final Boolean isMessageType;
+    private final SchemaFieldDataType schemaFieldDataType;
+
+    public OneofDescriptorProto oneOfProto() {
+        if (fieldProto.hasOneofIndex()) {
+            return protobufMessage.messageProto().getOneofDecl(fieldProto.getOneofIndex());
+        }
+        return null;
+    }
+
+    @Override
+    public FileDescriptorProto fileProto() {
+        return protobufMessage.fileProto();
+    }
+
+    @Override
+    public DescriptorProto messageProto() {
+        return protobufMessage.messageProto();
+    }
+
+    public String parentMessageName() {
+        return protobufMessage.fullName();
+    }
+
+    @Override
+    public String name() {
+        return fieldProto.getName();
+    }
+
+    @Override
+    public String fullName() {
+        return String.join(".", parentMessageName(), name());
+    }
+
+    @Override
+    public String nativeType() {
+        return Optional.ofNullable(nativeType).orElseGet(() -> {
+            if (fieldProto.getTypeName().isEmpty()) {
+                return fieldProto.getType().name().split("_")[1].toLowerCase();
+            } else {
+                return fieldProto.getTypeName().replaceFirst("^[.]", "");
+            }
+        });
+    }
+
+    @Override
+    public String fieldPathType() {
+        return Optional.ofNullable(fieldPathType).orElseGet(() -> {
+            final String pathType;
+
+            switch (fieldProto.getType()) {
+                case TYPE_DOUBLE:
+                    pathType = "double";
+                    break;
+                case TYPE_FLOAT:
+                    pathType = "float";
+                    break;
+                case TYPE_SFIXED64:
+                case TYPE_FIXED64:
+                case TYPE_UINT64:
+                case TYPE_INT64:
+                case TYPE_SINT64:
+                    pathType = "long";
+                    break;
+                case TYPE_FIXED32:
+                case TYPE_SFIXED32:
+                case TYPE_INT32:
+                case TYPE_UINT32:
+                case TYPE_SINT32:
+                    pathType = "int";
+                    break;
+                case TYPE_BYTES:
+                    pathType = "bytes";
+                    break;
+                case TYPE_ENUM:
+                    pathType = "enum";
+                    break;
+                case TYPE_BOOL:
+                    pathType = "boolean";
+                    break;
+                case TYPE_STRING:
+                    pathType = "string";
+                    break;
+                case TYPE_GROUP:
+                case TYPE_MESSAGE:
+                    pathType = nativeType().replace(".", "_");
+                    break;
+                default:
+                    throw new IllegalStateException(String.format("Unexpected FieldDescriptorProto => FieldPathType %s", fieldProto.getType()));
+            }
+
+            StringArray fieldPath = new StringArray();
+
+            if (schemaFieldDataType().getType().isArrayType()) {
+                fieldPath.add("[type=array]");
+            }
+
+            fieldPath.add(String.format("[type=%s]", pathType));
+
+            return String.join(".", fieldPath);
+        });
+    }
+
+    public boolean isMessage() {
+        return Optional.ofNullable(isMessageType).orElseGet(() ->
+                    fieldProto.getType().equals(FieldDescriptorProto.Type.TYPE_MESSAGE));
+    }
+
+    public int sortWeight() {
+        return messageProto().getFieldList().indexOf(fieldProto) + 1;
+    }
+
+    public SchemaFieldDataType schemaFieldDataType() throws IllegalStateException {
+        return Optional.ofNullable(schemaFieldDataType).orElseGet(() -> {
+            final SchemaFieldDataType.Type fieldType;
+
+            switch (fieldProto.getType()) {
+                case TYPE_DOUBLE:
+                case TYPE_FLOAT:
+                case TYPE_INT64:
+                case TYPE_UINT64:
+                case TYPE_INT32:
+                case TYPE_UINT32:
+                case TYPE_SINT32:
+                case TYPE_SINT64:
+                    fieldType = SchemaFieldDataType.Type.create(new NumberType());
+                    break;
+                case TYPE_GROUP:
+                case TYPE_MESSAGE:
+                    fieldType = SchemaFieldDataType.Type.create(new RecordType());
+                    break;
+                case TYPE_BYTES:
+                    fieldType = SchemaFieldDataType.Type.create(new BytesType());
+                    break;
+                case TYPE_ENUM:
+                    fieldType = SchemaFieldDataType.Type.create(new EnumType());
+                    break;
+                case TYPE_BOOL:
+                    fieldType = SchemaFieldDataType.Type.create(new BooleanType());
+                    break;
+                case TYPE_STRING:
+                    fieldType = SchemaFieldDataType.Type.create(new StringType());
+                    break;
+                case TYPE_FIXED64:
+                case TYPE_FIXED32:
+                case TYPE_SFIXED32:
+                case TYPE_SFIXED64:
+                    fieldType = SchemaFieldDataType.Type.create(new FixedType());
+                    break;
+                default:
+                    throw new IllegalStateException(String.format("Unexpected FieldDescriptorProto => SchemaFieldDataType: %s", fieldProto.getType()));
+            }
+
+            if (fieldProto.getLabel().equals(FieldDescriptorProto.Label.LABEL_REPEATED)) {
+                return new SchemaFieldDataType().setType(SchemaFieldDataType.Type.create(new ArrayType()
+                        .setNestedType(new StringArray())));
+            }
+
+            return new SchemaFieldDataType().setType(fieldType);
+        });
+    }
+
+    @Override
+    public String comment() {
+        return messageLocations()
+                .filter(loc -> loc.getPathCount() > 3
+                        && loc.getPath(2) == DescriptorProto.FIELD_FIELD_NUMBER
+                        && fieldProto == messageProto().getField(loc.getPath(3)))
+                .map(ProtobufUtils::collapseLocationComments)
+                .collect(Collectors.joining("\n"))
+                .trim();
+    }
+
+    @Override
+    public <T> Stream<T> accept(ProtobufModelVisitor<T> visitor, VisitContext context) {
+        return visitor.visitField(this, context);
+    }
+
+    @Override
+    public String toString() {
+        return String.format("ProtobufField[%s]", fullName());
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) {
+            return true;
+        }
+        if (o == null || getClass() != o.getClass()) {
+            return false;
+        }
+
+        ProtobufElement that = (ProtobufElement) o;
+
+        return fullName().equals(that.fullName());
+    }
+
+    @Override
+    public int hashCode() {
+        return fullName().hashCode();
+    }
+}
diff --git a/metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/model/ProtobufGraph.java b/metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/model/ProtobufGraph.java
new file mode 100644
index 0000000000000..a3394a0a579fb
--- /dev/null
+++ b/metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/model/ProtobufGraph.java
@@ -0,0 +1,351 @@
+package datahub.protobuf.model;
+
+import com.google.protobuf.DescriptorProtos;
+import com.google.protobuf.ExtensionRegistry;
+import com.google.protobuf.InvalidProtocolBufferException;
+import datahub.protobuf.ProtobufUtils;
+import datahub.protobuf.visitors.ProtobufModelVisitor;
+import datahub.protobuf.visitors.VisitContext;
+import org.jgrapht.GraphPath;
+import org.jgrapht.alg.shortestpath.AllDirectedPaths;
+import org.jgrapht.graph.DefaultDirectedGraph;
+
+import java.util.Collection;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Set;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+
+
+public class ProtobufGraph extends DefaultDirectedGraph<ProtobufElement, FieldTypeEdge> {
+    private final transient ProtobufMessage rootProtobufMessage;
+    private final transient AllDirectedPaths<ProtobufElement, FieldTypeEdge> directedPaths;
+    private final transient ExtensionRegistry registry;
+
+    public ProtobufGraph(DescriptorProtos.FileDescriptorSet fileSet) throws InvalidProtocolBufferException {
+        this(fileSet, null, true);
+    }
+
+    public ProtobufGraph(DescriptorProtos.FileDescriptorSet fileSet, String messageName) throws InvalidProtocolBufferException {
+        this(fileSet, messageName, true);
+    }
+
+    public ProtobufGraph(DescriptorProtos.FileDescriptorSet fileSet, String messageName, boolean flattenGoogleWrapped) throws InvalidProtocolBufferException {
+        super(FieldTypeEdge.class);
+        this.registry = ProtobufUtils.buildRegistry(fileSet);
+        DescriptorProtos.FileDescriptorSet fileSetExtended = DescriptorProtos.FileDescriptorSet
+                .parseFrom(fileSet.toByteArray(), this.registry);
+        buildProtobufGraph(fileSetExtended);
+        if (flattenGoogleWrapped) {
+            flattenGoogleWrapped();
+        }
+
+        if (messageName != null) {
+            this.rootProtobufMessage = findMessage(messageName);
+        } else {
+            DescriptorProtos.FileDescriptorProto lastFile = fileSetExtended.getFile(fileSetExtended.getFileCount() - 1);
+            this.rootProtobufMessage = autodetectRootMessage(lastFile);
+        }
+
+        this.directedPaths = new AllDirectedPaths<>(this);
+    }
+
+    public List<GraphPath<ProtobufElement, FieldTypeEdge>> getAllPaths(ProtobufElement a, ProtobufElement b) {
+        return directedPaths.getAllPaths(a, b, true, null);
+    }
+
+    public ExtensionRegistry getRegistry() {
+        return registry;
+    }
+
+    public String getFullName() {
+        return rootProtobufMessage.fullName();
+    }
+
+    public int getMajorVersion() {
+        return rootProtobufMessage.majorVersion();
+    }
+
+    public String getComment() {
+        return rootProtobufMessage.comment();
+    }
+
+    public ProtobufMessage root() {
+        return rootProtobufMessage;
+    }
+
+
+    public <T, V extends ProtobufModelVisitor<T>> Stream<T> accept(VisitContext.VisitContextBuilder contextBuilder, Collection<V> visitors) {
+        VisitContext context = Optional.ofNullable(contextBuilder).orElse(VisitContext.builder()).graph(this).build();
+        return accept(context, visitors);
+    }
+
+    public <T, V extends ProtobufModelVisitor<T>> Stream<T> accept(VisitContext context, Collection<V> visitors) {
+        return Stream.concat(
+                visitors.stream().flatMap(visitor -> visitor.visitGraph(context)),
+                vertexSet().stream().flatMap(vertex -> visitors.stream().flatMap(visitor -> vertex.accept(visitor, context)))
+        );
+    }
+
+    protected ProtobufMessage autodetectRootMessage(DescriptorProtos.FileDescriptorProto lastFile) throws IllegalArgumentException {
+        return (ProtobufMessage) vertexSet().stream()
+                .filter(v -> // incoming edges of fields
+                        lastFile.equals(v.fileProto())
+                                && v instanceof ProtobufMessage
+                                && incomingEdgesOf(v).isEmpty()
+                                && outgoingEdgesOf(v).stream()
+                                .flatMap(e -> incomingEdgesOf(e.getEdgeTarget()).stream())
+                                .allMatch(e -> e.getEdgeSource().equals(v))) // all the incoming edges on the child vertices should be self
+                .findFirst().orElseThrow(() ->
+                        new IllegalArgumentException("Cannot autodetect protobuf Message.")
+                );
+    }
+
+    public ProtobufMessage findMessage(String messageName) throws IllegalArgumentException {
+        return (ProtobufMessage) vertexSet().stream()
+                .filter(v -> v instanceof ProtobufMessage && messageName.equals(v.fullName()))
+                .findFirst().orElseThrow(() -> new IllegalArgumentException(String.format("Cannot find protobuf Message %s", messageName)));
+    }
+
+    private void buildProtobufGraph(DescriptorProtos.FileDescriptorSet fileSet) {
+        // Attach fields to messages
+        Map<String, List<ProtobufField>> messageFieldMap = fileSet.getFileList().stream().flatMap(fileProto ->
+                fileProto.getMessageTypeList().stream().flatMap(messageProto -> {
+
+                    ProtobufMessage messageVertex = ProtobufMessage.builder()
+                            .fileProto(fileProto)
+                            .messageProto(messageProto)
+                            .build();
+                    addVertex(messageVertex);
+
+                    // Handle nested fields
+                    Stream<ProtobufField> nestedFields = addNestedMessage(fileProto, messageProto);
+
+                    // Add enum types
+                    addEnum(fileProto, messageProto);
+
+                    // handle normal fields and oneofs
+                    Stream<ProtobufField> fields = messageProto.getFieldList().stream().flatMap(fieldProto -> {
+                        ProtobufField fieldVertex = ProtobufField.builder()
+                                .protobufMessage(messageVertex)
+                                .fieldProto(fieldProto)
+                                .build();
+
+                        // Add field vertex
+                        addVertex(fieldVertex);
+
+                        if (fieldVertex.oneOfProto() != null) {
+                            // Handle oneOf
+                            return addOneOf(messageVertex, fieldVertex);
+                        } else {
+                            // Add schema to field edge
+                            return linkMessageToField(messageVertex, fieldVertex);
+                        }
+                    });
+
+                    return Stream.concat(nestedFields, fields);
+                })
+        ).collect(Collectors.groupingBy(ProtobufField::parentMessageName));
+
+        attachMessagesToFields(messageFieldMap);
+    }
+
+    private void attachMessagesToFields(Map<String, List<ProtobufField>> messageFieldMap) {
+        // Connect field to Message
+        List<FieldTypeEdge> messageFieldEdges = edgeSet().stream()
+                .filter(FieldTypeEdge::isMessageType)
+                .collect(Collectors.toList());
+
+        messageFieldEdges.forEach(e -> {
+            ProtobufField source = (ProtobufField) e.getEdgeTarget();
+
+            List<ProtobufField> targetFields = messageFieldMap.get(source.nativeType());
+            if (targetFields != null) {
+                targetFields.forEach(target ->
+                        FieldTypeEdge.builder()
+                                .edgeSource(source)
+                                .edgeTarget(target)
+                                .type(target.fieldPathType())
+                                .isMessageType(target.isMessage())
+                                .build().inGraph(this)
+                );
+            }
+        });
+    }
+
+    private void addEnum(DescriptorProtos.FileDescriptorProto fileProto, DescriptorProtos.DescriptorProto messageProto) {
+        messageProto.getEnumTypeList().forEach(enumProto -> {
+            ProtobufEnum enumVertex = ProtobufEnum.enumBuilder()
+                    .fileProto(fileProto)
+                    .messageProto(messageProto)
+                    .enumProto(enumProto)
+                    .build();
+            addVertex(enumVertex);
+        });
+    }
+
+    private Stream<ProtobufField> addNestedMessage(DescriptorProtos.FileDescriptorProto fileProto, DescriptorProtos.DescriptorProto messageProto) {
+        return messageProto.getNestedTypeList().stream().flatMap(nestedMessageProto -> {
+            ProtobufMessage nestedMessageVertex = ProtobufMessage.builder()
+                    .fileProto(fileProto)
+                    .parentMessageProto(messageProto)
+                    .messageProto(nestedMessageProto)
+                    .build();
+            addVertex(nestedMessageVertex);
+
+            return nestedMessageProto.getFieldList().stream().map(nestedFieldProto -> {
+                ProtobufField field = ProtobufField.builder()
+                        .protobufMessage(nestedMessageVertex)
+                        .fieldProto(nestedFieldProto)
+                        .build();
+
+                // Add field vertex
+                addVertex(field);
+
+                // Add schema to field edge
+                if (!field.isMessage()) {
+                    FieldTypeEdge.builder()
+                            .edgeSource(nestedMessageVertex)
+                            .edgeTarget(field)
+                            .type(field.fieldPathType())
+                            .build().inGraph(this);
+                }
+
+                return field;
+            });
+        });
+    }
+
+    private Stream<ProtobufField> addOneOf(ProtobufMessage messageVertex, ProtobufField fieldVertex) {
+        // Handle oneOf
+        ProtobufField oneOfVertex = ProtobufOneOfField.oneOfBuilder()
+                .protobufMessage(messageVertex)
+                .fieldProto(fieldVertex.getFieldProto())
+                .build();
+        addVertex(oneOfVertex);
+
+        FieldTypeEdge.builder()
+                .edgeSource(messageVertex)
+                .edgeTarget(oneOfVertex)
+                .type(oneOfVertex.fieldPathType())
+                .build().inGraph(this);
+
+        // Add oneOf field to field edge
+        FieldTypeEdge.builder()
+                .edgeSource(oneOfVertex)
+                .edgeTarget(fieldVertex)
+                .type(fieldVertex.fieldPathType())
+                .isMessageType(fieldVertex.isMessage())
+                .build().inGraph(this);
+
+        return Stream.of(oneOfVertex);
+    }
+
+    private Stream<ProtobufField> linkMessageToField(ProtobufMessage messageVertex, ProtobufField fieldVertex) {
+        FieldTypeEdge.builder()
+                .edgeSource(messageVertex)
+                .edgeTarget(fieldVertex)
+                .type(fieldVertex.fieldPathType())
+                .isMessageType(fieldVertex.isMessage())
+                .build().inGraph(this);
+
+        return Stream.of(fieldVertex);
+    }
+
+    private void flattenGoogleWrapped() {
+        HashSet<ProtobufElement> removeVertices = new HashSet<>();
+        HashSet<FieldTypeEdge> removeEdges = new HashSet<>();
+        HashSet<ProtobufElement> addVertices = new HashSet<>();
+        HashSet<FieldTypeEdge> addEdges = new HashSet<>();
+
+        Set<ProtobufElement> googleWrapped = vertexSet().stream()
+                .filter(v -> v instanceof ProtobufMessage
+                        && "google/protobuf/wrappers.proto".equals(v.fileProto().getName()))
+                .collect(Collectors.toSet());
+        removeVertices.addAll(googleWrapped);
+
+        Set<ProtobufField> wrappedPrimitiveFields = googleWrapped.stream()
+                .flatMap(wrapped -> outgoingEdgesOf(wrapped).stream())
+                .map(FieldTypeEdge::getEdgeTarget)
+                .map(ProtobufField.class::cast)
+                .collect(Collectors.toSet());
+        removeVertices.addAll(wrappedPrimitiveFields);
+
+        wrappedPrimitiveFields.forEach(primitiveField -> {
+            // remove incoming old edges to primitive
+            removeEdges.addAll(incomingEdgesOf(primitiveField));
+
+            Set<ProtobufField> originatingFields = incomingEdgesOf(primitiveField).stream()
+                    .map(FieldTypeEdge::getEdgeSource)
+                    .filter(edgeSource -> !googleWrapped.contains(edgeSource))
+                    .map(ProtobufField.class::cast)
+                    .collect(Collectors.toSet());
+            removeVertices.addAll(originatingFields);
+
+            originatingFields.forEach(originatingField -> {
+                // Replacement Field
+                ProtobufElement fieldVertex = originatingField.toBuilder()
+                        .fieldPathType(primitiveField.fieldPathType())
+                        .schemaFieldDataType(primitiveField.schemaFieldDataType())
+                        .isMessageType(false)
+                        .build();
+                addVertices.add(fieldVertex);
+
+                // link source field parent directly to primitive
+                Set<FieldTypeEdge> incomingEdges = incomingEdgesOf(originatingField);
+                removeEdges.addAll(incomingEdgesOf(originatingField));
+                addEdges.addAll(incomingEdges.stream().map(oldEdge ->
+                        // Replace old edge with new edge to primitive
+                        FieldTypeEdge.builder()
+                                .edgeSource(oldEdge.getEdgeSource())
+                                .edgeTarget(fieldVertex)
+                                .type(primitiveField.fieldPathType())
+                                .isMessageType(false) // known primitive
+                                .build()).collect(Collectors.toSet()));
+            });
+
+            // remove old fields
+            removeVertices.addAll(originatingFields);
+        });
+
+        // Remove edges
+        removeAllEdges(removeEdges);
+        // Remove vertices
+        removeAllVertices(removeVertices);
+        // Add vertices
+        addVertices.forEach(this::addVertex);
+        // Add edges
+        addEdges.forEach(e -> e.inGraph(this));
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) {
+            return true;
+        }
+        if (o == null || getClass() != o.getClass()) {
+            return false;
+        }
+        if (!super.equals(o)) {
+            return false;
+        }
+
+        ProtobufGraph that = (ProtobufGraph) o;
+
+        return rootProtobufMessage.equals(that.rootProtobufMessage);
+    }
+
+    @Override
+    public int hashCode() {
+        int result = super.hashCode();
+        result = 31 * result + rootProtobufMessage.hashCode();
+        return result;
+    }
+
+    public String getHash() {
+        return String.valueOf(super.hashCode());
+    }
+}
diff --git a/metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/model/ProtobufMessage.java b/metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/model/ProtobufMessage.java
new file mode 100644
index 0000000000000..6b46b11231623
--- /dev/null
+++ b/metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/model/ProtobufMessage.java
@@ -0,0 +1,122 @@
+package datahub.protobuf.model;
+
+import com.google.protobuf.DescriptorProtos.DescriptorProto;
+import com.google.protobuf.DescriptorProtos.FileDescriptorProto;
+import com.linkedin.schema.MapType;
+import com.linkedin.schema.RecordType;
+import lombok.AllArgsConstructor;
+import lombok.Builder;
+import java.util.Arrays;
+import java.util.stream.Stream;
+import com.linkedin.schema.SchemaFieldDataType;
+
+import datahub.protobuf.ProtobufUtils;
+import datahub.protobuf.visitors.ProtobufModelVisitor;
+import datahub.protobuf.visitors.VisitContext;
+
+
+
+@Builder
+@AllArgsConstructor
+public class ProtobufMessage implements ProtobufElement {
+    private final DescriptorProto messageProto;
+    private final DescriptorProto parentMessageProto;
+    private final FileDescriptorProto fileProto;
+
+    @Override
+    public String name() {
+        return messageProto.getName();
+    }
+
+    @Override
+    public String fullName() {
+        if (parentMessageProto  != null) {
+            return String.join(".", fileProto.getPackage(), parentMessageProto.getName(), name());
+        }
+        return String.join(".", fileProto.getPackage(), name());
+    }
+
+    @Override
+    public String nativeType() {
+        return fullName();
+    }
+
+    @Override
+    public String fieldPathType() {
+        return String.format("[type=%s]", nativeType().replace(".", "_"));
+    }
+
+    @Override
+    public FileDescriptorProto fileProto() {
+        return fileProto;
+    }
+
+    @Override
+    public DescriptorProto messageProto() {
+        return messageProto;
+    }
+
+    public SchemaFieldDataType schemaFieldDataType() {
+        if (parentMessageProto != null && messageProto.getName().equals("MapFieldEntry")) {
+            return new SchemaFieldDataType().setType(SchemaFieldDataType.Type.create(new MapType()));
+        }
+        return new SchemaFieldDataType().setType(SchemaFieldDataType.Type.create(new RecordType()));
+    }
+
+    public int majorVersion() {
+        return Integer.parseInt(Arrays.stream(fileProto.getName().split("/"))
+                .filter(p -> p.matches("^v[0-9]+$"))
+                .findFirst()
+                .map(p -> p.replace("v", ""))
+                .orElse("1"));
+    }
+
+    @Override
+    public String comment() {
+        return messageLocations()
+                .map(ProtobufUtils::collapseLocationComments)
+                .findFirst().orElse("");
+    }
+
+    @Override
+    public <T> Stream<T> accept(ProtobufModelVisitor<T> visitor, VisitContext context) {
+        return visitor.visitMessage(this, context);
+    }
+
+    @Override
+    public String toString() {
+        return String.format("ProtobufMessage[%s]", fullName());
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) {
+            return true;
+        }
+        if (o == null || getClass() != o.getClass()) {
+            return false;
+        }
+
+        ProtobufMessage that = (ProtobufMessage) o;
+
+        if (!fullName().equals(that.fullName())) {
+            return false;
+        }
+        if (!messageProto.equals(that.messageProto)) {
+            return false;
+        }
+        if (parentMessageProto != null ? !parentMessageProto.equals(that.parentMessageProto) : that.parentMessageProto != null) {
+            return false;
+        }
+        return fileProto.equals(that.fileProto);
+    }
+
+    @Override
+    public int hashCode() {
+        int result = messageProto.hashCode();
+        result = 31 * result + (parentMessageProto != null ? parentMessageProto.hashCode() : 0);
+        result = 31 * result + fileProto.hashCode();
+        result = 31 * result + fullName().hashCode();
+        return result;
+    }
+}
diff --git a/metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/model/ProtobufOneOfField.java b/metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/model/ProtobufOneOfField.java
new file mode 100644
index 0000000000000..1b4db1c0a8c8b
--- /dev/null
+++ b/metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/model/ProtobufOneOfField.java
@@ -0,0 +1,63 @@
+package datahub.protobuf.model;
+
+import com.google.protobuf.DescriptorProtos.DescriptorProto;
+import com.google.protobuf.DescriptorProtos.FieldDescriptorProto;
+import com.linkedin.schema.SchemaFieldDataType;
+import com.linkedin.schema.UnionType;
+import datahub.protobuf.ProtobufUtils;
+import lombok.Builder;
+import lombok.Getter;
+
+import java.util.stream.Collectors;
+
+
+@Getter
+public class ProtobufOneOfField extends ProtobufField {
+
+    @Builder(builderMethodName = "oneOfBuilder")
+    public ProtobufOneOfField(ProtobufMessage protobufMessage,
+                              FieldDescriptorProto fieldProto) {
+        super(protobufMessage, fieldProto, null, null, null, null);
+    }
+
+    @Override
+    public String name() {
+        return oneOfProto().getName();
+    }
+
+    @Override
+    public String fieldPathType() {
+        return "[type=union]";
+    }
+
+    @Override
+    public String nativeType() {
+        return "oneof";
+    }
+
+    @Override
+    public boolean isMessage() {
+        return false;
+    }
+
+    @Override
+    public SchemaFieldDataType schemaFieldDataType() throws IllegalStateException {
+        return new SchemaFieldDataType().setType(SchemaFieldDataType.Type.create(new UnionType()));
+    }
+
+    @Override
+    public String comment() {
+        return messageLocations()
+                .filter(loc -> loc.getPathCount() > 3
+                        && loc.getPath(2) == DescriptorProto.ONEOF_DECL_FIELD_NUMBER
+                        && oneOfProto() == messageProto().getOneofDecl(loc.getPath(3)))
+                .map(ProtobufUtils::collapseLocationComments)
+                .collect(Collectors.joining("\n"))
+                .trim();
+    }
+
+    @Override
+    public String toString() {
+        return String.format("ProtobufOneOf[%s]", fullName());
+    }
+}
diff --git a/metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/visitors/ProtobufExtensionUtil.java b/metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/visitors/ProtobufExtensionUtil.java
new file mode 100644
index 0000000000000..943363e736a62
--- /dev/null
+++ b/metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/visitors/ProtobufExtensionUtil.java
@@ -0,0 +1,109 @@
+package datahub.protobuf.visitors;
+
+import com.google.protobuf.DescriptorProtos;
+import com.google.protobuf.Descriptors;
+import com.google.protobuf.ExtensionRegistry;
+import com.google.protobuf.InvalidProtocolBufferException;
+import com.linkedin.common.GlossaryTermAssociation;
+import com.linkedin.common.urn.GlossaryTermUrn;
+import com.linkedin.tag.TagProperties;
+
+import java.util.Collection;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+
+public class ProtobufExtensionUtil {
+
+    private ProtobufExtensionUtil() { }
+
+    public static DescriptorProtos.FieldDescriptorProto extendProto(DescriptorProtos.FieldDescriptorProto proto, ExtensionRegistry registry) {
+        try {
+            return DescriptorProtos.FieldDescriptorProto.parseFrom(proto.toByteArray(), registry);
+        } catch (InvalidProtocolBufferException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    public enum DataHubMetadataType {
+        PROPERTY, TAG, TERM;
+
+        public static final String PROTOBUF_TYPE = "DataHubMetadataType";
+    }
+
+    public static Map<Descriptors.FieldDescriptor, Object> filterByDataHubType(Map<Descriptors.FieldDescriptor, Object> options,
+                                                                               ExtensionRegistry registry, DataHubMetadataType filterType) {
+        return options.entrySet().stream()
+                .filter(entry -> {
+                    DescriptorProtos.FieldDescriptorProto extendedProtoOptions = extendProto(entry.getKey().toProto(), registry);
+                    Optional<DataHubMetadataType> dataHubMetadataType = extendedProtoOptions.getOptions().getAllFields().entrySet().stream()
+                            .filter(extEntry -> extEntry.getKey().getJavaType() == Descriptors.FieldDescriptor.JavaType.ENUM)
+                            .flatMap(extEntry -> {
+                                if (extEntry.getKey().isRepeated()) {
+                                    return ((Collection<Descriptors.EnumValueDescriptor>) extEntry.getValue()).stream();
+                                } else {
+                                    return Stream.of((Descriptors.EnumValueDescriptor) extEntry.getValue());
+                                }
+                            })
+                            .filter(enumDesc -> enumDesc.getType().getFullName().endsWith("." + DataHubMetadataType.PROTOBUF_TYPE))
+                            .map(enumDesc -> DataHubMetadataType.valueOf(enumDesc.getName()))
+                            .filter(dhmt -> dhmt.equals(filterType))
+                            .findFirst();
+
+                    return filterType.equals(dataHubMetadataType.orElse(DataHubMetadataType.PROPERTY));
+                })
+                .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
+    }
+
+    public static Stream<TagProperties> extractTagPropertiesFromOptions(Map<Descriptors.FieldDescriptor, Object> options, ExtensionRegistry registry) {
+        return filterByDataHubType(options, registry, DataHubMetadataType.TAG).entrySet().stream()
+                .filter(e -> e.getKey().isExtension())
+                .map(entry -> {
+                    switch (entry.getKey().getJavaType()) {
+                        case STRING:
+                            return  new TagProperties()
+                                    .setName(String.format("%s.%s", entry.getKey().getName(), entry.getValue().toString()))
+                                    .setDescription(entry.getKey().getFullName());
+                        case BOOLEAN:
+                            if ((boolean) entry.getValue()) {
+                                return new TagProperties()
+                                        .setName(entry.getKey().getName())
+                                        .setDescription(String.format("%s is true.", entry.getKey().getFullName()));
+                            }
+                            return null;
+                        case ENUM:
+                            Descriptors.EnumValueDescriptor desc = (Descriptors.EnumValueDescriptor) entry.getValue();
+                            String name = String.format("%s.%s", desc.getType().getName(), desc.getName());
+                            String others = entry.getKey().getEnumType().getValues().stream()
+                                    .map(Descriptors.EnumValueDescriptor::getName).collect(Collectors.joining(", "));
+                            return new TagProperties()
+                                    .setName(name)
+                                    .setDescription(String.format("Enum %s of {%s}", name, others));
+                        default:
+                            return null;
+                    }
+                }).filter(Objects::nonNull);
+    }
+
+    public static Stream<GlossaryTermAssociation> extractTermAssociationsFromOptions(Map<Descriptors.FieldDescriptor, Object> options,
+                                                                                     ExtensionRegistry registry) {
+        return filterByDataHubType(options, registry, DataHubMetadataType.TERM).entrySet().stream()
+                .filter(e -> e.getKey().isExtension())
+                .map(entry -> {
+                    switch (entry.getKey().getJavaType()) {
+                        case STRING:
+                            return new GlossaryTermAssociation()
+                                    .setUrn(new GlossaryTermUrn(entry.getValue().toString()));
+                        case ENUM:
+                            Descriptors.EnumValueDescriptor desc = (Descriptors.EnumValueDescriptor) entry.getValue();
+                            String name = String.format("%s.%s", desc.getType().getName(), desc.getName());
+                            return new GlossaryTermAssociation()
+                                    .setUrn(new GlossaryTermUrn(name));
+                        default:
+                            return null;
+                    }
+                }).filter(Objects::nonNull);
+    }
+}
diff --git a/metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/visitors/ProtobufModelVisitor.java b/metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/visitors/ProtobufModelVisitor.java
new file mode 100644
index 0000000000000..b5c630302d946
--- /dev/null
+++ b/metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/visitors/ProtobufModelVisitor.java
@@ -0,0 +1,24 @@
+package datahub.protobuf.visitors;
+
+
+import datahub.protobuf.model.ProtobufElement;
+import datahub.protobuf.model.ProtobufField;
+import datahub.protobuf.model.ProtobufMessage;
+
+import java.util.stream.Stream;
+
+public interface ProtobufModelVisitor<T> {
+    default Stream<T> visitField(ProtobufField field, VisitContext context) {
+        return visitElement(field, context);
+    }
+    default Stream<T> visitMessage(ProtobufMessage message, VisitContext context) {
+        return visitElement(message, context);
+    }
+
+    default Stream<T> visitElement(ProtobufElement element, VisitContext context) {
+        return Stream.of();
+    }
+    default Stream<T> visitGraph(VisitContext context) {
+        return Stream.of();
+    }
+}
diff --git a/metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/visitors/VisitContext.java b/metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/visitors/VisitContext.java
new file mode 100644
index 0000000000000..984215fa9e389
--- /dev/null
+++ b/metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/visitors/VisitContext.java
@@ -0,0 +1,58 @@
+package datahub.protobuf.visitors;
+
+import com.linkedin.common.AuditStamp;
+import com.linkedin.common.urn.DatasetUrn;
+import datahub.protobuf.model.FieldTypeEdge;
+import datahub.protobuf.model.ProtobufElement;
+import datahub.protobuf.model.ProtobufField;
+import datahub.protobuf.model.ProtobufGraph;
+import datahub.protobuf.model.ProtobufMessage;
+import lombok.Builder;
+import lombok.Getter;
+import org.jgrapht.GraphPath;
+
+import java.util.List;
+import java.util.Optional;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+import java.util.stream.Stream;
+
+@Builder
+@Getter
+public class VisitContext {
+    public static final String FIELD_PATH_VERSION = "[version=2.0]";
+
+    private final DatasetUrn datasetUrn;
+    private final ProtobufGraph graph;
+    private final AuditStamp auditStamp;
+
+    public ProtobufMessage root() {
+        return graph.root();
+    }
+
+    public Optional<GraphPath<ProtobufElement, FieldTypeEdge>> getFirstFieldPath(ProtobufField field) {
+        return graph.getAllPaths(root(), field).stream().findFirst();
+    }
+
+    public String getFieldPath(GraphPath<ProtobufElement, FieldTypeEdge> path) {
+        String fieldPathString = path.getEdgeList().stream()
+                .flatMap(e -> Stream.of(e.getType(), e.getEdgeTarget().name()))
+                .collect(Collectors.joining("."));
+        return String.join(".", FIELD_PATH_VERSION, root().fieldPathType(), fieldPathString);
+    }
+
+    // This is because order matters for the frontend. Both for matching the protobuf field order
+    // and also the nested struct's fieldPaths
+    public Double calculateSortOrder(GraphPath<ProtobufElement, FieldTypeEdge> path, ProtobufField field) {
+        List<Integer> weights = path.getEdgeList().stream()
+                .map(FieldTypeEdge::getEdgeTarget)
+                .filter(f -> f instanceof ProtobufField)
+                .map(f -> ((ProtobufField) f).sortWeight())
+                .collect(Collectors.toList());
+
+        return IntStream.range(0, weights.size())
+                .mapToDouble(i -> weights.get(i) * (1.0 / (i + 1)))
+                .reduce(Double::sum)
+                .orElse(0);
+    }
+}
diff --git a/metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/visitors/dataset/DatasetVisitor.java b/metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/visitors/dataset/DatasetVisitor.java
new file mode 100644
index 0000000000000..4ae92a6da5d90
--- /dev/null
+++ b/metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/visitors/dataset/DatasetVisitor.java
@@ -0,0 +1,79 @@
+package datahub.protobuf.visitors.dataset;
+
+import com.linkedin.common.GlobalTags;
+import com.linkedin.common.GlossaryTermAssociation;
+import com.linkedin.common.GlossaryTermAssociationArray;
+import com.linkedin.common.GlossaryTerms;
+import com.linkedin.common.InstitutionalMemory;
+import com.linkedin.common.InstitutionalMemoryMetadata;
+import com.linkedin.common.InstitutionalMemoryMetadataArray;
+import com.linkedin.common.TagAssociation;
+import com.linkedin.common.TagAssociationArray;
+import com.linkedin.common.urn.DatasetUrn;
+import com.linkedin.data.template.RecordTemplate;
+import com.linkedin.data.template.StringMap;
+import com.linkedin.dataset.DatasetProperties;
+import com.linkedin.events.metadata.ChangeType;
+import datahub.protobuf.model.ProtobufGraph;
+import datahub.protobuf.visitors.ProtobufModelVisitor;
+import datahub.protobuf.visitors.VisitContext;
+import datahub.event.MetadataChangeProposalWrapper;
+import lombok.AllArgsConstructor;
+import lombok.Builder;
+
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+
+@Builder
+@AllArgsConstructor
+public class DatasetVisitor implements ProtobufModelVisitor<MetadataChangeProposalWrapper<? extends RecordTemplate>> {
+    @Builder.Default
+    private final List<ProtobufModelVisitor<InstitutionalMemoryMetadata>> institutionalMemoryMetadataVisitors = List.of();
+    @Builder.Default
+    private final List<ProtobufModelVisitor<DatasetProperties>> datasetPropertyVisitors = List.of();
+    @Builder.Default
+    private final List<ProtobufModelVisitor<TagAssociation>> tagAssociationVisitors = List.of();
+    @Builder.Default
+    private final List<ProtobufModelVisitor<GlossaryTermAssociation>> termAssociationVisitors = List.of();
+    @Builder.Default
+    private final String protocBase64 = "";
+    @Builder.Default
+    private final ProtobufModelVisitor<String> descriptionVisitor = new DescriptionVisitor();
+
+    @Override
+    public Stream<MetadataChangeProposalWrapper<? extends RecordTemplate>> visitGraph(VisitContext context) {
+        final String datasetUrn = context.getDatasetUrn().toString();
+        final ProtobufGraph g = context.getGraph();
+
+        return Stream.of(
+                new MetadataChangeProposalWrapper<>(DatasetUrn.ENTITY_TYPE, datasetUrn, ChangeType.UPSERT, new DatasetProperties()
+                        .setDescription(g.accept(context, List.of(descriptionVisitor)).collect(Collectors.joining("\n")))
+                        .setCustomProperties(new StringMap(
+                                Stream.concat(
+                                                Stream.of(Map.entry("protoc", protocBase64)),
+                                                g.accept(context, datasetPropertyVisitors).flatMap(props -> props.getCustomProperties().entrySet().stream()))
+                                        .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue))
+                        )), "datasetProperties"),
+                new MetadataChangeProposalWrapper<>(DatasetUrn.ENTITY_TYPE, datasetUrn, ChangeType.UPSERT, new InstitutionalMemory().setElements(
+                        new InstitutionalMemoryMetadataArray(
+                                g.accept(context, institutionalMemoryMetadataVisitors)
+                                        .map(inst -> inst.setCreateStamp(context.getAuditStamp()))
+                                        .collect(Collectors.toMap(InstitutionalMemoryMetadata::getUrl, Function.identity(),
+                                                (a1, a2) -> a1, LinkedHashMap::new))
+                                        .values()
+                    )), "institutionalMemory"),
+                new MetadataChangeProposalWrapper<>(DatasetUrn.ENTITY_TYPE, datasetUrn, ChangeType.UPSERT,
+                        new GlobalTags().setTags(new TagAssociationArray(
+                                g.accept(context, tagAssociationVisitors).collect(Collectors.toList())
+                    )), "globalTags"),
+                new MetadataChangeProposalWrapper<>(DatasetUrn.ENTITY_TYPE, datasetUrn, ChangeType.UPSERT,
+                        new GlossaryTerms().setTerms(new GlossaryTermAssociationArray(
+                                g.accept(context, termAssociationVisitors).collect(Collectors.toList())
+                    )).setAuditStamp(context.getAuditStamp()), "glossaryTerms")
+        );
+    }
+}
diff --git a/metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/visitors/dataset/DescriptionVisitor.java b/metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/visitors/dataset/DescriptionVisitor.java
new file mode 100644
index 0000000000000..802c7e0c05408
--- /dev/null
+++ b/metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/visitors/dataset/DescriptionVisitor.java
@@ -0,0 +1,14 @@
+package datahub.protobuf.visitors.dataset;
+
+import datahub.protobuf.visitors.ProtobufModelVisitor;
+import datahub.protobuf.visitors.VisitContext;
+
+import java.util.stream.Stream;
+
+public class DescriptionVisitor implements ProtobufModelVisitor<String> {
+
+    @Override
+    public Stream<String> visitGraph(VisitContext context) {
+        return Stream.of(context.root().comment());
+    }
+}
diff --git a/metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/visitors/dataset/InstitutionalMemoryVisitor.java b/metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/visitors/dataset/InstitutionalMemoryVisitor.java
new file mode 100644
index 0000000000000..b6f52fe01c109
--- /dev/null
+++ b/metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/visitors/dataset/InstitutionalMemoryVisitor.java
@@ -0,0 +1,140 @@
+package datahub.protobuf.visitors.dataset;
+
+import com.linkedin.common.InstitutionalMemoryMetadata;
+import com.linkedin.common.url.Url;
+import datahub.protobuf.model.ProtobufField;
+import datahub.protobuf.visitors.ProtobufModelVisitor;
+import datahub.protobuf.visitors.VisitContext;
+
+import javax.annotation.Nullable;
+import java.util.LinkedList;
+import java.util.List;
+import java.util.Optional;
+import java.util.Spliterator;
+import java.util.Spliterators;
+import java.util.function.Consumer;
+import java.util.regex.MatchResult;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+import java.util.stream.Stream;
+import java.util.stream.StreamSupport;
+
+public class InstitutionalMemoryVisitor implements ProtobufModelVisitor<InstitutionalMemoryMetadata> {
+    public static final String TEAM_DESC =  "Github Team";
+    public static final String SLACK_CHAN_DESC = "Slack Channel";
+
+    private static final Pattern SLACK_CHANNEL_REGEX = Pattern.compile("(?si).*#([a-z0-9-]+).*");
+    private static final Pattern LINK_REGEX = Pattern.compile("(?s)(\\b(https?|ftp|file)://[-a-zA-Z0-9+&@#/%?=~_|!:,.;]*[-a-zA-Z0-9+&@#/%=~_|])");
+    private final String githubOrganization;
+    private final Pattern githubTeamRegex;
+    private final String slackTeamId;
+
+    public InstitutionalMemoryVisitor(@Nullable String slackTeamId, @Nullable String githubOrganization) {
+        this.slackTeamId = slackTeamId;
+        this.githubOrganization = githubOrganization;
+        if (githubOrganization != null) {
+            this.githubTeamRegex = Pattern.compile(String.format("(?si).*@%s/([a-z-]+).*", githubOrganization));
+        } else {
+            this.githubTeamRegex = null;
+        }
+    }
+
+    //  https://slack.com/app_redirect?channel=fdn-analytics-data-catalog&team=T024F4EL1
+    private Optional<Url> slackLink(String text) {
+        return Optional.ofNullable(slackTeamId).map(teamId -> {
+            Matcher m = SLACK_CHANNEL_REGEX.matcher(text);
+            if (m.matches()) {
+                return new Url(String.format("https://slack.com/app_redirect?channel=%s&team=%s", m.group(1), slackTeamId));
+            } else {
+                return null;
+            }
+        });
+    }
+
+    private Optional<Url> teamLink(String text) {
+        return Optional.ofNullable(githubTeamRegex).map(regex -> {
+            Matcher m = regex.matcher(text);
+            if (m.matches()) {
+                return new Url(String.format("https://github.com/orgs/%s/teams/%s", githubOrganization, m.group(1)));
+            } else {
+                return null;
+            }
+        });
+    }
+
+    @Override
+    public Stream<InstitutionalMemoryMetadata> visitGraph(VisitContext context) {
+        List<InstitutionalMemoryMetadata> institutionalMemoryMetadata = new LinkedList<>();
+
+        teamLink(context.root().comment()).ifPresent(url ->
+                institutionalMemoryMetadata.add(new InstitutionalMemoryMetadata()
+                        .setCreateStamp(context.getAuditStamp())
+                        .setDescription(TEAM_DESC)
+                        .setUrl(url)));
+
+
+        slackLink(context.root().comment()).ifPresent(url ->
+                institutionalMemoryMetadata.add(new InstitutionalMemoryMetadata()
+                        .setCreateStamp(context.getAuditStamp())
+                        .setDescription(SLACK_CHAN_DESC)
+                        .setUrl(url)));
+
+        final int[] cnt = {0};
+        MatcherStream.findMatches(LINK_REGEX, context.root().comment()).forEach(match -> {
+            cnt[0] += 1;
+            institutionalMemoryMetadata.add(new InstitutionalMemoryMetadata()
+                    .setCreateStamp(context.getAuditStamp())
+                    .setDescription(String.format("%s Reference %d", context.root().name(), cnt[0]))
+                    .setUrl(new Url(match.group(1))));
+        });
+
+        return institutionalMemoryMetadata.stream();
+    }
+
+    @Override
+    public Stream<InstitutionalMemoryMetadata> visitField(ProtobufField field, VisitContext context) {
+        List<InstitutionalMemoryMetadata> institutionalMemoryMetadata = new LinkedList<>();
+
+        if (field.messageProto().equals(context.getGraph().root().messageProto())) {
+            final int[] cnt = {0};
+            MatcherStream.findMatches(LINK_REGEX, field.comment()).forEach(match -> {
+                cnt[0] += 1;
+                institutionalMemoryMetadata.add(new InstitutionalMemoryMetadata()
+                        .setCreateStamp(context.getAuditStamp())
+                        .setDescription(String.format("%s.%s Reference %d",
+                                field.getProtobufMessage().name(),
+                                field.getFieldProto().getName(),
+                                cnt[0]))
+                        .setUrl(new Url(match.group(1))));
+            });
+        }
+
+        return institutionalMemoryMetadata.stream();
+    }
+
+    private static class MatcherStream {
+        private MatcherStream() { }
+
+        public static Stream<String> find(Pattern pattern, CharSequence input) {
+            return findMatches(pattern, input).map(MatchResult::group);
+        }
+
+        public static Stream<MatchResult> findMatches(
+                Pattern pattern, CharSequence input) {
+            Matcher matcher = pattern.matcher(input);
+
+            Spliterator<MatchResult> spliterator = new Spliterators.AbstractSpliterator<MatchResult>(
+                    Long.MAX_VALUE, Spliterator.ORDERED | Spliterator.NONNULL) {
+                @Override
+                public boolean tryAdvance(Consumer<? super MatchResult> action) {
+                    if (!matcher.find()) {
+                        return false;
+                    }
+                    action.accept(matcher.toMatchResult());
+                    return true;
+                } };
+
+            return StreamSupport.stream(spliterator, false);
+        }
+    }
+}
diff --git a/metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/visitors/dataset/KafkaTopicPropertyVisitor.java b/metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/visitors/dataset/KafkaTopicPropertyVisitor.java
new file mode 100644
index 0000000000000..57ec38611d47f
--- /dev/null
+++ b/metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/visitors/dataset/KafkaTopicPropertyVisitor.java
@@ -0,0 +1,31 @@
+package datahub.protobuf.visitors.dataset;
+
+import com.linkedin.data.template.StringMap;
+import com.linkedin.dataset.DatasetProperties;
+import datahub.protobuf.visitors.ProtobufModelVisitor;
+import datahub.protobuf.visitors.VisitContext;
+
+import java.util.Map;
+import java.util.Optional;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+import java.util.stream.Stream;
+
+
+public class KafkaTopicPropertyVisitor implements ProtobufModelVisitor<DatasetProperties> {
+
+    @Override
+    public Stream<DatasetProperties> visitGraph(VisitContext context) {
+        return getKafkaTopic(context.root().comment()).stream().map(kafkaTopic ->
+            new DatasetProperties()
+                    .setCustomProperties(new StringMap(Map.of("kafka_topic", kafkaTopic)))
+        );
+    }
+
+    private static final Pattern TOPIC_NAME_REGEX = Pattern.compile("(?si).*kafka.+topic.+[`]([a-z._-]+)[`].*");
+
+    private static Optional<String> getKafkaTopic(String text) {
+        Matcher m = TOPIC_NAME_REGEX.matcher(text);
+        return m.matches() ? Optional.of(m.group(1)) : Optional.empty();
+    }
+}
diff --git a/metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/visitors/dataset/ProtobufExtensionPropertyVisitor.java b/metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/visitors/dataset/ProtobufExtensionPropertyVisitor.java
new file mode 100644
index 0000000000000..325b7e0c3908a
--- /dev/null
+++ b/metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/visitors/dataset/ProtobufExtensionPropertyVisitor.java
@@ -0,0 +1,43 @@
+package datahub.protobuf.visitors.dataset;
+
+import com.google.protobuf.ByteString;
+import com.google.protobuf.DescriptorProtos;
+import com.google.protobuf.Descriptors;
+import com.linkedin.data.template.StringMap;
+import com.linkedin.dataset.DatasetProperties;
+import datahub.protobuf.visitors.ProtobufModelVisitor;
+import datahub.protobuf.visitors.ProtobufExtensionUtil;
+import datahub.protobuf.visitors.VisitContext;
+
+import java.util.Map;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+
+
+public class ProtobufExtensionPropertyVisitor implements ProtobufModelVisitor<DatasetProperties> {
+
+    @Override
+    public Stream<DatasetProperties> visitGraph(VisitContext context) {
+        Map<String, String> properties = ProtobufExtensionUtil.filterByDataHubType(context.root().messageProto()
+                        .getOptions().getAllFields(), context.getGraph().getRegistry(), ProtobufExtensionUtil.DataHubMetadataType.PROPERTY)
+                .entrySet().stream().flatMap(fd -> {
+                    if (fd.getKey().getJavaType() != Descriptors.FieldDescriptor.JavaType.MESSAGE) {
+                        return Stream.of(Map.entry(fd.getKey().getName(), fd.getValue().toString()));
+                    } else {
+                        Descriptors.FieldDescriptor field = fd.getKey();
+                        DescriptorProtos.DescriptorProto value = (DescriptorProtos.DescriptorProto) fd.getValue();
+                        return getProperties(field, value);
+                    }
+                }).collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
+
+        return Stream.of(new DatasetProperties().setCustomProperties(new StringMap(properties)));
+    }
+
+    private static Stream<Map.Entry<String, String>> getProperties(Descriptors.FieldDescriptor field, DescriptorProtos.DescriptorProto value)  {
+       return value.getUnknownFields().asMap().entrySet().stream().map(unknown -> {
+            Descriptors.FieldDescriptor fieldDesc = field.getMessageType().findFieldByNumber(unknown.getKey());
+            String fieldValue = unknown.getValue().getLengthDelimitedList().stream().map(ByteString::toStringUtf8).collect(Collectors.joining(""));
+            return Map.entry(String.join(".", field.getFullName(), fieldDesc.getName()), fieldValue);
+        });
+    }
+}
diff --git a/metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/visitors/dataset/ProtobufExtensionTagAssocVisitor.java b/metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/visitors/dataset/ProtobufExtensionTagAssocVisitor.java
new file mode 100644
index 0000000000000..a19fc48099a7d
--- /dev/null
+++ b/metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/visitors/dataset/ProtobufExtensionTagAssocVisitor.java
@@ -0,0 +1,20 @@
+package datahub.protobuf.visitors.dataset;
+
+import com.linkedin.common.TagAssociation;
+import com.linkedin.common.urn.TagUrn;
+import datahub.protobuf.visitors.ProtobufModelVisitor;
+import datahub.protobuf.visitors.ProtobufExtensionUtil;
+import datahub.protobuf.visitors.VisitContext;
+
+import java.util.stream.Stream;
+
+
+public class ProtobufExtensionTagAssocVisitor implements ProtobufModelVisitor<TagAssociation> {
+
+    @Override
+    public Stream<TagAssociation> visitGraph(VisitContext context) {
+        return ProtobufExtensionUtil.extractTagPropertiesFromOptions(context.root().messageProto().getOptions()
+                        .getAllFields(), context.getGraph().getRegistry())
+                .map(tag -> new TagAssociation().setTag(new TagUrn(tag.getName())));
+    }
+}
diff --git a/metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/visitors/dataset/ProtobufExtensionTermAssocVisitor.java b/metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/visitors/dataset/ProtobufExtensionTermAssocVisitor.java
new file mode 100644
index 0000000000000..00227954f870d
--- /dev/null
+++ b/metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/visitors/dataset/ProtobufExtensionTermAssocVisitor.java
@@ -0,0 +1,17 @@
+package datahub.protobuf.visitors.dataset;
+
+import com.linkedin.common.GlossaryTermAssociation;
+import datahub.protobuf.visitors.ProtobufModelVisitor;
+import datahub.protobuf.visitors.ProtobufExtensionUtil;
+import datahub.protobuf.visitors.VisitContext;
+
+import java.util.stream.Stream;
+
+public class ProtobufExtensionTermAssocVisitor implements ProtobufModelVisitor<GlossaryTermAssociation> {
+
+    @Override
+    public Stream<GlossaryTermAssociation> visitGraph(VisitContext context) {
+        return ProtobufExtensionUtil.extractTermAssociationsFromOptions(context.root().messageProto().getOptions().getAllFields(),
+                context.getGraph().getRegistry());
+    }
+}
diff --git a/metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/visitors/field/ProtobufExtensionFieldVisitor.java b/metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/visitors/field/ProtobufExtensionFieldVisitor.java
new file mode 100644
index 0000000000000..97ef223097f14
--- /dev/null
+++ b/metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/visitors/field/ProtobufExtensionFieldVisitor.java
@@ -0,0 +1,91 @@
+package datahub.protobuf.visitors.field;
+
+import com.linkedin.common.GlobalTags;
+import com.linkedin.common.GlossaryTermAssociation;
+import com.linkedin.common.GlossaryTermAssociationArray;
+import com.linkedin.common.GlossaryTerms;
+import com.linkedin.common.TagAssociation;
+import com.linkedin.common.TagAssociationArray;
+import com.linkedin.common.urn.TagUrn;
+import com.linkedin.schema.SchemaField;
+import com.linkedin.tag.TagProperties;
+import com.linkedin.util.Pair;
+import datahub.protobuf.model.ProtobufField;
+import datahub.protobuf.visitors.ProtobufExtensionUtil;
+import datahub.protobuf.visitors.VisitContext;
+
+import java.util.Comparator;
+import java.util.List;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+
+public class ProtobufExtensionFieldVisitor extends SchemaFieldVisitor {
+
+    @Override
+    public Stream<Pair<SchemaField, Double>> visitField(ProtobufField field, VisitContext context) {
+        boolean isPrimaryKey = field.getFieldProto().getOptions().getAllFields().keySet().stream()
+                .anyMatch(fieldDesc -> fieldDesc.getName().matches("(?i).*primary_?key"));
+
+        List<TagAssociation> tags = Stream.concat(
+                ProtobufExtensionUtil.extractTagPropertiesFromOptions(
+                        field.getFieldProto().getOptions().getAllFields(),
+                        context.getGraph().getRegistry()),
+                        promotedTags(field, context))
+                .distinct().map(tag -> new TagAssociation().setTag(new TagUrn(tag.getName())))
+                .sorted(Comparator.comparing(t -> t.getTag().getName()))
+                .collect(Collectors.toList());
+
+        List<GlossaryTermAssociation> terms =  Stream.concat(
+                ProtobufExtensionUtil.extractTermAssociationsFromOptions(
+                                field.getFieldProto().getOptions().getAllFields(), context.getGraph().getRegistry()),
+                promotedTerms(field, context))
+                .distinct()
+                .sorted(Comparator.comparing(a -> a.getUrn().getNameEntity()))
+                .collect(Collectors.toList());
+
+        return context.getFirstFieldPath(field).map(path -> Pair.of(
+                new SchemaField()
+                        .setFieldPath(context.getFieldPath(path))
+                        .setNullable(!isPrimaryKey)
+                        .setIsPartOfKey(isPrimaryKey)
+                        .setDescription(field.comment())
+                        .setNativeDataType(field.nativeType())
+                        .setType(field.schemaFieldDataType())
+                        .setGlobalTags(new GlobalTags().setTags(new TagAssociationArray(tags)))
+                        .setGlossaryTerms(new GlossaryTerms()
+                                .setTerms(new GlossaryTermAssociationArray(terms))
+                                .setAuditStamp(context.getAuditStamp())),
+                context.calculateSortOrder(path, field))).stream();
+    }
+
+    /**
+     * Promote tags from nested message to field.
+     * @return tags
+     */
+    private Stream<TagProperties> promotedTags(ProtobufField field, VisitContext context) {
+        if (field.isMessage()) {
+            return context.getGraph().outgoingEdgesOf(field).stream().flatMap(e ->
+                    ProtobufExtensionUtil.extractTagPropertiesFromOptions(e.getEdgeTarget().messageProto()
+                            .getOptions().getAllFields(), context.getGraph().getRegistry())
+            ).distinct();
+        } else {
+            return Stream.of();
+        }
+    }
+
+    /**
+     * Promote terms from nested message to field.
+     * @return terms
+     */
+    private Stream<GlossaryTermAssociation> promotedTerms(ProtobufField field, VisitContext context) {
+        if (field.isMessage()) {
+            return context.getGraph().outgoingEdgesOf(field).stream().flatMap(e ->
+                    ProtobufExtensionUtil.extractTermAssociationsFromOptions(e.getEdgeTarget().messageProto()
+                            .getOptions().getAllFields(), context.getGraph().getRegistry())
+            ).distinct();
+        } else {
+            return Stream.of();
+        }
+    }
+
+}
diff --git a/metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/visitors/field/SchemaFieldVisitor.java b/metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/visitors/field/SchemaFieldVisitor.java
new file mode 100644
index 0000000000000..9a3333ea4d660
--- /dev/null
+++ b/metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/visitors/field/SchemaFieldVisitor.java
@@ -0,0 +1,25 @@
+package datahub.protobuf.visitors.field;
+
+import com.linkedin.schema.SchemaField;
+import com.linkedin.util.Pair;
+import datahub.protobuf.model.ProtobufField;
+import datahub.protobuf.visitors.ProtobufModelVisitor;
+import datahub.protobuf.visitors.VisitContext;
+
+import java.util.stream.Stream;
+
+public class SchemaFieldVisitor implements ProtobufModelVisitor<Pair<SchemaField, Double>> {
+
+    @Override
+    public Stream<Pair<SchemaField, Double>> visitField(ProtobufField field, VisitContext context) {
+        return context.getFirstFieldPath(field).map(path ->
+                Pair.of(
+                        new SchemaField()
+                                .setFieldPath(context.getFieldPath(path))
+                                .setNullable(true)
+                                .setDescription(field.comment())
+                                .setNativeDataType(field.nativeType())
+                                .setType(field.schemaFieldDataType()),
+                        context.calculateSortOrder(path, field))).stream();
+    }
+}
diff --git a/metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/visitors/tags/ProtobufExtensionTagVisitor.java b/metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/visitors/tags/ProtobufExtensionTagVisitor.java
new file mode 100644
index 0000000000000..07f2931a1af7b
--- /dev/null
+++ b/metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/visitors/tags/ProtobufExtensionTagVisitor.java
@@ -0,0 +1,40 @@
+package datahub.protobuf.visitors.tags;
+
+import com.linkedin.common.urn.TagUrn;
+import com.linkedin.data.template.RecordTemplate;
+import com.linkedin.events.metadata.ChangeType;
+import com.linkedin.tag.TagProperties;
+import datahub.protobuf.model.ProtobufField;
+import datahub.protobuf.visitors.ProtobufModelVisitor;
+import datahub.protobuf.visitors.ProtobufExtensionUtil;
+import datahub.protobuf.visitors.VisitContext;
+import datahub.event.MetadataChangeProposalWrapper;
+
+import java.util.stream.Stream;
+
+public class ProtobufExtensionTagVisitor implements ProtobufModelVisitor<MetadataChangeProposalWrapper<? extends RecordTemplate>> {
+    private static final String TAG_PROPERTIES_ASPECT = "tagProperties";
+
+    @Override
+    public Stream<MetadataChangeProposalWrapper<? extends RecordTemplate>> visitGraph(VisitContext context) {
+        return ProtobufExtensionUtil.extractTagPropertiesFromOptions(context.root().messageProto().getOptions()
+                        .getAllFields(), context.getGraph().getRegistry())
+                .map(ProtobufExtensionTagVisitor::wrapTagProperty);
+    }
+
+    @Override
+    public Stream<MetadataChangeProposalWrapper<? extends RecordTemplate>> visitField(ProtobufField field, VisitContext context) {
+        return ProtobufExtensionUtil.extractTagPropertiesFromOptions(field.getFieldProto().getOptions().getAllFields(),
+                        context.getGraph().getRegistry())
+                .map(ProtobufExtensionTagVisitor::wrapTagProperty);
+    }
+
+    private static MetadataChangeProposalWrapper<TagProperties> wrapTagProperty(TagProperties tagProperty) {
+        return new MetadataChangeProposalWrapper<>(
+                "tag",
+                new TagUrn(tagProperty.getName()).toString(),
+                ChangeType.UPSERT,
+                tagProperty,
+                TAG_PROPERTIES_ASPECT);
+    }
+}
diff --git a/metadata-integration/java/datahub-protobuf/src/main/resources/protobuf/meta.proto b/metadata-integration/java/datahub-protobuf/src/main/resources/protobuf/meta.proto
new file mode 100644
index 0000000000000..d3ffed747206c
--- /dev/null
+++ b/metadata-integration/java/datahub-protobuf/src/main/resources/protobuf/meta.proto
@@ -0,0 +1,38 @@
+syntax = "proto3";
+
+package meta;
+
+/*
+   This is assigned to metadata fields. It describes how the metadata field should be represented
+   in DataHub. This enum must be used in the `meta` package. Multiple can be used for the same
+   metadata annotation. This allows a single piece of information to be captured in DataHub
+   as a property, tag and/or term.
+
+   Tags can be strings, enums, or booleans
+   Terms can be strings or enums
+   Properties should be strings
+
+*/
+enum DataHubMetadataType {
+  PROPERTY = 0; // Datahub Custom Property
+  TAG      = 1; // Datahub Tag
+  TERM     = 2; // Datahub Term
+}
+
+/*
+   Example below: The following is not required for annotation processing. This is an example
+   of creating an annotation using an enum.
+ */
+
+enum MetaEnumExample {
+  UNKNOWN = 0;
+  ENTITY = 1;
+  EVENT = 2;
+}
+
+// Assuming Glossary Term defined from bootstrap example
+enum Classification {
+  HighlyConfidential = 0;
+  Confidential = 1;
+  Sensitive = 2;
+}
diff --git a/metadata-integration/java/datahub-protobuf/src/main/resources/protobuf/meta_field_options.proto b/metadata-integration/java/datahub-protobuf/src/main/resources/protobuf/meta_field_options.proto
new file mode 100644
index 0000000000000..c2a6683a623b5
--- /dev/null
+++ b/metadata-integration/java/datahub-protobuf/src/main/resources/protobuf/meta_field_options.proto
@@ -0,0 +1,30 @@
+syntax = "proto3";
+
+package meta.fld;
+
+import "google/protobuf/descriptor.proto";
+
+import "protobuf/meta.proto";
+
+
+extend google.protobuf.FieldOptions {
+    // Required: Mark option field with how to export to DataHub in one or more places.
+    repeated meta.DataHubMetadataType type = 6000;
+
+    /*
+       Examples below: The following is not required for annotation processing.
+     */
+
+    // Set true if the field is a primary key. This works for any boolean with `primary_key` in it.
+    bool is_primary_key = 6010;
+
+    // Extract classification field option as a Term, either works
+    string classification                   = 6001 [(meta.fld.type) = TERM];
+    meta.Classification classification_enum = 6002 [(meta.fld.type) = TERM];
+
+    // Expose this option as a tag on the field.
+    string product_type                    = 70004 [(meta.fld.type) = TAG];
+    bool product_type_bool                 = 70005 [(meta.fld.type) = TAG];
+    meta.MetaEnumExample product_type_enum = 70006 [(meta.fld.type) = TAG];
+}
+
diff --git a/metadata-integration/java/datahub-protobuf/src/main/resources/protobuf/meta_message_options.proto b/metadata-integration/java/datahub-protobuf/src/main/resources/protobuf/meta_message_options.proto
new file mode 100644
index 0000000000000..d3a55bf92e6f4
--- /dev/null
+++ b/metadata-integration/java/datahub-protobuf/src/main/resources/protobuf/meta_message_options.proto
@@ -0,0 +1,29 @@
+syntax = "proto3";
+
+package meta.msg;
+
+import "google/protobuf/descriptor.proto";
+
+import "protobuf/meta.proto";
+import "protobuf/meta_field_options.proto";
+
+
+extend google.protobuf.MessageOptions {
+    /*
+       Examples below: The following is not required for annotation processing.
+     */
+
+    // Place the classification term at the Message/Dataset level, either string or enum is supported
+    string classification                   = 4000 [(meta.fld.type) = TERM, (meta.fld.type) = PROPERTY];
+    meta.Classification classification_enum = 4001 [(meta.fld.type) = TERM, (meta.fld.type) = PROPERTY];
+
+    // Attach these Message/Dataset options as a tag and property.
+    string product            = 5001 [(meta.fld.type) = TAG, (meta.fld.type) = PROPERTY];
+    string project            = 5002 [(meta.fld.type) = TAG, (meta.fld.type) = PROPERTY];
+    string team               = 5003 [(meta.fld.type) = TAG, (meta.fld.type) = PROPERTY];
+
+    string domain             = 60003 [(meta.fld.type) = TAG, (meta.fld.type) = PROPERTY];
+    meta.MetaEnumExample type = 60004 [(meta.fld.type) = TAG, (meta.fld.type) = PROPERTY];
+    bool bool_feature         = 60005 [(meta.fld.type) = TAG];
+    string alert_channel      = 60007 [(meta.fld.type) = PROPERTY];
+}
\ No newline at end of file
diff --git a/metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/ProtobufDatasetTest.java b/metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/ProtobufDatasetTest.java
new file mode 100644
index 0000000000000..6239802e318e4
--- /dev/null
+++ b/metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/ProtobufDatasetTest.java
@@ -0,0 +1,382 @@
+package datahub.protobuf;
+
+import com.linkedin.common.GlobalTags;
+import com.linkedin.common.GlossaryTermAssociationArray;
+import com.linkedin.common.GlossaryTerms;
+import com.linkedin.common.InstitutionalMemory;
+import com.linkedin.common.InstitutionalMemoryMetadata;
+import com.linkedin.common.InstitutionalMemoryMetadataArray;
+import com.linkedin.common.Status;
+import com.linkedin.common.TagAssociationArray;
+import com.linkedin.common.url.Url;
+import com.linkedin.data.template.StringArray;
+import com.linkedin.schema.ArrayType;
+import com.linkedin.schema.BooleanType;
+import com.linkedin.schema.BytesType;
+import com.linkedin.schema.NumberType;
+import com.linkedin.schema.RecordType;
+import com.linkedin.schema.SchemaField;
+import com.linkedin.schema.SchemaFieldDataType;
+import com.linkedin.schema.SchemaMetadata;
+import com.linkedin.schema.StringType;
+import com.linkedin.schema.UnionType;
+import com.linkedin.common.urn.DataPlatformUrn;
+import com.linkedin.common.FabricType;
+import com.linkedin.util.Pair;
+import datahub.protobuf.model.ProtobufField;
+import datahub.protobuf.visitors.ProtobufModelVisitor;
+import datahub.protobuf.visitors.VisitContext;
+import org.junit.Test;
+
+import java.io.IOException;
+import java.util.stream.Stream;
+
+import static datahub.protobuf.TestFixtures.*;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+
+
+public class ProtobufDatasetTest {
+
+    @Test
+    public void noSchemaTest() throws IOException {
+        ProtobufDataset dataset = ProtobufDataset.builder()
+                .setDataPlatformUrn(new DataPlatformUrn("kafka"))
+                .setProtocIn(getTestProtoc("protobuf", "messageA"))
+                .setAuditStamp(TEST_AUDIT_STAMP)
+                .setFabricType(FabricType.DEV)
+                .build();
+
+        assertNotNull(dataset);
+        assertEquals(2, dataset.getAllMetadataChangeProposals().count());
+        assertEquals(6, dataset.getDatasetMCPs().size());
+        assertEquals(0, dataset.getVisitorMCPs().size());
+    }
+
+    @Test
+    public void platformSchemaTest() throws IOException {
+        assertEquals(getTestProtoSource("protobuf", "messageA"),
+                extractDocumentSchema(getTestProtobufDataset("protobuf", "messageA")));
+    }
+
+    @Test
+    public void messageA() throws IOException {
+        ProtobufDataset test = getTestProtobufDataset("protobuf", "messageA");
+
+        assertEquals("urn:li:dataset:(urn:li:dataPlatform:kafka,protobuf.MessageA,TEST)",
+                test.getDatasetUrn().toString());
+
+        SchemaMetadata testMetadata = test.getSchemaMetadata();
+
+        assertEquals(1, testMetadata.getVersion());
+        assertEquals(9, testMetadata.getFields().size());
+
+
+        assertEquals("platform.topic", extractCustomProperty(test.getDatasetMCPs().get(0), "kafka_topic"));
+
+        assertEquals(new InstitutionalMemory().setElements(new InstitutionalMemoryMetadataArray(
+                        new InstitutionalMemoryMetadata()
+                                .setDescription("Github Team")
+                                .setCreateStamp(TEST_AUDIT_STAMP)
+                                .setUrl(new Url("https://github.com/orgs/myOrg/teams/teama")),
+                        new InstitutionalMemoryMetadata()
+                                .setDescription("Slack Channel")
+                                .setCreateStamp(TEST_AUDIT_STAMP)
+                                .setUrl(new Url("https://slack.com/app_redirect?channel=test-slack&team=SLACK123")),
+                        new InstitutionalMemoryMetadata()
+                                .setCreateStamp(TEST_AUDIT_STAMP)
+                                .setDescription("MessageA Reference 1")
+                                .setUrl(new Url("https://some/link")),
+                        new InstitutionalMemoryMetadata()
+                                .setCreateStamp(TEST_AUDIT_STAMP)
+                                .setDescription("MessageA Reference 2")
+                                .setUrl(new Url("https://www.google.com/search?q=protobuf+messages")),
+                        new InstitutionalMemoryMetadata()
+                                .setCreateStamp(TEST_AUDIT_STAMP)
+                                .setDescription("MessageA Reference 3")
+                                .setUrl(new Url("https://github.com/apache/kafka")),
+                        new InstitutionalMemoryMetadata()
+                                .setCreateStamp(TEST_AUDIT_STAMP)
+                                .setDescription("MessageA.map_field Reference 1")
+                                .setUrl(new Url("https://developers.google.com/protocol-buffers/docs/proto3#maps")))).data(),
+                test.getDatasetMCPs().get(1).getAspect().data());
+
+        assertEquals(new Status().setRemoved(false).data(), test.getDatasetMCPs().get(test.getDatasetMCPs().size() - 1).getAspect().data());
+
+        assertEquals(new SchemaField()
+                        .setFieldPath("[version=2.0].[type=protobuf_MessageA].[type=bytes].sequence_id")
+                        .setType(new SchemaFieldDataType().setType(SchemaFieldDataType.Type.create(new BytesType())))
+                        .setNativeDataType("bytes")
+                        .setNullable(true)
+                        .setIsPartOfKey(false)
+                        .setDescription("Leading single line comment")
+                        .setGlobalTags(new GlobalTags().setTags(new TagAssociationArray()))
+                        .setGlossaryTerms(new GlossaryTerms().setTerms(new GlossaryTermAssociationArray()).setAuditStamp(TEST_AUDIT_STAMP)),
+                testMetadata.getFields().stream().filter(f -> f.getFieldPath()
+                        .equals("[version=2.0].[type=protobuf_MessageA].[type=bytes].sequence_id")).findFirst().orElseThrow());
+
+        assertEquals(new SchemaField()
+                        .setFieldPath("[version=2.0].[type=protobuf_MessageA].[type=int].position")
+                        .setType(new SchemaFieldDataType().setType(SchemaFieldDataType.Type.create(new NumberType())))
+                        .setNativeDataType("uint32")
+                        .setNullable(true)
+                        .setIsPartOfKey(false)
+                        .setDescription("Leading multiline comment\nSecond line of leading multiline comment")
+                        .setGlobalTags(new GlobalTags().setTags(new TagAssociationArray()))
+                        .setGlossaryTerms(new GlossaryTerms().setTerms(new GlossaryTermAssociationArray()).setAuditStamp(TEST_AUDIT_STAMP)),
+                testMetadata.getFields().stream().filter(f -> f.getFieldPath()
+                        .equals("[version=2.0].[type=protobuf_MessageA].[type=int].position")).findFirst().orElseThrow());
+
+        assertEquals(new SchemaField()
+                        .setFieldPath("[version=2.0].[type=protobuf_MessageA].[type=int].total")
+                        .setType(new SchemaFieldDataType().setType(SchemaFieldDataType.Type.create(new NumberType())))
+                        .setNativeDataType("uint32")
+                        .setNullable(true)
+                        .setIsPartOfKey(false)
+                        .setDescription("Detached comment")
+                        .setGlobalTags(new GlobalTags().setTags(new TagAssociationArray()))
+                        .setGlossaryTerms(new GlossaryTerms().setTerms(new GlossaryTermAssociationArray()).setAuditStamp(TEST_AUDIT_STAMP)),
+                testMetadata.getFields().stream().filter(f -> f.getFieldPath()
+                        .equals("[version=2.0].[type=protobuf_MessageA].[type=int].total")).findFirst().orElseThrow());
+
+        assertEquals(new SchemaField()
+                        .setFieldPath("[version=2.0].[type=protobuf_MessageA].[type=array].[type=long].repeated_num")
+                        .setType(new SchemaFieldDataType().setType(SchemaFieldDataType.Type.create(new ArrayType().setNestedType(new StringArray()))))
+                        .setNativeDataType("uint64")
+                        .setNullable(true)
+                        .setIsPartOfKey(false)
+                        .setDescription("Test repeated and trailing comment")
+                        .setGlobalTags(new GlobalTags().setTags(new TagAssociationArray()))
+                        .setGlossaryTerms(new GlossaryTerms().setTerms(new GlossaryTermAssociationArray()).setAuditStamp(TEST_AUDIT_STAMP)),
+                testMetadata.getFields().stream().filter(f -> f.getFieldPath()
+                        .equals("[version=2.0].[type=protobuf_MessageA].[type=array].[type=long].repeated_num")).findFirst().orElseThrow());
+
+        assertEquals(new SchemaField()
+                        .setFieldPath("[version=2.0].[type=protobuf_MessageA].[type=array].[type=string].repeated_str")
+                        .setType(new SchemaFieldDataType().setType(SchemaFieldDataType.Type.create(new ArrayType().setNestedType(new StringArray()))))
+                        .setNativeDataType("string")
+                        .setNullable(true)
+                        .setIsPartOfKey(false)
+                        .setDescription("")
+                        .setGlobalTags(new GlobalTags().setTags(new TagAssociationArray()))
+                        .setGlossaryTerms(new GlossaryTerms().setTerms(new GlossaryTermAssociationArray()).setAuditStamp(TEST_AUDIT_STAMP)),
+                testMetadata.getFields().stream().filter(f -> f.getFieldPath()
+                        .equals("[version=2.0].[type=protobuf_MessageA].[type=array].[type=string].repeated_str")).findFirst().orElseThrow());
+
+    }
+
+    @Test
+    public void messageB() throws IOException {
+        ProtobufDataset test = getTestProtobufDataset("protobuf", "messageB");
+
+        assertEquals("urn:li:dataset:(urn:li:dataPlatform:kafka,protobuf.MessageB,TEST)",
+                test.getDatasetUrn().toString());
+
+        SchemaMetadata testMetadata = test.getSchemaMetadata();
+
+        assertEquals(1, testMetadata.getVersion());
+        assertEquals(13, testMetadata.getFields().size());
+
+        assertEquals(new SchemaField()
+                        .setFieldPath("[version=2.0].[type=protobuf_MessageB].[type=long].id")
+                        .setNullable(true)
+                        .setIsPartOfKey(false)
+                        .setType(new SchemaFieldDataType().setType(SchemaFieldDataType.Type.create(new NumberType())))
+                        .setNativeDataType("google.protobuf.Int64Value")
+                        .setDescription("wrapped int64")
+                        .setGlobalTags(new GlobalTags().setTags(new TagAssociationArray()))
+                        .setGlossaryTerms(new GlossaryTerms().setTerms(new GlossaryTermAssociationArray()).setAuditStamp(test.getAuditStamp())),
+                testMetadata.getFields().stream().filter(f -> f.getFieldPath()
+                        .equals("[version=2.0].[type=protobuf_MessageB].[type=long].id")).findFirst().orElseThrow());
+
+        assertEquals(new SchemaField()
+                        .setFieldPath("[version=2.0].[type=protobuf_MessageB].[type=boolean].hot")
+                        .setNullable(true)
+                        .setIsPartOfKey(false)
+                        .setType(new SchemaFieldDataType().setType(SchemaFieldDataType.Type.create(new BooleanType())))
+                        .setNativeDataType("google.protobuf.BoolValue")
+                        .setDescription("Indicator")
+                        .setGlobalTags(new GlobalTags().setTags(new TagAssociationArray()))
+                        .setGlossaryTerms(new GlossaryTerms().setTerms(new GlossaryTermAssociationArray()).setAuditStamp(test.getAuditStamp())),
+                testMetadata.getFields().stream().filter(f -> f.getFieldPath()
+                        .equals("[version=2.0].[type=protobuf_MessageB].[type=boolean].hot")).findFirst().orElseThrow());
+
+
+        assertEquals(new SchemaField()
+                        .setNullable(true)
+                        .setIsPartOfKey(false)
+                        .setFieldPath("[version=2.0].[type=protobuf_MessageB].[type=string].value")
+                        .setType(new SchemaFieldDataType().setType(SchemaFieldDataType.Type.create(new StringType())))
+                        .setNativeDataType("string")
+                        .setDescription("message value")
+                        .setGlobalTags(new GlobalTags().setTags(new TagAssociationArray()))
+                        .setGlossaryTerms(new GlossaryTerms().setTerms(new GlossaryTermAssociationArray()).setAuditStamp(test.getAuditStamp())),
+                testMetadata.getFields().stream().filter(f -> f.getFieldPath()
+                        .equals("[version=2.0].[type=protobuf_MessageB].[type=string].value")).findFirst().orElseThrow());
+
+    }
+
+    @Test
+    public void messageC() throws IOException {
+        ProtobufDataset test = getTestProtobufDataset("protobuf", "messageC");
+
+
+        assertEquals("urn:li:dataset:(urn:li:dataPlatform:kafka,protobuf.MessageC,TEST)",
+                test.getDatasetUrn().toString());
+
+        SchemaMetadata testMetadata = test.getSchemaMetadata();
+
+        assertEquals(1, testMetadata.getVersion());
+        assertEquals(4, testMetadata.getFields().size());
+
+        assertEquals(new SchemaField()
+                        .setFieldPath("[version=2.0].[type=protobuf_MessageC].[type=union].one_of_field")
+                        .setNullable(true)
+                        .setIsPartOfKey(false)
+                        .setType(new SchemaFieldDataType().setType(SchemaFieldDataType.Type.create(new UnionType())))
+                        .setNativeDataType("oneof")
+                        .setDescription("one of field comment")
+                        .setGlobalTags(new GlobalTags().setTags(new TagAssociationArray()))
+                        .setGlossaryTerms(new GlossaryTerms().setTerms(new GlossaryTermAssociationArray()).setAuditStamp(test.getAuditStamp())),
+                testMetadata.getFields().stream().filter(f -> f.getFieldPath()
+                        .equals("[version=2.0].[type=protobuf_MessageC].[type=union].one_of_field")).findFirst().orElseThrow());
+
+        assertEquals(new SchemaField()
+                        .setNullable(true)
+                        .setIsPartOfKey(false)
+                        .setFieldPath("[version=2.0].[type=protobuf_MessageC].[type=union].one_of_field.[type=string].one_of_string")
+                        .setType(new SchemaFieldDataType().setType(SchemaFieldDataType.Type.create(new StringType())))
+                        .setNativeDataType("string")
+                        .setDescription("one of string comment")
+                        .setGlobalTags(new GlobalTags().setTags(new TagAssociationArray()))
+                        .setGlossaryTerms(new GlossaryTerms().setTerms(new GlossaryTermAssociationArray()).setAuditStamp(test.getAuditStamp())),
+                testMetadata.getFields().stream().filter(f -> f.getFieldPath()
+                        .equals("[version=2.0].[type=protobuf_MessageC].[type=union].one_of_field.[type=string].one_of_string")).findFirst().orElseThrow());
+
+        assertEquals(new SchemaField()
+                        .setNullable(true)
+                        .setIsPartOfKey(false)
+                        .setFieldPath("[version=2.0].[type=protobuf_MessageC].[type=union].one_of_field.[type=int].one_of_int")
+                        .setType(new SchemaFieldDataType().setType(SchemaFieldDataType.Type.create(new NumberType())))
+                        .setNativeDataType("int32")
+                        .setDescription("one of int comment")
+                        .setGlobalTags(new GlobalTags().setTags(new TagAssociationArray()))
+                        .setGlossaryTerms(new GlossaryTerms().setTerms(new GlossaryTermAssociationArray()).setAuditStamp(test.getAuditStamp())),
+                testMetadata.getFields().stream().filter(f -> f.getFieldPath()
+                        .equals("[version=2.0].[type=protobuf_MessageC].[type=union].one_of_field.[type=int].one_of_int")).findFirst().orElseThrow());
+    }
+
+    @Test
+    @SuppressWarnings("LineLength")
+    public void messageC2NestedOneOf() throws IOException {
+        ProtobufDataset test = getTestProtobufDataset("protobuf", "messageC2");
+
+
+        assertEquals("urn:li:dataset:(urn:li:dataPlatform:kafka,protobuf.MessageC1,TEST)",
+                test.getDatasetUrn().toString());
+
+        SchemaMetadata testMetadata = test.getSchemaMetadata();
+
+        assertEquals(1, testMetadata.getVersion());
+        assertEquals(6, testMetadata.getFields().size());
+
+        assertEquals(new SchemaField()
+                        .setFieldPath("[version=2.0].[type=protobuf_MessageC1].[type=protobuf_MessageC2].messageList")
+                        .setNullable(true)
+                        .setIsPartOfKey(false)
+                        .setType(new SchemaFieldDataType().setType(SchemaFieldDataType.Type.create(new RecordType())))
+                        .setNativeDataType("protobuf.MessageC2")
+                        .setDescription("")
+                        .setGlobalTags(new GlobalTags().setTags(new TagAssociationArray()))
+                        .setGlossaryTerms(new GlossaryTerms().setTerms(new GlossaryTermAssociationArray()).setAuditStamp(test.getAuditStamp())),
+                testMetadata.getFields().stream().filter(f -> f.getFieldPath()
+                        .equals("[version=2.0].[type=protobuf_MessageC1].[type=protobuf_MessageC2].messageList")).findFirst().orElseThrow());
+
+        assertEquals(new SchemaField()
+                        .setFieldPath("[version=2.0].[type=protobuf_MessageC1].[type=protobuf_MessageC2].messageList.[type=array].[type=protobuf_MessageC3].list")
+                        .setNullable(true)
+                        .setIsPartOfKey(false)
+                        .setType(new SchemaFieldDataType().setType(SchemaFieldDataType.Type.create(new ArrayType().setNestedType(new StringArray()))))
+                        .setNativeDataType("protobuf.MessageC3")
+                        .setDescription("")
+                        .setGlobalTags(new GlobalTags().setTags(new TagAssociationArray()))
+                        .setGlossaryTerms(new GlossaryTerms().setTerms(new GlossaryTermAssociationArray()).setAuditStamp(test.getAuditStamp())),
+                testMetadata.getFields().stream().filter(f -> f.getFieldPath()
+                        .equals("[version=2.0].[type=protobuf_MessageC1].[type=protobuf_MessageC2].messageList.[type=array].[type=protobuf_MessageC3].list")).findFirst().orElseThrow());
+
+        assertEquals(new SchemaField()
+                        .setFieldPath("[version=2.0].[type=protobuf_MessageC1].[type=protobuf_MessageC2].messageList.[type=array].[type=protobuf_MessageC3].list.[type=string].normal")
+                        .setNullable(true)
+                        .setIsPartOfKey(false)
+                        .setType(new SchemaFieldDataType().setType(SchemaFieldDataType.Type.create(new StringType())))
+                        .setNativeDataType("string")
+                        .setDescription("")
+                        .setGlobalTags(new GlobalTags().setTags(new TagAssociationArray()))
+                        .setGlossaryTerms(new GlossaryTerms().setTerms(new GlossaryTermAssociationArray()).setAuditStamp(test.getAuditStamp())),
+                testMetadata.getFields().stream().filter(f -> f.getFieldPath()
+                        .equals("[version=2.0].[type=protobuf_MessageC1].[type=protobuf_MessageC2].messageList.[type=array].[type=protobuf_MessageC3].list.[type=string].normal")).findFirst().orElseThrow());
+
+        assertEquals(new SchemaField()
+                        .setFieldPath("[version=2.0].[type=protobuf_MessageC1].[type=protobuf_MessageC2].messageList.[type=array].[type=protobuf_MessageC3].list.[type=union].one_of_field")
+                        .setNullable(true)
+                        .setIsPartOfKey(false)
+                        .setType(new SchemaFieldDataType().setType(SchemaFieldDataType.Type.create(new UnionType())))
+                        .setNativeDataType("oneof")
+                        .setDescription("one of field comment")
+                        .setGlobalTags(new GlobalTags().setTags(new TagAssociationArray()))
+                        .setGlossaryTerms(new GlossaryTerms().setTerms(new GlossaryTermAssociationArray()).setAuditStamp(test.getAuditStamp())),
+                testMetadata.getFields().stream().filter(f -> f.getFieldPath()
+                        .equals("[version=2.0].[type=protobuf_MessageC1].[type=protobuf_MessageC2].messageList.[type=array].[type=protobuf_MessageC3].list.[type=union].one_of_field")).findFirst().orElseThrow());
+
+        assertEquals(new SchemaField()
+                        .setNullable(true)
+                        .setIsPartOfKey(false)
+                        .setFieldPath("[version=2.0].[type=protobuf_MessageC1].[type=protobuf_MessageC2].messageList.[type=array].[type=protobuf_MessageC3].list.[type=union].one_of_field.[type=string].one_of_string")
+                        .setType(new SchemaFieldDataType().setType(SchemaFieldDataType.Type.create(new StringType())))
+                        .setNativeDataType("string")
+                        .setDescription("one of string comment")
+                        .setGlobalTags(new GlobalTags().setTags(new TagAssociationArray()))
+                        .setGlossaryTerms(new GlossaryTerms().setTerms(new GlossaryTermAssociationArray()).setAuditStamp(test.getAuditStamp())),
+                testMetadata.getFields().stream().filter(f -> f.getFieldPath()
+                        .equals("[version=2.0].[type=protobuf_MessageC1].[type=protobuf_MessageC2].messageList.[type=array].[type=protobuf_MessageC3].list.[type=union].one_of_field.[type=string].one_of_string")).findFirst().orElseThrow());
+
+        assertEquals(new SchemaField()
+                        .setNullable(true)
+                        .setIsPartOfKey(false)
+                        .setFieldPath("[version=2.0].[type=protobuf_MessageC1].[type=protobuf_MessageC2].messageList.[type=array].[type=protobuf_MessageC3].list.[type=union].one_of_field.[type=int].one_of_int")
+                        .setType(new SchemaFieldDataType().setType(SchemaFieldDataType.Type.create(new NumberType())))
+                        .setNativeDataType("int32")
+                        .setDescription("one of int comment")
+                        .setGlobalTags(new GlobalTags().setTags(new TagAssociationArray()))
+                        .setGlossaryTerms(new GlossaryTerms().setTerms(new GlossaryTermAssociationArray()).setAuditStamp(test.getAuditStamp())),
+                testMetadata.getFields().stream().filter(f -> f.getFieldPath()
+                        .equals("[version=2.0].[type=protobuf_MessageC1].[type=protobuf_MessageC2].messageList.[type=array].[type=protobuf_MessageC3].list.[type=union].one_of_field.[type=int].one_of_int")).findFirst().orElseThrow());
+    }
+
+    @Test
+    public void customFieldVisitors() throws IOException {
+        ProtobufDataset test = getTestProtobufDataset("protobuf", "messageA");
+
+        test.setFieldVisitor(new ProtobufModelVisitor<Pair<SchemaField, Double>>() {
+            @Override
+            public Stream<Pair<SchemaField, Double>> visitField(ProtobufField field, VisitContext context) {
+                if (field.fullName().equals("protobuf.MessageA.sequence_id")) {
+                    return Stream.of(Pair.of(
+                            new SchemaField()
+                                    .setDescription("my comment")
+                                    .setNativeDataType("my type")
+                                    .setType(new SchemaFieldDataType().setType(SchemaFieldDataType.Type.create(new BytesType()))),
+                            0d));
+                } else {
+                    return Stream.of();
+                }
+            }
+        });
+        assertEquals(1, test.getSchemaMetadata().getFields().size());
+        assertEquals(new SchemaField()
+                        .setDescription("my comment")
+                        .setNativeDataType("my type")
+                        .setType(new SchemaFieldDataType().setType(SchemaFieldDataType.Type.create(new BytesType()))),
+                test.getSchemaMetadata().getFields().get(0));
+    }
+}
diff --git a/metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/ProtobufUtilsTest.java b/metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/ProtobufUtilsTest.java
new file mode 100644
index 0000000000000..a967157b500f8
--- /dev/null
+++ b/metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/ProtobufUtilsTest.java
@@ -0,0 +1,37 @@
+package datahub.protobuf;
+
+import com.google.protobuf.DescriptorProtos;
+import com.google.protobuf.ExtensionRegistry;
+import datahub.protobuf.model.ProtobufGraph;
+import org.junit.Test;
+
+import java.io.IOException;
+
+import static datahub.protobuf.TestFixtures.getTestProtobufFileSet;
+import static datahub.protobuf.TestFixtures.getTestProtoc;
+import static org.junit.jupiter.api.Assertions.*;
+
+
+public class ProtobufUtilsTest {
+
+    @Test
+    public void registryTest() throws IOException, IllegalArgumentException {
+        byte[] protocBytes = getTestProtoc("extended_protobuf", "messageA").readAllBytes();
+        DescriptorProtos.FileDescriptorSet fileSet = getTestProtobufFileSet("extended_protobuf", "messageA");
+        ExtensionRegistry registry = ProtobufUtils.buildRegistry(fileSet);
+        DescriptorProtos.FileDescriptorSet fileSetWithRegistry = DescriptorProtos.FileDescriptorSet.parseFrom(protocBytes, registry);
+
+        assertNotEquals(fileSet, fileSetWithRegistry);
+
+        /*
+         *
+         * Without the ExtensionRegistry we get field numbers instead of the names.
+         */
+        ProtobufGraph graph = new ProtobufGraph(fileSet, null);
+        assertEquals("[meta.msg.classification_enum]: HighlyConfidential\n"
+                + "[meta.msg.team]: \"TeamB\"\n"
+                + "[meta.msg.type]: ENTITY\n"
+                + "[meta.msg.bool_feature]: true\n"
+                + "[meta.msg.alert_channel]: \"#alerts\"\n", graph.root().messageProto().getOptions().toString());
+    }
+}
diff --git a/metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/TestFixtures.java b/metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/TestFixtures.java
new file mode 100644
index 0000000000000..b485ceb87d586
--- /dev/null
+++ b/metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/TestFixtures.java
@@ -0,0 +1,78 @@
+package datahub.protobuf;
+
+import com.google.protobuf.DescriptorProtos;
+import com.linkedin.common.AuditStamp;
+import com.linkedin.common.FabricType;
+import com.linkedin.common.urn.CorpuserUrn;
+import com.linkedin.common.urn.DataPlatformUrn;
+import com.linkedin.common.urn.DatasetUrn;
+import com.linkedin.data.DataMap;
+import com.linkedin.data.template.RecordTemplate;
+import datahub.event.MetadataChangeProposalWrapper;
+import datahub.protobuf.model.ProtobufGraph;
+import datahub.protobuf.visitors.VisitContext;
+
+import java.io.IOException;
+import java.io.InputStream;
+import java.nio.charset.StandardCharsets;
+import java.util.Objects;
+
+public class TestFixtures {
+    private TestFixtures() { }
+
+    public static final DataPlatformUrn TEST_DATA_PLATFORM = new DataPlatformUrn("kafka");
+    public static final AuditStamp TEST_AUDIT_STAMP = new AuditStamp()
+            .setTime(System.currentTimeMillis())
+            .setActor(new CorpuserUrn("datahub"));
+
+    public static InputStream getTestProtoc(String protoPackage, String filename) {
+        return Objects.requireNonNull(TestFixtures.class.getClassLoader()
+                .getResourceAsStream(String.format("%s/%s.protoc", protoPackage, filename)));
+    }
+
+    public static String getTestProtoSource(String protoPackage, String filename) throws IOException {
+        return new String(Objects.requireNonNull(TestFixtures.class.getClassLoader()
+                .getResourceAsStream(String.format("%s/%s.proto", protoPackage, filename))).readAllBytes(),
+                StandardCharsets.UTF_8);
+    }
+
+    public static ProtobufDataset getTestProtobufDataset(String protoPackage, String filename) throws IOException {
+        return ProtobufDataset.builder()
+                .setDataPlatformUrn(TEST_DATA_PLATFORM)
+                .setSchema(getTestProtoSource(protoPackage, filename))
+                .setProtocIn(getTestProtoc(protoPackage, filename))
+                .setAuditStamp(TEST_AUDIT_STAMP)
+                .setFabricType(FabricType.TEST)
+                .setGithubOrganization("myOrg")
+                .setSlackTeamId("SLACK123")
+                .build();
+    }
+
+    public static DescriptorProtos.FileDescriptorSet getTestProtobufFileSet(String protoPackage, String filename) throws IOException {
+        return DescriptorProtos.FileDescriptorSet
+                .parseFrom(getTestProtoc(protoPackage, filename).readAllBytes());
+    }
+
+    public static VisitContext.VisitContextBuilder getVisitContextBuilder(String message) {
+        return VisitContext.builder()
+                .datasetUrn(new DatasetUrn(TEST_DATA_PLATFORM, message, FabricType.TEST))
+                .auditStamp(TEST_AUDIT_STAMP);
+    }
+
+    public static ProtobufGraph getTestProtobufGraph(String protoPackage, String filename) throws IOException {
+        return new ProtobufGraph(getTestProtobufFileSet(protoPackage, filename));
+    }
+
+    public static Object extractAspect(MetadataChangeProposalWrapper<? extends RecordTemplate> mcp, String aspect) {
+        return mcp.getAspect().data().get(aspect);
+    }
+
+    public static Object extractCustomProperty(MetadataChangeProposalWrapper<? extends RecordTemplate> mcp, String key) {
+        return ((DataMap) extractAspect(mcp, "customProperties")).get(key);
+    }
+
+    public static String extractDocumentSchema(ProtobufDataset protobufDataset) {
+        return String.valueOf(((DataMap) ((DataMap) protobufDataset.getSchemaMetadata().getPlatformSchema().data())
+                .get("com.linkedin.schema.KafkaSchema")).get("documentSchema"));
+    }
+}
diff --git a/metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/model/ProtobufEnumTest.java b/metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/model/ProtobufEnumTest.java
new file mode 100644
index 0000000000000..3696f5795e1f9
--- /dev/null
+++ b/metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/model/ProtobufEnumTest.java
@@ -0,0 +1,80 @@
+package datahub.protobuf.model;
+
+import com.google.protobuf.DescriptorProtos.DescriptorProto;
+import com.google.protobuf.DescriptorProtos.EnumDescriptorProto;
+import com.google.protobuf.DescriptorProtos.FileDescriptorProto;
+import com.linkedin.schema.EnumType;
+import com.linkedin.schema.SchemaFieldDataType;
+import org.junit.Test;
+
+import java.util.List;
+import java.util.Set;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+
+public class ProtobufEnumTest {
+
+    @Test
+    public void enumTest() {
+        EnumDescriptorProto expectedEnum =  EnumDescriptorProto.newBuilder()
+                .setName("enum1")
+                .build();
+        DescriptorProto expectedMessage = DescriptorProto.newBuilder().build();
+        FileDescriptorProto expectedFile = FileDescriptorProto.newBuilder()
+                .addMessageType(expectedMessage)
+                .setPackage("protobuf")
+                .addEnumType(expectedEnum)
+                .build();
+
+        ProtobufEnum test = ProtobufEnum.enumBuilder()
+                .enumProto(expectedEnum)
+                .messageProto(expectedMessage)
+                .fileProto(expectedFile)
+                .build();
+
+        assertEquals("enum1", test.name());
+        assertEquals("protobuf.enum1", test.fullName());
+        assertEquals("[type=enum]", test.fieldPathType());
+        assertEquals("enum", test.nativeType());
+        assertEquals(expectedMessage, test.messageProto());
+        assertEquals(expectedFile, test.fileProto());
+        assertEquals(new SchemaFieldDataType().setType(SchemaFieldDataType.Type.create(new EnumType())), test.schemaFieldDataType());
+        assertEquals("ProtobufEnum[protobuf.enum1]", test.toString());
+        assertEquals("", test.comment());
+    }
+
+    @Test
+    public void enumEqualityTest() {
+        EnumDescriptorProto enum1 =  EnumDescriptorProto.newBuilder().setName("enum1").build();
+        EnumDescriptorProto enum2 =  EnumDescriptorProto.newBuilder().setName("enum2").build();
+        EnumDescriptorProto enum1Dup =  EnumDescriptorProto.newBuilder().setName("enum1").build();
+
+        DescriptorProto expectedMessage = DescriptorProto.newBuilder().build();
+        FileDescriptorProto expectedFile = FileDescriptorProto.newBuilder()
+                .addMessageType(expectedMessage)
+                .setPackage("protobuf")
+                .addAllEnumType(List.of(enum1, enum2, enum1Dup))
+                .build();
+
+        ProtobufEnum test1 = ProtobufEnum.enumBuilder().enumProto(enum1)
+                .messageProto(expectedMessage)
+                .fileProto(expectedFile)
+                .build();
+        ProtobufEnum test2 = ProtobufEnum.enumBuilder().enumProto(enum2)
+                .messageProto(expectedMessage)
+                .fileProto(expectedFile)
+                .build();
+        ProtobufEnum test1Dup = ProtobufEnum.enumBuilder().enumProto(enum1Dup)
+                .messageProto(expectedMessage)
+                .fileProto(expectedFile)
+                .build();
+
+        assertEquals(test1, test1Dup);
+        assertNotEquals(test1, test2);
+        assertEquals(Set.of(test1, test2), Stream.of(test1, test2, test1Dup).collect(Collectors.toSet()));
+    }
+
+}
diff --git a/metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/model/ProtobufFieldTest.java b/metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/model/ProtobufFieldTest.java
new file mode 100644
index 0000000000000..7ec0456c9c740
--- /dev/null
+++ b/metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/model/ProtobufFieldTest.java
@@ -0,0 +1,230 @@
+package datahub.protobuf.model;
+
+import com.google.protobuf.DescriptorProtos.DescriptorProto;
+import com.google.protobuf.DescriptorProtos.FieldDescriptorProto;
+import com.google.protobuf.DescriptorProtos.FileDescriptorProto;
+import com.linkedin.data.template.StringArray;
+import com.linkedin.schema.ArrayType;
+import com.linkedin.schema.BooleanType;
+import com.linkedin.schema.BytesType;
+import com.linkedin.schema.EnumType;
+import com.linkedin.schema.FixedType;
+import com.linkedin.schema.NumberType;
+import com.linkedin.schema.RecordType;
+import com.linkedin.schema.SchemaFieldDataType;
+import com.linkedin.schema.StringType;
+import org.junit.Test;
+
+import java.util.Arrays;
+import java.util.Set;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+
+public class ProtobufFieldTest {
+    private static final DescriptorProto EXPECTED_MESSAGE_PROTO = DescriptorProto.newBuilder()
+            .setName("message1")
+            .build();
+    private static final FileDescriptorProto EXPECTED_FILE_PROTO = FileDescriptorProto.newBuilder()
+            .addMessageType(EXPECTED_MESSAGE_PROTO)
+            .setPackage("protobuf")
+            .build();
+    private static final ProtobufMessage EXPECTED_MESSAGE = ProtobufMessage.builder()
+            .messageProto(EXPECTED_MESSAGE_PROTO)
+            .fileProto(EXPECTED_FILE_PROTO)
+            .build();
+
+
+    @Test
+    public void fieldTest() {
+        FieldDescriptorProto expectedField =  FieldDescriptorProto.newBuilder()
+                .setName("field1")
+                .setNumber(1)
+                .setType(FieldDescriptorProto.Type.TYPE_BYTES)
+                .build();
+        DescriptorProto expectedMessage1 = DescriptorProto.newBuilder()
+                .setName("message1")
+                .addField(expectedField)
+                .build();
+        FileDescriptorProto expectedFile = FileDescriptorProto.newBuilder()
+                .addMessageType(expectedMessage1)
+                .setPackage("protobuf")
+                .build();
+        ProtobufMessage expectedMessage = ProtobufMessage.builder()
+                .messageProto(expectedMessage1)
+                .fileProto(expectedFile)
+                .build();
+
+        ProtobufField test = ProtobufField.builder()
+                .fieldProto(expectedField)
+                .protobufMessage(expectedMessage)
+                .build();
+
+        assertEquals("field1", test.name());
+        assertEquals("protobuf.message1.field1", test.fullName());
+        assertEquals("[type=bytes]", test.fieldPathType());
+        assertEquals("protobuf.message1", test.parentMessageName());
+        assertEquals(expectedMessage1, test.messageProto());
+        assertEquals(expectedFile, test.fileProto());
+        assertNull(test.oneOfProto());
+        assertEquals("bytes", test.nativeType());
+        assertFalse(test.isMessage());
+        assertEquals(1, test.sortWeight());
+        assertEquals(new SchemaFieldDataType().setType(SchemaFieldDataType.Type.create(new BytesType())), test.schemaFieldDataType());
+        assertEquals("ProtobufField[protobuf.message1.field1]", test.toString());
+    }
+
+    @Test
+    public void fieldPathTypeTest() {
+        Arrays.stream(FieldDescriptorProto.Type.values()).forEach(type -> {
+            final  FieldDescriptorProto expectedField;
+            if (type == FieldDescriptorProto.Type.TYPE_MESSAGE) {
+                expectedField = FieldDescriptorProto.newBuilder()
+                        .setName("field1")
+                        .setNumber(1)
+                        .setTypeName(EXPECTED_MESSAGE.fullName())
+                        .setType(type)
+                        .build();
+            } else {
+                expectedField = FieldDescriptorProto.newBuilder()
+                        .setName("field1")
+                        .setNumber(1)
+                        .setType(type)
+                        .build();
+            }
+
+            ProtobufField test = ProtobufField.builder()
+                    .fieldProto(expectedField)
+                    .protobufMessage(EXPECTED_MESSAGE)
+                    .build();
+
+            if (type.equals(FieldDescriptorProto.Type.TYPE_MESSAGE)) {
+                assertEquals("[type=protobuf_message1]", test.fieldPathType());
+            } else if (type.name().endsWith("64")) {
+                assertEquals("[type=long]", test.fieldPathType());
+            } else if (type.name().endsWith("32")) {
+                assertEquals("[type=int]", test.fieldPathType());
+            } else if (type.name().endsWith("BOOL")) {
+                assertEquals("[type=boolean]", test.fieldPathType());
+            } else {
+                assertEquals(String.format("[type=%s]", type.name().split("_")[1].toLowerCase()), test.fieldPathType());
+            }
+        });
+    }
+
+    @Test
+    public void fieldPathTypeArrayTest() {
+        Arrays.stream(FieldDescriptorProto.Type.values()).forEach(type -> {
+            final  FieldDescriptorProto expectedField;
+
+            if (type == FieldDescriptorProto.Type.TYPE_MESSAGE) {
+                expectedField = FieldDescriptorProto.newBuilder()
+                        .setName("field1")
+                        .setNumber(1)
+                        .setTypeName(EXPECTED_MESSAGE.fullName())
+                        .setType(type)
+                        .setLabel(FieldDescriptorProto.Label.LABEL_REPEATED)
+                        .build();
+            } else {
+                expectedField = FieldDescriptorProto.newBuilder()
+                        .setName("field1")
+                        .setNumber(1)
+                        .setType(type)
+                        .setLabel(FieldDescriptorProto.Label.LABEL_REPEATED)
+                        .build();
+            }
+
+            ProtobufField test = ProtobufField.builder()
+                    .fieldProto(expectedField)
+                    .protobufMessage(EXPECTED_MESSAGE)
+                    .build();
+
+            if (type.equals(FieldDescriptorProto.Type.TYPE_MESSAGE)) {
+                assertEquals("[type=array].[type=protobuf_message1]", test.fieldPathType());
+            } else if (type.name().endsWith("64")) {
+                assertEquals("[type=array].[type=long]", test.fieldPathType());
+            } else if (type.name().endsWith("32")) {
+                assertEquals("[type=array].[type=int]", test.fieldPathType());
+            } else if (type.name().endsWith("BOOL")) {
+                assertEquals("[type=array].[type=boolean]", test.fieldPathType());
+            } else {
+                assertEquals(String.format("[type=array].[type=%s]", type.name().split("_")[1].toLowerCase()), test.fieldPathType());
+            }
+        });
+    }
+
+    @Test
+    public void schemaFieldTypeTest() {
+        Arrays.stream(FieldDescriptorProto.Type.values()).forEach(type -> {
+            final  FieldDescriptorProto expectedField;
+            if (type == FieldDescriptorProto.Type.TYPE_MESSAGE) {
+                expectedField = FieldDescriptorProto.newBuilder()
+                        .setName("field1")
+                        .setNumber(1)
+                        .setTypeName(EXPECTED_MESSAGE.fullName())
+                        .setType(type)
+                        .build();
+            } else {
+                expectedField = FieldDescriptorProto.newBuilder()
+                        .setName("field1")
+                        .setNumber(1)
+                        .setType(type)
+                        .build();
+            }
+
+            ProtobufField test = ProtobufField.builder()
+                    .fieldProto(expectedField)
+                    .protobufMessage(EXPECTED_MESSAGE)
+                    .build();
+
+            if (Set.of("TYPE_MESSAGE", "TYPE_GROUP").contains(type.name())) {
+                assertEquals(new SchemaFieldDataType().setType(SchemaFieldDataType.Type.create(new RecordType())), test.schemaFieldDataType());
+            } else if (type.name().contains("FIXED")) {
+                assertEquals(new SchemaFieldDataType().setType(SchemaFieldDataType.Type.create(new FixedType())), test.schemaFieldDataType());
+            } else if (type.name().endsWith("64") || type.name().endsWith("32") || Set.of("TYPE_DOUBLE", "TYPE_FLOAT").contains(type.name())) {
+                assertEquals(new SchemaFieldDataType().setType(SchemaFieldDataType.Type.create(new NumberType())), test.schemaFieldDataType());
+            } else if (type.name().endsWith("BOOL")) {
+                assertEquals(new SchemaFieldDataType().setType(SchemaFieldDataType.Type.create(new BooleanType())), test.schemaFieldDataType());
+            } else if (type.name().endsWith("STRING")) {
+                assertEquals(new SchemaFieldDataType().setType(SchemaFieldDataType.Type.create(new StringType())), test.schemaFieldDataType());
+            } else if (type.name().endsWith("ENUM")) {
+                assertEquals(new SchemaFieldDataType().setType(SchemaFieldDataType.Type.create(new EnumType())), test.schemaFieldDataType());
+            } else if (type.name().endsWith("BYTES")) {
+                assertEquals(new SchemaFieldDataType().setType(SchemaFieldDataType.Type.create(new BytesType())), test.schemaFieldDataType());
+            } else {
+                fail(String.format("Add test case for %s", type));
+            }
+        });
+    }
+
+    @Test
+    public void schemaFieldTypeArrayTest() {
+        Arrays.stream(FieldDescriptorProto.Type.values()).forEach(type -> {
+            final  FieldDescriptorProto expectedField;
+            if (type == FieldDescriptorProto.Type.TYPE_MESSAGE) {
+                expectedField = FieldDescriptorProto.newBuilder()
+                        .setName("field1")
+                        .setNumber(1)
+                        .setTypeName(EXPECTED_MESSAGE.fullName())
+                        .setType(type)
+                        .setLabel(FieldDescriptorProto.Label.LABEL_REPEATED)
+                        .build();
+            } else {
+                expectedField = FieldDescriptorProto.newBuilder()
+                        .setName("field1")
+                        .setNumber(1)
+                        .setType(type)
+                        .setLabel(FieldDescriptorProto.Label.LABEL_REPEATED)
+                        .build();
+            }
+
+            ProtobufField test = ProtobufField.builder()
+                    .fieldProto(expectedField)
+                    .protobufMessage(EXPECTED_MESSAGE)
+                    .build();
+
+            assertEquals(new SchemaFieldDataType().setType(SchemaFieldDataType.Type.create(new ArrayType()
+                    .setNestedType(new StringArray()))), test.schemaFieldDataType());
+        });
+    }
+}
diff --git a/metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/model/ProtobufGraphTest.java b/metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/model/ProtobufGraphTest.java
new file mode 100644
index 0000000000000..cfefaf1e04635
--- /dev/null
+++ b/metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/model/ProtobufGraphTest.java
@@ -0,0 +1,71 @@
+package datahub.protobuf.model;
+
+import com.google.protobuf.DescriptorProtos.FileDescriptorSet;
+import org.junit.Test;
+
+import java.io.IOException;
+import java.util.HashSet;
+
+import static datahub.protobuf.TestFixtures.getTestProtobufFileSet;
+import static datahub.protobuf.TestFixtures.getTestProtobufGraph;
+import static org.junit.jupiter.api.Assertions.*;
+
+public class ProtobufGraphTest {
+
+    @Test
+    public void autodetectRootMessageTest() throws IOException {
+        FileDescriptorSet fileset = getTestProtobufFileSet("protobuf", "messageB");
+        ProtobufGraph test = getTestProtobufGraph("protobuf", "messageB");
+
+        assertEquals("MessageB", test.autodetectRootMessage(
+                fileset.getFile(2)).messageProto().getName());
+
+        assertEquals("MessageA", test.autodetectRootMessage(
+                fileset.getFile(1)).messageProto().getName());
+    }
+
+    @Test
+    public void autodetectRootMessageFailureTest() throws IOException {
+        FileDescriptorSet empty = getTestProtobufFileSet("protobuf", "messageEmpty");
+        assertThrows(IllegalArgumentException.class, () -> new ProtobufGraph(empty));
+    }
+
+    @Test
+    public void findMessageTest() throws IOException {
+        FileDescriptorSet fileset = getTestProtobufFileSet("protobuf", "messageB");
+        ProtobufGraph test = getTestProtobufGraph("protobuf", "messageB");
+
+        assertEquals("MessageA",
+                test.findMessage("protobuf.MessageA").messageProto().getName());
+        assertEquals("MessageB",
+                test.findMessage("protobuf.MessageB").messageProto().getName());
+
+        assertThrows(IllegalArgumentException.class, () -> test.findMessage("not found"));
+        assertThrows(IllegalArgumentException.class, () -> new ProtobufGraph(fileset, "not found"));
+        assertEquals(test, new ProtobufGraph(fileset, "protobuf.MessageB"));
+    }
+
+    @Test
+    public void commentTest() throws IOException {
+        ProtobufGraph test = getTestProtobufGraph("protobuf", "messageC");
+        assertEquals("Test for one of", test.getComment());
+    }
+
+    @Test
+    public void equalityHashCodeTest() throws IOException {
+        ProtobufGraph testA = getTestProtobufGraph("protobuf", "messageA");
+        ProtobufGraph testB = getTestProtobufGraph("protobuf", "messageB");
+        FileDescriptorSet filesetB = getTestProtobufFileSet("protobuf", "messageB");
+
+        assertEquals(testB, new ProtobufGraph(filesetB));
+        assertNotEquals(testA, new ProtobufGraph(filesetB));
+        assertEquals(testA, testA);
+        assertNotEquals(testA, testB);
+
+        HashSet<ProtobufGraph> graphs = new HashSet<>();
+        graphs.add(testA);
+        graphs.add(testB);
+        graphs.add(new ProtobufGraph(filesetB));
+        assertEquals(2, graphs.size());
+    }
+}
diff --git a/metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/model/ProtobufMessageTest.java b/metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/model/ProtobufMessageTest.java
new file mode 100644
index 0000000000000..035c16552aeb5
--- /dev/null
+++ b/metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/model/ProtobufMessageTest.java
@@ -0,0 +1,180 @@
+package datahub.protobuf.model;
+
+import com.google.protobuf.DescriptorProtos.DescriptorProto;
+import com.google.protobuf.DescriptorProtos.FileDescriptorProto;
+import com.linkedin.schema.MapType;
+import com.linkedin.schema.RecordType;
+import com.linkedin.schema.SchemaFieldDataType;
+import org.junit.Test;
+
+import java.util.List;
+import java.util.Set;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+
+public class ProtobufMessageTest {
+
+    @Test
+    public void messageTest() {
+        DescriptorProto expectedMessage = DescriptorProto.newBuilder()
+                .setName("message1")
+                .build();
+        DescriptorProto expectedParentMessage1 = DescriptorProto.newBuilder()
+                .setName("messageParent1")
+                .addNestedType(expectedMessage)
+                .build();
+
+        FileDescriptorProto expectedFile = FileDescriptorProto.newBuilder()
+                .addMessageType(expectedMessage)
+                .setPackage("protobuf")
+                .build();
+
+        ProtobufMessage testParent = ProtobufMessage.builder()
+                .messageProto(expectedParentMessage1)
+                .fileProto(expectedFile)
+                .build();
+        ProtobufMessage test = ProtobufMessage.builder()
+                .messageProto(expectedMessage)
+                .parentMessageProto(expectedParentMessage1)
+                .fileProto(expectedFile)
+                .build();
+
+        assertEquals("messageParent1", testParent.name());
+        assertEquals("protobuf.messageParent1", testParent.fullName());
+        assertEquals("protobuf.messageParent1", testParent.nativeType());
+        assertEquals("[type=protobuf_messageParent1]", testParent.fieldPathType());
+        assertEquals(expectedFile, testParent.fileProto());
+        assertEquals(expectedParentMessage1, testParent.messageProto());
+        assertEquals(new SchemaFieldDataType().setType(SchemaFieldDataType.Type.create(new RecordType())), testParent.schemaFieldDataType());
+        assertEquals("ProtobufMessage[protobuf.messageParent1]", testParent.toString());
+
+        assertEquals("message1", test.name());
+        assertEquals("protobuf.messageParent1.message1", test.fullName());
+        assertEquals("protobuf.messageParent1.message1", test.nativeType());
+        assertEquals("[type=protobuf_messageParent1_message1]", test.fieldPathType());
+        assertEquals(expectedFile, test.fileProto());
+        assertEquals(expectedMessage, test.messageProto());
+        assertEquals(new SchemaFieldDataType().setType(SchemaFieldDataType.Type.create(new RecordType())), test.schemaFieldDataType());
+        assertEquals("ProtobufMessage[protobuf.messageParent1.message1]", test.toString());
+    }
+
+    @Test
+    public void mapTest() {
+        DescriptorProto expectedMap = DescriptorProto.newBuilder()
+                .setName("MapFieldEntry")
+                .build();
+        DescriptorProto expectedParentMessage1 = DescriptorProto.newBuilder()
+                .setName("messageParent1")
+                .addNestedType(expectedMap)
+                .build();
+
+        FileDescriptorProto expectedFile = FileDescriptorProto.newBuilder()
+                .addMessageType(expectedMap)
+                .setPackage("protobuf")
+                .build();
+
+        ProtobufMessage testParent = ProtobufMessage.builder()
+                .messageProto(expectedParentMessage1)
+                .fileProto(expectedFile)
+                .build();
+        ProtobufMessage testMap = ProtobufMessage.builder()
+                .messageProto(expectedMap)
+                .parentMessageProto(expectedParentMessage1)
+                .fileProto(expectedFile)
+                .build();
+
+        assertEquals("messageParent1", testParent.name());
+        assertEquals("protobuf.messageParent1", testParent.fullName());
+        assertEquals("protobuf.messageParent1", testParent.nativeType());
+        assertEquals("[type=protobuf_messageParent1]", testParent.fieldPathType());
+        assertEquals(expectedFile, testParent.fileProto());
+        assertEquals(expectedParentMessage1, testParent.messageProto());
+        assertEquals(new SchemaFieldDataType().setType(SchemaFieldDataType.Type.create(new RecordType())), testParent.schemaFieldDataType());
+        assertEquals("ProtobufMessage[protobuf.messageParent1]", testParent.toString());
+
+        assertEquals("MapFieldEntry", testMap.name());
+        assertEquals("protobuf.messageParent1.MapFieldEntry", testMap.fullName());
+        assertEquals("protobuf.messageParent1.MapFieldEntry", testMap.nativeType());
+        assertEquals("[type=protobuf_messageParent1_MapFieldEntry]", testMap.fieldPathType());
+        assertEquals(expectedFile, testMap.fileProto());
+        assertEquals(expectedMap, testMap.messageProto());
+        assertEquals(new SchemaFieldDataType().setType(SchemaFieldDataType.Type.create(new MapType())), testMap.schemaFieldDataType());
+        assertEquals("ProtobufMessage[protobuf.messageParent1.MapFieldEntry]", testMap.toString());
+    }
+
+    @Test
+    public void messageEqualityTest() {
+        DescriptorProto expectedMessage1 = DescriptorProto.newBuilder()
+                .setName("message1")
+                .build();
+        DescriptorProto expectedMessage2 = DescriptorProto.newBuilder()
+                .setName("message2")
+                .build();
+        DescriptorProto expectedMessage1Dup = DescriptorProto.newBuilder()
+                .setName("message1")
+                .build();
+
+        FileDescriptorProto expectedFile = FileDescriptorProto.newBuilder()
+                .addAllMessageType(List.of(expectedMessage1, expectedMessage2, expectedMessage1Dup))
+                .setPackage("protobuf")
+                .build();
+
+
+        ProtobufMessage test1 = ProtobufMessage.builder()
+                .messageProto(expectedMessage1)
+                .fileProto(expectedFile)
+                .build();
+        ProtobufMessage test2 = ProtobufMessage.builder()
+                .messageProto(expectedMessage2)
+                .fileProto(expectedFile)
+                .build();
+        ProtobufMessage test1Dup = ProtobufMessage.builder()
+                .messageProto(expectedMessage1Dup)
+                .fileProto(expectedFile)
+                .build();
+
+        assertEquals(test1, test1Dup);
+        assertNotEquals(test1, test2);
+        assertEquals(Set.of(test1, test2), Stream.of(test1, test2, test1Dup).collect(Collectors.toSet()));
+    }
+
+    @Test
+    public void majorVersionTest() {
+        DescriptorProto expectedMessage1 = DescriptorProto.newBuilder()
+                .setName("message1")
+                .build();
+
+        FileDescriptorProto expectedFile1 = FileDescriptorProto.newBuilder()
+                .setName("zendesk/v1/platform/test.proto")
+                .setPackage("protobuf")
+                .build();
+        ProtobufMessage test1 = ProtobufMessage.builder()
+                .messageProto(expectedMessage1)
+                .fileProto(expectedFile1)
+                .build();
+        assertEquals(1, test1.majorVersion());
+
+        FileDescriptorProto expectedFile2 = FileDescriptorProto.newBuilder()
+                .setName("zendesk/v2/platform/test.proto")
+                .setPackage("protobuf")
+                .build();
+        ProtobufMessage test2 = ProtobufMessage.builder()
+                .messageProto(expectedMessage1)
+                .fileProto(expectedFile2)
+                .build();
+        assertEquals(2, test2.majorVersion());
+
+        FileDescriptorProto expectedFile3 = FileDescriptorProto.newBuilder()
+                .setName("zendesk/platform/test.proto")
+                .setPackage("protobuf")
+                .build();
+        ProtobufMessage test3 = ProtobufMessage.builder()
+                .messageProto(expectedMessage1)
+                .fileProto(expectedFile3)
+                .build();
+        assertEquals(1, test3.majorVersion());
+    }
+}
diff --git a/metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/model/ProtobufOneOfFieldTest.java b/metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/model/ProtobufOneOfFieldTest.java
new file mode 100644
index 0000000000000..f9b168437643b
--- /dev/null
+++ b/metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/model/ProtobufOneOfFieldTest.java
@@ -0,0 +1,121 @@
+package datahub.protobuf.model;
+
+import com.google.protobuf.DescriptorProtos.DescriptorProto;
+import com.google.protobuf.DescriptorProtos.FieldDescriptorProto;
+import com.google.protobuf.DescriptorProtos.FileDescriptorProto;
+import com.google.protobuf.DescriptorProtos.OneofDescriptorProto;
+import com.linkedin.schema.SchemaFieldDataType;
+import com.linkedin.schema.UnionType;
+import org.junit.Test;
+
+import java.util.List;
+import java.util.Set;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+
+public class ProtobufOneOfFieldTest {
+
+    @Test
+    public void oneOfTest() {
+        OneofDescriptorProto expectedOneOf =  OneofDescriptorProto.newBuilder()
+                .setName("oneof1")
+                .build();
+        FieldDescriptorProto expectedField = FieldDescriptorProto.newBuilder()
+                .setName("field1")
+                .setOneofIndex(0)
+                .build();
+        DescriptorProto expectedMessage = DescriptorProto.newBuilder()
+                .setName("message1")
+                .addOneofDecl(expectedOneOf)
+                .addField(expectedField)
+                .build();
+        FileDescriptorProto expectedFile = FileDescriptorProto.newBuilder()
+                .addMessageType(expectedMessage)
+                .setPackage("protobuf")
+                .build();
+
+        ProtobufOneOfField test = ProtobufOneOfField.oneOfBuilder()
+                .fieldProto(expectedField)
+                .protobufMessage(ProtobufMessage.builder().fileProto(expectedFile).messageProto(expectedMessage).build())
+                .build();
+
+        assertEquals("oneof1", test.name());
+        assertEquals("protobuf.message1.oneof1", test.fullName());
+        assertEquals("[type=union]", test.fieldPathType());
+        assertEquals("oneof", test.nativeType());
+        assertEquals(expectedOneOf, test.oneOfProto());
+        assertEquals(expectedMessage, test.messageProto());
+        assertEquals(expectedFile, test.fileProto());
+        assertFalse(test.isMessage());
+        assertEquals(new SchemaFieldDataType().setType(SchemaFieldDataType.Type.create(new UnionType())), test.schemaFieldDataType());
+        assertEquals("ProtobufOneOf[protobuf.message1.oneof1]", test.toString());
+    }
+
+    @Test
+    public void oneOfEqualityTest() {
+        OneofDescriptorProto oneof1Message1 =  OneofDescriptorProto.newBuilder().setName("oneof1").build();
+        OneofDescriptorProto oneof2Message1 =  OneofDescriptorProto.newBuilder().setName("oneof2").build();
+        OneofDescriptorProto oneof1Message2 =  OneofDescriptorProto.newBuilder().setName("oneof1").build();
+        OneofDescriptorProto oneof1Message1Dup =  OneofDescriptorProto.newBuilder().setName("oneof1").build();
+
+        FieldDescriptorProto expectedField1 = FieldDescriptorProto.newBuilder()
+                .setName("field1")
+                .setOneofIndex(0)
+                .build();
+        FieldDescriptorProto expectedField2 = FieldDescriptorProto.newBuilder()
+                .setName("field2")
+                .setOneofIndex(1)
+                .build();
+        FieldDescriptorProto expectedField1Dup = FieldDescriptorProto.newBuilder()
+                .setName("field3")
+                .setOneofIndex(3)
+                .build();
+        DescriptorProto expectedMessage1 = DescriptorProto.newBuilder()
+                .setName("message1")
+                .addAllOneofDecl(List.of(oneof1Message1, oneof2Message1, oneof1Message1Dup))
+                .addField(expectedField1)
+                .addField(expectedField2)
+                .addField(expectedField1Dup)
+                .build();
+
+        FieldDescriptorProto expectedField3 = FieldDescriptorProto.newBuilder()
+                .setName("field3")
+                .setOneofIndex(0)
+                .build();
+        DescriptorProto expectedMessage2 = DescriptorProto.newBuilder()
+                .setName("message2")
+                .addAllOneofDecl(List.of(oneof1Message2))
+                .addField(expectedField3)
+                .build();
+
+        FileDescriptorProto expectedFile = FileDescriptorProto.newBuilder()
+                .addAllMessageType(List.of(expectedMessage1, expectedMessage2))
+                .setPackage("protobuf")
+                .build();
+
+        ProtobufOneOfField test1 = ProtobufOneOfField.oneOfBuilder()
+                .fieldProto(expectedField1)
+                .protobufMessage(ProtobufMessage.builder().fileProto(expectedFile).messageProto(expectedMessage1).build())
+                .build();
+        ProtobufOneOfField test1Dup = ProtobufOneOfField.oneOfBuilder()
+                .fieldProto(expectedField1)
+                .protobufMessage(ProtobufMessage.builder().fileProto(expectedFile).messageProto(expectedMessage1).build())
+                .build();
+        ProtobufOneOfField test2 = ProtobufOneOfField.oneOfBuilder()
+                .fieldProto(expectedField2)
+                .protobufMessage(ProtobufMessage.builder().fileProto(expectedFile).messageProto(expectedMessage1).build())
+                .build();
+        ProtobufOneOfField test3 = ProtobufOneOfField.oneOfBuilder()
+                .fieldProto(expectedField3)
+                .protobufMessage(ProtobufMessage.builder().fileProto(expectedFile).messageProto(expectedMessage2).build())
+                .build();
+
+        assertEquals(test1, test1Dup);
+        assertNotEquals(test1, test3);
+        assertNotEquals(test1, test2);
+        assertEquals(Set.of(test1, test2, test3), Stream.of(test1, test2, test3, test1Dup).collect(Collectors.toSet()));
+    }
+}
diff --git a/metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/visitors/dataset/DatasetVisitorTest.java b/metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/visitors/dataset/DatasetVisitorTest.java
new file mode 100644
index 0000000000000..165823d8e4925
--- /dev/null
+++ b/metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/visitors/dataset/DatasetVisitorTest.java
@@ -0,0 +1,56 @@
+package datahub.protobuf.visitors.dataset;
+
+import com.linkedin.common.urn.DatasetUrn;
+import com.linkedin.data.template.RecordTemplate;
+import org.junit.Test;
+
+import java.io.IOException;
+import java.net.URISyntaxException;
+import java.util.List;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+
+import datahub.protobuf.ProtobufDataset;
+import datahub.protobuf.visitors.ProtobufModelVisitor;
+import datahub.protobuf.visitors.VisitContext;
+import datahub.event.MetadataChangeProposalWrapper;
+
+import static datahub.protobuf.TestFixtures.*;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+
+public class DatasetVisitorTest {
+
+    @Test
+    public void protocBase64Test() throws URISyntaxException, IOException {
+        String expected = "23454345452345233455";
+        DatasetVisitor test = DatasetVisitor.builder().protocBase64(expected).build();
+
+        List<MetadataChangeProposalWrapper<? extends RecordTemplate>> changes =
+                test.visitGraph(
+                        VisitContext.builder()
+                                .auditStamp(TEST_AUDIT_STAMP)
+                                .datasetUrn(DatasetUrn.createFromString("urn:li:dataset:(urn:li:dataPlatform:kafka,protobuf.MessageA,TEST)"))
+                                .graph(getTestProtobufGraph("protobuf", "messageA")).build()
+                ).collect(Collectors.toList());
+
+        assertEquals(expected, extractCustomProperty(changes.get(0), "protoc"));
+    }
+
+    @Test
+    public void customDescriptionVisitors() throws IOException {
+        ProtobufDataset testDataset = getTestProtobufDataset("protobuf", "messageA");
+
+        DatasetVisitor test = DatasetVisitor.builder()
+                .descriptionVisitor(new ProtobufModelVisitor<String>() {
+                    @Override
+                    public Stream<String> visitGraph(VisitContext context) {
+                        return Stream.of("Test Description");
+                    }
+                })
+                .build();
+        testDataset.setDatasetVisitor(test);
+
+        assertEquals("Test Description", extractAspect(testDataset.getDatasetMCPs().get(0), "description"));
+    }
+}
diff --git a/metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/visitors/dataset/DescriptionVisitorTest.java b/metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/visitors/dataset/DescriptionVisitorTest.java
new file mode 100644
index 0000000000000..4a02c1089fe2d
--- /dev/null
+++ b/metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/visitors/dataset/DescriptionVisitorTest.java
@@ -0,0 +1,27 @@
+package datahub.protobuf.visitors.dataset;
+
+import datahub.protobuf.model.ProtobufGraph;
+import org.junit.Test;
+
+import java.io.IOException;
+import java.util.List;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+import static datahub.protobuf.TestFixtures.*;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+
+public class DescriptionVisitorTest {
+
+    @Test
+    public void visitorTest() throws IOException {
+        ProtobufGraph graph = getTestProtobufGraph("protobuf", "messageB");
+
+        DescriptionVisitor test = new DescriptionVisitor();
+
+        assertEquals(Set.of("This contains nested types.\n\nOwned by TeamB"),
+                getTestProtobufGraph("protobuf", "messageB")
+                        .accept(getVisitContextBuilder("protobuf.MessageB"), List.of(test)).collect(Collectors.toSet()));
+    }
+}
diff --git a/metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/visitors/dataset/InstitutionalMemoryVisitorTest.java b/metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/visitors/dataset/InstitutionalMemoryVisitorTest.java
new file mode 100644
index 0000000000000..a313681c5a5a0
--- /dev/null
+++ b/metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/visitors/dataset/InstitutionalMemoryVisitorTest.java
@@ -0,0 +1,68 @@
+package datahub.protobuf.visitors.dataset;
+
+import com.linkedin.common.InstitutionalMemoryMetadata;
+import com.linkedin.common.url.Url;
+import org.junit.Test;
+
+import java.io.IOException;
+import java.util.List;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+import static datahub.protobuf.TestFixtures.*;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+
+public class InstitutionalMemoryVisitorTest {
+
+    @Test
+    public void messageATest() throws IOException {
+        InstitutionalMemoryVisitor test = new InstitutionalMemoryVisitor("SLACK123", "myOrg");
+        assertEquals(Set.of(new InstitutionalMemoryMetadata()
+                                .setCreateStamp(TEST_AUDIT_STAMP)
+                                .setDescription("Slack Channel")
+                                .setUrl(new Url("https://slack.com/app_redirect?channel=test-slack&team=SLACK123")),
+                        new InstitutionalMemoryMetadata()
+                                .setCreateStamp(TEST_AUDIT_STAMP)
+                                .setDescription("Github Team")
+                                .setUrl(new Url("https://github.com/orgs/myOrg/teams/teama")),
+                        new InstitutionalMemoryMetadata()
+                                .setCreateStamp(TEST_AUDIT_STAMP)
+                                .setDescription("MessageA Reference 1")
+                                .setUrl(new Url("https://some/link")),
+                        new InstitutionalMemoryMetadata()
+                                .setCreateStamp(TEST_AUDIT_STAMP)
+                                .setDescription("MessageA Reference 2")
+                                .setUrl(new Url("https://www.google.com/search?q=protobuf+messages")),
+                        new InstitutionalMemoryMetadata()
+                                .setCreateStamp(TEST_AUDIT_STAMP)
+                                .setDescription("MessageA Reference 3")
+                                .setUrl(new Url("https://github.com/apache/kafka")),
+                        new InstitutionalMemoryMetadata()
+                                .setCreateStamp(TEST_AUDIT_STAMP)
+                                .setDescription("MessageA.map_field Reference 1")
+                                .setUrl(new Url("https://developers.google.com/protocol-buffers/docs/proto3#maps"))
+                        ),
+
+                getTestProtobufGraph("protobuf", "messageA")
+                        .accept(getVisitContextBuilder("protobuf.MessageA"),
+                        List.of(test)).collect(Collectors.toSet()));
+    }
+
+    @Test
+    public void messageBTest() throws IOException {
+        InstitutionalMemoryVisitor test = new InstitutionalMemoryVisitor("SLACK123", "myOrg");
+        assertEquals(Set.of(),
+                getTestProtobufGraph("protobuf", "messageB")
+                        .accept(getVisitContextBuilder("protobuf.MessageB"),
+                                List.of(test)).collect(Collectors.toSet()));
+    }
+
+    @Test
+    public void messageCTest() throws IOException {
+        InstitutionalMemoryVisitor test = new InstitutionalMemoryVisitor("SLACK123", "myOrg");
+        assertEquals(Set.of(), getTestProtobufGraph("protobuf", "messageC")
+                .accept(getVisitContextBuilder("protobuf.MessageC"),
+                        List.of(test)).collect(Collectors.toSet()));
+    }
+}
diff --git a/metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/visitors/dataset/KafkaTopicPropertyVisitorTest.java b/metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/visitors/dataset/KafkaTopicPropertyVisitorTest.java
new file mode 100644
index 0000000000000..84e7eb19f893b
--- /dev/null
+++ b/metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/visitors/dataset/KafkaTopicPropertyVisitorTest.java
@@ -0,0 +1,36 @@
+package datahub.protobuf.visitors.dataset;
+
+import com.linkedin.data.template.StringMap;
+import com.linkedin.dataset.DatasetProperties;
+import org.junit.Test;
+
+import java.io.IOException;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+import static datahub.protobuf.TestFixtures.getTestProtobufGraph;
+import static datahub.protobuf.TestFixtures.getVisitContextBuilder;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+
+public class KafkaTopicPropertyVisitorTest {
+
+    @Test
+    public void visitorTest() throws IOException {
+        KafkaTopicPropertyVisitor test = new KafkaTopicPropertyVisitor();
+        assertEquals(List.of(new DatasetProperties()
+                        .setCustomProperties(new StringMap(Map.of("kafka_topic", "platform.topic")))),
+                getTestProtobufGraph("protobuf", "messageA")
+                        .accept(getVisitContextBuilder("MessageB"),
+                                List.of(test)).collect(Collectors.toList()));
+    }
+
+    @Test
+    public void visitorEmptyTest() throws IOException {
+        KafkaTopicPropertyVisitor test = new KafkaTopicPropertyVisitor();
+        assertEquals(Set.of(), getTestProtobufGraph("protobuf", "messageB")
+                .accept(getVisitContextBuilder("MessageB"), List.of(test)).collect(Collectors.toSet()));
+    }
+}
diff --git a/metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/visitors/dataset/ProtobufExtensionPropertyVisitorTest.java b/metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/visitors/dataset/ProtobufExtensionPropertyVisitorTest.java
new file mode 100644
index 0000000000000..d00c81be47956
--- /dev/null
+++ b/metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/visitors/dataset/ProtobufExtensionPropertyVisitorTest.java
@@ -0,0 +1,46 @@
+package datahub.protobuf.visitors.dataset;
+
+import com.linkedin.data.template.StringMap;
+import com.linkedin.dataset.DatasetProperties;
+import org.junit.Test;
+
+import java.io.IOException;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+
+import static datahub.protobuf.TestFixtures.getTestProtobufGraph;
+import static datahub.protobuf.TestFixtures.getVisitContextBuilder;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+
+public class ProtobufExtensionPropertyVisitorTest {
+
+    @Test
+    public void extendedMessageTest() throws IOException {
+        ProtobufExtensionPropertyVisitor test = new ProtobufExtensionPropertyVisitor();
+
+        List<DatasetProperties> actual = getTestProtobufGraph("extended_protobuf", "messageA")
+                .accept(getVisitContextBuilder("extended_protobuf.Person"),
+                        List.of(test)).collect(Collectors.toList());
+
+        assertEquals(List.of(
+                new DatasetProperties().setCustomProperties(new StringMap(Map.of("classification_enum", "HighlyConfidential",
+                        "bool_feature", "true",
+                        "alert_channel", "#alerts",
+                        "team", "TeamB",
+                        "type", "ENTITY")))),
+                actual);
+    }
+
+    @Test
+    public void extendedFieldTest() throws IOException {
+        ProtobufExtensionPropertyVisitor test = new ProtobufExtensionPropertyVisitor();
+        List<DatasetProperties> actual = getTestProtobufGraph("extended_protobuf", "messageB")
+                .accept(getVisitContextBuilder("extended_protobuf.Person"),
+                        List.of(test)).collect(Collectors.toList());
+
+        assertEquals(List.of(new DatasetProperties()
+                        .setCustomProperties(new StringMap(Map.of()))), actual);
+    }
+}
diff --git a/metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/visitors/dataset/ProtobufExtensionTermAssocVisitorTest.java b/metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/visitors/dataset/ProtobufExtensionTermAssocVisitorTest.java
new file mode 100644
index 0000000000000..ac8600a20c5e9
--- /dev/null
+++ b/metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/visitors/dataset/ProtobufExtensionTermAssocVisitorTest.java
@@ -0,0 +1,38 @@
+package datahub.protobuf.visitors.dataset;
+
+import com.linkedin.common.GlossaryTermAssociation;
+import com.linkedin.common.urn.GlossaryTermUrn;
+import org.junit.Test;
+
+import java.io.IOException;
+import java.util.List;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+import static datahub.protobuf.TestFixtures.getTestProtobufGraph;
+import static datahub.protobuf.TestFixtures.getVisitContextBuilder;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+
+public class ProtobufExtensionTermAssocVisitorTest {
+
+    @Test
+    public void extendedMessageTest() throws IOException {
+        ProtobufExtensionTermAssocVisitor test = new ProtobufExtensionTermAssocVisitor();
+        assertEquals(Set.of(
+                        new GlossaryTermAssociation().setUrn(new GlossaryTermUrn("Classification.HighlyConfidential"))
+                ),
+                getTestProtobufGraph("extended_protobuf", "messageA")
+                        .accept(getVisitContextBuilder("extended_protobuf.Person"),
+                                List.of(test)).collect(Collectors.toSet()));
+    }
+
+    @Test
+    public void extendedFieldTest() throws IOException {
+        ProtobufExtensionTermAssocVisitor test = new ProtobufExtensionTermAssocVisitor();
+        assertEquals(
+                Set.of(),
+                getTestProtobufGraph("extended_protobuf", "messageB").
+                        accept(getVisitContextBuilder("extended_protobuf.Person"), List.of(test)).collect(Collectors.toSet()));
+    }
+}
diff --git a/metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/visitors/field/ProtobufExtensionFieldVisitorTest.java b/metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/visitors/field/ProtobufExtensionFieldVisitorTest.java
new file mode 100644
index 0000000000000..0945fd11d6794
--- /dev/null
+++ b/metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/visitors/field/ProtobufExtensionFieldVisitorTest.java
@@ -0,0 +1,218 @@
+package datahub.protobuf.visitors.field;
+
+import com.linkedin.common.GlobalTags;
+import com.linkedin.common.GlossaryTermAssociation;
+import com.linkedin.common.GlossaryTermAssociationArray;
+import com.linkedin.common.GlossaryTerms;
+import com.linkedin.common.TagAssociation;
+import com.linkedin.common.TagAssociationArray;
+import com.linkedin.common.urn.GlossaryTermUrn;
+import com.linkedin.common.urn.TagUrn;
+import com.linkedin.schema.NumberType;
+import com.linkedin.schema.RecordType;
+import com.linkedin.schema.SchemaField;
+import com.linkedin.schema.SchemaFieldDataType;
+import com.linkedin.schema.StringType;
+import com.linkedin.util.Pair;
+import datahub.protobuf.ProtobufDataset;
+import org.junit.Test;
+
+import java.io.IOException;
+import java.net.URISyntaxException;
+import java.util.List;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+
+import static datahub.protobuf.TestFixtures.*;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+
+public class ProtobufExtensionFieldVisitorTest {
+
+    @Test
+    public void extendedMessageTest() throws IOException, URISyntaxException {
+        ProtobufExtensionFieldVisitor test = new ProtobufExtensionFieldVisitor();
+        List<SchemaField> actual = getTestProtobufGraph("extended_protobuf", "messageA")
+                .accept(getVisitContextBuilder("extended_protobuf.Person"), List.of(test))
+                .sorted(ProtobufDataset.COMPARE_BY_ROOT_MESSAGE_FIELD_WEIGHT.thenComparing(ProtobufDataset.COMPARE_BY_FIELD_PATH))
+                .map(Pair::getFirst)
+                .collect(Collectors.toList());
+
+        List<SchemaField> expected = Stream.of(
+                Pair.of(
+                        new SchemaField()
+                                .setFieldPath("[version=2.0].[type=extended_protobuf_Person].[type=string].name")
+                                .setNullable(true)
+                                .setIsPartOfKey(false)
+                                .setDescription("")
+                                .setNativeDataType("string")
+                                .setType(new SchemaFieldDataType().setType(SchemaFieldDataType.Type.create(new StringType())))
+                                .setGlobalTags(new GlobalTags().setTags(new TagAssociationArray()))
+                                .setGlossaryTerms(new GlossaryTerms().setTerms(new GlossaryTermAssociationArray()).setAuditStamp(TEST_AUDIT_STAMP)),
+                        1),
+                Pair.of(
+                        new SchemaField()
+                                .setFieldPath("[version=2.0].[type=extended_protobuf_Person].[type=int].id")
+                                .setNullable(true)
+                                .setIsPartOfKey(false)
+                                .setDescription("")
+                                .setNativeDataType("int32")
+                                .setType(new SchemaFieldDataType().setType(SchemaFieldDataType.Type.create(new NumberType())))
+                                .setGlobalTags(new GlobalTags().setTags(new TagAssociationArray()))
+                                .setGlossaryTerms(new GlossaryTerms().setTerms(new GlossaryTermAssociationArray()).setAuditStamp(TEST_AUDIT_STAMP)),
+                        2),
+                Pair.of(
+                        new SchemaField()
+                                .setFieldPath("[version=2.0].[type=extended_protobuf_Person].[type=string].email")
+                                .setNullable(true)
+                                .setIsPartOfKey(false)
+                                .setDescription("")
+                                .setNativeDataType("string")
+                                .setType(new SchemaFieldDataType().setType(SchemaFieldDataType.Type.create(new StringType())))
+                                .setGlobalTags(new GlobalTags().setTags(new TagAssociationArray()))
+                                .setGlossaryTerms(new GlossaryTerms().setTerms(new GlossaryTermAssociationArray()).setAuditStamp(TEST_AUDIT_STAMP)),
+                        3),
+                Pair.of(
+                        new SchemaField()
+                                .setFieldPath("[version=2.0].[type=extended_protobuf_Person].[type=extended_protobuf_Department].dept")
+                                .setNullable(true)
+                                .setIsPartOfKey(false)
+                                .setDescription("")
+                                .setNativeDataType("extended_protobuf.Department")
+                                .setType(new SchemaFieldDataType().setType(SchemaFieldDataType.Type.create(new RecordType())))
+                                .setGlobalTags(new GlobalTags().setTags(new TagAssociationArray(
+                                        new TagAssociation().setTag(new TagUrn("MetaEnumExample.ENTITY")),
+                                        new TagAssociation().setTag(new TagUrn("team.TeamA"))
+                                )))
+                                .setGlossaryTerms(new GlossaryTerms().setTerms(new GlossaryTermAssociationArray(
+                                        new GlossaryTermAssociation().setUrn(new GlossaryTermUrn("Classification.Sensitive"))
+                                )).setAuditStamp(TEST_AUDIT_STAMP)),
+                        4),
+                Pair.of(
+                        new SchemaField()
+                                .setFieldPath("[version=2.0].[type=extended_protobuf_Person].[type=extended_protobuf_Department].dept.[type=int].id")
+                                .setNullable(true)
+                                .setIsPartOfKey(false)
+                                .setDescription("")
+                                .setNativeDataType("int32")
+                                .setType(new SchemaFieldDataType().setType(SchemaFieldDataType.Type.create(new NumberType())))
+                                .setGlobalTags(new GlobalTags().setTags(new TagAssociationArray()))
+                                .setGlossaryTerms(new GlossaryTerms().setTerms(new GlossaryTermAssociationArray()).setAuditStamp(TEST_AUDIT_STAMP)),
+                        4),
+                Pair.of(
+                        new SchemaField()
+                                .setFieldPath("[version=2.0].[type=extended_protobuf_Person].[type=extended_protobuf_Department].dept.[type=string].name")
+                                .setNullable(true)
+                                .setIsPartOfKey(false)
+                                .setDescription("")
+                                .setNativeDataType("string")
+                                .setType(new SchemaFieldDataType().setType(SchemaFieldDataType.Type.create(new StringType())))
+                                .setGlobalTags(new GlobalTags().setTags(new TagAssociationArray()))
+                                .setGlossaryTerms(new GlossaryTerms().setTerms(new GlossaryTermAssociationArray()).setAuditStamp(TEST_AUDIT_STAMP)),
+                        4)
+        ).map(Pair::getFirst).collect(Collectors.toList());
+
+
+        assertEquals(expected, actual);
+    }
+
+    @Test
+    public void extendedFieldTest() throws IOException {
+        ProtobufExtensionFieldVisitor test = new ProtobufExtensionFieldVisitor();
+        List<SchemaField> actual = getTestProtobufGraph("extended_protobuf", "messageB")
+                .accept(getVisitContextBuilder("extended_protobuf.Person"), List.of(test))
+                .sorted(ProtobufDataset.COMPARE_BY_ROOT_MESSAGE_FIELD_WEIGHT.thenComparing(ProtobufDataset.COMPARE_BY_FIELD_PATH))
+                .map(Pair::getFirst)
+                .collect(Collectors.toList());
+
+        List<SchemaField> expected = Stream.of(
+                Pair.of(
+                        new SchemaField()
+                                .setFieldPath("[version=2.0].[type=extended_protobuf_Person].[type=string].name")
+                                .setNullable(true)
+                                .setIsPartOfKey(false)
+                                .setDescription("person name")
+                                .setNativeDataType("string")
+                                .setType(new SchemaFieldDataType().setType(SchemaFieldDataType.Type.create(new StringType())))
+                                .setGlobalTags(new GlobalTags().setTags(new TagAssociationArray()))
+                                .setGlossaryTerms(new GlossaryTerms().setTerms(new GlossaryTermAssociationArray(
+                                        new GlossaryTermAssociation().setUrn(new GlossaryTermUrn("Classification.HighlyConfidential"))
+                                )).setAuditStamp(TEST_AUDIT_STAMP)),
+                        1),
+                Pair.of(
+                        new SchemaField()
+                                .setFieldPath("[version=2.0].[type=extended_protobuf_Person].[type=int].id")
+                                .setNullable(false)
+                                .setIsPartOfKey(true)
+                                .setDescription("unique identifier for a given person")
+                                .setNativeDataType("int32")
+                                .setType(new SchemaFieldDataType().setType(SchemaFieldDataType.Type.create(new NumberType())))
+                                .setGlobalTags(new GlobalTags().setTags(new TagAssociationArray()))
+                                .setGlossaryTerms(new GlossaryTerms().setTerms(new GlossaryTermAssociationArray()).setAuditStamp(TEST_AUDIT_STAMP)),
+                        2),
+                Pair.of(
+                        new SchemaField()
+                                .setFieldPath("[version=2.0].[type=extended_protobuf_Person].[type=string].email")
+                                .setNullable(true)
+                                .setIsPartOfKey(false)
+                                .setDescription("official email address")
+                                .setNativeDataType("string")
+                                .setType(new SchemaFieldDataType().setType(SchemaFieldDataType.Type.create(new StringType())))
+                                .setGlobalTags(new GlobalTags().setTags(new TagAssociationArray()))
+                                .setGlossaryTerms(new GlossaryTerms().setTerms(new GlossaryTermAssociationArray(
+                                        new GlossaryTermAssociation().setUrn(new GlossaryTermUrn("Classification.HighlyConfidential"))
+                                )).setAuditStamp(TEST_AUDIT_STAMP)),
+                        3),
+                Pair.of(
+                        new SchemaField()
+                                .setFieldPath("[version=2.0].[type=extended_protobuf_Person].[type=extended_protobuf_Department].dept")
+                                .setNullable(true)
+                                .setIsPartOfKey(false)
+                                .setDescription("department name of the person")
+                                .setNativeDataType("extended_protobuf.Department")
+                                .setType(new SchemaFieldDataType().setType(SchemaFieldDataType.Type.create(new RecordType())))
+                                .setGlobalTags(new GlobalTags().setTags(new TagAssociationArray()))
+                                .setGlossaryTerms(new GlossaryTerms().setTerms(new GlossaryTermAssociationArray()).setAuditStamp(TEST_AUDIT_STAMP)),
+                        4),
+                Pair.of(
+                        new SchemaField()
+                                .setFieldPath("[version=2.0].[type=extended_protobuf_Person].[type=extended_protobuf_Department].dept.[type=int].id")
+                                .setNullable(false)
+                                .setIsPartOfKey(true)
+                                .setDescription("")
+                                .setNativeDataType("int32")
+                                .setType(new SchemaFieldDataType().setType(SchemaFieldDataType.Type.create(new NumberType())))
+                                .setGlobalTags(new GlobalTags().setTags(new TagAssociationArray()))
+                                .setGlossaryTerms(new GlossaryTerms().setTerms(new GlossaryTermAssociationArray()).setAuditStamp(TEST_AUDIT_STAMP)),
+                        4),
+                Pair.of(
+                        new SchemaField()
+                                .setFieldPath("[version=2.0].[type=extended_protobuf_Person].[type=extended_protobuf_Department].dept.[type=string].name")
+                                .setNullable(true)
+                                .setIsPartOfKey(false)
+                                .setDescription("")
+                                .setNativeDataType("string")
+                                .setType(new SchemaFieldDataType().setType(SchemaFieldDataType.Type.create(new StringType())))
+                                .setGlobalTags(new GlobalTags().setTags(new TagAssociationArray()))
+                                .setGlossaryTerms(new GlossaryTerms().setTerms(new GlossaryTermAssociationArray()).setAuditStamp(TEST_AUDIT_STAMP)),
+                        4),
+                Pair.of(
+                        new SchemaField()
+                                .setFieldPath("[version=2.0].[type=extended_protobuf_Person].[type=string].test_coverage")
+                                .setNullable(true)
+                                .setIsPartOfKey(false)
+                                .setDescription("")
+                                .setNativeDataType("string")
+                                .setType(new SchemaFieldDataType().setType(SchemaFieldDataType.Type.create(new StringType())))
+                                .setGlobalTags(new GlobalTags().setTags(new TagAssociationArray(
+                                        new TagAssociation().setTag(new TagUrn("MetaEnumExample.EVENT")),
+                                        new TagAssociation().setTag(new TagUrn("product_type.my type")),
+                                        new TagAssociation().setTag(new TagUrn("product_type_bool"))
+                                )))
+                                .setGlossaryTerms(new GlossaryTerms().setTerms(new GlossaryTermAssociationArray()).setAuditStamp(TEST_AUDIT_STAMP)),
+                        5)
+        ).map(Pair::getFirst).collect(Collectors.toList());
+
+        assertEquals(expected, actual);
+    }
+}
diff --git a/metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/visitors/field/SchemaFieldVisitorTest.java b/metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/visitors/field/SchemaFieldVisitorTest.java
new file mode 100644
index 0000000000000..6c855e70d7f37
--- /dev/null
+++ b/metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/visitors/field/SchemaFieldVisitorTest.java
@@ -0,0 +1,68 @@
+package datahub.protobuf.visitors.field;
+
+import com.linkedin.schema.NumberType;
+import com.linkedin.schema.SchemaField;
+import com.linkedin.schema.SchemaFieldDataType;
+import com.linkedin.schema.StringType;
+import com.linkedin.schema.UnionType;
+import com.linkedin.util.Pair;
+import datahub.protobuf.ProtobufDataset;
+import org.junit.Test;
+
+import java.io.IOException;
+import java.util.List;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+
+import static datahub.protobuf.TestFixtures.getTestProtobufGraph;
+import static datahub.protobuf.TestFixtures.getVisitContextBuilder;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+
+public class SchemaFieldVisitorTest {
+
+    @Test
+    public void visitorTest() throws IOException {
+        List<SchemaField> expected = Stream.of(
+                Pair.of(
+                        new SchemaField()
+                                .setFieldPath("[version=2.0].[type=protobuf_MessageC].[type=union].one_of_field")
+                                .setNullable(true)
+                                .setDescription("one of field comment")
+                                .setNativeDataType("oneof")
+                                .setType(new SchemaFieldDataType().setType(SchemaFieldDataType.Type.create(new UnionType()))),
+                        1),
+                Pair.of(
+                        new SchemaField()
+                                .setFieldPath("[version=2.0].[type=protobuf_MessageC].[type=union].one_of_field.[type=string].one_of_string")
+                                .setNullable(true)
+                                .setDescription("one of string comment")
+                                .setNativeDataType("string")
+                                .setType(new SchemaFieldDataType().setType(SchemaFieldDataType.Type.create(new StringType()))),
+                        1),
+                Pair.of(
+                        new SchemaField()
+                                .setFieldPath("[version=2.0].[type=protobuf_MessageC].[type=union].one_of_field.[type=int].one_of_int")
+                                .setNullable(true)
+                                .setDescription("one of int comment")
+                                .setNativeDataType("int32")
+                                .setType(new SchemaFieldDataType().setType(SchemaFieldDataType.Type.create(new NumberType()))),
+                        2),
+                Pair.of(
+                        new SchemaField()
+                                .setFieldPath("[version=2.0].[type=protobuf_MessageC].[type=string].normal")
+                                .setNullable(true)
+                                .setDescription("")
+                                .setNativeDataType("string")
+                                .setType(new SchemaFieldDataType().setType(SchemaFieldDataType.Type.create(new StringType()))),
+                        4)
+        ).map(Pair::getFirst).collect(Collectors.toList());
+
+        SchemaFieldVisitor test = new SchemaFieldVisitor();
+        assertEquals(expected, getTestProtobufGraph("protobuf", "messageC")
+                .accept(getVisitContextBuilder("protobuf.MessageC"), List.of(test))
+                .sorted(ProtobufDataset.COMPARE_BY_ROOT_MESSAGE_FIELD_WEIGHT.thenComparing(ProtobufDataset.COMPARE_BY_FIELD_PATH))
+                .map(Pair::getFirst)
+                .collect(Collectors.toList()));
+    }
+}
diff --git a/metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/visitors/tag/ProtobufExtensionTagVisitorTest.java b/metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/visitors/tag/ProtobufExtensionTagVisitorTest.java
new file mode 100644
index 0000000000000..09f0425136194
--- /dev/null
+++ b/metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/visitors/tag/ProtobufExtensionTagVisitorTest.java
@@ -0,0 +1,59 @@
+package datahub.protobuf.visitors.tag;
+
+import com.linkedin.tag.TagProperties;
+import datahub.protobuf.visitors.tags.ProtobufExtensionTagVisitor;
+import datahub.event.MetadataChangeProposalWrapper;
+import org.junit.Test;
+
+import java.io.IOException;
+import java.util.List;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+import static datahub.protobuf.TestFixtures.getTestProtobufGraph;
+import static datahub.protobuf.TestFixtures.getVisitContextBuilder;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+
+public class ProtobufExtensionTagVisitorTest {
+
+    @Test
+    public void extendedMessageTest() throws IOException {
+        ProtobufExtensionTagVisitor test = new ProtobufExtensionTagVisitor();
+        assertEquals(Set.of(
+                new TagProperties()
+                        .setName("team.TeamB")
+                        .setDescription("meta.msg.team"),
+                new TagProperties()
+                        .setName("bool_feature")
+                        .setDescription("meta.msg.bool_feature is true."),
+                new TagProperties()
+                        .setName("MetaEnumExample.ENTITY")
+                        .setDescription("Enum MetaEnumExample.ENTITY of {UNKNOWN, ENTITY, EVENT}")
+        ), getTestProtobufGraph("extended_protobuf", "messageA")
+                .accept(getVisitContextBuilder("extended_protobuf.Person"), List.of(test))
+                .map(MetadataChangeProposalWrapper::getAspect)
+                .collect(Collectors.toSet()));
+    }
+
+    @Test
+    public void extendedFieldTest() throws IOException {
+        Set<TagProperties> expectedTagProperties = Set.of(
+                new TagProperties()
+                        .setName("product_type_bool")
+                        .setDescription("meta.fld.product_type_bool is true."),
+                new TagProperties()
+                        .setName("product_type.my type")
+                        .setDescription("meta.fld.product_type"),
+                new TagProperties()
+                        .setName("MetaEnumExample.EVENT")
+                        .setDescription("Enum MetaEnumExample.EVENT of {UNKNOWN, ENTITY, EVENT}")
+        );
+
+        assertEquals(expectedTagProperties,
+                getTestProtobufGraph("extended_protobuf", "messageB")
+                        .accept(getVisitContextBuilder("extended_protobuf.Person"), List.of(new ProtobufExtensionTagVisitor()))
+                        .map(MetadataChangeProposalWrapper::getAspect)
+                        .collect(Collectors.toSet()));
+    }
+}
\ No newline at end of file
diff --git a/metadata-integration/java/datahub-protobuf/src/test/resources/extended_protobuf/messageA.proto b/metadata-integration/java/datahub-protobuf/src/test/resources/extended_protobuf/messageA.proto
new file mode 100644
index 0000000000000..8fe9d6d5dcb10
--- /dev/null
+++ b/metadata-integration/java/datahub-protobuf/src/test/resources/extended_protobuf/messageA.proto
@@ -0,0 +1,28 @@
+syntax = "proto3";
+package extended_protobuf;
+
+import "extended_protobuf/meta/meta.proto";
+import "extended_protobuf/meta/meta_field_options.proto";
+import "extended_protobuf/meta/meta_message_options.proto";
+
+message Department {
+  option(meta.msg.type) = ENTITY;
+  option(meta.msg.classification) = "Classification.Sensitive";
+  option(meta.msg.team) = "TeamA";
+
+  int32 id = 1;
+  string name = 2;
+}
+
+message Person {
+  option(meta.msg.type) = ENTITY;
+  option(meta.msg.classification_enum) = HighlyConfidential;
+  option(meta.msg.team) = "TeamB";
+  option(meta.msg.bool_feature) = true;
+  option(meta.msg.alert_channel) = "#alerts";
+
+  string name = 1;
+  int32 id = 2;
+  string email = 3;
+  Department dept = 4;
+}
diff --git a/metadata-integration/java/datahub-protobuf/src/test/resources/extended_protobuf/messageA.protoc b/metadata-integration/java/datahub-protobuf/src/test/resources/extended_protobuf/messageA.protoc
new file mode 100644
index 0000000000000..45444ef366cec
Binary files /dev/null and b/metadata-integration/java/datahub-protobuf/src/test/resources/extended_protobuf/messageA.protoc differ
diff --git a/metadata-integration/java/datahub-protobuf/src/test/resources/extended_protobuf/messageB.proto b/metadata-integration/java/datahub-protobuf/src/test/resources/extended_protobuf/messageB.proto
new file mode 100644
index 0000000000000..0f9afe3f5bc8f
--- /dev/null
+++ b/metadata-integration/java/datahub-protobuf/src/test/resources/extended_protobuf/messageB.proto
@@ -0,0 +1,33 @@
+syntax = "proto3";
+package extended_protobuf;
+
+import "extended_protobuf/meta/meta.proto";
+import "extended_protobuf/meta/meta_field_options.proto";
+import "extended_protobuf/meta/meta_message_options.proto";
+
+message Department {
+  int32 id = 1 [(meta.fld.is_primary_key) = true];
+  string name = 2;
+}
+
+/**
+  This is the description of the users table
+
+  The comment added after thought
+*/
+message Person {
+  string name = 1 [(meta.fld.classification) = "Classification.HighlyConfidential"]; // person name
+
+  // unique identifier for a given person
+  int32 id = 2
+  [(meta.fld.is_primary_key) = true];
+
+  // official email address
+  string email = 3
+  [(meta.fld.classification_enum) = HighlyConfidential];
+
+  Department dept = 4; // department name of the person
+
+  string test_coverage = 5
+  [(meta.fld.product_type_bool) = true, (meta.fld.product_type) = "my type", (meta.fld.product_type_enum) = EVENT];
+}
diff --git a/metadata-integration/java/datahub-protobuf/src/test/resources/extended_protobuf/messageB.protoc b/metadata-integration/java/datahub-protobuf/src/test/resources/extended_protobuf/messageB.protoc
new file mode 100644
index 0000000000000..a6b89272dd2b3
Binary files /dev/null and b/metadata-integration/java/datahub-protobuf/src/test/resources/extended_protobuf/messageB.protoc differ
diff --git a/metadata-integration/java/datahub-protobuf/src/test/resources/extended_protobuf/meta/meta.proto b/metadata-integration/java/datahub-protobuf/src/test/resources/extended_protobuf/meta/meta.proto
new file mode 100644
index 0000000000000..d3ffed747206c
--- /dev/null
+++ b/metadata-integration/java/datahub-protobuf/src/test/resources/extended_protobuf/meta/meta.proto
@@ -0,0 +1,38 @@
+syntax = "proto3";
+
+package meta;
+
+/*
+   This is assigned to metadata fields. It describes how the metadata field should be represented
+   in DataHub. This enum must be used in the `meta` package. Multiple can be used for the same
+   metadata annotation. This allows a single piece of information to be captured in DataHub
+   as a property, tag and/or term.
+
+   Tags can be strings, enums, or booleans
+   Terms can be strings or enums
+   Properties should be strings
+
+*/
+enum DataHubMetadataType {
+  PROPERTY = 0; // Datahub Custom Property
+  TAG      = 1; // Datahub Tag
+  TERM     = 2; // Datahub Term
+}
+
+/*
+   Example below: The following is not required for annotation processing. This is an example
+   of creating an annotation using an enum.
+ */
+
+enum MetaEnumExample {
+  UNKNOWN = 0;
+  ENTITY = 1;
+  EVENT = 2;
+}
+
+// Assuming Glossary Term defined from bootstrap example
+enum Classification {
+  HighlyConfidential = 0;
+  Confidential = 1;
+  Sensitive = 2;
+}
diff --git a/metadata-integration/java/datahub-protobuf/src/test/resources/extended_protobuf/meta/meta_field_options.proto b/metadata-integration/java/datahub-protobuf/src/test/resources/extended_protobuf/meta/meta_field_options.proto
new file mode 100644
index 0000000000000..22ee592516458
--- /dev/null
+++ b/metadata-integration/java/datahub-protobuf/src/test/resources/extended_protobuf/meta/meta_field_options.proto
@@ -0,0 +1,28 @@
+syntax = "proto3";
+
+package meta.fld;
+
+import "google/protobuf/descriptor.proto";
+import "extended_protobuf/meta/meta.proto";
+
+
+extend google.protobuf.FieldOptions {
+    // Required: Mark option field with how to export to DataHub in one or more places.
+    repeated meta.DataHubMetadataType type = 6000;
+
+    /*
+       Examples below: The following is not required for annotation processing.
+     */
+
+    // Set true if the field is a primary key. This works for any boolean with `primary_key` in it.
+    bool is_primary_key = 6010;
+
+    // Extract classification field option as a Term, either works
+    string classification                   = 6001 [(meta.fld.type) = TERM];
+    meta.Classification classification_enum = 6002 [(meta.fld.type) = TERM];
+
+    // Expose this option as a tag on the field.
+    string product_type                    = 70004 [(meta.fld.type) = TAG];
+    bool product_type_bool                 = 70005 [(meta.fld.type) = TAG];
+    meta.MetaEnumExample product_type_enum = 70006 [(meta.fld.type) = TAG];
+}
diff --git a/metadata-integration/java/datahub-protobuf/src/test/resources/extended_protobuf/meta/meta_message_options.proto b/metadata-integration/java/datahub-protobuf/src/test/resources/extended_protobuf/meta/meta_message_options.proto
new file mode 100644
index 0000000000000..572a24cb20332
--- /dev/null
+++ b/metadata-integration/java/datahub-protobuf/src/test/resources/extended_protobuf/meta/meta_message_options.proto
@@ -0,0 +1,28 @@
+syntax = "proto3";
+
+package meta.msg;
+
+import "google/protobuf/descriptor.proto";
+import "extended_protobuf/meta/meta.proto";
+import "extended_protobuf/meta/meta_field_options.proto";
+
+
+extend google.protobuf.MessageOptions {
+    /*
+       Examples below: The following is not required for annotation processing.
+     */
+
+    // Place the classification term at the Message/Dataset level, either string or enum is supported
+    string classification                   = 4000 [(meta.fld.type) = TERM, (meta.fld.type) = PROPERTY];
+    meta.Classification classification_enum = 4001 [(meta.fld.type) = TERM, (meta.fld.type) = PROPERTY];
+
+    // Attach these Message/Dataset options as a tag and property.
+    string product            = 5001 [(meta.fld.type) = TAG, (meta.fld.type) = PROPERTY];
+    string project            = 5002 [(meta.fld.type) = TAG, (meta.fld.type) = PROPERTY];
+    string team               = 5003 [(meta.fld.type) = TAG, (meta.fld.type) = PROPERTY];
+
+    string domain             = 60003 [(meta.fld.type) = TAG, (meta.fld.type) = PROPERTY];
+    meta.MetaEnumExample type = 60004 [(meta.fld.type) = TAG, (meta.fld.type) = PROPERTY];
+    bool bool_feature         = 60005 [(meta.fld.type) = TAG];
+    string alert_channel      = 60007 [(meta.fld.type) = PROPERTY];
+}
\ No newline at end of file
diff --git a/metadata-integration/java/datahub-protobuf/src/test/resources/protobuf/messageA.proto b/metadata-integration/java/datahub-protobuf/src/test/resources/protobuf/messageA.proto
new file mode 100644
index 0000000000000..2ebed30b3473f
--- /dev/null
+++ b/metadata-integration/java/datahub-protobuf/src/test/resources/protobuf/messageA.proto
@@ -0,0 +1,46 @@
+syntax = "proto3";
+package protobuf;
+
+/*
+   Multiline message description
+   https://some/link
+
+   #test-slack
+   @myOrg/teama
+
+   Kafka topic: `platform.topic`
+
+   References:
+   https://www.google.com/search?q=protobuf+messages
+   https://github.com/apache/kafka
+*/
+message MessageA {
+  // Leading single line comment
+  bytes sequence_id = 1;
+
+  // Leading multiline comment
+  // Second line of leading multiline comment
+  uint32 position = 2;
+
+  // Detached comment
+
+  uint32 total = 3;
+
+  repeated uint64 repeated_num = 4; // Test repeated and trailing comment
+  repeated string repeated_str = 5;
+
+  // Enum example
+  enum Corpus {
+    UNIVERSAL = 0;
+    WEB = 1;
+    IMAGES = 2;
+    LOCAL = 3;
+    NEWS = 4;
+    PRODUCTS = 5;
+    VIDEO = 6;
+  }
+
+  Corpus corpus_field = 6; // enum field see
+
+  map<string, uint32> map_field = 7; // https://developers.google.com/protocol-buffers/docs/proto3#maps
+}
\ No newline at end of file
diff --git a/metadata-integration/java/datahub-protobuf/src/test/resources/protobuf/messageA.protoc b/metadata-integration/java/datahub-protobuf/src/test/resources/protobuf/messageA.protoc
new file mode 100644
index 0000000000000..6ccad0f42772f
Binary files /dev/null and b/metadata-integration/java/datahub-protobuf/src/test/resources/protobuf/messageA.protoc differ
diff --git a/metadata-integration/java/datahub-protobuf/src/test/resources/protobuf/messageB.proto b/metadata-integration/java/datahub-protobuf/src/test/resources/protobuf/messageB.proto
new file mode 100644
index 0000000000000..0b9f6d0685de9
--- /dev/null
+++ b/metadata-integration/java/datahub-protobuf/src/test/resources/protobuf/messageB.proto
@@ -0,0 +1,17 @@
+syntax = "proto3";
+package protobuf;
+
+import "google/protobuf/wrappers.proto";
+import "protobuf/messageA.proto";
+
+/**
+ * This contains nested types.
+ *
+ * Owned by TeamB
+ */
+message MessageB {
+  google.protobuf.Int64Value id = 1; // wrapped int64
+  string value = 2; // message value
+  protobuf.MessageA nested = 3; // nested message a
+  google.protobuf.BoolValue hot = 6; // Indicator
+}
\ No newline at end of file
diff --git a/metadata-integration/java/datahub-protobuf/src/test/resources/protobuf/messageB.protoc b/metadata-integration/java/datahub-protobuf/src/test/resources/protobuf/messageB.protoc
new file mode 100644
index 0000000000000..d6305e5234c1e
Binary files /dev/null and b/metadata-integration/java/datahub-protobuf/src/test/resources/protobuf/messageB.protoc differ
diff --git a/metadata-integration/java/datahub-protobuf/src/test/resources/protobuf/messageC.proto b/metadata-integration/java/datahub-protobuf/src/test/resources/protobuf/messageC.proto
new file mode 100644
index 0000000000000..5a8f41a0540e1
--- /dev/null
+++ b/metadata-integration/java/datahub-protobuf/src/test/resources/protobuf/messageC.proto
@@ -0,0 +1,14 @@
+syntax = "proto3";
+package protobuf;
+
+/*
+   Test for one of
+*/
+message MessageC {
+  // one of field comment
+  oneof one_of_field {
+    string one_of_string = 1; // one of string comment
+    int32 one_of_int = 2; // one of int comment
+  }
+  string normal = 4;
+}
\ No newline at end of file
diff --git a/metadata-integration/java/datahub-protobuf/src/test/resources/protobuf/messageC.protoc b/metadata-integration/java/datahub-protobuf/src/test/resources/protobuf/messageC.protoc
new file mode 100644
index 0000000000000..4cfbbd9498f9f
Binary files /dev/null and b/metadata-integration/java/datahub-protobuf/src/test/resources/protobuf/messageC.protoc differ
diff --git a/metadata-integration/java/datahub-protobuf/src/test/resources/protobuf/messageC2.proto b/metadata-integration/java/datahub-protobuf/src/test/resources/protobuf/messageC2.proto
new file mode 100644
index 0000000000000..532ea18ba362c
--- /dev/null
+++ b/metadata-integration/java/datahub-protobuf/src/test/resources/protobuf/messageC2.proto
@@ -0,0 +1,22 @@
+syntax = "proto3";
+package protobuf;
+
+/*
+   Test for one of nested array
+*/
+message MessageC1 {
+  MessageC2 messageList = 1;
+}
+
+message MessageC2 {
+  repeated MessageC3 list = 1;
+}
+
+message MessageC3 {
+  // one of field comment
+  oneof one_of_field {
+    string one_of_string = 1; // one of string comment
+    int32 one_of_int = 2; // one of int comment
+  }
+  string normal = 4;
+}
diff --git a/metadata-integration/java/datahub-protobuf/src/test/resources/protobuf/messageC2.protoc b/metadata-integration/java/datahub-protobuf/src/test/resources/protobuf/messageC2.protoc
new file mode 100644
index 0000000000000..d2988d2882a26
Binary files /dev/null and b/metadata-integration/java/datahub-protobuf/src/test/resources/protobuf/messageC2.protoc differ
diff --git a/metadata-integration/java/datahub-protobuf/src/test/resources/protobuf/messageEmpty.proto b/metadata-integration/java/datahub-protobuf/src/test/resources/protobuf/messageEmpty.proto
new file mode 100644
index 0000000000000..a8cb2fe9e80f6
--- /dev/null
+++ b/metadata-integration/java/datahub-protobuf/src/test/resources/protobuf/messageEmpty.proto
@@ -0,0 +1,6 @@
+syntax = "proto3";
+package protobuf;
+
+/*
+   Test for empty
+*/
diff --git a/metadata-integration/java/datahub-protobuf/src/test/resources/protobuf/messageEmpty.protoc b/metadata-integration/java/datahub-protobuf/src/test/resources/protobuf/messageEmpty.protoc
new file mode 100644
index 0000000000000..6de6ec602a2b8
Binary files /dev/null and b/metadata-integration/java/datahub-protobuf/src/test/resources/protobuf/messageEmpty.protoc differ
diff --git a/metadata-integration/java/spark-lineage/README.md b/metadata-integration/java/spark-lineage/README.md
index 61933a94605b7..404cec40a28d7 100644
--- a/metadata-integration/java/spark-lineage/README.md
+++ b/metadata-integration/java/spark-lineage/README.md
@@ -47,14 +47,19 @@ spark = SparkSession.builder()
         .enableHiveSupport()
         .getOrCreate();
  ```
- 
-### Enable https and authentication token
-Add below config in spark config
 
-```
-spark.datahub.rest.server                   https://<server URL>
-spark.datahub.rest.token                    <token>
-```
+### Configuration details
+
+| Field                                           | Required | Default | Description                                                             |
+|-------------------------------------------------|----------|---------|-------------------------------------------------------------------------|
+| spark.jars.packages                              | ✅        |         | Set with latest/required version  io.acryl:datahub-spark-lineage:0.8.23 |
+| spark.extraListeners                             | ✅        |         | datahub.spark.DatahubSparkListener                                      |
+| spark.datahub.rest.server                        | ✅        |         | Datahub server url  eg:http://localhost:8080                            |
+| spark.datahub.rest.token                         |          |         | Authentication token.                         |
+| spark.datahub.metadata.pipeline.platformInstance|          |         | Pipeline level platform instance                                        |
+| spark.datahub.metadata.dataset.platformInstance|          |         | dataset level platform instance                                        |
+| spark.datahub.metadata.dataset.env              |          | PROD    | [Supported values](https://datahubproject.io/docs/graphql/enums#fabrictype). In all other cases, will fallback to PROD           |
+
 
 ## What to Expect: The Metadata Model
 
diff --git a/metadata-integration/java/spark-lineage/build.gradle b/metadata-integration/java/spark-lineage/build.gradle
index 1ce6916c5f326..0c26476303d44 100644
--- a/metadata-integration/java/spark-lineage/build.gradle
+++ b/metadata-integration/java/spark-lineage/build.gradle
@@ -88,6 +88,7 @@ dependencies {
   annotationProcessor externalDependency.lombok
   
   implementation externalDependency.typesafeConfig
+  implementation externalDependency.opentracingJdbc
 
   implementation project(path: ':metadata-integration:java:datahub-client', configuration: 'shadow')
 
@@ -146,6 +147,7 @@ shadowJar {
   relocate 'org.apache.commons.codec', 'datahub.spark2.shaded.o.a.c.codec'
   relocate 'mozilla', 'datahub.spark2.shaded.mozilla'
   relocate 'com.typesafe','datahub.spark2.shaded.typesafe'
+  relocate 'io.opentracing','datahub.spark2.shaded.io.opentracing'
   finalizedBy checkShadowJar
 }
 
diff --git a/metadata-integration/java/spark-lineage/src/main/java/datahub/spark/DatahubSparkListener.java b/metadata-integration/java/spark-lineage/src/main/java/datahub/spark/DatahubSparkListener.java
index f0b066d1d22c3..9d693f3926443 100644
--- a/metadata-integration/java/spark-lineage/src/main/java/datahub/spark/DatahubSparkListener.java
+++ b/metadata-integration/java/spark-lineage/src/main/java/datahub/spark/DatahubSparkListener.java
@@ -1,20 +1,9 @@
 package datahub.spark;
 
-import com.google.common.base.Splitter;
-import datahub.spark.model.LineageUtils;
-import datahub.spark.model.AppEndEvent;
-import datahub.spark.model.AppStartEvent;
-import datahub.spark.model.DatasetLineage;
-import datahub.spark.model.LineageConsumer;
-import datahub.spark.model.SQLQueryExecEndEvent;
-import datahub.spark.model.SQLQueryExecStartEvent;
-import datahub.spark.model.dataset.SparkDataset;
-import datahub.spark.consumer.impl.McpEmitter;
 import java.io.IOException;
 import java.io.PrintWriter;
 import java.io.StringWriter;
 import java.util.ArrayList;
-import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
@@ -26,7 +15,7 @@
 import java.util.concurrent.Executors;
 import java.util.stream.Collectors;
 import java.util.stream.StreamSupport;
-import lombok.extern.slf4j.Slf4j;
+
 import org.apache.spark.SparkConf;
 import org.apache.spark.SparkContext;
 import org.apache.spark.SparkEnv;
@@ -41,13 +30,25 @@
 import org.apache.spark.sql.execution.SQLExecution;
 import org.apache.spark.sql.execution.ui.SparkListenerSQLExecutionEnd;
 import org.apache.spark.sql.execution.ui.SparkListenerSQLExecutionStart;
+
+import com.google.common.base.Splitter;
+import com.google.common.util.concurrent.ThreadFactoryBuilder;
+import com.typesafe.config.Config;
+
+import datahub.spark.consumer.impl.McpEmitter;
+import datahub.spark.model.AppEndEvent;
+import datahub.spark.model.AppStartEvent;
+import datahub.spark.model.DatasetLineage;
+import datahub.spark.model.LineageConsumer;
+import datahub.spark.model.LineageUtils;
+import datahub.spark.model.SQLQueryExecEndEvent;
+import datahub.spark.model.SQLQueryExecStartEvent;
+import datahub.spark.model.dataset.SparkDataset;
+import lombok.extern.slf4j.Slf4j;
 import scala.collection.JavaConversions;
 import scala.runtime.AbstractFunction1;
 import scala.runtime.AbstractPartialFunction;
 
-import com.typesafe.config.Config;
-import com.typesafe.config.ConfigFactory;
-
 
 
 @Slf4j
@@ -56,11 +57,15 @@ public class DatahubSparkListener extends SparkListener {
   private static final int THREAD_CNT = 16;
   public static final String CONSUMER_TYPE_KEY = "spark.datahub.lineage.consumerTypes";
   public static final String DATAHUB_EMITTER = "mcpEmitter";
-
+  public static final String  DATABRICKS_CLUSTER_KEY = "databricks.cluster";
+  public static final String PIPELINE_KEY = "metadata.pipeline"; 
+  public static final String PIPELINE_PLATFORM_INSTANCE_KEY = PIPELINE_KEY + ".platformInstance";
+  
   private final Map<String, AppStartEvent> appDetails = new ConcurrentHashMap<>();
   private final Map<String, Map<Long, SQLQueryExecStartEvent>> appSqlDetails = new ConcurrentHashMap<>();
   private final Map<String, ExecutorService> appPoolDetails = new ConcurrentHashMap<>();
   private final Map<String, McpEmitter> appEmitters = new ConcurrentHashMap<>();
+  private final Map<String, Config> appConfig = new ConcurrentHashMap<>();
   
   public DatahubSparkListener() {
     log.info("DatahubSparkListener initialised.");
@@ -80,17 +85,16 @@ public SqlStartTask(SparkListenerSQLExecutionStart sqlStart, LogicalPlan plan, S
 
     @Override
     public void run() {
-      appSqlDetails.get(ctx.appName())
+      appSqlDetails.get(ctx.applicationId())
           .put(sqlStart.executionId(),
-              new SQLQueryExecStartEvent(ctx.conf().get("spark.master"), ctx.appName(), ctx.applicationId(),
+              new SQLQueryExecStartEvent(ctx.conf().get("spark.master"), getPipelineName(ctx), ctx.applicationId(),
                   sqlStart.time(), sqlStart.executionId(), null));
-      log.debug("PLAN for execution id: " + ctx.appName() + ":" + sqlStart.executionId() + "\n");
+      log.debug("PLAN for execution id: " + getPipelineName(ctx) + ":" + sqlStart.executionId() + "\n");
       log.debug(plan.toString());
 
-      DatasetExtractor extractor = new DatasetExtractor();
-      Optional<? extends SparkDataset> outputDS = extractor.asDataset(plan, ctx, true);
+      Optional<? extends SparkDataset> outputDS = DatasetExtractor.asDataset(plan, ctx, true);
       if (!outputDS.isPresent()) {
-        log.debug("Skipping execution as no output dataset present for execution id: " + ctx.appName() + ":"
+        log.debug("Skipping execution as no output dataset present for execution id: " + ctx.applicationId() + ":"
             + sqlStart.executionId());
         return;
       }
@@ -103,7 +107,7 @@ public void run() {
         @Override
         public Void apply(LogicalPlan plan) {
           log.debug("CHILD " + plan.getClass() + "\n" + plan + "\n-------------\n");
-          Optional<? extends SparkDataset> inputDS = extractor.asDataset(plan, ctx, false);
+          Optional<? extends SparkDataset> inputDS = DatasetExtractor.asDataset(plan, ctx, false);
           inputDS.ifPresent(x -> lineage.addSource(x));
           allInners.addAll(JavaConversions.asJavaCollection(plan.innerChildren()));
           return null;
@@ -126,7 +130,7 @@ public boolean isDefinedAt(LogicalPlan x) {
           @Override
           public Void apply(LogicalPlan plan) {
             log.debug("INNER CHILD " + plan.getClass() + "\n" + plan + "\n-------------\n");
-            Optional<? extends SparkDataset> inputDS = extractor.asDataset(plan, ctx, false);
+            Optional<? extends SparkDataset> inputDS = DatasetExtractor.asDataset(plan, ctx, false);
             inputDS.ifPresent(
                 x -> log.debug("source added for " + ctx.appName() + "/" + sqlStart.executionId() + ": " + x));
             inputDS.ifPresent(x -> lineage.addSource(x));
@@ -141,12 +145,12 @@ public boolean isDefinedAt(LogicalPlan x) {
       }
 
       SQLQueryExecStartEvent evt =
-          new SQLQueryExecStartEvent(ctx.conf().get("spark.master"), ctx.appName(), ctx.applicationId(),
+          new SQLQueryExecStartEvent(ctx.conf().get("spark.master"), getPipelineName(ctx), ctx.applicationId(),
               sqlStart.time(), sqlStart.executionId(), lineage);
 
-      appSqlDetails.get(ctx.appName()).put(sqlStart.executionId(), evt);
+      appSqlDetails.get(ctx.applicationId()).put(sqlStart.executionId(), evt);
 
-      McpEmitter emitter = appEmitters.get(ctx.appName());
+      McpEmitter emitter = appEmitters.get(ctx.applicationId());
       if (emitter != null) {
         emitter.accept(evt);
       }
@@ -157,14 +161,6 @@ public boolean isDefinedAt(LogicalPlan x) {
     }
   }
   
-  private Config parseSparkConfig() {
-    SparkConf conf = SparkEnv.get().conf();
-    String propertiesString = Arrays.stream(conf.getAllWithPrefix("spark.datahub."))
-            .map(tup -> tup._1 + "= \"" + tup._2 + "\"")
-            .collect(Collectors.joining("\n"));
-    return ConfigFactory.parseString(propertiesString);
-  }
-
   @Override
   public void onApplicationStart(SparkListenerApplicationStart applicationStart) {
     try {
@@ -173,18 +169,7 @@ public void onApplicationStart(SparkListenerApplicationStart applicationStart) {
 
         @Override
         public Void apply(SparkContext sc) {
-          String appId = applicationStart.appId().isDefined() ? applicationStart.appId().get() : "";
-          AppStartEvent evt =
-              new AppStartEvent(LineageUtils.getMaster(sc), applicationStart.appName(), appId, applicationStart.time(),
-                  applicationStart.sparkUser());
-          Config datahubConf = parseSparkConfig();
-          appEmitters.computeIfAbsent(applicationStart.appName(), s -> new McpEmitter(datahubConf)).accept(evt);
-          consumers().forEach(c -> c.accept(evt));
-
-          appDetails.put(applicationStart.appName(), evt);
-          appSqlDetails.put(applicationStart.appName(), new ConcurrentHashMap<>());
-          ExecutorService pool = Executors.newFixedThreadPool(THREAD_CNT);
-          appPoolDetails.put(applicationStart.appName(), pool);
+          getOrCreateApplicationSetup(sc);
           return null;
         }
       });
@@ -207,21 +192,21 @@ public void onApplicationEnd(SparkListenerApplicationEnd applicationEnd) {
         @Override
         public Void apply(SparkContext sc) {
           log.info("Application ended : {} {}", sc.appName(), sc.applicationId());
-          AppStartEvent start = appDetails.remove(sc.appName());
-          appPoolDetails.remove(sc.appName()).shutdown();
-          appSqlDetails.remove(sc.appName());
+          AppStartEvent start = appDetails.remove(sc.applicationId());
+          appPoolDetails.remove(sc.applicationId()).shutdown();
+          appSqlDetails.remove(sc.applicationId());
           if (start == null) {
             log.error("Application end event received, but start event missing for appId " + sc.applicationId());
           } else {
-            AppEndEvent evt = new AppEndEvent(LineageUtils.getMaster(sc), sc.appName(), sc.applicationId(),
+            AppEndEvent evt = new AppEndEvent(LineageUtils.getMaster(sc), getPipelineName(sc), sc.applicationId(),
                 applicationEnd.time(), start);
 
-            McpEmitter emitter = appEmitters.get(sc.appName());
+            McpEmitter emitter = appEmitters.get(sc.applicationId());
             if (emitter != null) {
               emitter.accept(evt);
               try {
                 emitter.close();
-                appEmitters.remove(sc.appName());
+                appEmitters.remove(sc.applicationId());
               } catch (Exception e) {
                 log.warn("Failed to close underlying emitter due to {}", e.getMessage());
               }
@@ -276,7 +261,7 @@ public void processExecutionEnd(SparkListenerSQLExecutionEnd sqlEnd) {
 
       @Override
       public Void apply(SparkContext sc) {
-        SQLQueryExecStartEvent start = appSqlDetails.get(sc.appName()).remove(sqlEnd.executionId());
+        SQLQueryExecStartEvent start = appSqlDetails.get(sc.applicationId()).remove(sqlEnd.executionId());
         if (start == null) {
           log.error(
               "Execution end event received, but start event missing for appId/sql exec Id " + sc.applicationId() + ":"
@@ -285,7 +270,7 @@ public Void apply(SparkContext sc) {
           SQLQueryExecEndEvent evt =
               new SQLQueryExecEndEvent(LineageUtils.getMaster(sc), sc.appName(), sc.applicationId(), sqlEnd.time(),
                   sqlEnd.executionId(), start);
-          McpEmitter emitter = appEmitters.get(sc.appName());
+          McpEmitter emitter = appEmitters.get(sc.applicationId());
           if (emitter != null) {
             emitter.accept(evt);
           }
@@ -294,8 +279,50 @@ public Void apply(SparkContext sc) {
       }
     });
   }
+  
+  private synchronized ExecutorService getOrCreateApplicationSetup(SparkContext ctx) {
+
+    ExecutorService pool = null;
+    String appId = ctx.applicationId();
+    Config datahubConfig = appConfig.get(appId);
+    if (datahubConfig == null) {
+      Config datahubConf = LineageUtils.parseSparkConfig();
+      appConfig.put(appId, datahubConf);
+      Config pipelineConfig = datahubConf.hasPath(PIPELINE_KEY) ? datahubConf.getConfig(PIPELINE_KEY) : com.typesafe.config.ConfigFactory.empty(); 
+      AppStartEvent evt = new AppStartEvent(LineageUtils.getMaster(ctx), getPipelineName(ctx), appId, ctx.startTime(),
+          ctx.sparkUser(), pipelineConfig);
+      
+      appEmitters.computeIfAbsent(appId, s -> new McpEmitter(datahubConf)).accept(evt);
+      consumers().forEach(c -> c.accept(evt));
+      appDetails.put(appId, evt);
+      appSqlDetails.put(appId, new ConcurrentHashMap<>());
+      pool = Executors.newFixedThreadPool(THREAD_CNT,
+          new ThreadFactoryBuilder().setNameFormat("datahub-emit-pool").build());
+      appPoolDetails.put(appId, pool);
+      log.debug("Execution thread pool initialised for {}", appId);
+    } else {
+      pool = appPoolDetails.get(appId);
+    }
+
+    return pool;
+
+  }
+
+  private String getPipelineName(SparkContext cx) {
+    Config datahubConfig = appConfig.computeIfAbsent(cx.applicationId(), s -> LineageUtils.parseSparkConfig());
+    String name = "";
+    if (datahubConfig.hasPath(DATABRICKS_CLUSTER_KEY)) {
+      name = datahubConfig.getString(DATABRICKS_CLUSTER_KEY) + "_" + cx.applicationId();
+    }
+    name = cx.appName();
+    //TODO: appending of platform instance needs to be done at central location like adding constructor to dataflowurl
+    if (datahubConfig.hasPath(PIPELINE_PLATFORM_INSTANCE_KEY)) {
+      name = datahubConfig.getString(PIPELINE_PLATFORM_INSTANCE_KEY) + "." + name;
+    }
+    return name;
+  }
+
 
-  // TODO sqlEvt.details() unused
   private void processExecution(SparkListenerSQLExecutionStart sqlStart) {
     QueryExecution queryExec = SQLExecution.getQueryExecution(sqlStart.executionId());
     if (queryExec == null) {
@@ -306,7 +333,7 @@ private void processExecution(SparkListenerSQLExecutionStart sqlStart) {
     LogicalPlan plan = queryExec.optimizedPlan();
     SparkSession sess = queryExec.sparkSession();
     SparkContext ctx = sess.sparkContext();
-    ExecutorService pool = appPoolDetails.get(ctx.appName());
+    ExecutorService pool = getOrCreateApplicationSetup(ctx);
     pool.execute(new SqlStartTask(sqlStart, plan, ctx));
   }
   private List<LineageConsumer> consumers() {
@@ -320,4 +347,4 @@ private List<LineageConsumer> consumers() {
       }
 
     }
-}
+}
\ No newline at end of file
diff --git a/metadata-integration/java/spark-lineage/src/main/java/datahub/spark/DatasetExtractor.java b/metadata-integration/java/spark-lineage/src/main/java/datahub/spark/DatasetExtractor.java
index 0cf801421918f..95787343b03e6 100644
--- a/metadata-integration/java/spark-lineage/src/main/java/datahub/spark/DatasetExtractor.java
+++ b/metadata-integration/java/spark-lineage/src/main/java/datahub/spark/DatasetExtractor.java
@@ -1,9 +1,12 @@
 package datahub.spark;
 
+import datahub.spark.model.LineageUtils;
 import datahub.spark.model.dataset.CatalogTableDataset;
 import datahub.spark.model.dataset.HdfsPathDataset;
 import datahub.spark.model.dataset.JdbcDataset;
 import datahub.spark.model.dataset.SparkDataset;
+import lombok.extern.slf4j.Slf4j;
+
 import java.io.IOException;
 import java.util.HashMap;
 import java.util.List;
@@ -29,55 +32,55 @@
 import org.apache.spark.sql.sources.BaseRelation;
 
 import com.google.common.collect.ImmutableSet;
+import com.linkedin.common.FabricType;
 
 import scala.Option;
 import scala.collection.JavaConversions;
 
+import com.typesafe.config.Config;
+
+@Slf4j
 public class DatasetExtractor {
+  
   private static final Map<Class<? extends LogicalPlan>, PlanToDataset> PLAN_TO_DATASET = new HashMap<>();
   private static final Map<Class<? extends BaseRelation>, RelationToDataset> REL_TO_DATASET = new HashMap<>();
-  private static final Set<Class<? extends LogicalPlan>> OUTPUT_CMD = ImmutableSet
-      .of(InsertIntoHadoopFsRelationCommand.class, SaveIntoDataSourceCommand.class,
-          CreateDataSourceTableAsSelectCommand.class, CreateHiveTableAsSelectCommand.class,
-          InsertIntoHiveTable.class);
+  private static final Set<Class<? extends LogicalPlan>> OUTPUT_CMD = ImmutableSet.of(
+      InsertIntoHadoopFsRelationCommand.class, SaveIntoDataSourceCommand.class,
+      CreateDataSourceTableAsSelectCommand.class, CreateHiveTableAsSelectCommand.class, InsertIntoHiveTable.class);
+  private static final String DATASET_ENV_KEY = "metadata.dataset.env";
+  private static final String DATASET_PLATFORM_INSTANCE_KEY = "metadata.dataset.platformInstance";
   // TODO InsertIntoHiveDirCommand, InsertIntoDataSourceDirCommand
 
+ private DatasetExtractor() {
+   
+ }
   private static interface PlanToDataset {
-    Optional<? extends SparkDataset> fromPlanNode(LogicalPlan plan, SparkContext ctx);
+    Optional<? extends SparkDataset> fromPlanNode(LogicalPlan plan, SparkContext ctx, Config datahubConfig);
   }
 
   private static interface RelationToDataset {
-    Optional<? extends SparkDataset> fromRelation(BaseRelation rel, SparkContext ctx);
+    Optional<? extends SparkDataset> fromRelation(BaseRelation rel, SparkContext ctx, Config datahubConfig);
   }
 
   static {
-    PLAN_TO_DATASET.put(InsertIntoHadoopFsRelationCommand.class, (p, ctx) -> {
+    PLAN_TO_DATASET.put(InsertIntoHadoopFsRelationCommand.class, (p, ctx, datahubConfig) -> {
       InsertIntoHadoopFsRelationCommand cmd = (InsertIntoHadoopFsRelationCommand) p;
       if (cmd.catalogTable().isDefined()) {
-        return Optional.of(new CatalogTableDataset(cmd.catalogTable().get()));
+        return Optional.of(new CatalogTableDataset(cmd.catalogTable().get(), getCommonPlatformInstance(datahubConfig), getCommonFabricType(datahubConfig)));
       }
-      return Optional.of(new HdfsPathDataset(cmd.outputPath()));
+      return Optional.of(new HdfsPathDataset(cmd.outputPath(), getCommonPlatformInstance(datahubConfig), getCommonFabricType(datahubConfig)));
     });
 
-    PLAN_TO_DATASET.put(LogicalRelation.class, (p, ctx) -> {
+    PLAN_TO_DATASET.put(LogicalRelation.class, (p, ctx, datahubConfig) -> {
       BaseRelation baseRel = ((LogicalRelation) p).relation();
       if (!REL_TO_DATASET.containsKey(baseRel.getClass())) {
         return Optional.empty();
       }
-      return REL_TO_DATASET.get(baseRel.getClass()).fromRelation(baseRel, ctx);
+      return REL_TO_DATASET.get(baseRel.getClass()).fromRelation(baseRel, ctx, datahubConfig);
     });
 
-    PLAN_TO_DATASET.put(SaveIntoDataSourceCommand.class, (p, ctx) -> {
-      /*
-       * BaseRelation relation; if (((SaveIntoDataSourceCommand) p).dataSource()
-       * instanceof RelationProvider) { RelationProvider relProvider =
-       * (RelationProvider) ((SaveIntoDataSourceCommand) p).dataSource(); relation =
-       * relProvider.createRelation(ctx, ((SaveIntoDataSourceCommand) p).options()); }
-       * else { SchemaRelationProvider relProvider = (SchemaRelationProvider)
-       * ((SaveIntoDataSourceCommand) p).dataSource(); relation =
-       * p.createRelation(ctx, ((SaveIntoDataSourceCommand) p).options(), p.schema());
-       * }
-       */
+    PLAN_TO_DATASET.put(SaveIntoDataSourceCommand.class, (p, ctx, datahubConfig) -> {
+
       SaveIntoDataSourceCommand cmd = (SaveIntoDataSourceCommand) p;
 
       Map<String, String> options = JavaConversions.mapAsJavaMap(cmd.options());
@@ -87,49 +90,47 @@ private static interface RelationToDataset {
       }
 
       String tbl = options.get("dbtable");
-      return Optional.of(new JdbcDataset(url, tbl));
+      return Optional.of(new JdbcDataset(url, tbl, getCommonPlatformInstance(datahubConfig), getCommonFabricType(datahubConfig)));
     });
 
-    PLAN_TO_DATASET.put(CreateDataSourceTableAsSelectCommand.class, (p, ctx) -> {
+    PLAN_TO_DATASET.put(CreateDataSourceTableAsSelectCommand.class, (p, ctx, datahubConfig) -> {
       CreateDataSourceTableAsSelectCommand cmd = (CreateDataSourceTableAsSelectCommand) p;
       // TODO what of cmd.mode()
-      return Optional.of(new CatalogTableDataset(cmd.table()));
+      return Optional.of(new CatalogTableDataset(cmd.table(), getCommonPlatformInstance(datahubConfig), getCommonFabricType(datahubConfig)));
     });
-    PLAN_TO_DATASET.put(CreateHiveTableAsSelectCommand.class, (p, ctx) -> {
+    PLAN_TO_DATASET.put(CreateHiveTableAsSelectCommand.class, (p, ctx, datahubConfig) -> {
       CreateHiveTableAsSelectCommand cmd = (CreateHiveTableAsSelectCommand) p;
-      return Optional.of(new CatalogTableDataset(cmd.tableDesc()));
+      return Optional.of(new CatalogTableDataset(cmd.tableDesc(), getCommonPlatformInstance(datahubConfig), getCommonFabricType(datahubConfig)));
     });
-    PLAN_TO_DATASET.put(InsertIntoHiveTable.class, (p, ctx) -> {
+    PLAN_TO_DATASET.put(InsertIntoHiveTable.class, (p, ctx, datahubConfig) -> {
       InsertIntoHiveTable cmd = (InsertIntoHiveTable) p;
-      return Optional.of(new CatalogTableDataset(cmd.table()));
+      return Optional.of(new CatalogTableDataset(cmd.table(), getCommonPlatformInstance(datahubConfig), getCommonFabricType(datahubConfig)));
     });
 
-    PLAN_TO_DATASET.put(HiveTableRelation.class, (p, ctx) -> {
+    PLAN_TO_DATASET.put(HiveTableRelation.class, (p, ctx, datahubConfig) -> {
       HiveTableRelation cmd = (HiveTableRelation) p;
-      return Optional.of(new CatalogTableDataset(cmd.tableMeta()));
+      return Optional.of(new CatalogTableDataset(cmd.tableMeta(), getCommonPlatformInstance(datahubConfig), getCommonFabricType(datahubConfig)));
     });
 
-    REL_TO_DATASET.put(HadoopFsRelation.class, (r, ctx) -> {
+    REL_TO_DATASET.put(HadoopFsRelation.class, (r, ctx, datahubConfig) -> {
       List<Path> res = JavaConversions.asJavaCollection(((HadoopFsRelation) r).location().rootPaths()).stream()
-          .map(p -> getDirectoryPath(p, ctx.hadoopConfiguration()))
-          .distinct()
-          .collect(Collectors.toList());
+          .map(p -> getDirectoryPath(p, ctx.hadoopConfiguration())).distinct().collect(Collectors.toList());
 
       // TODO mapping to URN TBD
-      return Optional.of(new HdfsPathDataset(res.get(0)));
+      return Optional.of(new HdfsPathDataset(res.get(0), getCommonPlatformInstance(datahubConfig), getCommonFabricType(datahubConfig)));
     });
-    REL_TO_DATASET.put(JDBCRelation.class, (r, ctx) -> {
+    REL_TO_DATASET.put(JDBCRelation.class, (r, ctx, datahubConfig) -> {
       JDBCRelation rel = (JDBCRelation) r;
       Option<String> tbl = rel.jdbcOptions().parameters().get(JDBCOptions.JDBC_TABLE_NAME());
       if (tbl.isEmpty()) {
         return Optional.empty();
       }
 
-      return Optional.of(new JdbcDataset(rel.jdbcOptions().url(), tbl.get()));
+      return Optional.of(new JdbcDataset(rel.jdbcOptions().url(), tbl.get(), getCommonPlatformInstance(datahubConfig), getCommonFabricType(datahubConfig)));
     });
   }
 
-  Optional<? extends SparkDataset> asDataset(LogicalPlan logicalPlan, SparkContext ctx, boolean outputNode) {
+  static Optional<? extends SparkDataset> asDataset(LogicalPlan logicalPlan, SparkContext ctx, boolean outputNode) {
     if (!outputNode && OUTPUT_CMD.contains(logicalPlan.getClass())) {
       return Optional.empty();
     }
@@ -137,8 +138,8 @@ Optional<? extends SparkDataset> asDataset(LogicalPlan logicalPlan, SparkContext
     if (!PLAN_TO_DATASET.containsKey(logicalPlan.getClass())) {
       return Optional.empty();
     }
-
-    return PLAN_TO_DATASET.get(logicalPlan.getClass()).fromPlanNode(logicalPlan, ctx);
+    Config datahubconfig = LineageUtils.parseSparkConfig();
+    return PLAN_TO_DATASET.get(logicalPlan.getClass()).fromPlanNode(logicalPlan, ctx, datahubconfig);
   }
 
   private static Path getDirectoryPath(Path p, Configuration hadoopConf) {
@@ -149,8 +150,26 @@ private static Path getDirectoryPath(Path p, Configuration hadoopConf) {
         return p;
       }
     } catch (IOException e) {
-      // log.warn("Unable to get file system for path ", e);
       return p;
     }
   }
+
+  private static FabricType getCommonFabricType(Config datahubConfig) {
+    String fabricTypeString = datahubConfig.hasPath(DATASET_ENV_KEY)
+        ? datahubConfig.getString(DATASET_ENV_KEY).toUpperCase()
+        : "PROD";
+    FabricType fabricType = null;
+    try {
+      fabricType = FabricType.valueOf(fabricTypeString);
+    } catch (IllegalArgumentException e) {
+      log.warn("Invalid env ({}). Setting env to default PROD", fabricTypeString);
+      fabricType = FabricType.PROD;
+    }
+    return fabricType;
+  }
+
+  private static String getCommonPlatformInstance(Config datahubConfig) {
+    return datahubConfig.hasPath(DATASET_PLATFORM_INSTANCE_KEY) ? datahubConfig.getString(DATASET_PLATFORM_INSTANCE_KEY)
+        : null;
+  }
 }
diff --git a/metadata-integration/java/spark-lineage/src/main/java/datahub/spark/model/AppStartEvent.java b/metadata-integration/java/spark-lineage/src/main/java/datahub/spark/model/AppStartEvent.java
index 8e6967e381e89..529e7ebae22be 100644
--- a/metadata-integration/java/spark-lineage/src/main/java/datahub/spark/model/AppStartEvent.java
+++ b/metadata-integration/java/spark-lineage/src/main/java/datahub/spark/model/AppStartEvent.java
@@ -1,34 +1,64 @@
 package datahub.spark.model;
 
+import java.io.PrintWriter;
+import java.io.StringWriter;
+import java.net.URISyntaxException;
+import java.util.ArrayList;
+import java.util.List;
+
+import com.linkedin.common.DataPlatformInstance;
 import com.linkedin.common.urn.DataFlowUrn;
+import com.linkedin.common.urn.DataPlatformUrn;
 import com.linkedin.data.template.StringMap;
 import com.linkedin.datajob.DataFlowInfo;
+import com.typesafe.config.Config;
+
 import datahub.event.MetadataChangeProposalWrapper;
-import java.util.Collections;
-import java.util.List;
 import lombok.Getter;
 import lombok.ToString;
+import lombok.extern.slf4j.Slf4j;
 
-
+@Slf4j
 @ToString
 @Getter
 public class AppStartEvent extends LineageEvent {
 
+  private static final String PLATFORM_INSTANCE_KEY = "platformInstance";
+  private static final String PLATFORM_SPARK = "spark";
   private final String sparkUser;
+  private Config pipelineConfig;
 
-  public AppStartEvent(String master, String appName, String appId, long time, String sparkUser) {
+  public AppStartEvent(String master, String appName, String appId, long time, String sparkUser,
+      Config pipelineConfig) {
     super(master, appName, appId, time);
     this.sparkUser = sparkUser;
+    this.pipelineConfig = pipelineConfig;
   }
 
   @Override
   public List<MetadataChangeProposalWrapper> asMetadataEvents() {
     DataFlowUrn flowUrn = LineageUtils.flowUrn(getMaster(), getAppName());
+    ArrayList<MetadataChangeProposalWrapper> mcps = new ArrayList<MetadataChangeProposalWrapper>();
 
+    if (this.pipelineConfig.hasPath(PLATFORM_INSTANCE_KEY)) {
+      try {
+        DataPlatformInstance dpi = new DataPlatformInstance().setPlatform(new DataPlatformUrn(PLATFORM_SPARK)).setInstance(
+            LineageUtils.dataPlatformInstanceUrn(PLATFORM_SPARK, this.pipelineConfig.getString(PLATFORM_INSTANCE_KEY)));
+        mcps.add(MetadataChangeProposalWrapper
+            .create(b -> b.entityType("dataFlow").entityUrn(flowUrn).upsert().aspect(dpi)));
+      } catch (URISyntaxException e) {
+        // log error, but don't impact thread
+        StringWriter s = new StringWriter();
+        PrintWriter p = new PrintWriter(s);
+        e.printStackTrace(p);
+        log.error(s.toString());
+        p.close();
+      }
+    }
     DataFlowInfo flowInfo = new DataFlowInfo().setName(getAppName()).setCustomProperties(customProps());
-
-    return Collections.singletonList(MetadataChangeProposalWrapper.create(
-        b -> b.entityType("dataFlow").entityUrn(flowUrn).upsert().aspect(flowInfo)));
+    mcps.add(MetadataChangeProposalWrapper
+        .create(b -> b.entityType("dataFlow").entityUrn(flowUrn).upsert().aspect(flowInfo)));
+    return mcps;
   }
 
   StringMap customProps() {
diff --git a/metadata-integration/java/spark-lineage/src/main/java/datahub/spark/model/LineageEvent.java b/metadata-integration/java/spark-lineage/src/main/java/datahub/spark/model/LineageEvent.java
index 238bf61e71167..f1b3d55f638f7 100644
--- a/metadata-integration/java/spark-lineage/src/main/java/datahub/spark/model/LineageEvent.java
+++ b/metadata-integration/java/spark-lineage/src/main/java/datahub/spark/model/LineageEvent.java
@@ -1,10 +1,10 @@
 package datahub.spark.model;
 
-import datahub.event.MetadataChangeProposalWrapper;
 import java.util.Date;
 import java.util.List;
-import lombok.Data;
 
+import datahub.event.MetadataChangeProposalWrapper;
+import lombok.Data;
 
 @Data
 public abstract class LineageEvent {
diff --git a/metadata-integration/java/spark-lineage/src/main/java/datahub/spark/model/LineageUtils.java b/metadata-integration/java/spark-lineage/src/main/java/datahub/spark/model/LineageUtils.java
index 3f5d3da4d351d..ad837f034ad64 100644
--- a/metadata-integration/java/spark-lineage/src/main/java/datahub/spark/model/LineageUtils.java
+++ b/metadata-integration/java/spark-lineage/src/main/java/datahub/spark/model/LineageUtils.java
@@ -1,17 +1,31 @@
 package datahub.spark.model;
 
+import com.linkedin.common.FabricType;
 import com.linkedin.common.urn.DataFlowUrn;
+import com.linkedin.common.urn.DataPlatformUrn;
+import com.linkedin.common.urn.DatasetUrn;
+import com.linkedin.common.urn.TupleKey;
+import com.linkedin.common.urn.Urn;
+import com.typesafe.config.Config;
+import com.typesafe.config.ConfigFactory;
+
+import java.net.URISyntaxException;
+import java.util.Arrays;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.stream.Collectors;
+
 import lombok.extern.slf4j.Slf4j;
+
+import org.apache.spark.SparkConf;
 import org.apache.spark.SparkContext;
 import org.apache.spark.SparkContext$;
+import org.apache.spark.SparkEnv;
 import org.apache.spark.sql.SparkSession;
 import scala.Option;
 import scala.runtime.AbstractFunction0;
 import scala.runtime.AbstractFunction1;
 
-
 @Slf4j
 public class LineageUtils {
 
@@ -25,6 +39,11 @@ private LineageUtils() {
 
   }
 
+  public static Urn dataPlatformInstanceUrn(String platform, String instance) throws URISyntaxException {
+    return new Urn("dataPlatformInstance",
+        new TupleKey(Arrays.asList(new DataPlatformUrn(platform).toString(), instance)));
+  }
+
   public static DataFlowUrn flowUrn(String master, String appName) {
     return new DataFlowUrn("spark", appName, master.replaceAll(":", "_").replaceAll("/", "_").replaceAll("[_]+", "_"));
   }
@@ -58,31 +77,38 @@ public static LineageConsumer getConsumer(String consumerType) {
     return consumers.get(consumerType);
   }
 
-
-  /* This is for generating urn from a hash of the plan */
-  
-/*
-  public static String scrubPlan(String plan) {
-    String s = plan.replaceAll("#[0-9]*", "");
-    s = s.replaceAll("JdbcRelationProvider@[0-9a-zA-Z]*,", "JdbcRelationProvider,");
-    s = s.replaceAll("InMemoryFileIndex@[0-9a-zA-Z]*,", "InMemoryFileIndex,");
-    s = s.replaceAll("Created Time:[^\n]+\n", "");
-    s = s.replaceAll("Last Access:[^\n]+\n", "");
-    s = s.replaceAll("Owner:[^\n]+\n", "");
-    s = s.replaceAll("Statistics:[^\n]+\n", "");
-    s = s.replaceAll("Table Properties:[^\n]+\n", "");
-    // System.out.println("CLEAN: " + s);
-    return s;
+  public static Config parseSparkConfig() {
+    SparkConf conf = SparkEnv.get().conf();
+    String propertiesString = Arrays.stream(conf.getAllWithPrefix("spark.datahub."))
+        .map(tup -> tup._1 + "= \"" + tup._2 + "\"").collect(Collectors.joining("\n"));
+    return ConfigFactory.parseString(propertiesString);
   }
 
-  public static void setPathReplacer(Function<String, String> replacer) {
-    PATH_REPLACER = replacer;
-  }
-  
-  public static String hash(String s) {
-    s = PATH_REPLACER.apply(s);
-    log.debug("PATH REPLACED " + s);
-    return Hashing.md5().hashString(s, Charset.forName("US-ASCII")).toString();
+  // TODO: URN creation with platform instance needs to be inside DatasetUrn class
+  public static DatasetUrn createDatasetUrn(String platform, String platformInstance, String name,
+      FabricType fabricType) {
+    String datasteName = platformInstance == null ? name : platformInstance + "." + name;
+    return new DatasetUrn(new DataPlatformUrn(platform), datasteName, fabricType);
   }
-  */
+
+  /* This is for generating urn from a hash of the plan */
+
+  /*
+   * public static String scrubPlan(String plan) { String s =
+   * plan.replaceAll("#[0-9]*", ""); s =
+   * s.replaceAll("JdbcRelationProvider@[0-9a-zA-Z]*,", "JdbcRelationProvider,");
+   * s = s.replaceAll("InMemoryFileIndex@[0-9a-zA-Z]*,", "InMemoryFileIndex,"); s
+   * = s.replaceAll("Created Time:[^\n]+\n", ""); s =
+   * s.replaceAll("Last Access:[^\n]+\n", ""); s = s.replaceAll("Owner:[^\n]+\n",
+   * ""); s = s.replaceAll("Statistics:[^\n]+\n", ""); s =
+   * s.replaceAll("Table Properties:[^\n]+\n", ""); //
+   * System.out.println("CLEAN: " + s); return s; }
+   * 
+   * public static void setPathReplacer(Function<String, String> replacer) {
+   * PATH_REPLACER = replacer; }
+   * 
+   * public static String hash(String s) { s = PATH_REPLACER.apply(s);
+   * log.debug("PATH REPLACED " + s); return Hashing.md5().hashString(s,
+   * Charset.forName("US-ASCII")).toString(); }
+   */
 }
diff --git a/metadata-integration/java/spark-lineage/src/main/java/datahub/spark/model/dataset/CatalogTableDataset.java b/metadata-integration/java/spark-lineage/src/main/java/datahub/spark/model/dataset/CatalogTableDataset.java
index 0722fecd80ae3..a04985017a025 100644
--- a/metadata-integration/java/spark-lineage/src/main/java/datahub/spark/model/dataset/CatalogTableDataset.java
+++ b/metadata-integration/java/spark-lineage/src/main/java/datahub/spark/model/dataset/CatalogTableDataset.java
@@ -3,27 +3,18 @@
 import org.apache.spark.sql.catalyst.catalog.CatalogTable;
 
 import com.linkedin.common.FabricType;
-import com.linkedin.common.urn.DataPlatformUrn;
-import com.linkedin.common.urn.DatasetUrn;
 
-import lombok.EqualsAndHashCode;
 import lombok.ToString;
 
-@EqualsAndHashCode
 @ToString
-public class CatalogTableDataset implements SparkDataset {
-  private final DatasetUrn urn;
+public class CatalogTableDataset extends SparkDataset {
 
-  public CatalogTableDataset(CatalogTable table) {
-    this(table.qualifiedName());
+  public CatalogTableDataset(CatalogTable table, String platformInstance, FabricType fabricType) {
+    this(table.qualifiedName(), platformInstance, fabricType);
   }
 
-  public CatalogTableDataset(String dsName) {
-    this.urn = new DatasetUrn(new DataPlatformUrn("hive"), dsName, FabricType.PROD);
+  public CatalogTableDataset(String dsName, String platformInstance, FabricType fabricType) {
+    super("hive", platformInstance, dsName, fabricType);
   }
 
-  @Override
-  public DatasetUrn urn() {
-    return this.urn;
-  }
 }
diff --git a/metadata-integration/java/spark-lineage/src/main/java/datahub/spark/model/dataset/HdfsPathDataset.java b/metadata-integration/java/spark-lineage/src/main/java/datahub/spark/model/dataset/HdfsPathDataset.java
index 775bbd7738584..307bbf2522573 100644
--- a/metadata-integration/java/spark-lineage/src/main/java/datahub/spark/model/dataset/HdfsPathDataset.java
+++ b/metadata-integration/java/spark-lineage/src/main/java/datahub/spark/model/dataset/HdfsPathDataset.java
@@ -3,29 +3,20 @@
 import org.apache.hadoop.fs.Path;
 
 import com.linkedin.common.FabricType;
-import com.linkedin.common.urn.DataPlatformUrn;
-import com.linkedin.common.urn.DatasetUrn;
 
-import lombok.EqualsAndHashCode;
 import lombok.ToString;
 
-@EqualsAndHashCode
 @ToString
-public class HdfsPathDataset implements SparkDataset {
-  private final DatasetUrn urn;
-
-  public HdfsPathDataset(Path path) {
+public class HdfsPathDataset extends SparkDataset {
+  
+  public HdfsPathDataset(Path path, String platformInstance, FabricType fabricType) {
     // TODO check static partitions?
-    this(path.toUri().toString());
+    this(path.toUri().toString(), platformInstance, fabricType);
   }
 
-  public HdfsPathDataset(String pathUri) {
+  public HdfsPathDataset(String pathUri, String platformInstance, FabricType fabricType) {
     // TODO check static partitions?
-    this.urn = new DatasetUrn(new DataPlatformUrn("hdfs"), pathUri, FabricType.PROD);
+    super("hdfs", platformInstance, pathUri, fabricType);
   }
 
-  @Override
-  public DatasetUrn urn() {
-    return this.urn;
-  }
 }
diff --git a/metadata-integration/java/spark-lineage/src/main/java/datahub/spark/model/dataset/JdbcDataset.java b/metadata-integration/java/spark-lineage/src/main/java/datahub/spark/model/dataset/JdbcDataset.java
index 20d8df495249e..ea156b49fada7 100644
--- a/metadata-integration/java/spark-lineage/src/main/java/datahub/spark/model/dataset/JdbcDataset.java
+++ b/metadata-integration/java/spark-lineage/src/main/java/datahub/spark/model/dataset/JdbcDataset.java
@@ -1,40 +1,31 @@
 package datahub.spark.model.dataset;
 
+import java.util.HashMap;
+import java.util.Map;
+
 import com.linkedin.common.FabricType;
-import com.linkedin.common.urn.DataPlatformUrn;
-import com.linkedin.common.urn.DatasetUrn;
 
-import lombok.EqualsAndHashCode;
+import io.opentracing.contrib.jdbc.parser.URLParser;
 import lombok.ToString;
 
-@EqualsAndHashCode
 @ToString
-public class JdbcDataset implements SparkDataset {
-  private final DatasetUrn urn;
-
-  public JdbcDataset(String url, String tbl) {
-    this.urn = new DatasetUrn(new DataPlatformUrn(platformName(url)), dsName(url, tbl), FabricType.PROD);
+public class JdbcDataset extends SparkDataset {
+  //TODO: Should map to the central location on datahub for platform names
+  private static final Map<String, String> PLATFORM_NAME_MAPPING = new HashMap<>();
+  static {
+    PLATFORM_NAME_MAPPING.put("postgresql", "postgres");
   }
 
-  @Override
-  public DatasetUrn urn() {
-    return this.urn;
+  public JdbcDataset(String url, String tbl, String platformInstance, FabricType fabricType) {
+    super(platformName(url), platformInstance, dsName(url, tbl), fabricType);
   }
 
   private static String platformName(String url) {
-    if (url.contains("postgres")) {
-      return "postgres";
-    }
-    return "unknownJdbc";
+    String dbType = URLParser.parse(url).getDbType();
+    return PLATFORM_NAME_MAPPING.getOrDefault(dbType, dbType);
   }
 
   private static String dsName(String url, String tbl) {
-    url = url.replaceFirst("jdbc:", "");
-    if (url.contains("postgres")) {
-      url = url.substring(url.lastIndexOf('/') + 1);
-      url = url.substring(0, url.indexOf('?'));
-    }
-    // TODO different DBs have different formats. TBD mapping to data source names
-    return url + "." + tbl;
+    return URLParser.parse(url).getDbInstance() + "." + tbl;
   }
 }
diff --git a/metadata-integration/java/spark-lineage/src/main/java/datahub/spark/model/dataset/SparkDataset.java b/metadata-integration/java/spark-lineage/src/main/java/datahub/spark/model/dataset/SparkDataset.java
index cdc30d1d54e54..546b737576e60 100644
--- a/metadata-integration/java/spark-lineage/src/main/java/datahub/spark/model/dataset/SparkDataset.java
+++ b/metadata-integration/java/spark-lineage/src/main/java/datahub/spark/model/dataset/SparkDataset.java
@@ -1,7 +1,22 @@
 package datahub.spark.model.dataset;
 
+import com.linkedin.common.FabricType;
 import com.linkedin.common.urn.DatasetUrn;
 
-public interface SparkDataset {
-  DatasetUrn urn();
+import datahub.spark.model.LineageUtils;
+import lombok.EqualsAndHashCode;
+
+@EqualsAndHashCode
+public abstract class SparkDataset {
+  
+  private DatasetUrn urn;
+  
+  public SparkDataset(String platform, String platformInstance, String name, FabricType fabricType) {
+    super();
+    this.urn = LineageUtils.createDatasetUrn(platform, platformInstance, name, fabricType);
+  }
+
+  public DatasetUrn urn() {
+    return urn;
+  }
 }
diff --git a/metadata-integration/java/spark-lineage/src/test/java/datahub/spark/TestSparkJobsLineage.java b/metadata-integration/java/spark-lineage/src/test/java/datahub/spark/TestSparkJobsLineage.java
index 45e9de18c6513..4fff88007c2b1 100644
--- a/metadata-integration/java/spark-lineage/src/test/java/datahub/spark/TestSparkJobsLineage.java
+++ b/metadata-integration/java/spark-lineage/src/test/java/datahub/spark/TestSparkJobsLineage.java
@@ -43,6 +43,8 @@
 import org.mockserver.verify.VerificationTimes;
 import org.testcontainers.containers.PostgreSQLContainer;
 
+import com.linkedin.common.FabricType;
+
 import static org.junit.Assert.*;
 import static org.mockserver.integration.ClientAndServer.*;
 import static org.mockserver.model.HttpRequest.*;
@@ -69,6 +71,11 @@ public class TestSparkJobsLineage {
   private static final int GMS_PORT = MOCK_GMS ? MOCK_PORT : 8080;
 
   private static final String EXPECTED_JSON_ROOT = "src/test/resources/expected/";
+  
+  private static final FabricType DATASET_ENV = FabricType.DEV;
+  private static final String PIPELINE_PLATFORM_INSTANCE = "test_machine";
+  private static final String DATASET_PLATFORM_INSTANCE = "test_dev_dataset";
+  
   @ClassRule
   public static PostgreSQLContainer<?> db =
       new PostgreSQLContainer<>("postgres:9.6.12").withDatabaseName("sparktestdb");
@@ -145,6 +152,9 @@ public static void setup() {
         .config("spark.extraListeners", "datahub.spark.DatahubSparkListener")
         .config("spark.datahub.lineage.consumerTypes", "accumulator")
         .config("spark.datahub.rest.server", "http://localhost:" + mockServer.getPort())
+        .config("spark.datahub.metadata.pipeline.platformInstance", PIPELINE_PLATFORM_INSTANCE)
+        .config("spark.datahub.metadata.dataset.platformInstance", DATASET_PLATFORM_INSTANCE)
+        .config("spark.datahub.metadata.dataset.env", DATASET_ENV.name())
         .config("spark.sql.warehouse.dir", new File(WAREHOUSE_LOC).getAbsolutePath())
         .enableHiveSupport()
         .getOrCreate();
@@ -156,7 +166,7 @@ public static void setup() {
     jdbcConnnProperties.put("password", db.getPassword());
 
     if (VERIFY_EXPECTED) {
-      verify(1);
+      verify(2);
       clear();
     }
   }
@@ -200,15 +210,15 @@ private static DatasetLineage dsl(String callSite, SparkDataset sink, SparkDatas
   }
 
   private static HdfsPathDataset hdfsDs(String fileName) {
-    return new HdfsPathDataset("file:" + abs(DATA_DIR + "/" + fileName));
+    return new HdfsPathDataset("file:" + abs(DATA_DIR + "/" + fileName), DATASET_PLATFORM_INSTANCE, DATASET_ENV);
   }
 
   private static JdbcDataset pgDs(String tbl) {
-    return new JdbcDataset(db.getJdbcUrl(), tbl);
+    return new JdbcDataset(db.getJdbcUrl(), tbl, DATASET_PLATFORM_INSTANCE, DATASET_ENV);
   }
 
   private static CatalogTableDataset catTblDs(String tbl) {
-    return new CatalogTableDataset(tbl(tbl));
+    return new CatalogTableDataset(tbl(tbl), DATASET_PLATFORM_INSTANCE, DATASET_ENV);
   }
 
   private static String tbl(String tbl) {
diff --git a/metadata-integration/java/spark-lineage/src/test/resources/expected/testHdfsInHiveOut.json b/metadata-integration/java/spark-lineage/src/test/resources/expected/testHdfsInHiveOut.json
index b327969620aea..282b1d3a3c1ec 100644
--- a/metadata-integration/java/spark-lineage/src/test/resources/expected/testHdfsInHiveOut.json
+++ b/metadata-integration/java/spark-lineage/src/test/resources/expected/testHdfsInHiveOut.json
@@ -1,3 +1,3 @@
-{  "proposal" :{"aspectName":"dataJobInputOutput","entityUrn":"urn:li:dataJob:(urn:li:dataFlow:(spark,sparkTestApp,local),QueryExecId_9)","entityType":"dataJob","aspect":{"value":"{\"inputDatasets\":[\"urn:li:dataset:(urn:li:dataPlatform:hdfs,file:/src/test/resources/data/in1.csv,PROD)\",\"urn:li:dataset:(urn:li:dataPlatform:hdfs,file:/src/test/resources/data/in2.csv,PROD)\"],\"outputDatasets\":[\"urn:li:dataset:(urn:li:dataPlatform:hive,sparktestdb.foo4,PROD)\"]}","contentType":"application/json"},"changeType":"UPSERT"}}
-{  "proposal" :{"aspectName":"dataJobInputOutput","entityUrn":"urn:li:dataJob:(urn:li:dataFlow:(spark,sparkTestApp,local),QueryExecId_10)","entityType":"dataJob","aspect":{"value":"{\"inputDatasets\":[\"urn:li:dataset:(urn:li:dataPlatform:hdfs,file:/src/test/resources/data/in1.csv,PROD)\",\"urn:li:dataset:(urn:li:dataPlatform:hdfs,file:/src/test/resources/data/in2.csv,PROD)\"],\"outputDatasets\":[\"urn:li:dataset:(urn:li:dataPlatform:hive,sparktestdb.foo4,PROD)\"]}","contentType":"application/json"},"changeType":"UPSERT"}}
-{  "proposal" :{"aspectName":"dataJobInputOutput","entityUrn":"urn:li:dataJob:(urn:li:dataFlow:(spark,sparkTestApp,local),QueryExecId_11)","entityType":"dataJob","aspect":{"value":"{\"inputDatasets\":[\"urn:li:dataset:(urn:li:dataPlatform:hdfs,file:/src/test/resources/data/in1.csv,PROD)\",\"urn:li:dataset:(urn:li:dataPlatform:hdfs,file:/src/test/resources/data/in2.csv,PROD)\"],\"outputDatasets\":[\"urn:li:dataset:(urn:li:dataPlatform:hive,sparktestdb.foo4,PROD)\"]}","contentType":"application/json"},"changeType":"UPSERT"}}
+{  "proposal" :{"aspectName":"dataJobInputOutput","entityUrn":"urn:li:dataJob:(urn:li:dataFlow:(spark,test_machine.sparkTestApp,local),QueryExecId_9)","entityType":"dataJob","aspect":{"value":"{\"inputDatasets\":[\"urn:li:dataset:(urn:li:dataPlatform:hdfs,test_dev_dataset.file:/src/test/resources/data/in1.csv,DEV)\",\"urn:li:dataset:(urn:li:dataPlatform:hdfs,test_dev_dataset.file:/src/test/resources/data/in2.csv,DEV)\"],\"outputDatasets\":[\"urn:li:dataset:(urn:li:dataPlatform:hive,test_dev_dataset.sparktestdb.foo4,DEV)\"]}","contentType":"application/json"},"changeType":"UPSERT"}}
+{  "proposal" :{"aspectName":"dataJobInputOutput","entityUrn":"urn:li:dataJob:(urn:li:dataFlow:(spark,test_machine.sparkTestApp,local),QueryExecId_10)","entityType":"dataJob","aspect":{"value":"{\"inputDatasets\":[\"urn:li:dataset:(urn:li:dataPlatform:hdfs,test_dev_dataset.file:/src/test/resources/data/in1.csv,DEV)\",\"urn:li:dataset:(urn:li:dataPlatform:hdfs,test_dev_dataset.file:/src/test/resources/data/in2.csv,DEV)\"],\"outputDatasets\":[\"urn:li:dataset:(urn:li:dataPlatform:hive,test_dev_dataset.sparktestdb.foo4,DEV)\"]}","contentType":"application/json"},"changeType":"UPSERT"}}
+{  "proposal" :{"aspectName":"dataJobInputOutput","entityUrn":"urn:li:dataJob:(urn:li:dataFlow:(spark,test_machine.sparkTestApp,local),QueryExecId_11)","entityType":"dataJob","aspect":{"value":"{\"inputDatasets\":[\"urn:li:dataset:(urn:li:dataPlatform:hdfs,test_dev_dataset.file:/src/test/resources/data/in1.csv,DEV)\",\"urn:li:dataset:(urn:li:dataPlatform:hdfs,test_dev_dataset.file:/src/test/resources/data/in2.csv,DEV)\"],\"outputDatasets\":[\"urn:li:dataset:(urn:li:dataPlatform:hive,test_dev_dataset.sparktestdb.foo4,DEV)\"]}","contentType":"application/json"},"changeType":"UPSERT"}}
diff --git a/metadata-integration/java/spark-lineage/src/test/resources/expected/testHdfsInJdbcOut.json b/metadata-integration/java/spark-lineage/src/test/resources/expected/testHdfsInJdbcOut.json
index e89db8b2b8239..58938ab61713f 100644
--- a/metadata-integration/java/spark-lineage/src/test/resources/expected/testHdfsInJdbcOut.json
+++ b/metadata-integration/java/spark-lineage/src/test/resources/expected/testHdfsInJdbcOut.json
@@ -1 +1 @@
-{  "proposal" :{"aspectName":"dataJobInputOutput","entityUrn":"urn:li:dataJob:(urn:li:dataFlow:(spark,sparkTestApp,local),QueryExecId_24)","entityType":"dataJob","aspect":{"value":"{\"inputDatasets\":[\"urn:li:dataset:(urn:li:dataPlatform:hdfs,file:/src/test/resources/data/in1.csv,PROD)\",\"urn:li:dataset:(urn:li:dataPlatform:hdfs,file:/src/test/resources/data/in2.csv,PROD)\"],\"outputDatasets\":[\"urn:li:dataset:(urn:li:dataPlatform:postgres,sparktestdb.foo1,PROD)\"]}","contentType":"application/json"},"changeType":"UPSERT"}}
+{  "proposal" :{"aspectName":"dataJobInputOutput","entityUrn":"urn:li:dataJob:(urn:li:dataFlow:(spark,test_machine.sparkTestApp,local),QueryExecId_24)","entityType":"dataJob","aspect":{"value":"{\"inputDatasets\":[\"urn:li:dataset:(urn:li:dataPlatform:hdfs,test_dev_dataset.file:/src/test/resources/data/in1.csv,DEV)\",\"urn:li:dataset:(urn:li:dataPlatform:hdfs,test_dev_dataset.file:/src/test/resources/data/in2.csv,DEV)\"],\"outputDatasets\":[\"urn:li:dataset:(urn:li:dataPlatform:postgres,test_dev_dataset.sparktestdb.foo1,DEV)\"]}","contentType":"application/json"},"changeType":"UPSERT"}}
diff --git a/metadata-integration/java/spark-lineage/src/test/resources/expected/testHdfsInOut.json b/metadata-integration/java/spark-lineage/src/test/resources/expected/testHdfsInOut.json
index 3da8ec9002878..ee1bf2394c979 100644
--- a/metadata-integration/java/spark-lineage/src/test/resources/expected/testHdfsInOut.json
+++ b/metadata-integration/java/spark-lineage/src/test/resources/expected/testHdfsInOut.json
@@ -1 +1 @@
-{  "proposal" :{"aspectName":"dataJobInputOutput","entityUrn":"urn:li:dataJob:(urn:li:dataFlow:(spark,sparkTestApp,local),QueryExecId_6)","entityType":"dataJob","aspect":{"value":"{\"inputDatasets\":[\"urn:li:dataset:(urn:li:dataPlatform:hdfs,file:/src/test/resources/data/in1.csv,PROD)\",\"urn:li:dataset:(urn:li:dataPlatform:hdfs,file:/src/test/resources/data/in2.csv,PROD)\"],\"outputDatasets\":[\"urn:li:dataset:(urn:li:dataPlatform:hdfs,file:/src/test/resources/data/out.csv,PROD)\"]}","contentType":"application/json"},"changeType":"UPSERT"}}
+{  "proposal" :{"aspectName":"dataJobInputOutput","entityUrn":"urn:li:dataJob:(urn:li:dataFlow:(spark,test_machine.sparkTestApp,local),QueryExecId_6)","entityType":"dataJob","aspect":{"value":"{\"inputDatasets\":[\"urn:li:dataset:(urn:li:dataPlatform:hdfs,test_dev_dataset.file:/src/test/resources/data/in1.csv,DEV)\",\"urn:li:dataset:(urn:li:dataPlatform:hdfs,test_dev_dataset.file:/src/test/resources/data/in2.csv,DEV)\"],\"outputDatasets\":[\"urn:li:dataset:(urn:li:dataPlatform:hdfs,test_dev_dataset.file:/src/test/resources/data/out.csv,DEV)\"]}","contentType":"application/json"},"changeType":"UPSERT"}}
diff --git a/metadata-integration/java/spark-lineage/src/test/resources/expected/testHdfsJdbcInJdbcOut.json b/metadata-integration/java/spark-lineage/src/test/resources/expected/testHdfsJdbcInJdbcOut.json
index 06ea771467070..b3ff74dd216ec 100644
--- a/metadata-integration/java/spark-lineage/src/test/resources/expected/testHdfsJdbcInJdbcOut.json
+++ b/metadata-integration/java/spark-lineage/src/test/resources/expected/testHdfsJdbcInJdbcOut.json
@@ -1 +1 @@
-{  "proposal" :{"aspectName":"dataJobInputOutput","entityUrn":"urn:li:dataJob:(urn:li:dataFlow:(spark,sparkTestApp,local),QueryExecId_13)","entityType":"dataJob","aspect":{"value":"{\"inputDatasets\":[\"urn:li:dataset:(urn:li:dataPlatform:hdfs,file:/src/test/resources/data/in1.csv,PROD)\",\"urn:li:dataset:(urn:li:dataPlatform:postgres,sparktestdb.foo2,PROD)\"],\"outputDatasets\":[\"urn:li:dataset:(urn:li:dataPlatform:postgres,sparktestdb.foo3,PROD)\"]}","contentType":"application/json"},"changeType":"UPSERT"}}
+{  "proposal" :{"aspectName":"dataJobInputOutput","entityUrn":"urn:li:dataJob:(urn:li:dataFlow:(spark,test_machine.sparkTestApp,local),QueryExecId_13)","entityType":"dataJob","aspect":{"value":"{\"inputDatasets\":[\"urn:li:dataset:(urn:li:dataPlatform:hdfs,test_dev_dataset.file:/src/test/resources/data/in1.csv,DEV)\",\"urn:li:dataset:(urn:li:dataPlatform:postgres,test_dev_dataset.sparktestdb.foo2,DEV)\"],\"outputDatasets\":[\"urn:li:dataset:(urn:li:dataPlatform:postgres,test_dev_dataset.sparktestdb.foo3,DEV)\"]}","contentType":"application/json"},"changeType":"UPSERT"}}
diff --git a/metadata-integration/java/spark-lineage/src/test/resources/expected/testHdfsJdbcInJdbcOutTwoLevel.json b/metadata-integration/java/spark-lineage/src/test/resources/expected/testHdfsJdbcInJdbcOutTwoLevel.json
index 763b3319a0511..11cc041d83ae4 100644
--- a/metadata-integration/java/spark-lineage/src/test/resources/expected/testHdfsJdbcInJdbcOutTwoLevel.json
+++ b/metadata-integration/java/spark-lineage/src/test/resources/expected/testHdfsJdbcInJdbcOutTwoLevel.json
@@ -1 +1 @@
-{  "proposal" :{"aspectName":"dataJobInputOutput","entityUrn":"urn:li:dataJob:(urn:li:dataFlow:(spark,sparkTestApp,local),QueryExecId_27)","entityType":"dataJob","aspect":{"value":"{\"inputDatasets\":[\"urn:li:dataset:(urn:li:dataPlatform:hdfs,file:/src/test/resources/data/in1.csv,PROD)\",\"urn:li:dataset:(urn:li:dataPlatform:hdfs,file:/src/test/resources/data/in2.csv,PROD)\",\"urn:li:dataset:(urn:li:dataPlatform:postgres,sparktestdb.foo6,PROD)\"],\"outputDatasets\":[\"urn:li:dataset:(urn:li:dataPlatform:postgres,sparktestdb.foo7,PROD)\"]}","contentType":"application/json"},"changeType":"UPSERT"}}
+{  "proposal" :{"aspectName":"dataJobInputOutput","entityUrn":"urn:li:dataJob:(urn:li:dataFlow:(spark,test_machine.sparkTestApp,local),QueryExecId_27)","entityType":"dataJob","aspect":{"value":"{\"inputDatasets\":[\"urn:li:dataset:(urn:li:dataPlatform:hdfs,test_dev_dataset.file:/src/test/resources/data/in1.csv,DEV)\",\"urn:li:dataset:(urn:li:dataPlatform:hdfs,test_dev_dataset.file:/src/test/resources/data/in2.csv,DEV)\",\"urn:li:dataset:(urn:li:dataPlatform:postgres,test_dev_dataset.sparktestdb.foo6,DEV)\"],\"outputDatasets\":[\"urn:li:dataset:(urn:li:dataPlatform:postgres,test_dev_dataset.sparktestdb.foo7,DEV)\"]}","contentType":"application/json"},"changeType":"UPSERT"}}
diff --git a/metadata-integration/java/spark-lineage/src/test/resources/expected/testHiveInHiveOut.json b/metadata-integration/java/spark-lineage/src/test/resources/expected/testHiveInHiveOut.json
index 9cd7268a20381..2c4da82bdc3ae 100644
--- a/metadata-integration/java/spark-lineage/src/test/resources/expected/testHiveInHiveOut.json
+++ b/metadata-integration/java/spark-lineage/src/test/resources/expected/testHiveInHiveOut.json
@@ -1,4 +1,4 @@
-{  "proposal" :{"aspectName":"dataJobInputOutput","entityUrn":"urn:li:dataJob:(urn:li:dataFlow:(spark,sparkTestApp,local),QueryExecId_18)","entityType":"dataJob","aspect":{"value":"{\"inputDatasets\":[\"urn:li:dataset:(urn:li:dataPlatform:hdfs,file:/src/test/resources/data/in1.csv,PROD)\",\"urn:li:dataset:(urn:li:dataPlatform:hdfs,file:/src/test/resources/data/in2.csv,PROD)\"],\"outputDatasets\":[\"urn:li:dataset:(urn:li:dataPlatform:hive,sparktestdb.foo5,PROD)\"]}","contentType":"application/json"},"changeType":"UPSERT"}}
-{  "proposal" :{"aspectName":"dataJobInputOutput","entityUrn":"urn:li:dataJob:(urn:li:dataFlow:(spark,sparkTestApp,local),QueryExecId_19)","entityType":"dataJob","aspect":{"value":"{\"inputDatasets\":[\"urn:li:dataset:(urn:li:dataPlatform:hive,sparktestdb.foo5,PROD)\"],\"outputDatasets\":[\"urn:li:dataset:(urn:li:dataPlatform:hive,sparktestdb.hivetab,PROD)\"]}","contentType":"application/json"},"changeType":"UPSERT"}}
-{  "proposal" :{"aspectName":"dataJobInputOutput","entityUrn":"urn:li:dataJob:(urn:li:dataFlow:(spark,sparkTestApp,local),QueryExecId_20)","entityType":"dataJob","aspect":{"value":"{\"inputDatasets\":[\"urn:li:dataset:(urn:li:dataPlatform:hive,sparktestdb.foo5,PROD)\"],\"outputDatasets\":[\"urn:li:dataset:(urn:li:dataPlatform:hive,sparktestdb.hivetab,PROD)\"]}","contentType":"application/json"},"changeType":"UPSERT"}}
-{  "proposal" :{"aspectName":"dataJobInputOutput","entityUrn":"urn:li:dataJob:(urn:li:dataFlow:(spark,sparkTestApp,local),QueryExecId_21)","entityType":"dataJob","aspect":{"value":"{\"inputDatasets\":[\"urn:li:dataset:(urn:li:dataPlatform:hive,sparktestdb.foo5,PROD)\"],\"outputDatasets\":[\"urn:li:dataset:(urn:li:dataPlatform:hive,sparktestdb.hivetab,PROD)\"]}","contentType":"application/json"},"changeType":"UPSERT"}}
+{  "proposal" :{"aspectName":"dataJobInputOutput","entityUrn":"urn:li:dataJob:(urn:li:dataFlow:(spark,test_machine.sparkTestApp,local),QueryExecId_18)","entityType":"dataJob","aspect":{"value":"{\"inputDatasets\":[\"urn:li:dataset:(urn:li:dataPlatform:hdfs,test_dev_dataset.file:/src/test/resources/data/in1.csv,DEV)\",\"urn:li:dataset:(urn:li:dataPlatform:hdfs,test_dev_dataset.file:/src/test/resources/data/in2.csv,DEV)\"],\"outputDatasets\":[\"urn:li:dataset:(urn:li:dataPlatform:hive,test_dev_dataset.sparktestdb.foo5,DEV)\"]}","contentType":"application/json"},"changeType":"UPSERT"}}
+{  "proposal" :{"aspectName":"dataJobInputOutput","entityUrn":"urn:li:dataJob:(urn:li:dataFlow:(spark,test_machine.sparkTestApp,local),QueryExecId_19)","entityType":"dataJob","aspect":{"value":"{\"inputDatasets\":[\"urn:li:dataset:(urn:li:dataPlatform:hive,test_dev_dataset.sparktestdb.foo5,DEV)\"],\"outputDatasets\":[\"urn:li:dataset:(urn:li:dataPlatform:hive,test_dev_dataset.sparktestdb.hivetab,DEV)\"]}","contentType":"application/json"},"changeType":"UPSERT"}}
+{  "proposal" :{"aspectName":"dataJobInputOutput","entityUrn":"urn:li:dataJob:(urn:li:dataFlow:(spark,test_machine.sparkTestApp,local),QueryExecId_20)","entityType":"dataJob","aspect":{"value":"{\"inputDatasets\":[\"urn:li:dataset:(urn:li:dataPlatform:hive,test_dev_dataset.sparktestdb.foo5,DEV)\"],\"outputDatasets\":[\"urn:li:dataset:(urn:li:dataPlatform:hive,test_dev_dataset.sparktestdb.hivetab,DEV)\"]}","contentType":"application/json"},"changeType":"UPSERT"}}
+{  "proposal" :{"aspectName":"dataJobInputOutput","entityUrn":"urn:li:dataJob:(urn:li:dataFlow:(spark,test_machine.sparkTestApp,local),QueryExecId_21)","entityType":"dataJob","aspect":{"value":"{\"inputDatasets\":[\"urn:li:dataset:(urn:li:dataPlatform:hive,test_dev_dataset.sparktestdb.foo5,DEV)\"],\"outputDatasets\":[\"urn:li:dataset:(urn:li:dataPlatform:hive,test_dev_dataset.sparktestdb.hivetab,DEV)\"]}","contentType":"application/json"},"changeType":"UPSERT"}}
diff --git a/metadata-io/src/main/java/com/linkedin/metadata/config/IngestionConfiguration.java b/metadata-io/src/main/java/com/linkedin/metadata/config/IngestionConfiguration.java
index 9c85c4398ad0b..5b10b59ff0c20 100644
--- a/metadata-io/src/main/java/com/linkedin/metadata/config/IngestionConfiguration.java
+++ b/metadata-io/src/main/java/com/linkedin/metadata/config/IngestionConfiguration.java
@@ -10,10 +10,10 @@ public class IngestionConfiguration {
   /**
    * Whether managed ingestion is enabled
    */
-  private boolean enabled;
+  public boolean enabled;
 
   /**
    * The default CLI version to use in managed ingestion
    */
-  private String defaultCliVersion;
+  public String defaultCliVersion;
 }
\ No newline at end of file
diff --git a/metadata-io/src/main/java/com/linkedin/metadata/datahubusage/DataHubUsageEventConstants.java b/metadata-io/src/main/java/com/linkedin/metadata/datahubusage/DataHubUsageEventConstants.java
index 2a132137a1336..b622fc5bb6af2 100644
--- a/metadata-io/src/main/java/com/linkedin/metadata/datahubusage/DataHubUsageEventConstants.java
+++ b/metadata-io/src/main/java/com/linkedin/metadata/datahubusage/DataHubUsageEventConstants.java
@@ -12,4 +12,5 @@ private DataHubUsageEventConstants() {
   // Event specific fields
   public static final String ENTITY_URN = "entityUrn";
   public static final String ENTITY_TYPE = "entityType";
+  public static final String QUERY = "query";
 }
diff --git a/metadata-io/src/main/java/com/linkedin/metadata/entity/EntityService.java b/metadata-io/src/main/java/com/linkedin/metadata/entity/EntityService.java
index 270a9f32a0a3f..2f45deb1d63cb 100644
--- a/metadata-io/src/main/java/com/linkedin/metadata/entity/EntityService.java
+++ b/metadata-io/src/main/java/com/linkedin/metadata/entity/EntityService.java
@@ -7,6 +7,7 @@
 import com.google.common.collect.Streams;
 import com.linkedin.common.AuditStamp;
 import com.linkedin.common.BrowsePaths;
+import com.linkedin.common.Status;
 import com.linkedin.common.urn.Urn;
 import com.linkedin.data.schema.RecordDataSchema;
 import com.linkedin.data.schema.TyperefDataSchema;
@@ -104,6 +105,7 @@ public abstract class EntityService {
   public static final String DEFAULT_RUN_ID = "no-run-id-provided";
   public static final String BROWSE_PATHS = "browsePaths";
   public static final String DATA_PLATFORM_INSTANCE = "dataPlatformInstance";
+  public static final String STATUS = "status";
 
   protected EntityService(@Nonnull final EntityEventProducer producer, @Nonnull final EntityRegistry entityRegistry) {
     _producer = producer;
@@ -658,6 +660,11 @@ public List<Pair<String, RecordTemplate>> generateDefaultAspectsIfMissing(@Nonnu
       aspectsToGet.add(DATA_PLATFORM_INSTANCE);
     }
 
+    boolean shouldHaveStatusSet = isAspectProvided(entityType, STATUS, includedAspects);
+    if (shouldHaveStatusSet) {
+      aspectsToGet.add(STATUS);
+    }
+
     List<Pair<String, RecordTemplate>> aspects = new ArrayList<>();
     final String keyAspectName = getKeyAspectName(urn);
     aspectsToGet.add(keyAspectName);
@@ -686,6 +693,12 @@ public List<Pair<String, RecordTemplate>> generateDefaultAspectsIfMissing(@Nonnu
           .ifPresent(aspect -> aspects.add(Pair.of(DATA_PLATFORM_INSTANCE, aspect)));
     }
 
+    if (shouldHaveStatusSet && latestAspects.get(STATUS) != null) {
+      Status status = new Status();
+      status.setRemoved(false);
+      aspects.add(Pair.of(STATUS, status));
+    }
+
     return aspects;
   }
 
@@ -821,12 +834,12 @@ public Set<String> getEntityAspectNames(final String entityName) {
 
   public abstract void setWritable(boolean canWrite);
 
-  public RollbackRunResult rollbackRun(List<AspectRowSummary> aspectRows, String runId) {
-    return rollbackWithConditions(aspectRows, Collections.singletonMap("runId", runId));
+  public RollbackRunResult rollbackRun(List<AspectRowSummary> aspectRows, String runId, boolean hardDelete) {
+    return rollbackWithConditions(aspectRows, Collections.singletonMap("runId", runId), hardDelete);
   }
 
   public abstract RollbackRunResult rollbackWithConditions(List<AspectRowSummary> aspectRows,
-      Map<String, String> conditions);
+                                                           Map<String, String> conditions, boolean hardDelete);
 
   public abstract RollbackRunResult deleteUrn(Urn urn);
 
diff --git a/metadata-io/src/main/java/com/linkedin/metadata/entity/EntityUtils.java b/metadata-io/src/main/java/com/linkedin/metadata/entity/EntityUtils.java
new file mode 100644
index 0000000000000..66a6e59ede4fa
--- /dev/null
+++ b/metadata-io/src/main/java/com/linkedin/metadata/entity/EntityUtils.java
@@ -0,0 +1,31 @@
+package com.linkedin.metadata.entity;
+
+import com.linkedin.common.Status;
+import com.linkedin.common.urn.Urn;
+import com.linkedin.entity.EnvelopedAspect;
+import lombok.extern.slf4j.Slf4j;
+
+
+@Slf4j
+public class EntityUtils {
+  private EntityUtils() {
+  }
+
+  /**
+   * Check if entity is removed (removed=true in Status aspect)
+   */
+  public static boolean checkIfRemoved(EntityService entityService, Urn entityUrn) {
+    try {
+      EnvelopedAspect statusAspect =
+          entityService.getLatestEnvelopedAspect(entityUrn.getEntityType(), entityUrn, "status");
+      if (statusAspect == null) {
+        return false;
+      }
+      Status status = new Status(statusAspect.getValue().data());
+      return status.isRemoved();
+    } catch (Exception e) {
+      log.error("Error while checking if {} is removed", entityUrn, e);
+      return false;
+    }
+  }
+}
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 880099708d679..2c28e5016e717 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
@@ -19,6 +19,7 @@
 import io.ebean.RawSql;
 import io.ebean.RawSqlBuilder;
 import io.ebean.Transaction;
+import io.ebean.TxScope;
 import io.ebean.annotation.TxIsolation;
 import java.net.URISyntaxException;
 import java.sql.Timestamp;
@@ -443,7 +444,7 @@ public <T> T runInTransactionWithRetry(@Nonnull final Supplier<T> block, final i
 
     T result = null;
     do {
-      try (Transaction transaction = _server.beginTransaction(TxIsolation.REPEATABLE_READ)) {
+      try (Transaction transaction = _server.beginTransaction(TxScope.requiresNew().setIsolation(TxIsolation.REPEATABLE_READ))) {
         transaction.setBatchMode(true);
         result = block.get();
         transaction.commit();
diff --git a/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/EbeanEntityService.java b/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/EbeanEntityService.java
index 4a41ece95013e..d878147f3c75d 100644
--- a/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/EbeanEntityService.java
+++ b/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/EbeanEntityService.java
@@ -6,6 +6,7 @@
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Iterators;
 import com.linkedin.common.AuditStamp;
+import com.linkedin.common.Status;
 import com.linkedin.common.UrnArray;
 import com.linkedin.common.urn.Urn;
 import com.linkedin.data.schema.RecordDataSchema;
@@ -14,6 +15,7 @@
 import com.linkedin.data.template.RecordTemplate;
 import com.linkedin.entity.EnvelopedAspect;
 import com.linkedin.events.metadata.ChangeType;
+import com.linkedin.metadata.Constants;
 import com.linkedin.metadata.aspect.Aspect;
 import com.linkedin.metadata.aspect.VersionedAspect;
 import com.linkedin.metadata.entity.EntityService;
@@ -26,11 +28,17 @@
 import com.linkedin.metadata.models.registry.EntityRegistry;
 import com.linkedin.metadata.query.ListUrnsResult;
 import com.linkedin.metadata.run.AspectRowSummary;
+import com.linkedin.metadata.utils.AuditStampUtils;
 import com.linkedin.metadata.utils.EntityKeyUtils;
 import com.linkedin.metadata.utils.PegasusUtils;
+import com.linkedin.metadata.utils.SystemMetadataUtils;
 import com.linkedin.mxe.MetadataAuditOperation;
 import com.linkedin.mxe.SystemMetadata;
 import com.linkedin.util.Pair;
+import lombok.extern.slf4j.Slf4j;
+
+import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
 import java.net.URISyntaxException;
 import java.sql.Timestamp;
 import java.util.ArrayList;
@@ -45,13 +53,11 @@
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.Function;
 import java.util.stream.Collectors;
-import javax.annotation.Nonnull;
-import javax.annotation.Nullable;
-import lombok.extern.slf4j.Slf4j;
 
-import static com.linkedin.metadata.Constants.*;
+import static com.linkedin.metadata.Constants.ASPECT_LATEST_VERSION;
+import static com.linkedin.metadata.Constants.SYSTEM_ACTOR;
 import static com.linkedin.metadata.entity.ebean.EbeanUtils.*;
-import static com.linkedin.metadata.utils.PegasusUtils.*;
+import static com.linkedin.metadata.utils.PegasusUtils.urnToEntityName;
 
 
 /**
@@ -72,8 +78,9 @@ public EbeanEntityService(@Nonnull final EbeanAspectDao entityDao, @Nonnull fina
     _entityDao = entityDao;
   }
 
-  @Nonnull Map<String, EbeanAspectV2> getLatestEbeanAspectForUrn(@Nonnull final Urn urn,
-    @Nonnull final Set<String> aspectNames) {
+  @Nonnull
+  Map<String, EbeanAspectV2> getLatestEbeanAspectForUrn(@Nonnull final Urn urn,
+                                                        @Nonnull final Set<String> aspectNames) {
       Set<Urn> urns = new HashSet<>();
       urns.add(urn);
 
@@ -469,11 +476,14 @@ private boolean filterMatch(SystemMetadata systemMetadata, Map<String, String> c
   }
 
   @Nullable
-  public RollbackResult deleteAspect(String urn, String aspectName, Map<String, String> conditions) {
+  public RollbackResult deleteAspect(String urn, String aspectName, Map<String, String> conditions, boolean hardDelete) {
     // Validate pre-conditions before running queries
+    Urn entityUrn;
+    EntitySpec entitySpec;
     try {
-      String entityName = PegasusUtils.urnToEntityName(Urn.createFromString(urn));
-      EntitySpec entitySpec = getEntityRegistry().getEntitySpec(entityName);
+      entityUrn = Urn.createFromString(urn);
+      String entityName = PegasusUtils.urnToEntityName(entityUrn);
+      entitySpec = getEntityRegistry().getEntitySpec(entityName);
       Preconditions.checkState(entitySpec != null, String.format("Could not find entity definition for %s", entityName));
       Preconditions.checkState(entitySpec.hasAspect(aspectName), String.format("Could not find aspect %s in definition for %s", aspectName, entityName));
     } catch (URISyntaxException uriSyntaxException) {
@@ -543,8 +553,18 @@ public RollbackResult deleteAspect(String urn, String aspectName, Map<String, St
         _entityDao.deleteAspect(survivingAspect);
       } else {
         if (isKeyAspect) {
-          // If this is the key aspect, delete the entity entirely.
-          additionalRowsDeleted = _entityDao.deleteUrn(urn);
+          if (hardDelete) {
+            // If this is the key aspect, delete the entity entirely.
+            additionalRowsDeleted = _entityDao.deleteUrn(urn);
+          } else if (entitySpec.hasAspect(Constants.STATUS_ASPECT_NAME)) {
+            // soft delete by setting status.removed=true (if applicable)
+            final Status statusAspect = new Status();
+            statusAspect.setRemoved(true);
+            final SystemMetadata systemMetadata = SystemMetadataUtils.createDefaultSystemMetadata();
+            final AuditStamp auditStamp = AuditStampUtils.createDefaultAuditStamp();
+
+            this.ingestAspect(entityUrn, Constants.STATUS_ASPECT_NAME, statusAspect, auditStamp, systemMetadata);
+          }
         } else {
           // Else, only delete the specific aspect.
           _entityDao.deleteAspect(latest);
@@ -562,6 +582,10 @@ public RollbackResult deleteAspect(String urn, String aspectName, Map<String, St
                 survivingAspect.getKey().getAspect(), previousMetadata, getEntityRegistry());
 
         final Urn urnObj = Urn.createFromString(urn);
+        // We are not deleting key aspect if hardDelete has not been set so do not return a rollback result
+        if (isKeyAspect && !hardDelete) {
+          return null;
+        }
         return new RollbackResult(urnObj, urnObj.getEntityType(), latest.getAspect(), latestValue,
             previousValue == null ? latestValue : previousValue, latestSystemMetadata,
             previousValue == null ? null : parseSystemMetadata(survivingAspect.getSystemMetadata()),
@@ -575,19 +599,19 @@ public RollbackResult deleteAspect(String urn, String aspectName, Map<String, St
   }
 
   @Override
-  public RollbackRunResult rollbackRun(List<AspectRowSummary> aspectRows, String runId) {
-    return rollbackWithConditions(aspectRows, Collections.singletonMap("runId", runId));
+  public RollbackRunResult rollbackRun(List<AspectRowSummary> aspectRows, String runId, boolean hardDelete) {
+    return rollbackWithConditions(aspectRows, Collections.singletonMap("runId", runId), hardDelete);
   }
 
   @Override
-  public RollbackRunResult rollbackWithConditions(List<AspectRowSummary> aspectRows, Map<String, String> conditions) {
+  public RollbackRunResult rollbackWithConditions(List<AspectRowSummary> aspectRows, Map<String, String> conditions, boolean hardDelete) {
     List<AspectRowSummary> removedAspects = new ArrayList<>();
     AtomicInteger rowsDeletedFromEntityDeletion = new AtomicInteger(0);
 
     aspectRows.forEach(aspectToRemove -> {
 
       RollbackResult result = deleteAspect(aspectToRemove.getUrn(), aspectToRemove.getAspectName(),
-          conditions);
+          conditions, hardDelete);
       if (result != null) {
         Optional<AspectSpec> aspectSpec = getAspectSpec(result.entityName, result.aspectName);
         if (!aspectSpec.isPresent()) {
@@ -621,7 +645,8 @@ public RollbackRunResult deleteUrn(Urn urn) {
     }
 
     SystemMetadata latestKeySystemMetadata = parseSystemMetadata(latestKey.getSystemMetadata());
-    RollbackResult result = deleteAspect(urn.toString(), keyAspectName, Collections.singletonMap("runId", latestKeySystemMetadata.getRunId()));
+    RollbackResult result = deleteAspect(urn.toString(), keyAspectName, Collections.singletonMap("runId",
+            latestKeySystemMetadata.getRunId()), true);
 
     if (result != null) {
       AspectRowSummary summary = new AspectRowSummary();
diff --git a/metadata-io/src/main/java/com/linkedin/metadata/graph/GraphService.java b/metadata-io/src/main/java/com/linkedin/metadata/graph/GraphService.java
index 124033e602c4f..b010995e0c4da 100644
--- a/metadata-io/src/main/java/com/linkedin/metadata/graph/GraphService.java
+++ b/metadata-io/src/main/java/com/linkedin/metadata/graph/GraphService.java
@@ -1,6 +1,7 @@
 package com.linkedin.metadata.graph;
 
 import com.linkedin.common.urn.Urn;
+import com.linkedin.metadata.models.registry.LineageRegistry;
 import com.linkedin.metadata.query.filter.Filter;
 import com.linkedin.metadata.query.filter.RelationshipDirection;
 import com.linkedin.metadata.query.filter.RelationshipFilter;
diff --git a/metadata-io/src/main/java/com/linkedin/metadata/graph/dgraph/DgraphGraphService.java b/metadata-io/src/main/java/com/linkedin/metadata/graph/dgraph/DgraphGraphService.java
index 7ee29d2f72e05..22b70156c51e8 100644
--- a/metadata-io/src/main/java/com/linkedin/metadata/graph/dgraph/DgraphGraphService.java
+++ b/metadata-io/src/main/java/com/linkedin/metadata/graph/dgraph/DgraphGraphService.java
@@ -6,7 +6,7 @@
 import com.linkedin.common.urn.Urn;
 import com.linkedin.metadata.graph.Edge;
 import com.linkedin.metadata.graph.GraphService;
-import com.linkedin.metadata.graph.LineageRegistry;
+import com.linkedin.metadata.models.registry.LineageRegistry;
 import com.linkedin.metadata.graph.RelatedEntitiesResult;
 import com.linkedin.metadata.graph.RelatedEntity;
 import com.linkedin.metadata.query.filter.Criterion;
diff --git a/metadata-io/src/main/java/com/linkedin/metadata/graph/elastic/ESGraphQueryDAO.java b/metadata-io/src/main/java/com/linkedin/metadata/graph/elastic/ESGraphQueryDAO.java
index cc2ab4fc8d079..bcba903ce1a03 100644
--- a/metadata-io/src/main/java/com/linkedin/metadata/graph/elastic/ESGraphQueryDAO.java
+++ b/metadata-io/src/main/java/com/linkedin/metadata/graph/elastic/ESGraphQueryDAO.java
@@ -7,8 +7,8 @@
 import com.linkedin.common.urn.Urn;
 import com.linkedin.common.urn.UrnUtils;
 import com.linkedin.metadata.graph.LineageDirection;
-import com.linkedin.metadata.graph.LineageRegistry;
-import com.linkedin.metadata.graph.LineageRegistry.EdgeInfo;
+import com.linkedin.metadata.models.registry.LineageRegistry;
+import com.linkedin.metadata.models.registry.LineageRegistry.EdgeInfo;
 import com.linkedin.metadata.graph.LineageRelationship;
 import com.linkedin.metadata.query.filter.Condition;
 import com.linkedin.metadata.query.filter.ConjunctiveCriterion;
diff --git a/metadata-io/src/main/java/com/linkedin/metadata/graph/elastic/ElasticSearchGraphService.java b/metadata-io/src/main/java/com/linkedin/metadata/graph/elastic/ElasticSearchGraphService.java
index bec46a9f66b17..2441ff83d8647 100644
--- a/metadata-io/src/main/java/com/linkedin/metadata/graph/elastic/ElasticSearchGraphService.java
+++ b/metadata-io/src/main/java/com/linkedin/metadata/graph/elastic/ElasticSearchGraphService.java
@@ -8,7 +8,7 @@
 import com.linkedin.metadata.graph.EntityLineageResult;
 import com.linkedin.metadata.graph.GraphService;
 import com.linkedin.metadata.graph.LineageDirection;
-import com.linkedin.metadata.graph.LineageRegistry;
+import com.linkedin.metadata.models.registry.LineageRegistry;
 import com.linkedin.metadata.graph.LineageRelationshipArray;
 import com.linkedin.metadata.graph.RelatedEntitiesResult;
 import com.linkedin.metadata.graph.RelatedEntity;
diff --git a/metadata-io/src/main/java/com/linkedin/metadata/graph/neo4j/Neo4jGraphService.java b/metadata-io/src/main/java/com/linkedin/metadata/graph/neo4j/Neo4jGraphService.java
index 49934950792f8..3476ddfa99dab 100644
--- a/metadata-io/src/main/java/com/linkedin/metadata/graph/neo4j/Neo4jGraphService.java
+++ b/metadata-io/src/main/java/com/linkedin/metadata/graph/neo4j/Neo4jGraphService.java
@@ -7,7 +7,7 @@
 import com.linkedin.common.urn.Urn;
 import com.linkedin.metadata.graph.Edge;
 import com.linkedin.metadata.graph.GraphService;
-import com.linkedin.metadata.graph.LineageRegistry;
+import com.linkedin.metadata.models.registry.LineageRegistry;
 import com.linkedin.metadata.graph.RelatedEntitiesResult;
 import com.linkedin.metadata.graph.RelatedEntity;
 import com.linkedin.metadata.query.filter.Condition;
diff --git a/metadata-io/src/main/java/com/linkedin/metadata/recommendation/candidatesource/MostPopularSource.java b/metadata-io/src/main/java/com/linkedin/metadata/recommendation/candidatesource/MostPopularSource.java
index c235dc19d6c7f..8222561757d68 100644
--- a/metadata-io/src/main/java/com/linkedin/metadata/recommendation/candidatesource/MostPopularSource.java
+++ b/metadata-io/src/main/java/com/linkedin/metadata/recommendation/candidatesource/MostPopularSource.java
@@ -3,8 +3,11 @@
 import com.codahale.metrics.Timer;
 import com.datahub.util.exception.ESQueryException;
 import com.linkedin.common.urn.Urn;
+import com.linkedin.common.urn.UrnUtils;
 import com.linkedin.metadata.datahubusage.DataHubUsageEventConstants;
 import com.linkedin.metadata.datahubusage.DataHubUsageEventType;
+import com.linkedin.metadata.entity.EntityService;
+import com.linkedin.metadata.entity.EntityUtils;
 import com.linkedin.metadata.recommendation.EntityProfileParams;
 import com.linkedin.metadata.recommendation.RecommendationContent;
 import com.linkedin.metadata.recommendation.RecommendationParams;
@@ -16,7 +19,6 @@
 import com.linkedin.metadata.utils.metrics.MetricUtils;
 import io.opentelemetry.extension.annotations.WithSpan;
 import java.io.IOException;
-import java.net.URISyntaxException;
 import java.util.List;
 import java.util.Optional;
 import java.util.stream.Collectors;
@@ -41,6 +43,7 @@
 public class MostPopularSource implements RecommendationSource {
   private final RestHighLevelClient _searchClient;
   private final IndexConvention _indexConvention;
+  private final EntityService _entityService;
 
   private static final String DATAHUB_USAGE_INDEX = "datahub_usage_event";
   private static final String ENTITY_AGG_NAME = "entity";
@@ -87,6 +90,7 @@ public List<RecommendationContent> getRecommendations(@Nonnull Urn userUrn,
           .map(bucket -> buildContent(bucket.getKeyAsString()))
           .filter(Optional::isPresent)
           .map(Optional::get)
+          .limit(MAX_CONTENT)
           .collect(Collectors.toList());
     } catch (Exception e) {
       log.error("Search query to get most popular entities failed", e);
@@ -106,7 +110,7 @@ private SearchRequest buildSearchRequest(@Nonnull Urn userUrn) {
     // Find the entities with the most views
     AggregationBuilder aggregation = AggregationBuilders.terms(ENTITY_AGG_NAME)
         .field(DataHubUsageEventConstants.ENTITY_URN + ESUtils.KEYWORD_SUFFIX)
-        .size(MAX_CONTENT);
+        .size(MAX_CONTENT * 2);
     source.aggregation(aggregation);
     source.size(0);
 
@@ -116,13 +120,11 @@ private SearchRequest buildSearchRequest(@Nonnull Urn userUrn) {
   }
 
   private Optional<RecommendationContent> buildContent(@Nonnull String entityUrn) {
-    Urn entity;
-    try {
-      entity = Urn.createFromString(entityUrn);
-    } catch (URISyntaxException e) {
-      log.error("Error decoding entity URN: {}", entityUrn, e);
+    Urn entity = UrnUtils.getUrn(entityUrn);
+    if (EntityUtils.checkIfRemoved(_entityService, entity)) {
       return Optional.empty();
     }
+
     return Optional.of(new RecommendationContent().setEntity(entity)
         .setValue(entityUrn)
         .setParams(new RecommendationParams().setEntityProfileParams(new EntityProfileParams().setUrn(entity))));
diff --git a/metadata-io/src/main/java/com/linkedin/metadata/recommendation/candidatesource/RecentlySearchedSource.java b/metadata-io/src/main/java/com/linkedin/metadata/recommendation/candidatesource/RecentlySearchedSource.java
new file mode 100644
index 0000000000000..b0b503c8894a3
--- /dev/null
+++ b/metadata-io/src/main/java/com/linkedin/metadata/recommendation/candidatesource/RecentlySearchedSource.java
@@ -0,0 +1,125 @@
+package com.linkedin.metadata.recommendation.candidatesource;
+
+import com.codahale.metrics.Timer;
+import com.datahub.util.exception.ESQueryException;
+import com.linkedin.common.urn.Urn;
+import com.linkedin.metadata.datahubusage.DataHubUsageEventConstants;
+import com.linkedin.metadata.datahubusage.DataHubUsageEventType;
+import com.linkedin.metadata.recommendation.RecommendationContent;
+import com.linkedin.metadata.recommendation.RecommendationParams;
+import com.linkedin.metadata.recommendation.RecommendationRenderType;
+import com.linkedin.metadata.recommendation.RecommendationRequestContext;
+import com.linkedin.metadata.recommendation.ScenarioType;
+import com.linkedin.metadata.recommendation.SearchParams;
+import com.linkedin.metadata.utils.elasticsearch.IndexConvention;
+import com.linkedin.metadata.utils.metrics.MetricUtils;
+import java.io.IOException;
+import java.util.List;
+import java.util.Optional;
+import java.util.stream.Collectors;
+import javax.annotation.Nonnull;
+import lombok.RequiredArgsConstructor;
+import lombok.extern.slf4j.Slf4j;
+import org.elasticsearch.action.search.SearchRequest;
+import org.elasticsearch.action.search.SearchResponse;
+import org.elasticsearch.client.RequestOptions;
+import org.elasticsearch.client.RestHighLevelClient;
+import org.elasticsearch.client.indices.GetIndexRequest;
+import org.elasticsearch.index.query.BoolQueryBuilder;
+import org.elasticsearch.index.query.QueryBuilders;
+import org.elasticsearch.search.aggregations.AggregationBuilder;
+import org.elasticsearch.search.aggregations.AggregationBuilders;
+import org.elasticsearch.search.aggregations.BucketOrder;
+import org.elasticsearch.search.aggregations.bucket.terms.ParsedTerms;
+import org.elasticsearch.search.builder.SearchSourceBuilder;
+
+
+@Slf4j
+@RequiredArgsConstructor
+public class RecentlySearchedSource implements RecommendationSource {
+  private final RestHighLevelClient _searchClient;
+  private final IndexConvention _indexConvention;
+
+  private static final String DATAHUB_USAGE_INDEX = "datahub_usage_event";
+  private static final String ENTITY_AGG_NAME = "entity";
+  private static final int MAX_CONTENT = 5;
+
+  @Override
+  public String getTitle() {
+    return "Recent Searches";
+  }
+
+  @Override
+  public String getModuleId() {
+    return "RecentSearches";
+  }
+
+  @Override
+  public RecommendationRenderType getRenderType() {
+    return RecommendationRenderType.SEARCH_QUERY_LIST;
+  }
+
+  @Override
+  public boolean isEligible(@Nonnull Urn userUrn, @Nonnull RecommendationRequestContext requestContext) {
+    boolean analyticsEnabled = false;
+    try {
+      analyticsEnabled = _searchClient.indices()
+          .exists(new GetIndexRequest(_indexConvention.getIndexName(DATAHUB_USAGE_INDEX)), RequestOptions.DEFAULT);
+    } catch (IOException e) {
+      log.error("Failed to check whether DataHub usage index exists");
+    }
+    return requestContext.getScenario() == ScenarioType.SEARCH_BAR && analyticsEnabled;
+  }
+
+  @Override
+  public List<RecommendationContent> getRecommendations(@Nonnull Urn userUrn,
+      @Nonnull RecommendationRequestContext requestContext) {
+    SearchRequest searchRequest = buildSearchRequest(userUrn);
+    try (Timer.Context ignored = MetricUtils.timer(this.getClass(), "getRecentlyViewed").time()) {
+      final SearchResponse searchResponse = _searchClient.search(searchRequest, RequestOptions.DEFAULT);
+      // extract results
+      ParsedTerms parsedTerms = searchResponse.getAggregations().get(ENTITY_AGG_NAME);
+      return parsedTerms.getBuckets()
+          .stream()
+          .map(bucket -> buildContent(bucket.getKeyAsString()))
+          .filter(Optional::isPresent)
+          .map(Optional::get)
+          .collect(Collectors.toList());
+    } catch (Exception e) {
+      log.error("Search query to get most recently viewed entities failed", e);
+      throw new ESQueryException("Search query failed:", e);
+    }
+  }
+
+  private SearchRequest buildSearchRequest(@Nonnull Urn userUrn) {
+    SearchRequest request = new SearchRequest();
+    SearchSourceBuilder source = new SearchSourceBuilder();
+    BoolQueryBuilder query = QueryBuilders.boolQuery();
+    // Filter for the entity view events of the user requesting recommendation
+    query.must(QueryBuilders.termQuery(DataHubUsageEventConstants.ACTOR_URN + ".keyword", userUrn.toString()));
+    query.must(QueryBuilders.termQuery(DataHubUsageEventConstants.TYPE,
+        DataHubUsageEventType.SEARCH_RESULTS_VIEW_EVENT.getType()));
+    query.must(QueryBuilders.rangeQuery("total").gt(0));
+    query.must(QueryBuilders.existsQuery(DataHubUsageEventConstants.QUERY));
+    source.query(query);
+
+    // Find the entity with the largest last viewed timestamp
+    String lastSearched = "last_searched";
+    AggregationBuilder aggregation = AggregationBuilders.terms(ENTITY_AGG_NAME)
+        .field(DataHubUsageEventConstants.QUERY + ".keyword")
+        .size(MAX_CONTENT)
+        .order(BucketOrder.aggregation(lastSearched, false))
+        .subAggregation(AggregationBuilders.max(lastSearched).field(DataHubUsageEventConstants.TIMESTAMP));
+    source.aggregation(aggregation);
+    source.size(0);
+
+    request.source(source);
+    request.indices(_indexConvention.getIndexName(DATAHUB_USAGE_INDEX));
+    return request;
+  }
+
+  private Optional<RecommendationContent> buildContent(@Nonnull String query) {
+    return Optional.of(new RecommendationContent().setValue(query)
+        .setParams(new RecommendationParams().setSearchParams(new SearchParams().setQuery(query))));
+  }
+}
diff --git a/metadata-io/src/main/java/com/linkedin/metadata/recommendation/candidatesource/RecentlyViewedSource.java b/metadata-io/src/main/java/com/linkedin/metadata/recommendation/candidatesource/RecentlyViewedSource.java
index e281fce25dea3..e22c569e48f16 100644
--- a/metadata-io/src/main/java/com/linkedin/metadata/recommendation/candidatesource/RecentlyViewedSource.java
+++ b/metadata-io/src/main/java/com/linkedin/metadata/recommendation/candidatesource/RecentlyViewedSource.java
@@ -3,8 +3,11 @@
 import com.codahale.metrics.Timer;
 import com.datahub.util.exception.ESQueryException;
 import com.linkedin.common.urn.Urn;
+import com.linkedin.common.urn.UrnUtils;
 import com.linkedin.metadata.datahubusage.DataHubUsageEventConstants;
 import com.linkedin.metadata.datahubusage.DataHubUsageEventType;
+import com.linkedin.metadata.entity.EntityService;
+import com.linkedin.metadata.entity.EntityUtils;
 import com.linkedin.metadata.recommendation.EntityProfileParams;
 import com.linkedin.metadata.recommendation.RecommendationContent;
 import com.linkedin.metadata.recommendation.RecommendationParams;
@@ -16,7 +19,6 @@
 import com.linkedin.metadata.utils.metrics.MetricUtils;
 import io.opentelemetry.extension.annotations.WithSpan;
 import java.io.IOException;
-import java.net.URISyntaxException;
 import java.util.List;
 import java.util.Optional;
 import java.util.stream.Collectors;
@@ -42,6 +44,7 @@
 public class RecentlyViewedSource implements RecommendationSource {
   private final RestHighLevelClient _searchClient;
   private final IndexConvention _indexConvention;
+  private final EntityService _entityService;
 
   private static final String DATAHUB_USAGE_INDEX = "datahub_usage_event";
   private static final String ENTITY_AGG_NAME = "entity";
@@ -87,7 +90,7 @@ public List<RecommendationContent> getRecommendations(@Nonnull Urn userUrn,
           .stream()
           .map(bucket -> buildContent(bucket.getKeyAsString()))
           .filter(Optional::isPresent)
-          .map(Optional::get)
+          .map(Optional::get).limit(MAX_CONTENT)
           .collect(Collectors.toList());
     } catch (Exception e) {
       log.error("Search query to get most recently viewed entities failed", e);
@@ -122,13 +125,11 @@ private SearchRequest buildSearchRequest(@Nonnull Urn userUrn) {
   }
 
   private Optional<RecommendationContent> buildContent(@Nonnull String entityUrn) {
-    Urn entity;
-    try {
-      entity = Urn.createFromString(entityUrn);
-    } catch (URISyntaxException e) {
-      log.error("Error decoding entity URN: {}", entityUrn, e);
+    Urn entity = UrnUtils.getUrn(entityUrn);
+    if (EntityUtils.checkIfRemoved(_entityService, entity)) {
       return Optional.empty();
     }
+
     return Optional.of(new RecommendationContent().setEntity(entity)
         .setValue(entityUrn)
         .setParams(new RecommendationParams().setEntityProfileParams(new EntityProfileParams().setUrn(entity))));
diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/EntitySearchService.java b/metadata-io/src/main/java/com/linkedin/metadata/search/EntitySearchService.java
index 280d72333b1c4..46bc682f7a8ac 100644
--- a/metadata-io/src/main/java/com/linkedin/metadata/search/EntitySearchService.java
+++ b/metadata-io/src/main/java/com/linkedin/metadata/search/EntitySearchService.java
@@ -27,6 +27,7 @@ public interface EntitySearchService {
    */
   long docCount(@Nonnull String entityName);
 
+
   /**
    * Updates or inserts the given search document.
    *
diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/SearchService.java b/metadata-io/src/main/java/com/linkedin/metadata/search/SearchService.java
index 3db6196c9242d..d2b6ef3bd3930 100644
--- a/metadata-io/src/main/java/com/linkedin/metadata/search/SearchService.java
+++ b/metadata-io/src/main/java/com/linkedin/metadata/search/SearchService.java
@@ -25,13 +25,15 @@ public class SearchService {
   private final AllEntitiesSearchAggregatorCache _allEntitiesSearchAggregatorCache;
 
   public SearchService(EntityRegistry entityRegistry, EntitySearchService entitySearchService,
-      SearchRanker searchRanker, CacheManager cacheManager, int batchSize) {
+      SearchRanker searchRanker, CacheManager cacheManager, int batchSize, boolean enableCache) {
     _entitySearchService = entitySearchService;
     _searchRanker = searchRanker;
     _aggregator =
-        new AllEntitiesSearchAggregator(entityRegistry, entitySearchService, searchRanker, cacheManager, batchSize);
-    _entitySearchServiceCache = new EntitySearchServiceCache(cacheManager, entitySearchService, batchSize);
-    _allEntitiesSearchAggregatorCache = new AllEntitiesSearchAggregatorCache(cacheManager, _aggregator, batchSize);
+        new AllEntitiesSearchAggregator(entityRegistry, entitySearchService, searchRanker, cacheManager, batchSize,
+            enableCache);
+    _entitySearchServiceCache = new EntitySearchServiceCache(cacheManager, entitySearchService, batchSize, enableCache);
+    _allEntitiesSearchAggregatorCache =
+        new AllEntitiesSearchAggregatorCache(cacheManager, _aggregator, batchSize, enableCache);
   }
 
   /**
diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/aggregator/AllEntitiesSearchAggregator.java b/metadata-io/src/main/java/com/linkedin/metadata/search/aggregator/AllEntitiesSearchAggregator.java
index e1de29a1828bf..98a7081e65015 100644
--- a/metadata-io/src/main/java/com/linkedin/metadata/search/aggregator/AllEntitiesSearchAggregator.java
+++ b/metadata-io/src/main/java/com/linkedin/metadata/search/aggregator/AllEntitiesSearchAggregator.java
@@ -47,13 +47,13 @@ public class AllEntitiesSearchAggregator {
   private final EntitySearchServiceCache _entitySearchServiceCache;
 
   public AllEntitiesSearchAggregator(EntityRegistry entityRegistry, EntitySearchService entitySearchService,
-      SearchRanker searchRanker, CacheManager cacheManager, int batchSize) {
+      SearchRanker searchRanker, CacheManager cacheManager, int batchSize, boolean enableCache) {
     _entityRegistry = entityRegistry;
     _entitySearchService = entitySearchService;
     _searchRanker = searchRanker;
     _cacheManager = cacheManager;
     _nonEmptyEntitiesCache = new NonEmptyEntitiesCache(entityRegistry, entitySearchService, cacheManager);
-    _entitySearchServiceCache = new EntitySearchServiceCache(cacheManager, entitySearchService, batchSize);
+    _entitySearchServiceCache = new EntitySearchServiceCache(cacheManager, entitySearchService, batchSize, enableCache);
   }
 
   @Nonnull
diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/cache/AllEntitiesSearchAggregatorCache.java b/metadata-io/src/main/java/com/linkedin/metadata/search/cache/AllEntitiesSearchAggregatorCache.java
index b436e8a1dc156..27ca66b7dadb8 100644
--- a/metadata-io/src/main/java/com/linkedin/metadata/search/cache/AllEntitiesSearchAggregatorCache.java
+++ b/metadata-io/src/main/java/com/linkedin/metadata/search/cache/AllEntitiesSearchAggregatorCache.java
@@ -19,12 +19,13 @@ public class AllEntitiesSearchAggregatorCache {
   private final CacheManager cacheManager;
   private final AllEntitiesSearchAggregator aggregator;
   private final int batchSize;
+  private final boolean enableCache;
 
   public CacheableSearcher<?> getSearcher(List<String> entities, @Nonnull String input, @Nullable Filter postFilters,
       @Nullable SortCriterion sortCriterion, @Nullable SearchFlags searchFlags) {
     return new CacheableSearcher<>(cacheManager.getCache(ALL_ENTITIES_SEARCH_AGGREGATOR_CACHE_NAME), batchSize,
         querySize -> aggregator.search(entities, input, postFilters, sortCriterion, querySize.getFrom(),
             querySize.getSize(), searchFlags),
-        querySize -> Quintet.with(entities, input, postFilters, sortCriterion, querySize), searchFlags);
+        querySize -> Quintet.with(entities, input, postFilters, sortCriterion, querySize), searchFlags, enableCache);
   }
 }
diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/cache/CacheableSearcher.java b/metadata-io/src/main/java/com/linkedin/metadata/search/cache/CacheableSearcher.java
index fb448d325d6bd..f0ed205d7b17d 100644
--- a/metadata-io/src/main/java/com/linkedin/metadata/search/cache/CacheableSearcher.java
+++ b/metadata-io/src/main/java/com/linkedin/metadata/search/cache/CacheableSearcher.java
@@ -28,6 +28,7 @@ public class CacheableSearcher<K> {
   private final Function<QueryPagination, K> cacheKeyGenerator;
   @Nullable
   private final SearchFlags searchFlags;
+  private final boolean enableCache;
 
   @Value
   public static class QueryPagination {
@@ -95,6 +96,6 @@ private SearchResult getBatch(int batchId) {
   }
 
   private boolean enableCache() {
-    return searchFlags == null || !searchFlags.isSkipCache();
+    return enableCache && (searchFlags == null || !searchFlags.isSkipCache());
   }
 }
diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/cache/EntitySearchServiceCache.java b/metadata-io/src/main/java/com/linkedin/metadata/search/cache/EntitySearchServiceCache.java
index 954529c135a15..cc5d8c91d4d1b 100644
--- a/metadata-io/src/main/java/com/linkedin/metadata/search/cache/EntitySearchServiceCache.java
+++ b/metadata-io/src/main/java/com/linkedin/metadata/search/cache/EntitySearchServiceCache.java
@@ -18,12 +18,13 @@ public class EntitySearchServiceCache {
   private final CacheManager cacheManager;
   private final EntitySearchService entitySearchService;
   private final int batchSize;
+  private final boolean enableCache;
 
   public CacheableSearcher<?> getSearcher(@Nonnull String entityName, @Nonnull String input,
       @Nullable Filter postFilters, @Nullable SortCriterion sortCriterion, @Nullable SearchFlags searchFlags) {
     return new CacheableSearcher<>(cacheManager.getCache(ENTITY_SEARCH_SERVICE_CACHE_NAME), batchSize,
         querySize -> entitySearchService.search(entityName, input, postFilters, sortCriterion, querySize.getFrom(),
             querySize.getSize()), querySize -> Quintet.with(entityName, input, postFilters, sortCriterion, querySize),
-        searchFlags);
+        searchFlags, enableCache);
   }
 }
diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/AutocompleteRequestHandler.java b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/AutocompleteRequestHandler.java
index 1883783c066af..d81f59eaf09f2 100644
--- a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/AutocompleteRequestHandler.java
+++ b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/AutocompleteRequestHandler.java
@@ -1,13 +1,18 @@
 package com.linkedin.metadata.search.elasticsearch.query.request;
 
 import com.google.common.collect.ImmutableList;
+import com.linkedin.common.urn.Urn;
 import com.linkedin.data.template.StringArray;
 import com.linkedin.metadata.models.EntitySpec;
 import com.linkedin.metadata.models.SearchableFieldSpec;
 import com.linkedin.metadata.models.annotation.SearchableAnnotation;
+import com.linkedin.metadata.query.AutoCompleteEntity;
+import com.linkedin.metadata.query.AutoCompleteEntityArray;
 import com.linkedin.metadata.query.AutoCompleteResult;
 import com.linkedin.metadata.query.filter.Filter;
 import com.linkedin.metadata.search.utils.ESUtils;
+import java.net.URISyntaxException;
+import java.util.HashSet;
 import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
@@ -101,18 +106,30 @@ private List<String> getAutocompleteFields(@Nullable String field) {
 
   public AutoCompleteResult extractResult(@Nonnull SearchResponse searchResponse, @Nonnull String input) {
     Set<String> results = new LinkedHashSet<>();
+    Set<AutoCompleteEntity> entityResults = new HashSet<>();
     for (SearchHit hit : searchResponse.getHits()) {
       Optional<String> matchedFieldValue = hit.getHighlightFields()
           .entrySet()
           .stream()
           .findFirst()
           .map(entry -> entry.getValue().getFragments()[0].string());
+      Optional<String> matchedUrn = Optional.ofNullable((String) hit.getSourceAsMap().get("urn"));
+      try {
+        if (matchedUrn.isPresent()) {
+          entityResults.add(new AutoCompleteEntity().setUrn(Urn.createFromString(matchedUrn.get())));
+        }
+      } catch (URISyntaxException e) {
+        throw new RuntimeException(String.format("Failed to create urn %s", matchedUrn.get()), e);
+      }
       if (matchedFieldValue.isPresent()) {
         results.add(matchedFieldValue.get());
       } else {
         log.info("No highlighted field for query {}, hit {}", input, hit);
       }
     }
-    return new AutoCompleteResult().setQuery(input).setSuggestions(new StringArray(results));
+    return new AutoCompleteResult()
+        .setQuery(input)
+        .setSuggestions(new StringArray(results))
+        .setEntities(new AutoCompleteEntityArray(entityResults));
   }
 }
diff --git a/metadata-io/src/main/java/com/linkedin/metadata/systemmetadata/ESSystemMetadataDAO.java b/metadata-io/src/main/java/com/linkedin/metadata/systemmetadata/ESSystemMetadataDAO.java
index d0cbc6d4992b4..955d1c43119ef 100644
--- a/metadata-io/src/main/java/com/linkedin/metadata/systemmetadata/ESSystemMetadataDAO.java
+++ b/metadata-io/src/main/java/com/linkedin/metadata/systemmetadata/ESSystemMetadataDAO.java
@@ -114,14 +114,21 @@ public BulkByScrollResponse deleteByUrnAspect(
     return null;
   }
 
-  public SearchResponse findByParams(Map<String, String> searchParams) {
+  public SearchResponse findByParams(Map<String, String> searchParams, boolean includeSoftDeleted) {
     SearchRequest searchRequest = new SearchRequest();
 
     SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
 
     BoolQueryBuilder finalQuery = QueryBuilders.boolQuery();
-    searchParams.entrySet()
-        .forEach(entry -> finalQuery.must(QueryBuilders.termQuery(entry.getKey(), entry.getValue())));
+
+    for (String key : searchParams.keySet()) {
+      finalQuery.must(QueryBuilders.termQuery(key, searchParams.get(key)));
+    }
+
+    if (!includeSoftDeleted) {
+      finalQuery.mustNot(QueryBuilders.termQuery("removed", "true"));
+    }
+
     searchSourceBuilder.query(finalQuery);
 
     // this is the max page size elastic will return
@@ -140,15 +147,15 @@ public SearchResponse findByParams(Map<String, String> searchParams) {
     return null;
   }
 
-  public SearchResponse findByRegistry(String registryName, String registryVersion) {
+  public SearchResponse findByRegistry(String registryName, String registryVersion, boolean includeSoftDeleted) {
     Map<String, String> params = new HashMap<>();
     params.put("registryName", registryName);
     params.put("registryVersion", registryVersion);
-    return findByParams(params);
+    return findByParams(params, includeSoftDeleted);
   }
 
-  public SearchResponse findByRunId(String runId) {
-    return findByParams(Collections.singletonMap("runId", runId));
+  public SearchResponse findByRunId(String runId, boolean includeSoftDeleted) {
+    return findByParams(Collections.singletonMap("runId", runId), includeSoftDeleted);
   }
 
   public SearchResponse findRuns(Integer pageOffset, Integer pageSize) {
@@ -170,7 +177,8 @@ public SearchResponse findRuns(Integer pageOffset, Integer pageSize) {
     TermsAggregationBuilder aggregation = AggregationBuilders.terms("runId")
         .field("runId")
         .subAggregation(AggregationBuilders.max("maxTimestamp").field("lastUpdated"))
-        .subAggregation(bucketSort);
+        .subAggregation(bucketSort)
+        .subAggregation(AggregationBuilders.filter("removed", QueryBuilders.termQuery("removed", "true")));
 
     searchSourceBuilder.aggregation(aggregation);
 
diff --git a/metadata-io/src/main/java/com/linkedin/metadata/systemmetadata/ElasticSearchSystemMetadataService.java b/metadata-io/src/main/java/com/linkedin/metadata/systemmetadata/ElasticSearchSystemMetadataService.java
index 744817460e6ad..a1fe2efa7f0c5 100644
--- a/metadata-io/src/main/java/com/linkedin/metadata/systemmetadata/ElasticSearchSystemMetadataService.java
+++ b/metadata-io/src/main/java/com/linkedin/metadata/systemmetadata/ElasticSearchSystemMetadataService.java
@@ -2,6 +2,7 @@
 
 import com.fasterxml.jackson.databind.node.JsonNodeFactory;
 import com.fasterxml.jackson.databind.node.ObjectNode;
+import com.google.common.collect.ImmutableMap;
 import com.linkedin.metadata.run.AspectRowSummary;
 import com.linkedin.metadata.run.IngestionRunSummary;
 import com.linkedin.metadata.search.elasticsearch.indexbuilder.ESIndexBuilder;
@@ -30,6 +31,7 @@
 import org.elasticsearch.index.query.QueryBuilders;
 import org.elasticsearch.index.reindex.DeleteByQueryRequest;
 import org.elasticsearch.search.SearchHits;
+import org.elasticsearch.search.aggregations.bucket.filter.ParsedFilter;
 import org.elasticsearch.search.aggregations.bucket.terms.ParsedStringTerms;
 import org.elasticsearch.search.aggregations.bucket.terms.Terms;
 import org.elasticsearch.search.aggregations.metrics.ParsedMax;
@@ -65,6 +67,7 @@ private String toDocument(SystemMetadata systemMetadata, String urn, String aspe
     document.put("lastUpdated", systemMetadata.getLastObserved());
     document.put("registryName", systemMetadata.getRegistryName());
     document.put("registryVersion", systemMetadata.getRegistryVersion());
+    document.put("removed", false);
     return document.toString();
   }
 
@@ -92,6 +95,21 @@ public void deleteUrn(String urn) {
     _esDAO.deleteByUrn(urn);
   }
 
+  @Override
+  public void setDocStatus(String urn, boolean removed) {
+    // searchBy findByParams
+    // If status.removed -> false (from removed to not removed) --> get soft deleted entities.
+    // If status.removed -> true (from not removed to removed) --> do not get soft deleted entities.
+    final List<AspectRowSummary> aspectList = findByParams(ImmutableMap.of("urn", urn), !removed);
+    // for each -> toDocId and set removed to true for all
+    aspectList.forEach(aspect -> {
+      final String docId = toDocId(aspect.getUrn(), aspect.getAspectName());
+      final ObjectNode document = JsonNodeFactory.instance.objectNode();
+      document.put("removed", removed);
+      _esDAO.upsertDocument(docId, document.toString());
+    });
+  }
+
   @Override
   public void insert(@Nullable SystemMetadata systemMetadata, String urn, String aspect) {
     if (systemMetadata == null) {
@@ -105,12 +123,12 @@ public void insert(@Nullable SystemMetadata systemMetadata, String urn, String a
   }
 
   @Override
-  public List<AspectRowSummary> findByRunId(String runId) {
-    return findByParams(Collections.singletonMap(FIELD_RUNID, runId));
+  public List<AspectRowSummary> findByRunId(String runId, boolean includeSoftDeleted) {
+    return findByParams(Collections.singletonMap(FIELD_RUNID, runId), includeSoftDeleted);
   }
 
-  private List<AspectRowSummary> findByParams(Map<String, String> systemMetaParams) {
-    SearchResponse searchResponse = _esDAO.findByParams(systemMetaParams);
+  private List<AspectRowSummary> findByParams(Map<String, String> systemMetaParams, boolean includeSoftDeleted) {
+    SearchResponse searchResponse = _esDAO.findByParams(systemMetaParams, includeSoftDeleted);
     if (searchResponse != null) {
       SearchHits hits = searchResponse.getHits();
       List<AspectRowSummary> summaries = Arrays.stream(hits.getHits()).map(hit -> {
@@ -135,17 +153,26 @@ private List<AspectRowSummary> findByParams(Map<String, String> systemMetaParams
   }
 
   @Override
-  public List<AspectRowSummary> findByRegistry(String registryName, String registryVersion) {
+  public List<AspectRowSummary> findByRegistry(String registryName, String registryVersion, boolean includeSoftDeleted) {
     Map<String, String> registryParams = new HashMap<>();
     registryParams.put(FIELD_REGISTRY_NAME, registryName);
     registryParams.put(FIELD_REGISTRY_VERSION, registryVersion);
-    return findByParams(registryParams);
+    return findByParams(registryParams, includeSoftDeleted);
   }
 
   @Override
-  public List<IngestionRunSummary> listRuns(Integer pageOffset, Integer pageSize) {
+  public List<IngestionRunSummary> listRuns(Integer pageOffset, Integer pageSize, boolean includeSoftDeleted) {
     SearchResponse response = _esDAO.findRuns(pageOffset, pageSize);
     List<? extends Terms.Bucket> buckets = ((ParsedStringTerms) response.getAggregations().get("runId")).getBuckets();
+
+    if (!includeSoftDeleted) {
+      buckets.removeIf(bucket -> {
+        long totalDocs = bucket.getDocCount();
+        long softDeletedDocs = ((ParsedFilter) bucket.getAggregations().get("removed")).getDocCount();
+        return totalDocs == softDeletedDocs;
+      });
+    }
+
     // TODO(gabe-lyons): add sample urns
     return buckets.stream().map(bucket -> {
       IngestionRunSummary entry = new IngestionRunSummary();
diff --git a/metadata-io/src/main/java/com/linkedin/metadata/systemmetadata/SystemMetadataService.java b/metadata-io/src/main/java/com/linkedin/metadata/systemmetadata/SystemMetadataService.java
index dbc7965ef87ed..065de46876146 100644
--- a/metadata-io/src/main/java/com/linkedin/metadata/systemmetadata/SystemMetadataService.java
+++ b/metadata-io/src/main/java/com/linkedin/metadata/systemmetadata/SystemMetadataService.java
@@ -18,15 +18,15 @@ public interface SystemMetadataService {
 
   void deleteUrn(String finalOldUrn);
 
+  void setDocStatus(String urn, boolean removed);
+
   void insert(@Nullable SystemMetadata systemMetadata, String urn, String aspect);
 
-  List<AspectRowSummary> findByRunId(String runId);
+  List<AspectRowSummary> findByRunId(String runId, boolean includeSoftDeleted);
 
-  List<AspectRowSummary> findByRegistry(String registryName, String registryVersion);
+  List<AspectRowSummary> findByRegistry(String registryName, String registryVersion, boolean includeSoftDeleted);
 
-  List<IngestionRunSummary> listRuns(
-      final Integer pageOffset,
-      final Integer pageSize);
+  List<IngestionRunSummary> listRuns(Integer pageOffset, Integer pageSize, boolean includeSoftDeleted);
 
   void configure();
 
diff --git a/metadata-io/src/test/java/com/linkedin/metadata/entity/EbeanEntityServiceTest.java b/metadata-io/src/test/java/com/linkedin/metadata/entity/EbeanEntityServiceTest.java
index ebc0d4e5fb85a..2e16167b4d231 100644
--- a/metadata-io/src/test/java/com/linkedin/metadata/entity/EbeanEntityServiceTest.java
+++ b/metadata-io/src/test/java/com/linkedin/metadata/entity/EbeanEntityServiceTest.java
@@ -1,5 +1,6 @@
 package com.linkedin.metadata.entity;
 
+import com.datahub.util.RecordUtils;
 import com.fasterxml.jackson.annotation.JsonInclude;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.common.collect.ImmutableList;
@@ -22,7 +23,6 @@
 import com.linkedin.metadata.aspect.CorpUserAspect;
 import com.linkedin.metadata.aspect.CorpUserAspectArray;
 import com.linkedin.metadata.aspect.VersionedAspect;
-import com.datahub.util.RecordUtils;
 import com.linkedin.metadata.entity.ebean.EbeanAspectDao;
 import com.linkedin.metadata.entity.ebean.EbeanAspectV2;
 import com.linkedin.metadata.entity.ebean.EbeanEntityService;
@@ -52,33 +52,27 @@
 import com.linkedin.util.Pair;
 import io.ebean.EbeanServer;
 import io.ebean.EbeanServerFactory;
+import io.ebean.Transaction;
+import io.ebean.TxScope;
+import io.ebean.annotation.TxIsolation;
 import io.ebean.config.ServerConfig;
 import io.ebean.datasource.DataSourceConfig;
-
-
-import javax.annotation.Nonnull;
 import org.mockito.ArgumentCaptor;
 import org.mockito.Mockito;
 import org.testng.annotations.BeforeMethod;
 import org.testng.annotations.Test;
 
+import javax.annotation.Nonnull;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 
-import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.reset;
-import static org.mockito.Mockito.times;
-import static org.mockito.Mockito.verify;
-import static org.mockito.Mockito.verifyNoMoreInteractions;
-import static org.testng.Assert.assertEquals;
-import static org.testng.Assert.assertFalse;
-import static org.testng.Assert.assertNotNull;
-import static org.testng.Assert.assertNull;
-import static org.testng.Assert.assertTrue;
+import static org.mockito.Mockito.*;
+import static org.testng.Assert.*;
 
 
 public class EbeanEntityServiceTest {
@@ -151,7 +145,7 @@ public void testIngestGetEntity() throws Exception {
     com.linkedin.entity.Entity readEntity = _entityService.getEntity(entityUrn, Collections.emptySet());
 
     // 3. Compare Entity Objects
-    assertEquals(2, readEntity.getValue().getCorpUserSnapshot().getAspects().size()); // Key + Info aspect.
+    assertEquals(readEntity.getValue().getCorpUserSnapshot().getAspects().size(), 2); // Key + Info aspect.
     assertTrue(DataTemplateUtil.areEqual(writeEntity.getValue().getCorpUserSnapshot().getAspects().get(0),
         readEntity.getValue().getCorpUserSnapshot().getAspects().get(1)));
     CorpUserKey expectedKey = new CorpUserKey();
@@ -190,7 +184,7 @@ public void testAddKey() throws Exception {
     com.linkedin.entity.Entity readEntity = _entityService.getEntity(entityUrn, Collections.emptySet());
 
     // 3. Compare Entity Objects
-    assertEquals(2, readEntity.getValue().getCorpUserSnapshot().getAspects().size()); // Key + Info aspect.
+    assertEquals(readEntity.getValue().getCorpUserSnapshot().getAspects().size(), 2); // Key + Info aspects.
     assertTrue(DataTemplateUtil.areEqual(writeEntity.getValue().getCorpUserSnapshot().getAspects().get(0),
         readEntity.getValue().getCorpUserSnapshot().getAspects().get(1)));
     CorpUserKey expectedKey = new CorpUserKey();
@@ -241,7 +235,7 @@ public void testIngestGetEntities() throws Exception {
 
     // Entity 1
     com.linkedin.entity.Entity readEntity1 = readEntities.get(entityUrn1);
-    assertEquals(2, readEntity1.getValue().getCorpUserSnapshot().getAspects().size()); // Key + Info aspect.
+    assertEquals(readEntity1.getValue().getCorpUserSnapshot().getAspects().size(), 2); // Key + Info aspect.
     assertTrue(DataTemplateUtil.areEqual(writeEntity1.getValue().getCorpUserSnapshot().getAspects().get(0),
         readEntity1.getValue().getCorpUserSnapshot().getAspects().get(1)));
     CorpUserKey expectedKey1 = new CorpUserKey();
@@ -251,9 +245,15 @@ public void testIngestGetEntities() throws Exception {
 
     // Entity 2
     com.linkedin.entity.Entity readEntity2 = readEntities.get(entityUrn2);
-    assertEquals(2, readEntity2.getValue().getCorpUserSnapshot().getAspects().size()); // Key + Info aspect.
-    assertTrue(DataTemplateUtil.areEqual(writeEntity2.getValue().getCorpUserSnapshot().getAspects().get(0),
-        readEntity2.getValue().getCorpUserSnapshot().getAspects().get(1)));
+    assertEquals(readEntity2.getValue().getCorpUserSnapshot().getAspects().size(), 2); // Key + Info aspect.
+    Optional<CorpUserAspect> writer2UserInfo = writeEntity2.getValue().getCorpUserSnapshot().getAspects()
+            .stream().filter(CorpUserAspect::isCorpUserInfo).findAny();
+    Optional<CorpUserAspect> reader2UserInfo = writeEntity2.getValue().getCorpUserSnapshot().getAspects()
+            .stream().filter(CorpUserAspect::isCorpUserInfo).findAny();
+
+    assertTrue(writer2UserInfo.isPresent(), "Writer2 user info exists");
+    assertTrue(reader2UserInfo.isPresent(), "Reader2 user info exists");
+    assertTrue(DataTemplateUtil.areEqual(writer2UserInfo.get(),  reader2UserInfo.get()), "UserInfo's are the same");
     CorpUserKey expectedKey2 = new CorpUserKey();
     expectedKey2.setUsername("tester2");
     assertTrue(DataTemplateUtil.areEqual(expectedKey2,
@@ -317,7 +317,7 @@ public void testIngestGetEntitiesV2() throws Exception {
 
     // Entity 1
     EntityResponse readEntityResponse1 = readEntities.get(entityUrn1);
-    assertEquals(2, readEntityResponse1.getAspects().size()); // Key + Info aspect.
+    assertEquals(readEntityResponse1.getAspects().size(), 2); // Key + Info aspect.
     EnvelopedAspect envelopedAspect1 = readEntityResponse1.getAspects().get(aspectName);
     assertEquals(envelopedAspect1.getName(), aspectName);
     assertTrue(
@@ -330,7 +330,7 @@ public void testIngestGetEntitiesV2() throws Exception {
 
     // Entity 2
     EntityResponse readEntityResponse2 = readEntities.get(entityUrn2);
-    assertEquals(2, readEntityResponse2.getAspects().size()); // Key + Info aspect.
+    assertEquals(readEntityResponse2.getAspects().size(), 2); // Key + Info aspect.
     EnvelopedAspect envelopedAspect2 = readEntityResponse2.getAspects().get(aspectName);
     assertEquals(envelopedAspect2.getName(), aspectName);
     assertTrue(
@@ -725,7 +725,7 @@ public void testRollbackAspect() throws Exception {
     rollbackOverwrittenAspect.setAspectName(aspectName);
     rollbackOverwrittenAspect.setUrn(entityUrn1.toString());
 
-    _entityService.rollbackRun(ImmutableList.of(rollbackOverwrittenAspect), "run-123");
+    _entityService.rollbackRun(ImmutableList.of(rollbackOverwrittenAspect), "run-123", true);
 
     // assert nothing was deleted
     RecordTemplate readAspectOriginal = _entityService.getAspect(entityUrn1, aspectName, 1);
@@ -740,7 +740,7 @@ public void testRollbackAspect() throws Exception {
     rollbackRecentAspect.setAspectName(aspectName);
     rollbackRecentAspect.setUrn(entityUrn1.toString());
 
-    _entityService.rollbackRun(ImmutableList.of(rollbackOverwrittenAspect), "run-456");
+    _entityService.rollbackRun(ImmutableList.of(rollbackOverwrittenAspect), "run-456", true);
 
     // assert the new most recent aspect is the original one
     RecordTemplate readNewRecentAspect = _entityService.getAspect(entityUrn1, aspectName, 0);
@@ -779,7 +779,7 @@ public void testRollbackKey() throws Exception {
     rollbackKeyWithWrongRunId.setAspectName("corpUserKey");
     rollbackKeyWithWrongRunId.setUrn(entityUrn1.toString());
 
-    _entityService.rollbackRun(ImmutableList.of(rollbackKeyWithWrongRunId), "run-456");
+    _entityService.rollbackRun(ImmutableList.of(rollbackKeyWithWrongRunId), "run-456", true);
 
     // assert nothing was deleted
     RecordTemplate readAspectOriginal = _entityService.getAspect(entityUrn1, aspectName, 1);
@@ -794,13 +794,36 @@ public void testRollbackKey() throws Exception {
     rollbackKeyWithCorrectRunId.setAspectName("corpUserKey");
     rollbackKeyWithCorrectRunId.setUrn(entityUrn1.toString());
 
-    _entityService.rollbackRun(ImmutableList.of(rollbackKeyWithCorrectRunId), "run-123");
+    _entityService.rollbackRun(ImmutableList.of(rollbackKeyWithCorrectRunId), "run-123", true);
 
     // assert the new most recent aspect is null
     RecordTemplate readNewRecentAspect = _entityService.getAspect(entityUrn1, aspectName, 0);
     assertTrue(DataTemplateUtil.areEqual(null, readNewRecentAspect));
   }
 
+  @Test
+  public void testNestedTransactions() throws Exception {
+    EbeanServer server = _aspectDao.getServer();
+
+    try (Transaction transaction = server.beginTransaction(TxScope.requiresNew()
+            .setIsolation(TxIsolation.REPEATABLE_READ))) {
+      transaction.setBatchMode(true);
+      // Work 1
+      try (Transaction transaction2 = server.beginTransaction(TxScope.requiresNew()
+              .setIsolation(TxIsolation.REPEATABLE_READ))) {
+        transaction2.setBatchMode(true);
+        // Work 2
+        transaction2.commit();
+      }
+      transaction.commit();
+      } catch (Exception e) {
+      System.out.printf("Top level catch %s%n", e);
+      e.printStackTrace();
+      throw e;
+    }
+    System.out.println("done");
+  }
+
   @Test
   public void testRollbackUrn() throws Exception {
     Urn entityUrn1 = Urn.createFromString("urn:li:corpuser:test1");
diff --git a/metadata-io/src/test/java/com/linkedin/metadata/graph/dgraph/DgraphGraphServiceTest.java b/metadata-io/src/test/java/com/linkedin/metadata/graph/dgraph/DgraphGraphServiceTest.java
index 70f1123ce74bd..3626fa2df2030 100644
--- a/metadata-io/src/test/java/com/linkedin/metadata/graph/dgraph/DgraphGraphServiceTest.java
+++ b/metadata-io/src/test/java/com/linkedin/metadata/graph/dgraph/DgraphGraphServiceTest.java
@@ -2,7 +2,7 @@
 
 import com.linkedin.metadata.graph.GraphService;
 import com.linkedin.metadata.graph.GraphServiceTestBase;
-import com.linkedin.metadata.graph.LineageRegistry;
+import com.linkedin.metadata.models.registry.LineageRegistry;
 import com.linkedin.metadata.graph.RelatedEntity;
 import com.linkedin.metadata.models.registry.SnapshotEntityRegistry;
 import com.linkedin.metadata.query.filter.RelationshipDirection;
diff --git a/metadata-io/src/test/java/com/linkedin/metadata/graph/elastic/ElasticSearchGraphServiceTest.java b/metadata-io/src/test/java/com/linkedin/metadata/graph/elastic/ElasticSearchGraphServiceTest.java
index 3f0f9a4f51349..8f016c38a3bcc 100644
--- a/metadata-io/src/test/java/com/linkedin/metadata/graph/elastic/ElasticSearchGraphServiceTest.java
+++ b/metadata-io/src/test/java/com/linkedin/metadata/graph/elastic/ElasticSearchGraphServiceTest.java
@@ -7,7 +7,7 @@
 import com.linkedin.metadata.graph.GraphService;
 import com.linkedin.metadata.graph.GraphServiceTestBase;
 import com.linkedin.metadata.graph.LineageDirection;
-import com.linkedin.metadata.graph.LineageRegistry;
+import com.linkedin.metadata.models.registry.LineageRegistry;
 import com.linkedin.metadata.graph.LineageRelationship;
 import com.linkedin.metadata.graph.RelatedEntitiesResult;
 import com.linkedin.metadata.graph.RelatedEntity;
diff --git a/metadata-io/src/test/java/com/linkedin/metadata/graph/neo4j/Neo4jGraphServiceTest.java b/metadata-io/src/test/java/com/linkedin/metadata/graph/neo4j/Neo4jGraphServiceTest.java
index 1c7065e6b11f6..458cb525998c7 100644
--- a/metadata-io/src/test/java/com/linkedin/metadata/graph/neo4j/Neo4jGraphServiceTest.java
+++ b/metadata-io/src/test/java/com/linkedin/metadata/graph/neo4j/Neo4jGraphServiceTest.java
@@ -2,7 +2,7 @@
 
 import com.linkedin.metadata.graph.GraphService;
 import com.linkedin.metadata.graph.GraphServiceTestBase;
-import com.linkedin.metadata.graph.LineageRegistry;
+import com.linkedin.metadata.models.registry.LineageRegistry;
 import com.linkedin.metadata.graph.RelatedEntitiesResult;
 import com.linkedin.metadata.graph.RelatedEntity;
 import com.linkedin.metadata.models.registry.SnapshotEntityRegistry;
diff --git a/metadata-io/src/test/java/com/linkedin/metadata/search/LineageSearchServiceTest.java b/metadata-io/src/test/java/com/linkedin/metadata/search/LineageSearchServiceTest.java
index 06bc775aa7d5d..d1ffecec52e2e 100644
--- a/metadata-io/src/test/java/com/linkedin/metadata/search/LineageSearchServiceTest.java
+++ b/metadata-io/src/test/java/com/linkedin/metadata/search/LineageSearchServiceTest.java
@@ -83,7 +83,7 @@ public void setup() {
     _cacheManager = new ConcurrentMapCacheManager();
     _graphService = mock(GraphService.class);
     _lineageSearchService = new LineageSearchService(
-        new SearchService(_entityRegistry, _elasticSearchService, new SimpleRanker(), _cacheManager, 100),
+        new SearchService(_entityRegistry, _elasticSearchService, new SimpleRanker(), _cacheManager, 100, true),
         _graphService, _cacheManager.getCache("test"));
   }
 
diff --git a/metadata-io/src/test/java/com/linkedin/metadata/search/SearchServiceTest.java b/metadata-io/src/test/java/com/linkedin/metadata/search/SearchServiceTest.java
index a0882c342dbd7..5bab4579870d8 100644
--- a/metadata-io/src/test/java/com/linkedin/metadata/search/SearchServiceTest.java
+++ b/metadata-io/src/test/java/com/linkedin/metadata/search/SearchServiceTest.java
@@ -60,7 +60,8 @@ public void setup() {
     _elasticSearchService = buildEntitySearchService();
     _elasticSearchService.configure();
     _cacheManager = new ConcurrentMapCacheManager();
-    _searchService = new SearchService(_entityRegistry, _elasticSearchService, new SimpleRanker(), _cacheManager, 100);
+    _searchService =
+        new SearchService(_entityRegistry, _elasticSearchService, new SimpleRanker(), _cacheManager, 100, true);
   }
 
   @BeforeMethod
diff --git a/metadata-io/src/test/java/com/linkedin/metadata/search/cache/CacheableSearcherTest.java b/metadata-io/src/test/java/com/linkedin/metadata/search/cache/CacheableSearcherTest.java
index 88d3693184f2f..0ad644649a88d 100644
--- a/metadata-io/src/test/java/com/linkedin/metadata/search/cache/CacheableSearcherTest.java
+++ b/metadata-io/src/test/java/com/linkedin/metadata/search/cache/CacheableSearcherTest.java
@@ -26,7 +26,7 @@ public class CacheableSearcherTest {
   public void testCacheableSearcherWhenEmpty() {
     CacheableSearcher<Integer> emptySearcher =
         new CacheableSearcher<>(cacheManager.getCache("emptySearcher"), 10, this::getEmptySearchResult,
-            CacheableSearcher.QueryPagination::getFrom, null);
+            CacheableSearcher.QueryPagination::getFrom, null, true);
     assertTrue(emptySearcher.getSearchResults(0, 0).getEntities().isEmpty());
     assertTrue(emptySearcher.getSearchResults(0, 10).getEntities().isEmpty());
     assertTrue(emptySearcher.getSearchResults(5, 10).getEntities().isEmpty());
@@ -36,7 +36,7 @@ public void testCacheableSearcherWhenEmpty() {
   public void testCacheableSearcherWithFixedNumResults() {
     CacheableSearcher<Integer> fixedBatchSearcher =
         new CacheableSearcher<>(cacheManager.getCache("fixedBatchSearcher"), 10, qs -> getSearchResult(qs, 10),
-            CacheableSearcher.QueryPagination::getFrom, null);
+            CacheableSearcher.QueryPagination::getFrom, null, true);
 
     SearchResult result = fixedBatchSearcher.getSearchResults(0, 0);
     assertTrue(result.getEntities().isEmpty());
@@ -59,7 +59,8 @@ public void testCacheableSearcherWithFixedNumResults() {
   public void testCacheableSearcherWithVariableNumResults() {
     CacheableSearcher<Integer> variableBatchSearcher =
         new CacheableSearcher<>(cacheManager.getCache("variableBatchSearcher"), 10,
-            qs -> getSearchResult(qs, qs.getFrom() + qs.getSize()), CacheableSearcher.QueryPagination::getFrom, null);
+            qs -> getSearchResult(qs, qs.getFrom() + qs.getSize()), CacheableSearcher.QueryPagination::getFrom, null,
+            true);
 
     SearchResult result = variableBatchSearcher.getSearchResults(0, 0);
     assertTrue(result.getEntities().isEmpty());
diff --git a/metadata-io/src/test/java/com/linkedin/metadata/systemmetadata/ElasticSearchSystemMetadataServiceTest.java b/metadata-io/src/test/java/com/linkedin/metadata/systemmetadata/ElasticSearchSystemMetadataServiceTest.java
index 94024f7401fcc..c3a2ae0b570dd 100644
--- a/metadata-io/src/test/java/com/linkedin/metadata/systemmetadata/ElasticSearchSystemMetadataServiceTest.java
+++ b/metadata-io/src/test/java/com/linkedin/metadata/systemmetadata/ElasticSearchSystemMetadataServiceTest.java
@@ -78,7 +78,7 @@ public void testListRuns() throws Exception {
 
     syncAfterWrite(_searchClient, _indexName);
 
-    List<IngestionRunSummary> runs = _client.listRuns(0, 20);
+    List<IngestionRunSummary> runs = _client.listRuns(0, 20, false);
 
     assertEquals(runs.size(), 2);
     assertEquals(runs.get(0).getRows(), Long.valueOf(2));
@@ -107,7 +107,7 @@ public void testOverwriteRuns() throws Exception {
 
     syncAfterWrite(_searchClient, _indexName);
 
-    List<IngestionRunSummary> runs = _client.listRuns(0, 20);
+    List<IngestionRunSummary> runs = _client.listRuns(0, 20, false);
 
     assertEquals(runs.size(), 2);
     assertEquals(runs.get(0).getRows(), Long.valueOf(4));
@@ -136,7 +136,7 @@ public void testFindByRunId() throws Exception {
 
     syncAfterWrite(_searchClient, _indexName);
 
-    List<AspectRowSummary> rows = _client.findByRunId("abc-456");
+    List<AspectRowSummary> rows = _client.findByRunId("abc-456", false);
 
     assertEquals(rows.size(), 4);
     rows.forEach(row -> assertEquals(row.getRunId(), "abc-456"));
@@ -168,7 +168,7 @@ public void testDelete() throws Exception {
 
     syncAfterWrite(_searchClient, _indexName);
 
-    List<AspectRowSummary> rows = _client.findByRunId("abc-456");
+    List<AspectRowSummary> rows = _client.findByRunId("abc-456", false);
 
     assertEquals(rows.size(), 2);
     rows.forEach(row -> assertEquals(row.getRunId(), "abc-456"));
@@ -180,7 +180,7 @@ public void testInsertNullData() throws Exception {
 
     syncAfterWrite(_searchClient, _indexName);
 
-    List<IngestionRunSummary> runs = _client.listRuns(0, 20);
+    List<IngestionRunSummary> runs = _client.listRuns(0, 20, false);
 
     assertEquals(runs.size(), 0);
   }
diff --git a/metadata-jobs/mae-consumer/src/main/java/com/linkedin/metadata/kafka/hook/UpdateIndicesHook.java b/metadata-jobs/mae-consumer/src/main/java/com/linkedin/metadata/kafka/hook/UpdateIndicesHook.java
index daf483b3d69ff..bb5193b9826d1 100644
--- a/metadata-jobs/mae-consumer/src/main/java/com/linkedin/metadata/kafka/hook/UpdateIndicesHook.java
+++ b/metadata-jobs/mae-consumer/src/main/java/com/linkedin/metadata/kafka/hook/UpdateIndicesHook.java
@@ -2,6 +2,7 @@
 
 import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.databind.JsonNode;
+import com.linkedin.common.Status;
 import com.linkedin.common.urn.Urn;
 import com.linkedin.data.template.RecordTemplate;
 import com.linkedin.events.metadata.ChangeType;
@@ -11,6 +12,7 @@
 import com.linkedin.gms.factory.search.EntitySearchServiceFactory;
 import com.linkedin.gms.factory.search.SearchDocumentTransformerFactory;
 import com.linkedin.gms.factory.timeseries.TimeseriesAspectServiceFactory;
+import com.linkedin.metadata.Constants;
 import com.linkedin.metadata.models.extractor.FieldExtractor;
 import com.linkedin.metadata.graph.Edge;
 import com.linkedin.metadata.graph.GraphService;
@@ -114,7 +116,7 @@ public void invoke(@Nonnull MetadataChangeLog event) {
       } else {
         updateSearchService(entitySpec.getName(), urn, aspectSpec, aspect);
         updateGraphService(urn, aspectSpec, aspect);
-        updateSystemMetadata(event.getSystemMetadata(), urn, aspectSpec);
+        updateSystemMetadata(event.getSystemMetadata(), urn, aspectSpec, aspect);
       }
     } else if (event.getChangeType() == ChangeType.DELETE) {
       if (!event.hasAspectName() || !event.hasAspect()) {
@@ -224,8 +226,13 @@ private void updateTimeseriesFields(String entityType, String aspectName, Urn ur
     });
   }
 
-  private void updateSystemMetadata(SystemMetadata systemMetadata, Urn urn, AspectSpec aspectSpec) {
+  private void updateSystemMetadata(SystemMetadata systemMetadata, Urn urn, AspectSpec aspectSpec, RecordTemplate aspect) {
     _systemMetadataService.insert(systemMetadata, urn.toString(), aspectSpec.getName());
+
+    // If processing status aspect update all aspects for this urn to removed
+    if (aspectSpec.getName().equals(Constants.STATUS_ASPECT_NAME)) {
+      _systemMetadataService.setDocStatus(urn.toString(), ((Status) aspect).isRemoved());
+    }
   }
 
   private void deleteSystemMetadata(Urn urn, AspectSpec aspectSpec, Boolean isKeyAspect) {
diff --git a/metadata-models-custom/README.md b/metadata-models-custom/README.md
index 9d4b95f7821dd..40895ed2c3396 100644
--- a/metadata-models-custom/README.md
+++ b/metadata-models-custom/README.md
@@ -4,7 +4,7 @@
 
 ### Caveats
 
-Currently, this project only supports adding new aspects to existing entities. You cannot add new entities to the metadata model yet.
+Currently, this project only supports aspects defined in PDL to existing or newly defined entities. You cannot add new aspects to the metadata model directly through yaml configuration yet.
 
 ## Pre-Requisites
 
@@ -15,7 +15,7 @@ Before proceeding further, make sure you understand the DataHub Metadata Model c
 Follow the regular process in creating a new aspect by adding it to the [`src/main/pegasus`](./src/main/pegasus) folder. e.g. This repository has an Aspect called `customDataQualityRules` hosted in the [`DataQualityRules.pdl`](./src/main/pegasus/com/mycompany/dq/DataQualityRules.pdl) file that you can follow.
 Once you've gone through this exercise, feel free to delete the sample aspects that are stored in this module.
 
-**_Tip_**: PDL requires that the name of the file must match the name of the class that is defined in it, so keep that in mind when you create your aspect pdl file.
+**_Tip_**: PDL requires that the name of the file must match the name of the class that is defined in it and the package path must also match the directory path, so keep that in mind when you create your aspect pdl file.
 
 ## Add your aspect(s) to the entity registry
 
@@ -34,8 +34,8 @@ entities:
 ```
 
 The entity registry has a few important fields to pay attention to: 
-- id: this is the name of your registry. This drives naming, artifact generation, so make sure you pick a unique name that will not conflict with other names you might create for other registries. 
-- entities: this is a list of entities with aspects attached to them that you are creating additional aspects for. In this example, we are adding the aspect `customDataQualityRules` to the `dataset` entity. 
+- id: The name of your registry. This drives naming, artifact generation, so make sure you pick a unique name that will not conflict with other names you might create for other registries. 
+- entities: A list of entities with aspects attached to them that you are creating additional aspects for as well as any new entities you wish to define. In this example, we are adding the aspect `customDataQualityRules` to the `dataset` entity. 
 
 ## Build your new model 
 
@@ -146,8 +146,10 @@ As you evolve the metadata model, you can publish new versions of the repository
 
 ## The Future
 
-Hopefully this repository shows you how easily you can extend and customize DataHub's metadata model. 
-We will be adding support for adding new entities soon, and programmatically generating Python classes to make it even easier to interact with the extended metadata model. 
+Hopefully this repository shows you how easily you can extend and customize DataHub's metadata model!
+
+We will be continuing to make the experience less reliant on core changes to DataHub and reducing the need to fork the main repository.
+
 
 
 
diff --git a/metadata-models/src/main/pegasus/com/linkedin/common/OwnershipType.pdl b/metadata-models/src/main/pegasus/com/linkedin/common/OwnershipType.pdl
index 43c456362b1c3..e47776d4428e3 100644
--- a/metadata-models/src/main/pegasus/com/linkedin/common/OwnershipType.pdl
+++ b/metadata-models/src/main/pegasus/com/linkedin/common/OwnershipType.pdl
@@ -1,37 +1,69 @@
 namespace com.linkedin.common
 
 /**
- * Owner category or owner role
+ * Asset owner types
  */
 enum OwnershipType {
 
+  /**
+   * person or group who is responsible for technical aspects of the asset.
+   */
+  TECHNICAL_OWNER
+
+  /**
+   *A person or group who is responsible for logical, or business related, aspects of the asset.
+   */
+  BUSINESS_OWNER
+
+  /**
+   * A steward, expert, or delegate responsible for the asset.
+   */
+  DATA_STEWARD
+
+  /**
+   * No specific type associated to the owner.
+   */
+  NONE
+
   /**
    * A person or group that is in charge of developing the code
+   * Deprecated! Use TECHNICAL_OWNER instead.
    */
+  @deprecated
   DEVELOPER
 
   /**
    * A person or group that is owning the data
+   * Deprecated! Use TECHNICAL_OWNER instead.
    */
+  @deprecated
   DATAOWNER
 
   /**
    * A person or a group that overseas the operation, e.g. a DBA or SRE.
+   * Deprecated! Use TECHNICAL_OWNER instead.
    */
+  @deprecated
   DELEGATE
 
   /**
    * A person, group, or service that produces/generates the data
+   * Deprecated! Use TECHNICAL_OWNER instead.
    */
+  @deprecated
   PRODUCER
 
   /**
    * A person, group, or service that consumes the data
+   * Deprecated! Use TECHNICAL_OWNER or BUSINESS_OWNER instead.
    */
+  @deprecated
   CONSUMER
 
   /**
    * A person or a group that has direct business interest
+   * Deprecated! Use TECHNICAL_OWNER, BUSINESS_OWNER, or STEWARD instead.
    */
+  @deprecated
   STAKEHOLDER
 }
\ No newline at end of file
diff --git a/metadata-models/src/main/pegasus/com/linkedin/metadata/query/AutoCompleteEntity.pdl b/metadata-models/src/main/pegasus/com/linkedin/metadata/query/AutoCompleteEntity.pdl
new file mode 100644
index 0000000000000..59f1fa1fa37b3
--- /dev/null
+++ b/metadata-models/src/main/pegasus/com/linkedin/metadata/query/AutoCompleteEntity.pdl
@@ -0,0 +1,13 @@
+namespace com.linkedin.metadata.query
+
+import com.linkedin.common.Urn
+
+/**
+ * Data model for an entity returned as an autocomplete suggestion
+ */
+record AutoCompleteEntity {
+  /**
+   * Urn of the entity
+   */
+  urn: Urn
+}
\ No newline at end of file
diff --git a/metadata-models/src/main/pegasus/com/linkedin/metadata/query/AutoCompleteResult.pdl b/metadata-models/src/main/pegasus/com/linkedin/metadata/query/AutoCompleteResult.pdl
index 3079e9dab3901..185e9e6076582 100644
--- a/metadata-models/src/main/pegasus/com/linkedin/metadata/query/AutoCompleteResult.pdl
+++ b/metadata-models/src/main/pegasus/com/linkedin/metadata/query/AutoCompleteResult.pdl
@@ -11,7 +11,12 @@ record AutoCompleteResult {
   query: string
 
   /**
-   * A list of typeahead suggestions
+   * A list of typeahead suggestions -- Just the text
    */
   suggestions: array[string]
+
+  /**
+   * Specific entities to recommend
+   */
+  entities: array[AutoCompleteEntity]
 }
\ No newline at end of file
diff --git a/metadata-models/src/main/pegasus/com/linkedin/metadata/recommendation/ScenarioType.pdl b/metadata-models/src/main/pegasus/com/linkedin/metadata/recommendation/ScenarioType.pdl
index 28c020e250665..e51300de20ead 100644
--- a/metadata-models/src/main/pegasus/com/linkedin/metadata/recommendation/ScenarioType.pdl
+++ b/metadata-models/src/main/pegasus/com/linkedin/metadata/recommendation/ScenarioType.pdl
@@ -18,5 +18,10 @@ enum ScenarioType {
   /**
    * Recommendations to show on an Entity Profile page
    */
-  ENTITY_PROFILE
+  ENTITY_PROFILE,
+
+  /**
+   * Recommendations to show on the search bar when clicked
+   */
+  SEARCH_BAR
 }
\ No newline at end of file
diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/common/ElasticSearchGraphServiceFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/common/ElasticSearchGraphServiceFactory.java
index 463fccd3bbfcd..12ede1ef940b5 100644
--- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/common/ElasticSearchGraphServiceFactory.java
+++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/common/ElasticSearchGraphServiceFactory.java
@@ -3,7 +3,7 @@
 import com.linkedin.gms.factory.entityregistry.EntityRegistryFactory;
 import com.linkedin.gms.factory.search.BaseElasticSearchComponentsFactory;
 import com.linkedin.gms.factory.spring.YamlPropertySourceFactory;
-import com.linkedin.metadata.graph.LineageRegistry;
+import com.linkedin.metadata.models.registry.LineageRegistry;
 import com.linkedin.metadata.graph.elastic.ESGraphQueryDAO;
 import com.linkedin.metadata.graph.elastic.ESGraphWriteDAO;
 import com.linkedin.metadata.graph.elastic.ElasticSearchGraphService;
diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/common/Neo4jGraphServiceFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/common/Neo4jGraphServiceFactory.java
index 86705c9b71ac6..e62dfd50f897d 100644
--- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/common/Neo4jGraphServiceFactory.java
+++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/common/Neo4jGraphServiceFactory.java
@@ -1,7 +1,7 @@
 package com.linkedin.gms.factory.common;
 
 import com.linkedin.gms.factory.entityregistry.EntityRegistryFactory;
-import com.linkedin.metadata.graph.LineageRegistry;
+import com.linkedin.metadata.models.registry.LineageRegistry;
 import com.linkedin.metadata.graph.neo4j.Neo4jGraphService;
 import com.linkedin.metadata.models.registry.EntityRegistry;
 import javax.annotation.Nonnull;
diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/config/ConfigurationProvider.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/config/ConfigurationProvider.java
index 039a62bee0534..2b075ec7ae233 100644
--- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/config/ConfigurationProvider.java
+++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/config/ConfigurationProvider.java
@@ -22,4 +22,9 @@ public class ConfigurationProvider {
    * Ingestion related configs
    */
   private IngestionConfiguration ingestion;
+
+  /**
+   * Telemetry related configs
+   */
+  private TelemetryConfiguration telemetry;
 }
\ No newline at end of file
diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/config/TelemetryConfiguration.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/config/TelemetryConfiguration.java
new file mode 100644
index 0000000000000..b78a1cbec5c85
--- /dev/null
+++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/config/TelemetryConfiguration.java
@@ -0,0 +1,17 @@
+package com.linkedin.gms.factory.config;
+
+import lombok.Data;
+/**
+ * POJO representing the "telemtry" configuration block in application.yml.
+ */
+@Data
+public class TelemetryConfiguration {
+    /**
+    * Whether cli telemtry is enabled
+    */
+    public boolean enabledCli;
+    /**
+    * Whether reporting telemetry is enabled
+    */
+    public boolean enabledIngestion;
+}
\ No newline at end of file
diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/recommendation/RecommendationServiceFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/recommendation/RecommendationServiceFactory.java
index d25eebe19e81f..96ab40b186968 100644
--- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/recommendation/RecommendationServiceFactory.java
+++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/recommendation/RecommendationServiceFactory.java
@@ -2,7 +2,7 @@
 
 import com.google.common.collect.ImmutableList;
 import com.linkedin.gms.factory.recommendation.candidatesource.DomainsCandidateSourceFactory;
-import com.linkedin.gms.factory.recommendation.candidatesource.HighUsageCandidateSourceFactory;
+import com.linkedin.gms.factory.recommendation.candidatesource.MostPopularCandidateSourceFactory;
 import com.linkedin.gms.factory.recommendation.candidatesource.RecentlyViewedCandidateSourceFactory;
 import com.linkedin.gms.factory.recommendation.candidatesource.TopPlatformsCandidateSourceFactory;
 import com.linkedin.gms.factory.recommendation.candidatesource.TopTagsCandidateSourceFactory;
@@ -10,6 +10,7 @@
 import com.linkedin.metadata.recommendation.RecommendationsService;
 import com.linkedin.metadata.recommendation.candidatesource.DomainsCandidateSource;
 import com.linkedin.metadata.recommendation.candidatesource.MostPopularSource;
+import com.linkedin.metadata.recommendation.candidatesource.RecentlySearchedSource;
 import com.linkedin.metadata.recommendation.candidatesource.RecentlyViewedSource;
 import com.linkedin.metadata.recommendation.candidatesource.RecommendationSource;
 import com.linkedin.metadata.recommendation.candidatesource.TopPlatformsSource;
@@ -27,7 +28,7 @@
 
 @Configuration
 @Import({TopPlatformsCandidateSourceFactory.class, RecentlyViewedCandidateSourceFactory.class,
-    HighUsageCandidateSourceFactory.class, TopTagsCandidateSourceFactory.class, TopTermsCandidateSourceFactory.class, DomainsCandidateSourceFactory.class})
+    MostPopularCandidateSourceFactory.class, TopTagsCandidateSourceFactory.class, TopTermsCandidateSourceFactory.class, DomainsCandidateSourceFactory.class})
 public class RecommendationServiceFactory {
 
   @Autowired
@@ -39,7 +40,7 @@ public class RecommendationServiceFactory {
   private RecentlyViewedSource recentlyViewedCandidateSource;
 
   @Autowired
-  @Qualifier("highUsageCandidateSource")
+  @Qualifier("mostPopularCandidateSource")
   private MostPopularSource _mostPopularCandidateSource;
 
   @Autowired
@@ -54,6 +55,10 @@ public class RecommendationServiceFactory {
   @Qualifier("domainsCandidateSource")
   private DomainsCandidateSource domainsCandidateSource;
 
+  @Autowired
+  @Qualifier("recentlySearchedCandidateSource")
+  private RecentlySearchedSource recentlySearchedCandidateSource;
+
   @Bean
   @Nonnull
   protected RecommendationsService getInstance() {
@@ -63,7 +68,7 @@ protected RecommendationsService getInstance() {
         topPlatformsCandidateSource,
         domainsCandidateSource,
         recentlyViewedCandidateSource, _mostPopularCandidateSource,
-        topTagsCandidateSource, topTermsCandidateSource);
+        topTagsCandidateSource, topTermsCandidateSource, recentlySearchedCandidateSource);
     return new RecommendationsService(candidateSources, new SimpleRecommendationRanker());
   }
 }
diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/recommendation/candidatesource/MostPopularCandidateSourceFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/recommendation/candidatesource/MostPopularCandidateSourceFactory.java
new file mode 100644
index 0000000000000..c74f5e11cadce
--- /dev/null
+++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/recommendation/candidatesource/MostPopularCandidateSourceFactory.java
@@ -0,0 +1,38 @@
+package com.linkedin.gms.factory.recommendation.candidatesource;
+
+import com.linkedin.gms.factory.common.IndexConventionFactory;
+import com.linkedin.gms.factory.common.RestHighLevelClientFactory;
+import com.linkedin.gms.factory.entity.EntityServiceFactory;
+import com.linkedin.metadata.entity.EntityService;
+import com.linkedin.metadata.recommendation.candidatesource.MostPopularSource;
+import com.linkedin.metadata.utils.elasticsearch.IndexConvention;
+import javax.annotation.Nonnull;
+import org.elasticsearch.client.RestHighLevelClient;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.beans.factory.annotation.Qualifier;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+import org.springframework.context.annotation.Import;
+
+
+@Configuration
+@Import({RestHighLevelClientFactory.class, IndexConventionFactory.class, EntityServiceFactory.class})
+public class MostPopularCandidateSourceFactory {
+  @Autowired
+  @Qualifier("elasticSearchRestHighLevelClient")
+  private RestHighLevelClient searchClient;
+
+  @Autowired
+  @Qualifier(IndexConventionFactory.INDEX_CONVENTION_BEAN)
+  private IndexConvention indexConvention;
+
+  @Autowired
+  @Qualifier("entityService")
+  private EntityService entityService;
+
+  @Bean(name = "mostPopularCandidateSource")
+  @Nonnull
+  protected MostPopularSource getInstance() {
+    return new MostPopularSource(searchClient, indexConvention, entityService);
+  }
+}
diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/recommendation/candidatesource/HighUsageCandidateSourceFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/recommendation/candidatesource/RecentlySearchedCandidateSourceFactory.java
similarity index 75%
rename from metadata-service/factories/src/main/java/com/linkedin/gms/factory/recommendation/candidatesource/HighUsageCandidateSourceFactory.java
rename to metadata-service/factories/src/main/java/com/linkedin/gms/factory/recommendation/candidatesource/RecentlySearchedCandidateSourceFactory.java
index 4b564c0ddc09f..b3779a132284f 100644
--- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/recommendation/candidatesource/HighUsageCandidateSourceFactory.java
+++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/recommendation/candidatesource/RecentlySearchedCandidateSourceFactory.java
@@ -2,7 +2,7 @@
 
 import com.linkedin.gms.factory.common.IndexConventionFactory;
 import com.linkedin.gms.factory.common.RestHighLevelClientFactory;
-import com.linkedin.metadata.recommendation.candidatesource.MostPopularSource;
+import com.linkedin.metadata.recommendation.candidatesource.RecentlySearchedSource;
 import com.linkedin.metadata.utils.elasticsearch.IndexConvention;
 import javax.annotation.Nonnull;
 import org.elasticsearch.client.RestHighLevelClient;
@@ -15,7 +15,7 @@
 
 @Configuration
 @Import({RestHighLevelClientFactory.class, IndexConventionFactory.class})
-public class HighUsageCandidateSourceFactory {
+public class RecentlySearchedCandidateSourceFactory {
   @Autowired
   @Qualifier("elasticSearchRestHighLevelClient")
   private RestHighLevelClient searchClient;
@@ -24,9 +24,9 @@ public class HighUsageCandidateSourceFactory {
   @Qualifier(IndexConventionFactory.INDEX_CONVENTION_BEAN)
   private IndexConvention indexConvention;
 
-  @Bean(name = "highUsageCandidateSource")
+  @Bean(name = "recentlySearchedCandidateSource")
   @Nonnull
-  protected MostPopularSource getInstance() {
-    return new MostPopularSource(searchClient, indexConvention);
+  protected RecentlySearchedSource getInstance() {
+    return new RecentlySearchedSource(searchClient, indexConvention);
   }
 }
diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/recommendation/candidatesource/RecentlyViewedCandidateSourceFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/recommendation/candidatesource/RecentlyViewedCandidateSourceFactory.java
index 44b5a671cf85e..d0505e8d2a3ea 100644
--- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/recommendation/candidatesource/RecentlyViewedCandidateSourceFactory.java
+++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/recommendation/candidatesource/RecentlyViewedCandidateSourceFactory.java
@@ -2,6 +2,8 @@
 
 import com.linkedin.gms.factory.common.IndexConventionFactory;
 import com.linkedin.gms.factory.common.RestHighLevelClientFactory;
+import com.linkedin.gms.factory.entity.EntityServiceFactory;
+import com.linkedin.metadata.entity.EntityService;
 import com.linkedin.metadata.recommendation.candidatesource.RecentlyViewedSource;
 import com.linkedin.metadata.utils.elasticsearch.IndexConvention;
 import javax.annotation.Nonnull;
@@ -14,7 +16,7 @@
 
 
 @Configuration
-@Import({RestHighLevelClientFactory.class, IndexConventionFactory.class})
+@Import({RestHighLevelClientFactory.class, IndexConventionFactory.class, EntityServiceFactory.class})
 public class RecentlyViewedCandidateSourceFactory {
   @Autowired
   @Qualifier("elasticSearchRestHighLevelClient")
@@ -24,9 +26,13 @@ public class RecentlyViewedCandidateSourceFactory {
   @Qualifier(IndexConventionFactory.INDEX_CONVENTION_BEAN)
   private IndexConvention indexConvention;
 
+  @Autowired
+  @Qualifier("entityService")
+  private EntityService entityService;
+
   @Bean(name = "recentlyViewedCandidateSource")
   @Nonnull
   protected RecentlyViewedSource getInstance() {
-    return new RecentlyViewedSource(searchClient, indexConvention);
+    return new RecentlyViewedSource(searchClient, indexConvention, entityService);
   }
 }
diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/search/SearchServiceFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/search/SearchServiceFactory.java
index a4945a2b12f40..8fd6641d47ee9 100644
--- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/search/SearchServiceFactory.java
+++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/search/SearchServiceFactory.java
@@ -38,10 +38,13 @@ public class SearchServiceFactory {
   @Value("${searchService.resultBatchSize}")
   private Integer batchSize;
 
+  @Value("${searchService.enableCache}")
+  private Boolean enableCache;
+
   @Bean(name = "searchService")
   @Primary
   @Nonnull
   protected SearchService getInstance() {
-    return new SearchService(entityRegistry, entitySearchService, searchRanker, cacheManager, batchSize);
+    return new SearchService(entityRegistry, entitySearchService, searchRanker, cacheManager, batchSize, enableCache);
   }
 }
diff --git a/metadata-service/factories/src/main/resources/application.yml b/metadata-service/factories/src/main/resources/application.yml
index 52ab6d5d72021..77da07d98b328 100644
--- a/metadata-service/factories/src/main/resources/application.yml
+++ b/metadata-service/factories/src/main/resources/application.yml
@@ -28,6 +28,10 @@ ingestion:
   enabled: ${UI_INGESTION_ENABLED:true}
   defaultCliVersion: '${UI_INGESTION_DEFAULT_CLI_VERSION:0.8.26.6}'
 
+telemetry:
+  enabledCli: ${CLI_TELEMETRY_ENABLED:true}
+  enabledIngestion: ${INGESTION_REPORTING_ENABLED:false}
+
 secretService:
   encryptionKey: ${SECRET_SERVICE_ENCRYPTION_KEY:ENCRYPTION_KEY}
 
@@ -54,6 +58,7 @@ graphService:
 
 searchService:
   resultBatchSize: ${SEARCH_SERVICE_BATCH_SIZE:100}
+  enableCache: ${SEARCH_SERVICE_ENABLE_CACHE:false}
 
 configEntityRegistry:
   # TODO: Change to read from resources on classpath.
diff --git a/metadata-service/restli-api/src/main/idl/com.linkedin.entity.runs.restspec.json b/metadata-service/restli-api/src/main/idl/com.linkedin.entity.runs.restspec.json
index 9501d0e7d5858..501b1cbab8f1a 100644
--- a/metadata-service/restli-api/src/main/idl/com.linkedin.entity.runs.restspec.json
+++ b/metadata-service/restli-api/src/main/idl/com.linkedin.entity.runs.restspec.json
@@ -21,6 +21,10 @@
         "name" : "pageSize",
         "type" : "int",
         "optional" : true
+      }, {
+        "name" : "includeSoft",
+        "type" : "boolean",
+        "optional" : true
       } ],
       "returns" : "{ \"type\" : \"array\", \"items\" : \"com.linkedin.metadata.run.IngestionRunSummary\" }"
     }, {
@@ -33,6 +37,10 @@
         "name" : "dryRun",
         "type" : "boolean",
         "optional" : true
+      }, {
+        "name" : "hardDelete",
+        "type" : "boolean",
+        "optional" : true
       } ],
       "returns" : "com.linkedin.metadata.run.RollbackResponse"
     } ],
diff --git a/metadata-service/restli-api/src/main/snapshot/com.linkedin.entity.aspects.snapshot.json b/metadata-service/restli-api/src/main/snapshot/com.linkedin.entity.aspects.snapshot.json
index 690b50734a6aa..8327bc5e12fc2 100644
--- a/metadata-service/restli-api/src/main/snapshot/com.linkedin.entity.aspects.snapshot.json
+++ b/metadata-service/restli-api/src/main/snapshot/com.linkedin.entity.aspects.snapshot.json
@@ -1022,15 +1022,27 @@
       "type" : {
         "type" : "enum",
         "name" : "OwnershipType",
-        "doc" : "Owner category or owner role",
-        "symbols" : [ "DEVELOPER", "DATAOWNER", "DELEGATE", "PRODUCER", "CONSUMER", "STAKEHOLDER" ],
+        "doc" : "Asset owner types",
+        "symbols" : [ "TECHNICAL_OWNER", "BUSINESS_OWNER", "DATA_STEWARD", "NONE", "DEVELOPER", "DATAOWNER", "DELEGATE", "PRODUCER", "CONSUMER", "STAKEHOLDER" ],
         "symbolDocs" : {
-          "CONSUMER" : "A person, group, or service that consumes the data",
-          "DATAOWNER" : "A person or group that is owning the data",
-          "DELEGATE" : "A person or a group that overseas the operation, e.g. a DBA or SRE.",
-          "DEVELOPER" : "A person or group that is in charge of developing the code",
-          "PRODUCER" : "A person, group, or service that produces/generates the data",
-          "STAKEHOLDER" : "A person or a group that has direct business interest"
+          "BUSINESS_OWNER" : "A person or group who is responsible for logical, or business related, aspects of the asset.",
+          "CONSUMER" : "A person, group, or service that consumes the data\nDeprecated! Use TECHNICAL_OWNER or BUSINESS_OWNER instead.",
+          "DATAOWNER" : "A person or group that is owning the data\nDeprecated! Use TECHNICAL_OWNER instead.",
+          "DATA_STEWARD" : "A steward, expert, or delegate responsible for the asset.",
+          "DELEGATE" : "A person or a group that overseas the operation, e.g. a DBA or SRE.\nDeprecated! Use TECHNICAL_OWNER instead.",
+          "DEVELOPER" : "A person or group that is in charge of developing the code\nDeprecated! Use TECHNICAL_OWNER instead.",
+          "NONE" : "No specific type associated to the owner.",
+          "PRODUCER" : "A person, group, or service that produces/generates the data\nDeprecated! Use TECHNICAL_OWNER instead.",
+          "STAKEHOLDER" : "A person or a group that has direct business interest\nDeprecated! Use TECHNICAL_OWNER, BUSINESS_OWNER, or STEWARD instead.",
+          "TECHNICAL_OWNER" : "person or group who is responsible for technical aspects of the asset."
+        },
+        "deprecatedSymbols" : {
+          "CONSUMER" : true,
+          "DATAOWNER" : true,
+          "DELEGATE" : true,
+          "DEVELOPER" : true,
+          "PRODUCER" : true,
+          "STAKEHOLDER" : true
         }
       },
       "doc" : "The type of the ownership"
@@ -1549,9 +1561,10 @@
       "doc" : "Fully-qualified name of the Dataset",
       "optional" : true,
       "Searchable" : {
+        "addToFilters" : false,
         "boostScore" : 10.0,
         "enableAutocomplete" : true,
-        "fieldType" : "TEXT_PARTIAL"
+        "fieldType" : "TEXT"
       }
     }, {
       "name" : "description",
diff --git a/metadata-service/restli-api/src/main/snapshot/com.linkedin.entity.entities.snapshot.json b/metadata-service/restli-api/src/main/snapshot/com.linkedin.entity.entities.snapshot.json
index f1b543d0023d6..b0172cb76ae8f 100644
--- a/metadata-service/restli-api/src/main/snapshot/com.linkedin.entity.entities.snapshot.json
+++ b/metadata-service/restli-api/src/main/snapshot/com.linkedin.entity.entities.snapshot.json
@@ -1041,15 +1041,27 @@
       "type" : {
         "type" : "enum",
         "name" : "OwnershipType",
-        "doc" : "Owner category or owner role",
-        "symbols" : [ "DEVELOPER", "DATAOWNER", "DELEGATE", "PRODUCER", "CONSUMER", "STAKEHOLDER" ],
+        "doc" : "Asset owner types",
+        "symbols" : [ "TECHNICAL_OWNER", "BUSINESS_OWNER", "DATA_STEWARD", "NONE", "DEVELOPER", "DATAOWNER", "DELEGATE", "PRODUCER", "CONSUMER", "STAKEHOLDER" ],
         "symbolDocs" : {
-          "CONSUMER" : "A person, group, or service that consumes the data",
-          "DATAOWNER" : "A person or group that is owning the data",
-          "DELEGATE" : "A person or a group that overseas the operation, e.g. a DBA or SRE.",
-          "DEVELOPER" : "A person or group that is in charge of developing the code",
-          "PRODUCER" : "A person, group, or service that produces/generates the data",
-          "STAKEHOLDER" : "A person or a group that has direct business interest"
+          "BUSINESS_OWNER" : "A person or group who is responsible for logical, or business related, aspects of the asset.",
+          "CONSUMER" : "A person, group, or service that consumes the data\nDeprecated! Use TECHNICAL_OWNER or BUSINESS_OWNER instead.",
+          "DATAOWNER" : "A person or group that is owning the data\nDeprecated! Use TECHNICAL_OWNER instead.",
+          "DATA_STEWARD" : "A steward, expert, or delegate responsible for the asset.",
+          "DELEGATE" : "A person or a group that overseas the operation, e.g. a DBA or SRE.\nDeprecated! Use TECHNICAL_OWNER instead.",
+          "DEVELOPER" : "A person or group that is in charge of developing the code\nDeprecated! Use TECHNICAL_OWNER instead.",
+          "NONE" : "No specific type associated to the owner.",
+          "PRODUCER" : "A person, group, or service that produces/generates the data\nDeprecated! Use TECHNICAL_OWNER instead.",
+          "STAKEHOLDER" : "A person or a group that has direct business interest\nDeprecated! Use TECHNICAL_OWNER, BUSINESS_OWNER, or STEWARD instead.",
+          "TECHNICAL_OWNER" : "person or group who is responsible for technical aspects of the asset."
+        },
+        "deprecatedSymbols" : {
+          "CONSUMER" : true,
+          "DATAOWNER" : true,
+          "DELEGATE" : true,
+          "DEVELOPER" : true,
+          "PRODUCER" : true,
+          "STAKEHOLDER" : true
         }
       },
       "doc" : "The type of the ownership"
@@ -1757,9 +1769,10 @@
       "doc" : "Fully-qualified name of the Dataset",
       "optional" : true,
       "Searchable" : {
+        "addToFilters" : false,
         "boostScore" : 10.0,
         "enableAutocomplete" : true,
-        "fieldType" : "TEXT_PARTIAL"
+        "fieldType" : "TEXT"
       }
     }, {
       "name" : "description",
@@ -2541,6 +2554,7 @@
                     "Searchable" : {
                       "boostScore" : 10.0,
                       "enableAutocomplete" : true,
+                      "fieldName" : "id",
                       "fieldType" : "TEXT_PARTIAL"
                     }
                   }, {
@@ -4874,6 +4888,16 @@
       "doc" : "The total number of elements (entities + groups) directly under queried path"
     } ]
   }, "com.linkedin.metadata.browse.BrowseResultEntity", "com.linkedin.metadata.browse.BrowseResultGroup", "com.linkedin.metadata.browse.BrowseResultMetadata", "com.linkedin.metadata.key.ChartKey", "com.linkedin.metadata.key.CorpGroupKey", "com.linkedin.metadata.key.CorpUserKey", "com.linkedin.metadata.key.DashboardKey", "com.linkedin.metadata.key.DataFlowKey", "com.linkedin.metadata.key.DataHubPolicyKey", "com.linkedin.metadata.key.DataHubRetentionKey", "com.linkedin.metadata.key.DataJobKey", "com.linkedin.metadata.key.DataPlatformKey", "com.linkedin.metadata.key.DataProcessKey", "com.linkedin.metadata.key.DatasetKey", "com.linkedin.metadata.key.GlossaryNodeKey", "com.linkedin.metadata.key.GlossaryTermKey", "com.linkedin.metadata.key.MLFeatureKey", "com.linkedin.metadata.key.MLFeatureTableKey", "com.linkedin.metadata.key.MLModelDeploymentKey", "com.linkedin.metadata.key.MLModelGroupKey", "com.linkedin.metadata.key.MLModelKey", "com.linkedin.metadata.key.MLPrimaryKeyKey", "com.linkedin.metadata.key.SchemaFieldKey", "com.linkedin.metadata.key.TagKey", {
+    "type" : "record",
+    "name" : "AutoCompleteEntity",
+    "namespace" : "com.linkedin.metadata.query",
+    "doc" : "Data model for an entity returned as an autocomplete suggestion",
+    "fields" : [ {
+      "name" : "urn",
+      "type" : "com.linkedin.common.Urn",
+      "doc" : "Urn of the entity"
+    } ]
+  }, {
     "type" : "record",
     "name" : "AutoCompleteResult",
     "namespace" : "com.linkedin.metadata.query",
@@ -4888,7 +4912,14 @@
         "type" : "array",
         "items" : "string"
       },
-      "doc" : "A list of typeahead suggestions"
+      "doc" : "A list of typeahead suggestions -- Just the text"
+    }, {
+      "name" : "entities",
+      "type" : {
+        "type" : "array",
+        "items" : "AutoCompleteEntity"
+      },
+      "doc" : "Specific entities to recommend"
     } ]
   }, {
     "type" : "record",
diff --git a/metadata-service/restli-api/src/main/snapshot/com.linkedin.entity.runs.snapshot.json b/metadata-service/restli-api/src/main/snapshot/com.linkedin.entity.runs.snapshot.json
index a5035a7bb57dc..a6f4a9292bb0a 100644
--- a/metadata-service/restli-api/src/main/snapshot/com.linkedin.entity.runs.snapshot.json
+++ b/metadata-service/restli-api/src/main/snapshot/com.linkedin.entity.runs.snapshot.json
@@ -784,15 +784,27 @@
       "type" : {
         "type" : "enum",
         "name" : "OwnershipType",
-        "doc" : "Owner category or owner role",
-        "symbols" : [ "DEVELOPER", "DATAOWNER", "DELEGATE", "PRODUCER", "CONSUMER", "STAKEHOLDER" ],
+        "doc" : "Asset owner types",
+        "symbols" : [ "TECHNICAL_OWNER", "BUSINESS_OWNER", "DATA_STEWARD", "NONE", "DEVELOPER", "DATAOWNER", "DELEGATE", "PRODUCER", "CONSUMER", "STAKEHOLDER" ],
         "symbolDocs" : {
-          "CONSUMER" : "A person, group, or service that consumes the data",
-          "DATAOWNER" : "A person or group that is owning the data",
-          "DELEGATE" : "A person or a group that overseas the operation, e.g. a DBA or SRE.",
-          "DEVELOPER" : "A person or group that is in charge of developing the code",
-          "PRODUCER" : "A person, group, or service that produces/generates the data",
-          "STAKEHOLDER" : "A person or a group that has direct business interest"
+          "BUSINESS_OWNER" : "A person or group who is responsible for logical, or business related, aspects of the asset.",
+          "CONSUMER" : "A person, group, or service that consumes the data\nDeprecated! Use TECHNICAL_OWNER or BUSINESS_OWNER instead.",
+          "DATAOWNER" : "A person or group that is owning the data\nDeprecated! Use TECHNICAL_OWNER instead.",
+          "DATA_STEWARD" : "A steward, expert, or delegate responsible for the asset.",
+          "DELEGATE" : "A person or a group that overseas the operation, e.g. a DBA or SRE.\nDeprecated! Use TECHNICAL_OWNER instead.",
+          "DEVELOPER" : "A person or group that is in charge of developing the code\nDeprecated! Use TECHNICAL_OWNER instead.",
+          "NONE" : "No specific type associated to the owner.",
+          "PRODUCER" : "A person, group, or service that produces/generates the data\nDeprecated! Use TECHNICAL_OWNER instead.",
+          "STAKEHOLDER" : "A person or a group that has direct business interest\nDeprecated! Use TECHNICAL_OWNER, BUSINESS_OWNER, or STEWARD instead.",
+          "TECHNICAL_OWNER" : "person or group who is responsible for technical aspects of the asset."
+        },
+        "deprecatedSymbols" : {
+          "CONSUMER" : true,
+          "DATAOWNER" : true,
+          "DELEGATE" : true,
+          "DEVELOPER" : true,
+          "PRODUCER" : true,
+          "STAKEHOLDER" : true
         }
       },
       "doc" : "The type of the ownership"
@@ -1311,9 +1323,10 @@
       "doc" : "Fully-qualified name of the Dataset",
       "optional" : true,
       "Searchable" : {
+        "addToFilters" : false,
         "boostScore" : 10.0,
         "enableAutocomplete" : true,
-        "fieldType" : "TEXT_PARTIAL"
+        "fieldType" : "TEXT"
       }
     }, {
       "name" : "description",
@@ -3213,6 +3226,10 @@
           "name" : "pageSize",
           "type" : "int",
           "optional" : true
+        }, {
+          "name" : "includeSoft",
+          "type" : "boolean",
+          "optional" : true
         } ],
         "returns" : "{ \"type\" : \"array\", \"items\" : \"com.linkedin.metadata.run.IngestionRunSummary\" }"
       }, {
@@ -3225,6 +3242,10 @@
           "name" : "dryRun",
           "type" : "boolean",
           "optional" : true
+        }, {
+          "name" : "hardDelete",
+          "type" : "boolean",
+          "optional" : true
         } ],
         "returns" : "com.linkedin.metadata.run.RollbackResponse"
       } ],
diff --git a/metadata-service/restli-servlet-impl/build.gradle b/metadata-service/restli-servlet-impl/build.gradle
index b790ae5c83eaf..eafe43f74e8e3 100644
--- a/metadata-service/restli-servlet-impl/build.gradle
+++ b/metadata-service/restli-servlet-impl/build.gradle
@@ -46,7 +46,6 @@ dependencies {
   annotationProcessor externalDependency.lombok
 
   testCompile project(':test-models')
-  testCompile externalDependency.parseqTest
   testCompile externalDependency.mockito
   testCompile externalDependency.testng
 
diff --git a/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/entity/BatchIngestionRunResource.java b/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/entity/BatchIngestionRunResource.java
index 8015d637aa990..90f1eb94dcd0c 100644
--- a/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/entity/BatchIngestionRunResource.java
+++ b/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/entity/BatchIngestionRunResource.java
@@ -7,6 +7,7 @@
 import com.linkedin.metadata.restli.RestliUtil;
 import com.linkedin.metadata.run.AspectRowSummary;
 import com.linkedin.metadata.run.AspectRowSummaryArray;
+import com.linkedin.metadata.run.IngestionRunSummary;
 import com.linkedin.metadata.run.IngestionRunSummaryArray;
 import com.linkedin.metadata.run.RollbackResponse;
 import com.linkedin.metadata.systemmetadata.SystemMetadataService;
@@ -36,6 +37,8 @@ public class BatchIngestionRunResource extends CollectionResourceTaskTemplate<St
 
   private static final Integer DEFAULT_OFFSET = 0;
   private static final Integer DEFAULT_PAGE_SIZE = 100;
+  private static final boolean DEFAULT_INCLUDE_SOFT_DELETED = false;
+  private static final boolean DEFAULT_HARD_DELETE = false;
   private static final Integer ELASTIC_MAX_PAGE_SIZE = 10000;
   private static final Integer ELASTIC_BATCH_DELETE_SLEEP_SEC = 5;
 
@@ -54,8 +57,11 @@ public class BatchIngestionRunResource extends CollectionResourceTaskTemplate<St
   @Nonnull
   @WithSpan
   public Task<RollbackResponse> rollback(@ActionParam("runId") @Nonnull String runId,
-      @ActionParam("dryRun") @Optional @Nullable Boolean dryRun) {
+      @ActionParam("dryRun") @Optional Boolean dryRun, @ActionParam("hardDelete") @Optional Boolean hardDelete) {
     log.info("ROLLBACK RUN runId: {} dry run: {}", runId, dryRun);
+
+    boolean doHardDelete = hardDelete != null ? hardDelete : DEFAULT_HARD_DELETE;
+
     return RestliUtil.toTask(() -> {
       if (runId.equals(EntityService.DEFAULT_RUN_ID)) {
         throw new IllegalArgumentException(String.format(
@@ -64,10 +70,15 @@ public Task<RollbackResponse> rollback(@ActionParam("runId") @Nonnull String run
       }
       RollbackResponse response = new RollbackResponse();
       List<AspectRowSummary> aspectRowsToDelete;
-      aspectRowsToDelete = _systemMetadataService.findByRunId(runId);
+      aspectRowsToDelete = _systemMetadataService.findByRunId(runId, doHardDelete);
 
       log.info("found {} rows to delete...", stringifyRowCount(aspectRowsToDelete.size()));
       if (dryRun) {
+
+        if (!doHardDelete) {
+          aspectRowsToDelete.removeIf(AspectRowSummary::isKeyAspect);
+        }
+
         response.setAspectsAffected(aspectRowsToDelete.size());
         response.setEntitiesAffected(
             aspectRowsToDelete.stream().collect(Collectors.groupingBy(AspectRowSummary::getUrn)).keySet().size());
@@ -77,17 +88,17 @@ public Task<RollbackResponse> rollback(@ActionParam("runId") @Nonnull String run
         return response;
       }
 
-      RollbackRunResult rollbackRunResult = _entityService.rollbackRun(aspectRowsToDelete, runId);
+      RollbackRunResult rollbackRunResult = _entityService.rollbackRun(aspectRowsToDelete, runId, doHardDelete);
       List<AspectRowSummary> deletedRows = rollbackRunResult.getRowsRolledBack();
       Integer rowsDeletedFromEntityDeletion = rollbackRunResult.getRowsDeletedFromEntityDeletion();
 
       // since elastic limits how many rows we can access at once, we need to iteratively delete
       while (aspectRowsToDelete.size() >= ELASTIC_MAX_PAGE_SIZE) {
         sleep(ELASTIC_BATCH_DELETE_SLEEP_SEC);
-        aspectRowsToDelete = _systemMetadataService.findByRunId(runId);
+        aspectRowsToDelete = _systemMetadataService.findByRunId(runId, doHardDelete);
         log.info("{} remaining rows to delete...", stringifyRowCount(aspectRowsToDelete.size()));
         log.info("deleting...");
-        rollbackRunResult = _entityService.rollbackRun(aspectRowsToDelete, runId);
+        rollbackRunResult = _entityService.rollbackRun(aspectRowsToDelete, runId, doHardDelete);
         deletedRows.addAll(rollbackRunResult.getRowsRolledBack());
         rowsDeletedFromEntityDeletion += rollbackRunResult.getRowsDeletedFromEntityDeletion();
       }
@@ -124,11 +135,17 @@ private void sleep(Integer seconds) {
   @Nonnull
   @WithSpan
   public Task<IngestionRunSummaryArray> list(@ActionParam("pageOffset") @Optional @Nullable Integer pageOffset,
-      @ActionParam("pageSize") @Optional @Nullable Integer pageSize) {
+      @ActionParam("pageSize") @Optional @Nullable Integer pageSize,
+      @ActionParam("includeSoft") @Optional @Nullable Boolean includeSoft) {
     log.info("LIST RUNS offset: {} size: {}", pageOffset, pageSize);
 
-    return RestliUtil.toTask(() -> new IngestionRunSummaryArray(
-        _systemMetadataService.listRuns(pageOffset != null ? pageOffset : DEFAULT_OFFSET,
-            pageSize != null ? pageSize : DEFAULT_PAGE_SIZE)), MetricRegistry.name(this.getClass(), "list"));
+    return RestliUtil.toTask(() -> {
+              List<IngestionRunSummary> summaries = _systemMetadataService.listRuns(
+                      pageOffset != null ? pageOffset : DEFAULT_OFFSET,
+                      pageSize != null ? pageSize : DEFAULT_PAGE_SIZE,
+                      includeSoft != null ? includeSoft : DEFAULT_INCLUDE_SOFT_DELETED);
+
+              return new IngestionRunSummaryArray(summaries);
+            }, MetricRegistry.name(this.getClass(), "list"));
   }
 }
diff --git a/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/entity/EntityResource.java b/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/entity/EntityResource.java
index c3203ebce8494..5790629e2900e 100644
--- a/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/entity/EntityResource.java
+++ b/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/entity/EntityResource.java
@@ -368,7 +368,7 @@ public Task<RollbackResponse> deleteEntities(@ActionParam("registryId") @Optiona
     return RestliUtil.toTask(() -> {
       RollbackResponse response = new RollbackResponse();
       List<AspectRowSummary> aspectRowsToDelete =
-          _systemMetadataService.findByRegistry(finalRegistryName, finalRegistryVersion.toString());
+          _systemMetadataService.findByRegistry(finalRegistryName, finalRegistryVersion.toString(), false);
       log.info("found {} rows to delete...", stringifyRowCount(aspectRowsToDelete.size()));
       response.setAspectsAffected(aspectRowsToDelete.size());
       response.setEntitiesAffected(
@@ -380,7 +380,7 @@ public Task<RollbackResponse> deleteEntities(@ActionParam("registryId") @Optiona
         Map<String, String> conditions = new HashMap();
         conditions.put("registryName", finalRegistryName1);
         conditions.put("registryVersion", finalRegistryVersion1.toString());
-        _entityService.rollbackWithConditions(aspectRowsToDelete, conditions);
+        _entityService.rollbackWithConditions(aspectRowsToDelete, conditions, false);
       }
       return response;
     }, MetricRegistry.name(this.getClass(), "deleteAll"));
diff --git a/metadata-service/servlet/build.gradle b/metadata-service/servlet/build.gradle
index 3cf2e62882f25..9242d21201886 100644
--- a/metadata-service/servlet/build.gradle
+++ b/metadata-service/servlet/build.gradle
@@ -9,4 +9,5 @@ dependencies {
   compile externalDependency.springWebMVC
   annotationProcessor externalDependency.lombok
   compile project(':entity-registry')
+  compile project(':metadata-service:factories')
 }
diff --git a/metadata-service/servlet/src/main/java/com/datahub/gms/servlet/Config.java b/metadata-service/servlet/src/main/java/com/datahub/gms/servlet/Config.java
index bef66f40b3159..c857b2b051b7d 100644
--- a/metadata-service/servlet/src/main/java/com/datahub/gms/servlet/Config.java
+++ b/metadata-service/servlet/src/main/java/com/datahub/gms/servlet/Config.java
@@ -20,6 +20,7 @@
 import org.apache.maven.artifact.versioning.ComparableVersion;
 import org.springframework.web.context.WebApplicationContext;
 import org.springframework.web.context.support.WebApplicationContextUtils;
+import com.linkedin.gms.factory.config.ConfigurationProvider;
 
 // Return a 200 for health checks
 
@@ -50,6 +51,10 @@ private Map<String, Map<ComparableVersion, EntityRegistryLoadResult>> getPluginM
     return patchDiagnostics;
   }
 
+  private ConfigurationProvider getConfigProvider(WebApplicationContext ctx) {
+    return (ConfigurationProvider) ctx.getBean("configurationProvider");
+  }
+
   private GitVersion getGitVersion(WebApplicationContext ctx) {
     return (GitVersion) ctx.getBean("gitVersion");
   }
@@ -71,6 +76,20 @@ protected void doGet(HttpServletRequest req, HttpServletResponse resp) throws IO
     versionConfig.put("linkedin/datahub", version.toConfig());
     config.put("versions", versionConfig);
 
+    ConfigurationProvider configProvider = getConfigProvider(ctx);
+
+    Map<String, Object> telemetryConfig = new HashMap<String, Object>() {{
+      put("enabledCli", configProvider.getTelemetry().enabledCli);
+      put("enabledIngestion", configProvider.getTelemetry().enabledIngestion);
+    }};
+    config.put("telemetry", telemetryConfig);
+
+    Map<String, Object> ingestionConfig = new HashMap<String, Object>() {{
+      put("enabled", configProvider.getIngestion().enabled);
+      put("defaultCliVersion", configProvider.getIngestion().defaultCliVersion);
+    }};
+    config.put("managedIngestion", ingestionConfig);
+
     resp.setContentType("application/json");
     PrintWriter out = resp.getWriter();
 
diff --git a/metadata-service/war/src/main/resources/boot/data_platforms.json b/metadata-service/war/src/main/resources/boot/data_platforms.json
index a614c9010706d..87db83fad45dd 100644
--- a/metadata-service/war/src/main/resources/boot/data_platforms.json
+++ b/metadata-service/war/src/main/resources/boot/data_platforms.json
@@ -405,5 +405,14 @@
       "type": "OTHERS",
       "logoUrl": "/assets/platforms/powerbilogo.png"
     }
+  },
+  {
+    "urn": "urn:li:dataPlatform:unknown",
+    "aspect": {
+      "datasetNameDelimiter": ".",
+      "name": "Unknown Platform",
+      "displayName": "N/A",
+      "type": "OTHERS"
+    }
   }
 ]
\ No newline at end of file
diff --git a/metadata-utils/src/main/java/com/linkedin/metadata/Constants.java b/metadata-utils/src/main/java/com/linkedin/metadata/Constants.java
index 4156fdcb5f333..b90e1884b5f01 100644
--- a/metadata-utils/src/main/java/com/linkedin/metadata/Constants.java
+++ b/metadata-utils/src/main/java/com/linkedin/metadata/Constants.java
@@ -11,6 +11,9 @@ public class Constants {
   public static final String SYSTEM_ACTOR = "urn:li:corpuser:__datahub_system"; // DataHub internal service principal.
   public static final String UNKNOWN_ACTOR = "urn:li:corpuser:UNKNOWN"; // Unknown principal.
   public static final Long ASPECT_LATEST_VERSION = 0L;
+  public static final String UNKNOWN_DATA_PLATFORM = "urn:li:dataPlatform:unknown";
+
+  public static final String DEFAULT_RUN_ID = "no-run-id-provided";
 
   /**
    * Entities
diff --git a/metadata-utils/src/main/java/com/linkedin/metadata/utils/AuditStampUtils.java b/metadata-utils/src/main/java/com/linkedin/metadata/utils/AuditStampUtils.java
new file mode 100644
index 0000000000000..913def2a040f4
--- /dev/null
+++ b/metadata-utils/src/main/java/com/linkedin/metadata/utils/AuditStampUtils.java
@@ -0,0 +1,21 @@
+package com.linkedin.metadata.utils;
+
+import com.linkedin.common.AuditStamp;
+import com.linkedin.common.urn.UrnUtils;
+import lombok.extern.slf4j.Slf4j;
+
+import java.time.Clock;
+
+import static com.linkedin.metadata.Constants.SYSTEM_ACTOR;
+
+@Slf4j
+public class AuditStampUtils {
+    private AuditStampUtils() {
+    }
+
+    public static AuditStamp createDefaultAuditStamp() {
+        return new AuditStamp()
+                .setActor(UrnUtils.getUrn(SYSTEM_ACTOR))
+                .setTime(Clock.systemUTC().millis());
+    }
+}
diff --git a/metadata-utils/src/main/java/com/linkedin/metadata/utils/SystemMetadataUtils.java b/metadata-utils/src/main/java/com/linkedin/metadata/utils/SystemMetadataUtils.java
new file mode 100644
index 0000000000000..f4be950575624
--- /dev/null
+++ b/metadata-utils/src/main/java/com/linkedin/metadata/utils/SystemMetadataUtils.java
@@ -0,0 +1,18 @@
+package com.linkedin.metadata.utils;
+
+import com.linkedin.metadata.Constants;
+import com.linkedin.mxe.SystemMetadata;
+import lombok.extern.slf4j.Slf4j;
+
+@Slf4j
+public class SystemMetadataUtils {
+
+    private SystemMetadataUtils() {
+    }
+
+    public static SystemMetadata createDefaultSystemMetadata() {
+        return new SystemMetadata()
+                .setRunId(Constants.DEFAULT_RUN_ID)
+                .setLastObserved(System.currentTimeMillis());
+    }
+}
diff --git a/settings.gradle b/settings.gradle
index 3ec3a2856239f..911c2774e9f61 100644
--- a/settings.gradle
+++ b/settings.gradle
@@ -42,4 +42,5 @@ include 'metadata-models-custom'
 include 'entity-registry:custom-test-model'
 include 'metadata-integration:java:spark-lineage'
 include 'metadata-integration:java:datahub-client'
+include 'metadata-integration:java:datahub-protobuf'
 include 'ingestion-scheduler'
diff --git a/smoke-test/tests/cli/__init__.py b/smoke-test/tests/cli/__init__.py
new file mode 100644
index 0000000000000..e69de29bb2d1d
diff --git a/smoke-test/tests/cli/cli_editable_test_data.json b/smoke-test/tests/cli/cli_editable_test_data.json
new file mode 100644
index 0000000000000..84b5b1c5e43e5
--- /dev/null
+++ b/smoke-test/tests/cli/cli_editable_test_data.json
@@ -0,0 +1,15 @@
+[
+  {
+    "auditHeader": null,
+    "entityType": "dataset",
+    "entityUrn":  "urn:li:dataset:(urn:li:dataPlatform:kafka,test-rollback,PROD)",
+    "entityKeyAspect": null,
+    "changeType": "UPSERT",
+    "aspectName": "editableDatasetProperties",
+    "aspect": {
+      "value": "{\"description\": \"Some Documentation\" }",
+      "contentType": "application/json"
+    },
+    "systemMetadata": null
+  }
+]
\ No newline at end of file
diff --git a/smoke-test/tests/cli/cli_test_data.json b/smoke-test/tests/cli/cli_test_data.json
new file mode 100644
index 0000000000000..07f7978f12d73
--- /dev/null
+++ b/smoke-test/tests/cli/cli_test_data.json
@@ -0,0 +1,143 @@
+[
+  {
+    "auditHeader": null,
+    "proposedSnapshot": {
+      "com.linkedin.pegasus2avro.metadata.snapshot.DatasetSnapshot": {
+        "urn": "urn:li:dataset:(urn:li:dataPlatform:kafka,test-rollback,PROD)",
+        "aspects": [
+          {
+            "com.linkedin.pegasus2avro.common.BrowsePaths": {
+              "paths": ["/prod/kafka/SampleKafkaDataset"]
+            }
+          },
+          {
+            "com.linkedin.pegasus2avro.dataset.DatasetProperties": {
+              "description": null,
+              "uri": null,
+              "tags": [],
+              "customProperties": {
+                "prop1": "fakeprop",
+                "prop2": "pikachu"
+              }
+            }
+          },
+          {
+            "com.linkedin.pegasus2avro.common.Ownership": {
+              "owners": [
+                {
+                  "owner": "urn:li:corpuser:jdoe",
+                  "type": "DATAOWNER",
+                  "source": null
+                },
+                {
+                  "owner": "urn:li:corpuser:datahub",
+                  "type": "DATAOWNER",
+                  "source": null
+                }
+              ],
+              "lastModified": {
+                "time": 1581407189000,
+                "actor": "urn:li:corpuser:jdoe",
+                "impersonator": null
+              }
+            }
+          },
+          {
+            "com.linkedin.pegasus2avro.common.InstitutionalMemory": {
+              "elements": [
+                {
+                  "url": "https://www.linkedin.com",
+                  "description": "Sample doc",
+                  "createStamp": {
+                    "time": 1581407189000,
+                    "actor": "urn:li:corpuser:jdoe",
+                    "impersonator": null
+                  }
+                }
+              ]
+            }
+          },
+          {
+            "com.linkedin.pegasus2avro.schema.SchemaMetadata": {
+              "schemaName": "SampleKafkaSchema",
+              "platform": "urn:li:dataPlatform:kafka",
+              "version": 0,
+              "created": {
+                "time": 1581407189000,
+                "actor": "urn:li:corpuser:jdoe",
+                "impersonator": null
+              },
+              "lastModified": {
+                "time": 1581407189000,
+                "actor": "urn:li:corpuser:jdoe",
+                "impersonator": null
+              },
+              "deleted": null,
+              "dataset": null,
+              "cluster": null,
+              "hash": "",
+              "platformSchema": {
+                "com.linkedin.pegasus2avro.schema.KafkaSchema": {
+                  "documentSchema": "{\"type\":\"record\",\"name\":\"SampleKafkaSchema\",\"namespace\":\"com.linkedin.dataset\",\"doc\":\"Sample Kafka dataset\",\"fields\":[{\"name\":\"field_foo\",\"type\":[\"string\"]},{\"name\":\"field_bar\",\"type\":[\"boolean\"]}]}"
+                }
+              },
+              "fields": [
+                {
+                  "fieldPath": "[version=2.0].[type=boolean].field_foo_2",
+                  "jsonPath": null,
+                  "nullable": false,
+                  "description": {
+                    "string": "Foo field description"
+                  },
+                  "type": {
+                    "type": {
+                      "com.linkedin.pegasus2avro.schema.BooleanType": {}
+                    }
+                  },
+                  "nativeDataType": "varchar(100)",
+                  "globalTags": {
+                    "tags": [{ "tag": "urn:li:tag:NeedsDocumentation" }]
+                  },
+                  "recursive": false
+                },
+                {
+                  "fieldPath": "[version=2.0].[type=boolean].field_bar",
+                  "jsonPath": null,
+                  "nullable": false,
+                  "description": {
+                    "string": "Bar field description"
+                  },
+                  "type": {
+                    "type": {
+                      "com.linkedin.pegasus2avro.schema.BooleanType": {}
+                    }
+                  },
+                  "nativeDataType": "boolean",
+                  "recursive": false
+                },
+                {
+                  "fieldPath": "[version=2.0].[key=True].[type=int].id",
+                  "jsonPath": null,
+                  "nullable": false,
+                  "description": {
+                    "string": "Id specifying which partition the message should go to"
+                  },
+                  "type": {
+                    "type": {
+                      "com.linkedin.pegasus2avro.schema.BooleanType": {}
+                    }
+                  },
+                  "nativeDataType": "boolean",
+                  "recursive": false
+                }
+              ],
+              "primaryKeys": null,
+              "foreignKeysSpecs": null
+            }
+          }
+        ]
+      }
+    },
+    "proposedDelta": null
+  }
+]
\ No newline at end of file
diff --git a/smoke-test/tests/cli/datahub-cli.py b/smoke-test/tests/cli/datahub-cli.py
new file mode 100644
index 0000000000000..910ad5f64ced5
--- /dev/null
+++ b/smoke-test/tests/cli/datahub-cli.py
@@ -0,0 +1,66 @@
+import json
+import pytest
+from time import sleep
+from datahub.cli import delete_cli, ingest_cli
+from datahub.cli.cli_utils import guess_entity_type, post_entity, get_aspects_for_entity
+from datahub.cli.ingest_cli import get_session_and_host
+from datahub.cli.delete_cli import guess_entity_type, delete_one_urn_cmd
+from tests.utils import ingest_file_via_rest, delete_urns_from_file
+
+def delete_by_urn(urn, session, host):
+    entity_type = guess_entity_type(urn=urn)
+    delete_one_urn_cmd(
+        urn,
+        soft=False,
+        dry_run=False,
+        entity_type=entity_type,
+        cached_session_host=(session, host),
+    )
+    sleep(5)
+    print()
+
+@pytest.mark.dependency()
+def test_rollback_editable():
+    platform = "urn:li:dataPlatform:kafka"
+    dataset_name = (
+        "test-rollback"
+    )
+    env = "PROD"
+    dataset_urn = f"urn:li:dataset:({platform},{dataset_name},{env})"
+
+    session, gms_host = get_session_and_host()
+
+    # Clean slate.
+    delete_by_urn(dataset_urn, session, gms_host)
+
+    assert "browsePaths" not in get_aspects_for_entity(entity_urn=dataset_urn, aspects=["browsePaths"], typed=False)
+    assert "editableDatasetProperties" not in get_aspects_for_entity(entity_urn=dataset_urn, aspects=["editableDatasetProperties"], typed=False)
+
+    # Ingest dataset
+    ingested_dataset_run_id = ingest_file_via_rest("tests/cli/cli_test_data.json").config.run_id
+    print("Ingested dataset id:", ingested_dataset_run_id)
+    # Assert that second data ingestion worked
+    assert "browsePaths" in get_aspects_for_entity(entity_urn=dataset_urn, aspects=["browsePaths"], typed=False)
+
+    # Sleep forces ingestion of files to have distinct run-ids.
+    sleep(1)
+
+    # Make editable change
+    ingested_editable_run_id = ingest_file_via_rest("tests/cli/cli_editable_test_data.json").config.run_id
+    print("ingested editable id:", ingested_editable_run_id)
+    # Assert that second data ingestion worked
+    assert "editableDatasetProperties" in get_aspects_for_entity(entity_urn=dataset_urn, aspects=["editableDatasetProperties"], typed=False)
+
+    # rollback ingestion 1
+    rollback_url = f"{gms_host}/runs?action=rollback"
+
+    session.post(rollback_url, data=json.dumps({"runId": ingested_dataset_run_id, "dryRun": False, "hardDelete": False}))
+
+    # Allow async MCP processor to handle ingestions & rollbacks
+    sleep(10)
+
+    # EditableDatasetProperties should still be part of the entity that was soft deleted.
+    assert "editableDatasetProperties" in get_aspects_for_entity(entity_urn=dataset_urn, aspects=["editableDatasetProperties"], typed=False)
+    # But first ingestion aspects should not be present
+    assert "browsePaths" not in get_aspects_for_entity(entity_urn=dataset_urn, aspects=["browsePaths"], typed=False)
+    pass
\ No newline at end of file
diff --git a/smoke-test/tests/utils.py b/smoke-test/tests/utils.py
index b16931898fce7..354a98c8e473d 100644
--- a/smoke-test/tests/utils.py
+++ b/smoke-test/tests/utils.py
@@ -1,13 +1,14 @@
 import json
 
 import requests
+from typing import Any
 from datahub.cli import cli_utils
 from datahub.ingestion.run.pipeline import Pipeline
 
 GMS_ENDPOINT = "http://localhost:8080"
 FRONTEND_ENDPOINT = "http://localhost:9002"
 
-def ingest_file_via_rest(filename: str) -> None:
+def ingest_file_via_rest(filename: str) -> Any:
     pipeline = Pipeline.create(
         {
             "source": {
@@ -23,6 +24,8 @@ def ingest_file_via_rest(filename: str) -> None:
     pipeline.run()
     pipeline.raise_from_status()
 
+    return pipeline
+
 
 def delete_urns_from_file(filename: str) -> None:
     session = requests.Session()