Merge branch 'datahub-project:master' into master

acryldata · Mar 28, 2024 · 90460a2 · 90460a2
2 parents f286107 + db668e8
commit 90460a2
Show file tree

Hide file tree

Showing 127 changed files with 21,239 additions and 848 deletions.
diff --git a/.github/workflows/airflow-plugin.yml b/.github/workflows/airflow-plugin.yml
@@ -40,16 +40,16 @@ jobs:
             extra_pip_requirements: "apache-airflow~=2.2.4"
             extra_pip_extras: plugin-v1
           - python-version: "3.10"
-            extra_pip_requirements: 'apache-airflow~=2.4.0 pluggy==1.0.0 "pendulum<3.0" "Flask-Session<0.6.0"'
-            extra_pip_extras: plugin-v2
+            extra_pip_requirements: "apache-airflow==2.4.3"
+            extra_pip_extras: plugin-v2,test-airflow24
           - python-version: "3.10"
-            extra_pip_requirements: 'apache-airflow~=2.6.0 "pendulum<3.0" "Flask-Session<0.6.0"'
+            extra_pip_requirements: 'apache-airflow==2.6.3 -c https://raw.githubusercontent.com/apache/airflow/constraints-2.6.3/constraints-3.10.txt'
             extra_pip_extras: plugin-v2
           - python-version: "3.10"
-            extra_pip_requirements: 'apache-airflow~=2.7.0 pydantic==2.4.2 "Flask-Session<0.6.0"'
+            extra_pip_requirements: 'apache-airflow==2.7.3 -c https://raw.githubusercontent.com/apache/airflow/constraints-2.7.3/constraints-3.10.txt'
             extra_pip_extras: plugin-v2
           - python-version: "3.10"
-            extra_pip_requirements: 'apache-airflow>=2.8.0 pydantic>=2.4.2 "Flask-Session<0.6.0"'
+            extra_pip_requirements: 'apache-airflow==2.8.1 -c https://raw.githubusercontent.com/apache/airflow/constraints-2.8.1/constraints-3.10.txt'
             extra_pip_extras: plugin-v2
       fail-fast: false
     steps:

diff --git a/build.gradle b/build.gradle
@@ -95,6 +95,7 @@ project.ext.spec = [
             'restliDocgen' : 'com.linkedin.pegasus:restli-docgen:' + pegasusVersion,
             'restliServer' : 'com.linkedin.pegasus:restli-server:' + pegasusVersion,
             'restliSpringBridge': 'com.linkedin.pegasus:restli-spring-bridge:' + pegasusVersion,
+            'restliTestUtils' : 'com.linkedin.pegasus:restli-client-testutils:' + pegasusVersion,
         ]
     ]
 ]

diff --git a/.../main/java/com/linkedin/datahub/graphql/resolvers/assertion/EntityAssertionsResolver.java b/.../main/java/com/linkedin/datahub/graphql/resolvers/assertion/EntityAssertionsResolver.java
@@ -4,6 +4,7 @@
 import com.linkedin.common.EntityRelationship;
 import com.linkedin.common.EntityRelationships;
 import com.linkedin.common.urn.Urn;
+import com.linkedin.common.urn.UrnUtils;
 import com.linkedin.datahub.graphql.QueryContext;
 import com.linkedin.datahub.graphql.generated.Assertion;
 import com.linkedin.datahub.graphql.generated.Entity;
@@ -25,8 +26,10 @@
 import java.util.Objects;
 import java.util.concurrent.CompletableFuture;
 import java.util.stream.Collectors;
+import lombok.extern.slf4j.Slf4j;
 
 /** GraphQL Resolver used for fetching the list of Assertions associated with an Entity. */
