From 9928e942bdbaa634907673831356244b0aa8ccbb Mon Sep 17 00:00:00 2001
From: Pedro Silva <pedro@acryl.io>
Date: Tue, 6 Sep 2022 19:10:56 +0100
Subject: [PATCH] feat(hooks): Add toggle for enabling/disabling platform event
 hook

---
 .../metadata/kafka/MetadataChangeLogProcessor.java    |  3 +++
 .../metadata/kafka/hook/MetadataChangeLogHook.java    |  7 +++++++
 .../hook/event/EntityChangeEventGeneratorHook.java    | 11 ++++++++++-
 .../event/EntityChangeEventGeneratorHookTest.java     |  3 ++-
 .../factories/src/main/resources/application.yml      |  3 +++
 5 files changed, 25 insertions(+), 2 deletions(-)

diff --git a/metadata-jobs/mae-consumer/src/main/java/com/linkedin/metadata/kafka/MetadataChangeLogProcessor.java b/metadata-jobs/mae-consumer/src/main/java/com/linkedin/metadata/kafka/MetadataChangeLogProcessor.java
index 02946464d5fab3..6700dec65a11d8 100644
--- a/metadata-jobs/mae-consumer/src/main/java/com/linkedin/metadata/kafka/MetadataChangeLogProcessor.java
+++ b/metadata-jobs/mae-consumer/src/main/java/com/linkedin/metadata/kafka/MetadataChangeLogProcessor.java
@@ -81,6 +81,9 @@ public void consume(final ConsumerRecord<String, GenericRecord> consumerRecord)
 
     // Here - plug in additional "custom processor hooks"
     for (MetadataChangeLogHook hook : this.hooks) {
+      if (!hook.isEnabled()) {
+        continue;
+      }
       try (Timer.Context ignored = MetricUtils.timer(this.getClass(), hook.getClass().getSimpleName() + "_latency")
           .time()) {
         hook.invoke(event);
diff --git a/metadata-jobs/mae-consumer/src/main/java/com/linkedin/metadata/kafka/hook/MetadataChangeLogHook.java b/metadata-jobs/mae-consumer/src/main/java/com/linkedin/metadata/kafka/hook/MetadataChangeLogHook.java
index bc5ab2a01d055e..c7857eb7baffc1 100644
--- a/metadata-jobs/mae-consumer/src/main/java/com/linkedin/metadata/kafka/hook/MetadataChangeLogHook.java
+++ b/metadata-jobs/mae-consumer/src/main/java/com/linkedin/metadata/kafka/hook/MetadataChangeLogHook.java
@@ -18,6 +18,13 @@ public interface MetadataChangeLogHook {
    */
   default void init() { }
 
+  /**
+   * Return whether the hook is enabled or not. If not enabled, the below invoke method is not triggered
+   */
+  default boolean isEnabled() {
+    return true;
+  }
+
   /**
    * Invoke the hook when a MetadataChangeLog is received
    */
diff --git a/metadata-jobs/mae-consumer/src/main/java/com/linkedin/metadata/kafka/hook/event/EntityChangeEventGeneratorHook.java b/metadata-jobs/mae-consumer/src/main/java/com/linkedin/metadata/kafka/hook/event/EntityChangeEventGeneratorHook.java
index c0603d07cb6e48..1b78f77687e33d 100644
--- a/metadata-jobs/mae-consumer/src/main/java/com/linkedin/metadata/kafka/hook/event/EntityChangeEventGeneratorHook.java
+++ b/metadata-jobs/mae-consumer/src/main/java/com/linkedin/metadata/kafka/hook/event/EntityChangeEventGeneratorHook.java
@@ -35,6 +35,7 @@
 import javax.annotation.Nullable;
 import lombok.extern.slf4j.Slf4j;
 import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.beans.factory.annotation.Value;
 import org.springframework.context.annotation.Import;
 import org.springframework.stereotype.Component;
 
@@ -91,17 +92,25 @@ public class EntityChangeEventGeneratorHook implements MetadataChangeLogHook {
   private final EntityClient _entityClient;
   private final Authentication _systemAuthentication;
   private final EntityRegistry _entityRegistry;
+  private final Boolean _isEnabled;
 
   @Autowired
   public EntityChangeEventGeneratorHook(
       @Nonnull final AspectDifferRegistry aspectDifferRegistry,
       @Nonnull final RestliEntityClient entityClient,
       @Nonnull final Authentication systemAuthentication,
-      @Nonnull final EntityRegistry entityRegistry) {
+      @Nonnull final EntityRegistry entityRegistry,
+      @Nonnull @Value("${entityChangeEvents.enabled:true}") Boolean isEnabled) {
     _aspectDifferRegistry = Objects.requireNonNull(aspectDifferRegistry);
     _entityClient = Objects.requireNonNull(entityClient);
     _systemAuthentication = Objects.requireNonNull(systemAuthentication);
     _entityRegistry = Objects.requireNonNull(entityRegistry);
+    _isEnabled = isEnabled;
+  }
+
+  @Override
+  public boolean isEnabled() {
+    return _isEnabled;
   }
 
   @Override
diff --git a/metadata-jobs/mae-consumer/src/test/java/com/linkedin/metadata/kafka/hook/event/EntityChangeEventGeneratorHookTest.java b/metadata-jobs/mae-consumer/src/test/java/com/linkedin/metadata/kafka/hook/event/EntityChangeEventGeneratorHookTest.java
index 4c87d2c93daf02..e2cecfbeb20d2f 100644
--- a/metadata-jobs/mae-consumer/src/test/java/com/linkedin/metadata/kafka/hook/event/EntityChangeEventGeneratorHookTest.java
+++ b/metadata-jobs/mae-consumer/src/test/java/com/linkedin/metadata/kafka/hook/event/EntityChangeEventGeneratorHookTest.java
@@ -79,7 +79,8 @@ public void setupTest() {
         differRegistry,
         _mockClient,
         mockAuthentication,
-        createMockEntityRegistry());
+        createMockEntityRegistry(),
+        true);
   }
 
   @Test
diff --git a/metadata-service/factories/src/main/resources/application.yml b/metadata-service/factories/src/main/resources/application.yml
index edddf4a0ab701f..be259a12724a53 100644
--- a/metadata-service/factories/src/main/resources/application.yml
+++ b/metadata-service/factories/src/main/resources/application.yml
@@ -199,3 +199,6 @@ siblings:
 
 featureFlags:
   showSimplifiedHomepageByDefault: ${SHOW_SIMPLIFIED_HOMEPAGE_BY_DEFAULT:false} # shows a simplified homepage with just datasets, charts and dashboards by default to users. this can be configured in user settings
+
+entityChangeEvents:
+  enabled: ${ENABLE_ENTITY_CHANGE_EVENTS_HOOK:true}