Merge branch 'master' into feature/add-access-management-aspect-to-da…

…shboards
datahub-project · Nov 13, 2024 · 3d0b56b · 3d0b56b
2 parents 49f6823 + 470c858
commit 3d0b56b
Show file tree

Hide file tree

Showing 124 changed files with 14,966 additions and 1,897 deletions.
diff --git a/.github/workflows/docker-unified.yml b/.github/workflows/docker-unified.yml
@@ -144,6 +144,11 @@ jobs:
     needs: setup
     if: ${{ needs.setup.outputs.backend_change == 'true' || needs.setup.outputs.publish == 'true' }}
     steps:
+      - name: Free up disk space
+        run: |
+          sudo apt-get remove 'dotnet-*' azure-cli || true
+          sudo rm -rf /usr/local/lib/android/ || true
+          sudo docker image prune -a -f || true
       - name: Set up JDK 17
         uses: actions/setup-java@v4
         with:
@@ -210,6 +215,11 @@ jobs:
     needs: setup
     if: ${{ needs.setup.outputs.backend_change == 'true' || needs.setup.outputs.publish == 'true' }}
     steps:
+      - name: Free up disk space
+        run: |
+          sudo apt-get remove 'dotnet-*' azure-cli || true
+          sudo rm -rf /usr/local/lib/android/ || true
+          sudo docker image prune -a -f || true
       - name: Set up JDK 17
         uses: actions/setup-java@v4
         with:
@@ -276,6 +286,11 @@ jobs:
     needs: setup
     if: ${{ needs.setup.outputs.backend_change == 'true' || needs.setup.outputs.publish == 'true' }}
     steps:
+      - name: Free up disk space
+        run: |
+          sudo apt-get remove 'dotnet-*' azure-cli || true
+          sudo rm -rf /usr/local/lib/android/ || true
+          sudo docker image prune -a -f || true
       - name: Set up JDK 17
         uses: actions/setup-java@v4
         with:
@@ -342,6 +357,11 @@ jobs:
     needs: setup
     if: ${{ needs.setup.outputs.backend_change == 'true' || needs.setup.outputs.publish == 'true' }}
     steps:
+      - name: Free up disk space
+        run: |
+          sudo apt-get remove 'dotnet-*' azure-cli || true
+          sudo rm -rf /usr/local/lib/android/ || true
+          sudo docker image prune -a -f || true
       - name: Set up JDK 17
         uses: actions/setup-java@v4
         with:
@@ -408,6 +428,11 @@ jobs:
     needs: setup
     if: ${{ needs.setup.outputs.frontend_change == 'true' || needs.setup.outputs.publish == 'true' || needs.setup.outputs.pr-publish == 'true'}}
     steps:
+      - name: Free up disk space
+        run: |
+          sudo apt-get remove 'dotnet-*' azure-cli || true
+          sudo rm -rf /usr/local/lib/android/ || true
+          sudo docker image prune -a -f || true
       - name: Set up JDK 17
         uses: actions/setup-java@v4
         with:
@@ -476,6 +501,11 @@ jobs:
     needs: setup
     if: ${{ needs.setup.outputs.kafka_setup_change == 'true' || (needs.setup.outputs.publish == 'true' || needs.setup.outputs.pr-publish == 'true') }}
     steps:
+      - name: Free up disk space
+        run: |
+          sudo apt-get remove 'dotnet-*' azure-cli || true
+          sudo rm -rf /usr/local/lib/android/ || true
+          sudo docker image prune -a -f || true
       - name: Check out the repo
         uses: acryldata/sane-checkout-action@v3
       - name: Build and push
@@ -532,6 +562,11 @@ jobs:
     needs: setup
     if: ${{ needs.setup.outputs.mysql_setup_change == 'true' || (needs.setup.outputs.publish == 'true' || needs.setup.outputs.pr-publish == 'true') }}
     steps:
+      - name: Free up disk space
+        run: |
+          sudo apt-get remove 'dotnet-*' azure-cli || true
+          sudo rm -rf /usr/local/lib/android/ || true
+          sudo docker image prune -a -f || true
       - name: Check out the repo
         uses: acryldata/sane-checkout-action@v3
       - name: Build and push
@@ -588,6 +623,11 @@ jobs:
     needs: setup
     if: ${{ needs.setup.outputs.elasticsearch_setup_change == 'true' || (needs.setup.outputs.publish == 'true' || needs.setup.outputs.pr-publish == 'true' ) }}
     steps:
