feat(template-mcps): allow further control for helm (#11816)

datahub-project · Nov 7, 2024 · b3d3b63 · b3d3b63
1 parent 33751bf
commit b3d3b63
Show file tree

Hide file tree

Showing 7 changed files with 75 additions and 6 deletions.
diff --git a/...ade/src/main/java/com/linkedin/datahub/upgrade/system/bootstrapmcps/BootstrapMCPUtil.java b/...ade/src/main/java/com/linkedin/datahub/upgrade/system/bootstrapmcps/BootstrapMCPUtil.java
@@ -54,6 +54,7 @@ static List<UpgradeStep> generateSteps(
             .getBootstrap()
             .getTemplates()
             .stream()
+            .map(cfg -> cfg.withOverride(opContext.getObjectMapper()))
             .filter(cfg -> cfg.isBlocking() == isBlocking)
             .map(cfg -> new BootstrapMCPStep(opContext, entityService, cfg))
             .collect(Collectors.toList());

diff --git a/.../java/com/linkedin/datahub/upgrade/system/bootstrapmcps/model/BootstrapMCPConfigFile.java b/.../java/com/linkedin/datahub/upgrade/system/bootstrapmcps/model/BootstrapMCPConfigFile.java
@@ -1,12 +1,15 @@
 package com.linkedin.datahub.upgrade.system.bootstrapmcps.model;
 
+import com.fasterxml.jackson.databind.ObjectMapper;
+import java.io.IOException;
 import java.util.List;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
 import lombok.AllArgsConstructor;
 import lombok.Builder;
 import lombok.Data;
 import lombok.NoArgsConstructor;
+import lombok.extern.slf4j.Slf4j;
 
 @AllArgsConstructor
 @NoArgsConstructor
@@ -23,6 +26,7 @@ public static class Bootstrap {
     private List<MCPTemplate> templates;
   }
 
+  @Slf4j
   @AllArgsConstructor
   @NoArgsConstructor
   @Data
@@ -36,5 +40,19 @@ public static class MCPTemplate {
     @Builder.Default private boolean optional = false;
     @Nonnull private String mcps_location;
     @Nullable private String values_env;
+    @Nullable private String revision_env;
+
+    public MCPTemplate withOverride(ObjectMapper objectMapper) {
+      if (revision_env != null) {
+        String overrideJson = System.getenv().getOrDefault(revision_env, "{}");
+        try {
+          return objectMapper.readerForUpdating(this).readValue(overrideJson);
+        } catch (IOException e) {
+          log.error("Error applying override {} to {}", overrideJson, this);
+          throw new RuntimeException(e);
+        }
+      }
+      return this;
+    }
   }
 }
diff --git a/...src/test/java/com/linkedin/datahub/upgrade/system/bootstrapmcps/BootstrapMCPUtilTest.java b/...src/test/java/com/linkedin/datahub/upgrade/system/bootstrapmcps/BootstrapMCPUtilTest.java
@@ -4,6 +4,7 @@
 import static org.testng.Assert.assertFalse;
 import static org.testng.Assert.assertTrue;
 
+import com.fasterxml.jackson.databind.ObjectMapper;
 import com.fasterxml.jackson.databind.node.ObjectNode;
 import com.linkedin.common.AuditStamp;
 import com.linkedin.common.urn.UrnUtils;
@@ -17,6 +18,7 @@
 import io.datahubproject.test.metadata.context.TestOperationContexts;
 import java.io.IOException;
 import java.util.List;
+import org.testng.annotations.BeforeMethod;
 import org.testng.annotations.Listeners;
 import org.testng.annotations.Test;
 import uk.org.webcompere.systemstubs.environment.EnvironmentVariables;
@@ -28,10 +30,17 @@ public class BootstrapMCPUtilTest {
   static final OperationContext OP_CONTEXT =
       TestOperationContexts.systemContextNoSearchAuthorization();
   private static final String DATAHUB_TEST_VALUES_ENV = "DATAHUB_TEST_VALUES_ENV";
+  private static final String DATAHUB_TEST_REVISION_ENV = "DATAHUB_TEST_REVISION_ENV";
   private static final AuditStamp TEST_AUDIT_STAMP = AuditStampUtils.createDefaultAuditStamp();
 
   @SystemStub private EnvironmentVariables environmentVariables;
 
+  @BeforeMethod
+  private void resetEnvironment() {
+    environmentVariables.remove(DATAHUB_TEST_VALUES_ENV);
+    environmentVariables.remove(DATAHUB_TEST_REVISION_ENV);
+  }
+
   @Test
   public void testResolveYamlConf() throws IOException {
     BootstrapMCPConfigFile initConfig =
@@ -51,9 +60,28 @@ public void testResolveYamlConf() throws IOException {
   }
 
   @Test
-  public void testResolveMCPTemplateDefaults() throws IOException {
-    environmentVariables.remove(DATAHUB_TEST_VALUES_ENV);
+  public void testResolveYamlConfOverride() throws IOException {
+    environmentVariables.set(DATAHUB_TEST_REVISION_ENV, "{\"version\":\"2024110600\"}");
+
+    BootstrapMCPConfigFile initConfig =
+        BootstrapMCPUtil.resolveYamlConf(
+            OP_CONTEXT, "bootstrapmcp/test.yaml", BootstrapMCPConfigFile.class);
+    assertEquals(initConfig.getBootstrap().getTemplates().size(), 1);
+
+    BootstrapMCPConfigFile.MCPTemplate template =
+        initConfig.getBootstrap().getTemplates().get(0).withOverride(new ObjectMapper());
+    assertEquals(template.getName(), "datahub-test");
+    assertEquals(template.getVersion(), "2024110600");
+    assertFalse(template.isForce());
+    assertFalse(template.isBlocking());
+    assertTrue(template.isAsync());
+    assertFalse(template.isOptional());
+    assertEquals(template.getMcps_location(), "bootstrapmcp/datahub-test-mcp.yaml");
+    assertEquals(template.getValues_env(), "DATAHUB_TEST_VALUES_ENV");
+  }
 
+  @Test
+  public void testResolveMCPTemplateDefaults() throws IOException {
     BootstrapMCPConfigFile.MCPTemplate template =
         BootstrapMCPUtil.resolveYamlConf(
                 OP_CONTEXT, "bootstrapmcp/test.yaml", BootstrapMCPConfigFile.class)
@@ -186,8 +214,6 @@ public void testResolveMCPTemplateOverride() throws IOException {
 
   @Test
   public void testMCPBatch() throws IOException {
-    environmentVariables.remove(DATAHUB_TEST_VALUES_ENV);
-
     BootstrapMCPConfigFile.MCPTemplate template =
         BootstrapMCPUtil.resolveYamlConf(
                 OP_CONTEXT, "bootstrapmcp/test.yaml", BootstrapMCPConfigFile.class)

diff --git a/datahub-upgrade/src/test/resources/bootstrapmcp/test.yaml b/datahub-upgrade/src/test/resources/bootstrapmcp/test.yaml
@@ -6,4 +6,5 @@ bootstrap:
       # blocking: false
       # async: true
       mcps_location: "bootstrapmcp/datahub-test-mcp.yaml"
-      values_env: "DATAHUB_TEST_VALUES_ENV"
+      values_env: "DATAHUB_TEST_VALUES_ENV"
+      revision_env: "DATAHUB_TEST_REVISION_ENV"
diff --git a/docs/advanced/bootstrap-mcps.md b/docs/advanced/bootstrap-mcps.md
@@ -149,6 +149,28 @@ to the required json structure and stored as a string.
       executorId: default
 ```
 
+## `bootstrap_mcps.yaml` Override
+
+Additionally, the `bootstrap_mcps.yaml` can be overridden.
+This might be useful for applying changes to the version when using helm defined template values.
+
+```yaml
+bootstrap:
+  templates:
+    - name: myMCPTemplate
+      version: v1
+      mcps_location: <classpath or file location>
+      values_env: <value environment variable>
+      revision_env: REVISION_ENV
+```
+
+In the above example, we've added a `revision_env` which allows overriding the MCP bootstrap definition itself (excluding `revision_env`).
+
+In this example we could configure `REVISION_ENV` to contain a timestamp or hash: `{"version":"2024060600"}` 
+This value can be changed/incremented each time the helm supplied template values change. This ensures the MCP is updated
+with the latest values during deployment.
+
+
 ## Known Limitations
 
 * Supported change types: 

diff --git a/metadata-service/configuration/src/main/resources/bootstrap_mcps.yaml b/metadata-service/configuration/src/main/resources/bootstrap_mcps.yaml
@@ -42,3 +42,4 @@ bootstrap:
       optional: false
       mcps_location: "bootstrap_mcps/ingestion-datahub-gc.yaml"
       values_env: "DATAHUB_GC_BOOTSTRAP_VALUES"
+      revision_env: "DATAHUB_GC_BOOTSTRAP_REVISION"
diff --git a/metadata-service/configuration/src/main/resources/bootstrap_mcps/ingestion-datahub-gc.yaml b/metadata-service/configuration/src/main/resources/bootstrap_mcps/ingestion-datahub-gc.yaml
@@ -12,7 +12,7 @@
       timezone: '{{schedule.timezone}}{{^schedule.timezone}}UTC{{/schedule.timezone}}'
       interval: '{{schedule.interval}}{{^schedule.interval}}0 1 * * *{{/schedule.interval}}'
     config:
-      version: '{{&ingestion.version}}{{^ingestion.version}}0.14.1.6{{/ingestion.version}}'
+      version: '{{&ingestion.version}}{{^ingestion.version}}0.14.1.7rc2{{/ingestion.version}}'
       recipe:
         source:
           type: 'datahub-gc'