feat: update mlflow connector

.github/workflows/airflow-plugin.yml

@@ -27,7 +27,6 @@ jobs:
   airflow-plugin:
     runs-on: ubuntu-latest
     env:
-      SPARK_VERSION: 3.0.3
       DATAHUB_TELEMETRY_ENABLED: false
     strategy:
       matrix:
@@ -69,7 +68,7 @@ jobs:
       - name: pip freeze show list installed
         if: always()
         run: source metadata-ingestion-modules/airflow-plugin/venv/bin/activate && uv pip freeze
-      - uses: actions/upload-artifact@v3
+      - uses: actions/upload-artifact@v4
         if: ${{ always() && matrix.python-version == '3.10' && matrix.extra_pip_requirements == 'apache-airflow>=2.7.0' }}
         with:
           name: Test Results (Airflow Plugin ${{ matrix.python-version}})
@@ -93,7 +92,7 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Upload
-        uses: actions/upload-artifact@v3
+        uses: actions/upload-artifact@v4
         with:
           name: Event File
           path: ${{ github.event_path }}

.github/workflows/dagster-plugin.yml

@@ -27,7 +27,6 @@ jobs:
   dagster-plugin:
     runs-on: ubuntu-latest
     env:
-      SPARK_VERSION: 3.0.3
       DATAHUB_TELEMETRY_ENABLED: false
     strategy:
       matrix:
@@ -44,7 +43,8 @@ jobs:
         with:
           distribution: "zulu"
           java-version: 17
-      - uses: actions/checkout@v4
+      - uses: gradle/actions/setup-gradle@v3
+      - uses: acryldata/sane-checkout-action@v3
       - uses: actions/setup-python@v5
         with:
           python-version: ${{ matrix.python-version }}
@@ -56,7 +56,7 @@ jobs:
       - name: pip freeze show list installed
         if: always()
         run: source metadata-ingestion-modules/dagster-plugin/venv/bin/activate && uv pip freeze
-      - uses: actions/upload-artifact@v3
+      - uses: actions/upload-artifact@v4
         if: ${{ always() && matrix.python-version == '3.10' && matrix.extraPythonRequirement == 'dagster>=1.3.3' }}
         with:
           name: Test Results (dagster Plugin ${{ matrix.python-version}})
@@ -79,7 +79,7 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Upload
-        uses: actions/upload-artifact@v3
+        uses: actions/upload-artifact@v4
         with:
           name: Event File
           path: ${{ github.event_path }}

.github/workflows/gx-plugin.yml

@@ -27,7 +27,6 @@ jobs:
   gx-plugin:
     runs-on: ubuntu-latest
     env:
-      SPARK_VERSION: 3.0.3
       DATAHUB_TELEMETRY_ENABLED: false
     strategy:
       matrix:
@@ -48,7 +47,8 @@ jobs:
         with:
           distribution: "zulu"
           java-version: 17
-      - uses: actions/checkout@v4
+      - uses: gradle/actions/setup-gradle@v3
+      - uses: acryldata/sane-checkout-action@v3
       - uses: actions/setup-python@v5
         with:
           python-version: ${{ matrix.python-version }}
@@ -60,7 +60,7 @@ jobs:
       - name: pip freeze show list installed
         if: always()
         run: source metadata-ingestion-modules/gx-plugin/venv/bin/activate && uv pip freeze
-      - uses: actions/upload-artifact@v3
+      - uses: actions/upload-artifact@v4
         if: ${{ always() && matrix.python-version == '3.11' && matrix.extraPythonRequirement == 'great-expectations~=0.17.0' }}
         with:
           name: Test Results (GX Plugin ${{ matrix.python-version}})
@@ -83,7 +83,7 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Upload
-        uses: actions/upload-artifact@v3
+        uses: actions/upload-artifact@v4
         with:
           name: Event File
           path: ${{ github.event_path }}

.github/workflows/metadata-ingestion.yml

@@ -28,7 +28,6 @@ jobs:
     runs-on: ubuntu-latest
     timeout-minutes: 40
     env:
-      SPARK_VERSION: 3.3.2
       DATAHUB_TELEMETRY_ENABLED: false
       # TODO: Enable this once the test is fixed.
       # DATAHUB_LOOKML_GIT_TEST_SSH_KEY: ${{ secrets.DATAHUB_LOOKML_GIT_TEST_SSH_KEY }}