+      - name: Free up disk space
+        run: |
+          sudo apt-get remove 'dotnet-*' azure-cli || true
+          sudo rm -rf /usr/local/lib/android/ || true
+          sudo docker image prune -a -f || true
       - name: Check out the repo
         uses: acryldata/sane-checkout-action@v3
       - name: Build and push
@@ -646,6 +686,11 @@ jobs:
     needs: setup
     if: ${{ needs.setup.outputs.ingestion_change == 'true' || needs.setup.outputs.publish == 'true' || needs.setup.outputs.pr-publish == 'true' }}
     steps:
+      - name: Free up disk space
+        run: |
+          sudo apt-get remove 'dotnet-*' azure-cli || true
+          sudo rm -rf /usr/local/lib/android/ || true
+          sudo docker image prune -a -f || true
       - name: Check out the repo
         uses: acryldata/sane-checkout-action@v3
       - name: Build and push Base Image
@@ -674,6 +719,11 @@ jobs:
     needs: [setup, datahub_ingestion_base_build]
     if: ${{ needs.setup.outputs.ingestion_change == 'true' || needs.setup.outputs.publish == 'true' || needs.setup.outputs.pr-publish == 'true' }}
     steps:
+      - name: Free up disk space
+        run: |
+          sudo apt-get remove 'dotnet-*' azure-cli || true
+          sudo rm -rf /usr/local/lib/android/ || true
+          sudo docker image prune -a -f || true
       - name: Check out the repo
         uses: acryldata/sane-checkout-action@v3
       - name: Download Base Image
@@ -717,6 +767,11 @@ jobs:
     needs: [setup, datahub_ingestion_base_build]
     if: ${{ needs.setup.outputs.ingestion_change == 'true' || needs.setup.outputs.publish == 'true' || needs.setup.outputs.pr-publish == 'true' }}
     steps:
+      - name: Free up disk space
+        run: |
+          sudo apt-get remove 'dotnet-*' azure-cli || true
+          sudo rm -rf /usr/local/lib/android/ || true
+          sudo docker image prune -a -f || true
       - name: Check out the repo
         uses: acryldata/sane-checkout-action@v3
       - name: Download Base Image
@@ -760,6 +815,11 @@ jobs:
     needs: [setup, datahub_ingestion_base_slim_build]
     if: ${{ needs.setup.outputs.ingestion_change == 'true' || needs.setup.outputs.publish == 'true' || needs.setup.outputs.pr-publish == 'true' }}
     steps:
+      - name: Free up disk space
+        run: |
+          sudo apt-get remove 'dotnet-*' azure-cli || true
+          sudo rm -rf /usr/local/lib/android/ || true
+          sudo docker image prune -a -f || true
       - name: Check out the repo
         uses: acryldata/sane-checkout-action@v3
       - uses: actions/setup-python@v5
@@ -856,6 +916,11 @@ jobs:
     needs: [setup, datahub_ingestion_base_full_build]
     if: ${{ needs.setup.outputs.ingestion_change == 'true' || needs.setup.outputs.publish == 'true' || needs.setup.outputs.pr-publish == 'true' }}
     steps:
+      - name: Free up disk space
+        run: |
+          sudo apt-get remove 'dotnet-*' azure-cli || true
+          sudo rm -rf /usr/local/lib/android/ || true
+          sudo docker image prune -a -f || true
       - name: Check out the repo
         uses: acryldata/sane-checkout-action@v3
       - uses: actions/setup-python@v5

diff --git a/.github/workflows/gx-plugin.yml b/.github/workflows/gx-plugin.yml
@@ -39,6 +39,8 @@ jobs:
             extraPythonRequirement: "great-expectations~=0.16.0 numpy~=1.26.0"
           - python-version: "3.11"
             extraPythonRequirement: "great-expectations~=0.17.0"
+          - python-version: "3.11"
+            extraPythonRequirement: "great-expectations~=0.18.0"
       fail-fast: false
     steps:
       - name: Set up JDK 17

diff --git a/datahub-frontend/app/auth/AuthModule.java b/datahub-frontend/app/auth/AuthModule.java
@@ -13,6 +13,7 @@
 import com.google.inject.Provides;
 import com.google.inject.Singleton;
 import com.google.inject.name.Named;