+@Slf4j
 public class EntityAssertionsResolver
     implements DataFetcher<CompletableFuture<EntityAssertionsResult>> {
 
@@ -49,6 +52,8 @@ public CompletableFuture<EntityAssertionsResult> get(DataFetchingEnvironment env
           final String entityUrn = ((Entity) environment.getSource()).getUrn();
           final Integer start = environment.getArgumentOrDefault("start", 0);
           final Integer count = environment.getArgumentOrDefault("count", 200);
+          final Boolean includeSoftDeleted =
+              environment.getArgumentOrDefault("includeSoftDeleted", false);
 
           try {
             // Step 1: Fetch set of assertions associated with the target entity from the Graph
@@ -84,6 +89,7 @@ public CompletableFuture<EntityAssertionsResult> get(DataFetchingEnvironment env
                 gmsResults.stream()
                     .filter(Objects::nonNull)
                     .map(r -> AssertionMapper.map(context, r))
+                    .filter(assertion -> assertionExists(assertion, includeSoftDeleted, context))
                     .collect(Collectors.toList());
 
             // Step 4: Package and return result
@@ -98,4 +104,17 @@ public CompletableFuture<EntityAssertionsResult> get(DataFetchingEnvironment env
           }
         });
   }
+
+  private boolean assertionExists(
+      Assertion assertion, Boolean includeSoftDeleted, QueryContext context) {
+    try {
+      return _entityClient.exists(
+          UrnUtils.getUrn(assertion.getUrn()), includeSoftDeleted, context.getAuthentication());
+    } catch (RemoteInvocationException e) {
+      log.error(
+          String.format("Unable to check if assertion %s exists, ignoring it", assertion.getUrn()),
+          e);
+      return false;
+    }
+  }
 }
diff --git a/...phql-core/src/main/java/com/linkedin/datahub/graphql/types/assertion/AssertionMapper.java b/...phql-core/src/main/java/com/linkedin/datahub/graphql/types/assertion/AssertionMapper.java
@@ -2,6 +2,7 @@
 
 import com.linkedin.assertion.AssertionInfo;
 import com.linkedin.common.DataPlatformInstance;
+import com.linkedin.common.Status;
 import com.linkedin.common.urn.Urn;
 import com.linkedin.data.DataMap;
 import com.linkedin.datahub.graphql.QueryContext;
@@ -56,6 +57,18 @@ public static Assertion map(@Nullable QueryContext context, final EntityResponse
       result.setPlatform(unknownPlatform);
     }
 
+    final EnvelopedAspect envelopedStatus = aspects.get(Constants.STATUS_ASPECT_NAME);
+    if (envelopedStatus != null) {
+      result.setStatus(mapStatus(new Status(envelopedStatus.getValue().data())));
+    }
+
+    return result;
+  }
+
+  private static com.linkedin.datahub.graphql.generated.Status mapStatus(Status status) {
+    final com.linkedin.datahub.graphql.generated.Status result =
+        new com.linkedin.datahub.graphql.generated.Status();
+    result.setRemoved(status.isRemoved());
     return result;
   }
 

diff --git a/...edin/datahub/graphql/types/dataprocessinst/mappers/DataProcessInstanceRunEventMapper.java b/...edin/datahub/graphql/types/dataprocessinst/mappers/DataProcessInstanceRunEventMapper.java
@@ -42,6 +42,9 @@ public com.linkedin.datahub.graphql.generated.DataProcessRunEvent apply(
     if (runEvent.hasResult()) {
       result.setResult(DataProcessInstanceRunResultMapper.map(context, runEvent.getResult()));
     }
+    if (runEvent.hasDurationMillis()) {
+      result.setDurationMillis(runEvent.getDurationMillis());
+    }
 
     return result;
   }

diff --git a/datahub-graphql-core/src/main/resources/entity.graphql b/datahub-graphql-core/src/main/resources/entity.graphql
@@ -1548,7 +1548,7 @@ type Dataset implements EntityWithRelationships & Entity & BrowsableEntity {
     """
     Assertions associated with the Dataset
     """
-    assertions(start: Int, count: Int): EntityAssertionsResult
+    assertions(start: Int, count: Int, includeSoftDeleted: Boolean): EntityAssertionsResult
 
     """
     Edges extending from this entity
@@ -6471,6 +6471,11 @@ type DataProcessRunEvent implements TimeSeriesAspect {
     The timestamp associated with the run event in milliseconds
     """
     timestampMillis: Long!
+
+    """
+    The duration of the run in milliseconds
+    """
+    durationMillis: Long
 }
 
 """
@@ -7209,6 +7214,11 @@ type Assertion implements EntityWithRelationships & Entity {
     """
     lineage(input: LineageInput!): EntityLineageResult
 