@@ -84,9 +83,9 @@ jobs:
           df -hl
           docker image ls
           docker system df
-      - uses: actions/upload-artifact@v3
+      - uses: actions/upload-artifact@v4
         with:
-          name: Test Results (metadata ingestion ${{ matrix.python-version }})
+          name: Test Results (metadata ingestion ${{ matrix.python-version }} ${{ matrix.command }})
           path: |
             **/build/reports/tests/test/**
             **/build/test-results/test/**
@@ -100,14 +99,14 @@ jobs:
           directory: ./build/coverage-reports/
           fail_ci_if_error: false
           flags: pytest-${{ matrix.command }}
-          name: pytest-${{ matrix.command }}
+          name: pytest-${{ matrix.python-version }}-${{ matrix.command }}
           verbose: true
 
   event-file:
     runs-on: ubuntu-latest
     steps:
       - name: Upload
-        uses: actions/upload-artifact@v3
+        uses: actions/upload-artifact@v4
         with:
           name: Event File
           path: ${{ github.event_path }}

.github/workflows/prefect-plugin.yml

@@ -27,25 +27,20 @@ jobs:
   prefect-plugin:
     runs-on: ubuntu-latest
     env:
-      SPARK_VERSION: 3.0.3
       DATAHUB_TELEMETRY_ENABLED: false
     strategy:
       matrix:
         python-version: ["3.8", "3.9", "3.10"]
-        include:
-          - python-version: "3.8"
-          - python-version: "3.9"
-          - python-version: "3.10"
       fail-fast: false
     steps:
       - name: Set up JDK 17
-        uses: actions/setup-java@v3
+        uses: actions/setup-java@v4
         with:
           distribution: "zulu"
           java-version: 17
       - uses: gradle/actions/setup-gradle@v3
-      - uses: actions/checkout@v3
-      - uses: actions/setup-python@v4
+      - uses: acryldata/sane-checkout-action@v3
+      - uses: actions/setup-python@v5
         with:
           python-version: ${{ matrix.python-version }}
           cache: "pip"
@@ -56,7 +51,7 @@ jobs:
       - name: pip freeze show list installed
         if: always()
         run: source metadata-ingestion-modules/prefect-plugin/venv/bin/activate && uv pip freeze
-      - uses: actions/upload-artifact@v3
+      - uses: actions/upload-artifact@v4
         if: ${{ always() && matrix.python-version == '3.10'}}
         with:
           name: Test Results (Prefect Plugin ${{ matrix.python-version}})
@@ -72,15 +67,15 @@ jobs:
           token: ${{ secrets.CODECOV_TOKEN }}
           directory: ./build/coverage-reports/
           fail_ci_if_error: false
-          flags: prefect,prefect-${{ matrix.extra_pip_extras }}
+          flags: prefect,prefect-${{ matrix.python-version }}
           name: pytest-prefect-${{ matrix.python-version }}
           verbose: true
 
   event-file:
     runs-on: ubuntu-latest
     steps:
       - name: Upload
-        uses: actions/upload-artifact@v3
+        uses: actions/upload-artifact@v4
         with:
           name: Event File
           path: ${{ github.event_path }}

.github/workflows/python-build-pages.yml

@@ -0,0 +1,64 @@
+name: Python Build
+on:
+  push:
+    branches:
+      - master
+    paths:
+      - ".github/workflows/python-build-pages.yml"
+      - "metadata-ingestion/**"
+      - "metadata-ingestion-modules/**"
+      - "metadata-models/**"
+  pull_request:
+    branches:
+      - "**"
+    paths:
+      - ".github/workflows/python-build-pages.yml"
+      - "metadata-ingestion/**"
+      - "metadata-ingestion-modules/**"
+      - "metadata-models/**"
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  deploy-pages:
+    runs-on: ubuntu-latest
+    if: ${{ vars.CLOUDFLARE_WHEELS_PROJECT_NAME != '' }}
+
+    name: Python Wheels
+    permissions:
+      contents: read
+      pull-requests: read
+      deployments: write
+    steps:
+      - name: Set up JDK 17
+        uses: actions/setup-java@v4
+        with:
+          distribution: "zulu"
+          java-version: 17
+      - uses: gradle/actions/setup-gradle@v3
+      - uses: acryldata/sane-checkout-action@v3
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.10"
+          cache: "pip"
+      - uses: actions/cache@v4
+        with:
+          path: |
+            ~/.cache/uv
+          key: ${{ runner.os }}-uv-${{ hashFiles('**/requirements.txt') }}
+      - name: Build Python wheel site
+        run: |
+          ./gradlew :python-build:buildSite
+        env:
+          GITHUB_TOKEN: ${{ github.token }}
+      - name: Publish
+        uses: cloudflare/pages-action@v1
+        with:
+          apiToken: ${{ secrets.CLOUDFLARE_API_TOKEN }}
+          accountId: ${{ secrets.CLOUDFLARE_ACCOUNT_ID }}
+          projectName: ${{ vars.CLOUDFLARE_WHEELS_PROJECT_NAME }}
+          workingDirectory: python-build
+          directory: site
+          gitHubToken: ${{ github.token }}

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/common/mappers/TimeStampMapper.java