+import com.linkedin.entity.client.EntityClientConfig;
 import com.linkedin.entity.client.SystemEntityClient;
 import com.linkedin.entity.client.SystemRestliEntityClient;
 import com.linkedin.metadata.models.registry.EmptyEntityRegistry;
@@ -213,11 +214,13 @@ protected SystemEntityClient provideEntityClient(
 
     return new SystemRestliEntityClient(
         buildRestliClient(),
-        new ExponentialBackoff(configs.getInt(ENTITY_CLIENT_RETRY_INTERVAL)),
-        configs.getInt(ENTITY_CLIENT_NUM_RETRIES),
-        configurationProvider.getCache().getClient().getEntityClient(),
-        Math.max(1, configs.getInt(ENTITY_CLIENT_RESTLI_GET_BATCH_SIZE)),
-        Math.max(1, configs.getInt(ENTITY_CLIENT_RESTLI_GET_BATCH_CONCURRENCY)));
+        EntityClientConfig.builder()
+            .backoffPolicy(new ExponentialBackoff(configs.getInt(ENTITY_CLIENT_RETRY_INTERVAL)))
+            .retryCount(configs.getInt(ENTITY_CLIENT_NUM_RETRIES))
+            .batchGetV2Size(configs.getInt(ENTITY_CLIENT_RESTLI_GET_BATCH_SIZE))
+            .batchGetV2Concurrency(2)
+            .build(),
+        configurationProvider.getCache().getClient().getEntityClient());
   }
 
   @Provides

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
@@ -31,6 +31,7 @@
 import java.util.Map;
 import java.util.stream.Collectors;
 import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
 import lombok.extern.slf4j.Slf4j;
 import org.apache.commons.io.IOUtils;
 import org.springframework.core.io.ClassPathResource;
@@ -53,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());
@@ -109,13 +111,29 @@ static List<ObjectNode> resolveMCPTemplate(
       AuditStamp auditStamp)
       throws IOException {
 
-    String template = loadTemplate(mcpTemplate.getMcps_location());
-    Mustache mustache = MUSTACHE_FACTORY.compile(new StringReader(template), mcpTemplate.getName());
+    final String template = loadTemplate(mcpTemplate.getMcps_location());
     Map<String, Object> scopeValues = resolveValues(opContext, mcpTemplate, auditStamp);
+
     StringWriter writer = new StringWriter();
-    mustache.execute(writer, scopeValues);
+    try {
+      Mustache mustache =
+          MUSTACHE_FACTORY.compile(new StringReader(template), mcpTemplate.getName());
+      mustache.execute(writer, scopeValues);
+    } catch (Exception e) {
+      log.error(
+          "Failed to apply mustache template. Template: {} Values: {}",
+          template,
+          resolveEnv(mcpTemplate));
+      throw e;
+    }
 
-    return opContext.getYamlMapper().readValue(writer.toString(), new TypeReference<>() {});
+    final String yaml = writer.toString();
+    try {
+      return opContext.getYamlMapper().readValue(yaml, new TypeReference<>() {});
+    } catch (Exception e) {
+      log.error("Failed to parse rendered MCP bootstrap yaml: {}", yaml);
+      throw e;
+    }
   }
 
   static Map<String, Object> resolveValues(
@@ -128,13 +146,21 @@ static Map<String, Object> resolveValues(
     // built-in
     scopeValues.put("auditStamp", RecordUtils.toJsonString(auditStamp));
 
+    String envValue = resolveEnv(mcpTemplate);
+    if (envValue != null) {
+      scopeValues.putAll(opContext.getObjectMapper().readValue(envValue, new TypeReference<>() {}));
+    }
+    return scopeValues;
+  }
+
+  @Nullable
+  private static String resolveEnv(BootstrapMCPConfigFile.MCPTemplate mcpTemplate) {
     if (mcpTemplate.getValues_env() != null
         && !mcpTemplate.getValues_env().isEmpty()
         && System.getenv().containsKey(mcpTemplate.getValues_env())) {
-      String envValue = System.getenv(mcpTemplate.getValues_env());
-      scopeValues.putAll(opContext.getObjectMapper().readValue(envValue, new TypeReference<>() {}));
+      return System.getenv(mcpTemplate.getValues_env());
     }
-    return scopeValues;
+    return null;
   }
 
   private static String loadTemplate(String source) throws IOException {

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"