+    """
+    Status metadata of the assertion
+    """
+    status: Status
+
     """
     Experimental API.
     For fetching extra aspects that do not have custom UI code yet

diff --git a/...t/java/com/linkedin/datahub/graphql/resolvers/assertion/EntityAssertionsResolverTest.java b/...t/java/com/linkedin/datahub/graphql/resolvers/assertion/EntityAssertionsResolverTest.java
@@ -119,6 +119,11 @@ public void testGetSuccess() throws Exception {
                     .setUrn(assertionUrn)
                     .setAspects(new EnvelopedAspectMap(assertionAspects))));
 
+    Mockito.when(
+            mockClient.exists(
+                Mockito.any(Urn.class), Mockito.eq(false), Mockito.any(Authentication.class)))
+        .thenReturn(true);
+
     EntityAssertionsResolver resolver = new EntityAssertionsResolver(mockClient, graphClient);
 
     // Execute resolver
@@ -128,6 +133,8 @@ public void testGetSuccess() throws Exception {
 
     Mockito.when(mockEnv.getArgumentOrDefault(Mockito.eq("start"), Mockito.eq(0))).thenReturn(0);
     Mockito.when(mockEnv.getArgumentOrDefault(Mockito.eq("count"), Mockito.eq(200))).thenReturn(10);
+    Mockito.when(mockEnv.getArgumentOrDefault(Mockito.eq("includeSoftDeleted"), Mockito.eq(false)))
+        .thenReturn(false);
     Mockito.when(mockEnv.getContext()).thenReturn(mockContext);
 
     Dataset parentEntity = new Dataset();
@@ -148,6 +155,9 @@ public void testGetSuccess() throws Exception {
     Mockito.verify(mockClient, Mockito.times(1))
         .batchGetV2(Mockito.any(), Mockito.any(), Mockito.any(), Mockito.any());
 
+    Mockito.verify(mockClient, Mockito.times(1))
+        .exists(Mockito.any(), Mockito.any(), Mockito.any());
+
     // Assert that GraphQL assertion run event matches expectations
     assertEquals(result.getStart(), 0);
     assertEquals(result.getCount(), 1);

diff --git a/datahub-web-react/.eslintrc.js b/datahub-web-react/.eslintrc.js
@@ -8,7 +8,7 @@ module.exports = {
         'plugin:vitest/recommended',
         'prettier',
     ],
-    plugins: ['@typescript-eslint'],
+    plugins: ['@typescript-eslint', 'react-refresh'],
     parserOptions: {
         ecmaVersion: 2020, // Allows for the parsing of modern ECMAScript features
         sourceType: 'module', // Allows for the use of imports
@@ -48,6 +48,7 @@ module.exports = {
         ],
         'vitest/prefer-to-be': 'off',
         '@typescript-eslint/no-use-before-define': ['error', { functions: false, classes: false }],
+        'react-refresh/only-export-components': ['warn', { 'allowConstantExport': true }],
     },
     settings: {
         react: {

diff --git a/datahub-web-react/package.json b/datahub-web-react/package.json
@@ -124,6 +124,7 @@
         "eslint-plugin-jsx-a11y": "^6.5.1",
         "eslint-plugin-react": "^7.28.0",
         "eslint-plugin-react-hooks": "^4.3.0",
+        "eslint-plugin-react-refresh": "^0.4.6",
         "eslint-plugin-vitest": "^0.3.17",
         "jsdom": "^22.1.0",
         "less": "^4.2.0",

diff --git a/datahub-web-react/src/Mocks.tsx b/datahub-web-react/src/Mocks.tsx
@@ -625,8 +625,7 @@ export const dataset3 = {
     health: [],
     assertions: null,
     status: null,
-    readRuns: null,
-    writeRuns: null,
+    runs: null,
     testResults: null,
     siblings: null,
     statsSummary: null,

diff --git a/datahub-web-react/src/app/entity/dataset/DatasetEntity.tsx b/datahub-web-react/src/app/entity/dataset/DatasetEntity.tsx
@@ -178,18 +178,15 @@ export class DatasetEntity implements Entity<Dataset> {
                     },
                 },
                 {
-                    name: 'Operations',
+                    name: 'Runs',
+                    // TODO: Rename this to DatasetRunsTab.
                     component: OperationsTab,
                     display: {
                         visible: (_, dataset: GetDatasetQuery) => {
-                            return (
-                                (dataset?.dataset?.readRuns?.total || 0) + (dataset?.dataset?.writeRuns?.total || 0) > 0
-                            );
+                            return (dataset?.dataset?.runs?.total || 0) > 0;
                         },
                         enabled: (_, dataset: GetDatasetQuery) => {
-                            return (
-                                (dataset?.dataset?.readRuns?.total || 0) + (dataset?.dataset?.writeRuns?.total || 0) > 0
-                            );
+                            return (dataset?.dataset?.runs?.total || 0) > 0;
                         },
                     },
                 },