@@ -0,0 +1,28 @@
+package com.linkedin.datahub.graphql.types.common.mappers;
+
+import com.linkedin.datahub.graphql.QueryContext;
+import com.linkedin.datahub.graphql.generated.TimeStamp;
+import com.linkedin.datahub.graphql.types.mappers.ModelMapper;
+import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
+
+public class TimeStampMapper implements ModelMapper<com.linkedin.common.TimeStamp, TimeStamp> {
+
+  public static final TimeStampMapper INSTANCE = new TimeStampMapper();
+
+  public static TimeStamp map(
+      @Nullable QueryContext context, @Nonnull final com.linkedin.common.TimeStamp timestamp) {
+    return INSTANCE.apply(context, timestamp);
+  }
+
+  @Override
+  public TimeStamp apply(
+      @Nullable QueryContext context, @Nonnull final com.linkedin.common.TimeStamp timestamp) {
+    final TimeStamp result = new TimeStamp();
+    result.setTime(timestamp.getTime());
+    if (timestamp.hasActor()) {
+      result.setActor(timestamp.getActor().toString());
+    }
+    return result;
+  }
+}

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mappers/EmbeddedModelMapper.java

@@ -0,0 +1,12 @@
+package com.linkedin.datahub.graphql.types.mappers;
+
+import com.linkedin.common.urn.Urn;
+import com.linkedin.datahub.graphql.QueryContext;
+import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
+
+/** Made for models that are embedded in other models and thus do not encode their own URN. */
+public interface EmbeddedModelMapper<I, O> {
+  O apply(
+      @Nullable final QueryContext context, @Nonnull final I input, @Nonnull final Urn entityUrn);
+}

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/mappers/MLFeatureMapper.java

@@ -75,7 +75,8 @@ public MLFeature apply(
             mlFeature.setOwnership(
                 OwnershipMapper.map(context, new Ownership(dataMap), entityUrn)));
     mappingHelper.mapToResult(
-        context, ML_FEATURE_PROPERTIES_ASPECT_NAME, MLFeatureMapper::mapMLFeatureProperties);
+        ML_FEATURE_PROPERTIES_ASPECT_NAME,
+        (entity, dataMap) -> mapMLFeatureProperties(context, entity, dataMap, entityUrn));
     mappingHelper.mapToResult(
         INSTITUTIONAL_MEMORY_ASPECT_NAME,
         (mlFeature, dataMap) ->
@@ -138,10 +139,13 @@ private static void mapMLFeatureKey(@Nonnull MLFeature mlFeature, @Nonnull DataM
   private static void mapMLFeatureProperties(
       @Nullable final QueryContext context,
       @Nonnull MLFeature mlFeature,
-      @Nonnull DataMap dataMap) {
+      @Nonnull DataMap dataMap,
+      @Nonnull Urn entityUrn) {
     MLFeatureProperties featureProperties = new MLFeatureProperties(dataMap);
-    mlFeature.setFeatureProperties(MLFeaturePropertiesMapper.map(context, featureProperties));
-    mlFeature.setProperties(MLFeaturePropertiesMapper.map(context, featureProperties));
+    com.linkedin.datahub.graphql.generated.MLFeatureProperties graphqlProperties =
+        MLFeaturePropertiesMapper.map(context, featureProperties, entityUrn);
+    mlFeature.setFeatureProperties(graphqlProperties);
+    mlFeature.setProperties(graphqlProperties);
     mlFeature.setDescription(featureProperties.getDescription());
     if (featureProperties.getDataType() != null) {
       mlFeature.setDataType(MLFeatureDataType.valueOf(featureProperties.getDataType().toString()));

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/mappers/MLFeaturePropertiesMapper.java

@@ -1,29 +1,34 @@
 package com.linkedin.datahub.graphql.types.mlmodel.mappers;
 
+import com.linkedin.common.urn.Urn;
 import com.linkedin.datahub.graphql.QueryContext;
 import com.linkedin.datahub.graphql.generated.Dataset;
 import com.linkedin.datahub.graphql.generated.MLFeatureDataType;
 import com.linkedin.datahub.graphql.generated.MLFeatureProperties;
-import com.linkedin.datahub.graphql.types.mappers.ModelMapper;
+import com.linkedin.datahub.graphql.types.common.mappers.CustomPropertiesMapper;
+import com.linkedin.datahub.graphql.types.mappers.EmbeddedModelMapper;
 import java.util.stream.Collectors;
+import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
-import lombok.NonNull;
 
 public class MLFeaturePropertiesMapper
-    implements ModelMapper<com.linkedin.ml.metadata.MLFeatureProperties, MLFeatureProperties> {
+    implements EmbeddedModelMapper<
+        com.linkedin.ml.metadata.MLFeatureProperties, MLFeatureProperties> {
 
   public static final MLFeaturePropertiesMapper INSTANCE = new MLFeaturePropertiesMapper();
 
   public static MLFeatureProperties map(
       @Nullable QueryContext context,
-      @NonNull final com.linkedin.ml.metadata.MLFeatureProperties mlFeatureProperties) {
-    return INSTANCE.apply(context, mlFeatureProperties);
+      @Nonnull final com.linkedin.ml.metadata.MLFeatureProperties mlFeatureProperties,
+      @Nonnull Urn entityUrn) {
+    return INSTANCE.apply(context, mlFeatureProperties, entityUrn);
   }
 
   @Override
   public MLFeatureProperties apply(
       @Nullable QueryContext context,
-      @NonNull final com.linkedin.ml.metadata.MLFeatureProperties mlFeatureProperties) {
+      @Nonnull final com.linkedin.ml.metadata.MLFeatureProperties mlFeatureProperties,
+      @Nonnull Urn entityUrn) {
     final MLFeatureProperties result = new MLFeatureProperties();
 
     result.setDescription(mlFeatureProperties.getDescription());
@@ -45,6 +50,9 @@ public MLFeatureProperties apply(
               .collect(Collectors.toList()));
     }
 
+    result.setCustomProperties(
+        CustomPropertiesMapper.map(mlFeatureProperties.getCustomProperties(), entityUrn));
+
     return result;
   }
 }

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/mappers/MLFeatureTableMapper.java

@@ -76,7 +76,7 @@ public MLFeatureTable apply(
     mappingHelper.mapToResult(ML_FEATURE_TABLE_KEY_ASPECT_NAME, this::mapMLFeatureTableKey);
     mappingHelper.mapToResult(
         ML_FEATURE_TABLE_PROPERTIES_ASPECT_NAME,
-        (entity, dataMap) -> this.mapMLFeatureTableProperties(context, entity, dataMap, entityUrn));
+        (entity, dataMap) -> mapMLFeatureTableProperties(context, entity, dataMap, entityUrn));
     mappingHelper.mapToResult(
         INSTITUTIONAL_MEMORY_ASPECT_NAME,
         (mlFeatureTable, dataMap) ->
@@ -146,10 +146,10 @@ private static void mapMLFeatureTableProperties(
       @Nonnull DataMap dataMap,
       Urn entityUrn) {
     MLFeatureTableProperties featureTableProperties = new MLFeatureTableProperties(dataMap);
-    mlFeatureTable.setFeatureTableProperties(
-        MLFeatureTablePropertiesMapper.map(context, featureTableProperties, entityUrn));
-    mlFeatureTable.setProperties(
-        MLFeatureTablePropertiesMapper.map(context, featureTableProperties, entityUrn));
+    com.linkedin.datahub.graphql.generated.MLFeatureTableProperties graphqlProperties =
+        MLFeatureTablePropertiesMapper.map(context, featureTableProperties, entityUrn);
+    mlFeatureTable.setFeatureTableProperties(graphqlProperties);
+    mlFeatureTable.setProperties(graphqlProperties);
     mlFeatureTable.setDescription(featureTableProperties.getDescription());
   }