diff --git a/.github/dependabot.yml b/.github/dependabot.yml
new file mode 100644
index 00000000000000..0d08e261a2ae9d
--- /dev/null
+++ b/.github/dependabot.yml
@@ -0,0 +1,11 @@
+# To get started with Dependabot version updates, you'll need to specify which
+# package ecosystems to update and where the package manifests are located.
+# Please see the documentation for all configuration options:
+# https://docs.github.com/code-security/dependabot/dependabot-version-updates/configuration-options-for-the-dependabot.yml-file
+
+version: 2
+updates:
+  - package-ecosystem: "github-actions" # See documentation for possible values
+    directory: "/" # Location of package manifests
+    schedule:
+      interval: "weekly"
diff --git a/.github/pr-labeler-config.yml b/.github/pr-labeler-config.yml
index 2625cb1cfdff68..e3fcc15b637e22 100644
--- a/.github/pr-labeler-config.yml
+++ b/.github/pr-labeler-config.yml
@@ -1,22 +1,32 @@
 ingestion:
-  - "metadata-ingestion/**/*"
-  - "metadata-ingestion-modules/**/*"
-  - "metadata-integration/**/*"
+- changed-files:
+  - any-glob-to-any-file:
+    - "metadata-ingestion/**/*"
+    - "metadata-ingestion-modules/**/*"
+    - "metadata-integration/**/*"
 
 devops:
-  - "docker/**/*"
-  - ".github/**/*"
-  - "perf-test/**/*"
-  - "metadata-service/**/*"
+- changed-files:
+  - any-glob-to-any-file:
+    - "docker/**/*"
+    - ".github/**/*"
+    - "perf-test/**/*"
+    - "metadata-service/**/*"
 
 product:
-  - "datahub-web-react/**/*"
-  - "datahub-frontend/**/*"
-  - "datahub-graphql-core/**/*"
-  - "metadata-io/**/*"
+- changed-files:
+  - any-glob-to-any-file:
+    - "datahub-web-react/**/*"
+    - "datahub-frontend/**/*"
+    - "datahub-graphql-core/**/*"
+    - "metadata-io/**/*"
 
 docs:
-  - "docs/**/*"
+- changed-files:
+  - any-glob-to-any-file:
+    - "docs/**/*"
 
 smoke_test:
-  - "smoke-test/**/*"
+- changed-files:
+  - any-glob-to-any-file:
+    - "smoke-test/**/*"
diff --git a/.github/scripts/pre-commit-override.yaml b/.github/scripts/pre-commit-override.yaml
index 961134bebe2c98..ecd3c97ad61efa 100644
--- a/.github/scripts/pre-commit-override.yaml
+++ b/.github/scripts/pre-commit-override.yaml
@@ -5,5 +5,5 @@ repos:
         name: smoke-test cypress Lint Fix
         entry: ./gradlew :smoke-test:cypressLintFix
         language: system
-        files: ^smoke-test/tests/cypress/.*$
+        files: ^smoke-test/tests/cypress/.*\.tsx$
         pass_filenames: false
\ No newline at end of file
diff --git a/.github/workflows/airflow-plugin.yml b/.github/workflows/airflow-plugin.yml
index e1e0fb0a85e977..31c77d754dc69d 100644
--- a/.github/workflows/airflow-plugin.yml
+++ b/.github/workflows/airflow-plugin.yml
@@ -56,7 +56,7 @@ jobs:
         with:
           distribution: "zulu"
           java-version: 17
-      - uses: gradle/actions/setup-gradle@v3
+      - uses: gradle/actions/setup-gradle@v4
       - uses: acryldata/sane-checkout-action@v3
       - uses: actions/setup-python@v5
         with:
diff --git a/.github/workflows/build-and-test.yml b/.github/workflows/build-and-test.yml
index 86545946d6afea..7f7dcb9c1f7548 100644
--- a/.github/workflows/build-and-test.yml
+++ b/.github/workflows/build-and-test.yml
@@ -63,7 +63,7 @@ jobs:
           sudo apt-get remove 'dotnet-*' azure-cli || true
           sudo rm -rf /usr/local/lib/android/ || true
           sudo docker image prune -a -f || true
-      - uses: szenius/set-timezone@v1.2
+      - uses: szenius/set-timezone@v2.0
         with:
           timezoneLinux: ${{ matrix.timezone }}
       - name: Check out the repo
@@ -83,7 +83,7 @@ jobs:
         with:
           distribution: "zulu"
           java-version: 17
-      - uses: gradle/actions/setup-gradle@v3
+      - uses: gradle/actions/setup-gradle@v4
       - name: Gradle build (and test) for NOT metadata ingestion
         if: ${{  matrix.command == 'except_metadata_ingestion' && needs.setup.outputs.backend_change == 'true' }}
         # datahub-schematron:cli excluded due to dependency on metadata-ingestion
diff --git a/.github/workflows/check-datahub-jars.yml b/.github/workflows/check-datahub-jars.yml
index dc770f7fc83a61..1b284578483f30 100644
--- a/.github/workflows/check-datahub-jars.yml
+++ b/.github/workflows/check-datahub-jars.yml
@@ -43,7 +43,7 @@ jobs:
         with:
           distribution: "zulu"
           java-version: 17
-      - uses: gradle/actions/setup-gradle@v3
+      - uses: gradle/actions/setup-gradle@v4
       - name: check ${{ matrix.command }} jar
         run: |
           ./gradlew :metadata-integration:java:${{ matrix.command }}:build --info
diff --git a/.github/workflows/dagster-plugin.yml b/.github/workflows/dagster-plugin.yml
index a2ac59d6989a9f..61fddaeb7770bb 100644
--- a/.github/workflows/dagster-plugin.yml
+++ b/.github/workflows/dagster-plugin.yml
@@ -44,7 +44,7 @@ jobs:
         with:
           distribution: "zulu"
           java-version: 17
-      - uses: gradle/actions/setup-gradle@v3
+      - uses: gradle/actions/setup-gradle@v4
       - uses: acryldata/sane-checkout-action@v3
       - uses: actions/setup-python@v5
         with:
diff --git a/.github/workflows/docker-unified.yml b/.github/workflows/docker-unified.yml
index e44e6b11c6d057..5f944c8e28769f 100644
--- a/.github/workflows/docker-unified.yml
+++ b/.github/workflows/docker-unified.yml
@@ -130,7 +130,7 @@ jobs:
         with:
           distribution: "zulu"
           java-version: 17
-      - uses: gradle/actions/setup-gradle@v3
+      - uses: gradle/actions/setup-gradle@v4
       - name: Run lint on smoke test
         if: ${{ steps.ci-optimize.outputs.smoke-test-change == 'true' }}
         run: |
@@ -154,7 +154,7 @@ jobs:
         with:
           distribution: "zulu"
           java-version: 17
-      - uses: gradle/actions/setup-gradle@v3
+      - uses: gradle/actions/setup-gradle@v4
       - name: Check out the repo
         uses: acryldata/sane-checkout-action@v3
       - name: Pre-build artifacts for docker image
@@ -191,7 +191,7 @@ jobs:
         with:
           image: ${{ env.DATAHUB_GMS_IMAGE }}:${{ needs.setup.outputs.unique_tag }}
       - name: Run Trivy vulnerability scanner
-        uses: aquasecurity/trivy-action@0.26.0
+        uses: aquasecurity/trivy-action@0.29.0
         env:
           TRIVY_OFFLINE_SCAN: true
           TRIVY_DB_REPOSITORY: public.ecr.aws/aquasecurity/trivy-db:2,ghcr.io/aquasecurity/trivy-db:2
@@ -205,7 +205,7 @@ jobs:
           ignore-unfixed: true
           vuln-type: "os,library"
       - name: Upload Trivy scan results to GitHub Security tab
-        uses: github/codeql-action/upload-sarif@v2
+        uses: github/codeql-action/upload-sarif@v3
         with:
           sarif_file: "trivy-results.sarif"
 
@@ -225,7 +225,7 @@ jobs:
         with:
           distribution: "zulu"
           java-version: 17
-      - uses: gradle/actions/setup-gradle@v3
+      - uses: gradle/actions/setup-gradle@v4
       - name: Check out the repo
         uses: acryldata/sane-checkout-action@v3
       - name: Pre-build artifacts for docker image
@@ -262,7 +262,7 @@ jobs:
         with:
           image: ${{ env.DATAHUB_MAE_CONSUMER_IMAGE }}:${{ needs.setup.outputs.unique_tag }}
       - name: Run Trivy vulnerability scanner
-        uses: aquasecurity/trivy-action@0.26.0
+        uses: aquasecurity/trivy-action@0.29.0
         env:
           TRIVY_OFFLINE_SCAN: true
           TRIVY_DB_REPOSITORY: public.ecr.aws/aquasecurity/trivy-db:2,ghcr.io/aquasecurity/trivy-db:2
@@ -276,7 +276,7 @@ jobs:
           ignore-unfixed: true
           vuln-type: "os,library"
       - name: Upload Trivy scan results to GitHub Security tab
-        uses: github/codeql-action/upload-sarif@v2
+        uses: github/codeql-action/upload-sarif@v3
         with:
           sarif_file: "trivy-results.sarif"
 
@@ -296,7 +296,7 @@ jobs:
         with:
           distribution: "zulu"
           java-version: 17
-      - uses: gradle/actions/setup-gradle@v3
+      - uses: gradle/actions/setup-gradle@v4
       - name: Check out the repo
         uses: acryldata/sane-checkout-action@v3
       - name: Pre-build artifacts for docker image
@@ -333,7 +333,7 @@ jobs:
         with:
           image: ${{ env.DATAHUB_MCE_CONSUMER_IMAGE }}:${{ needs.setup.outputs.unique_tag }}
       - name: Run Trivy vulnerability scanner
-        uses: aquasecurity/trivy-action@0.26.0
+        uses: aquasecurity/trivy-action@0.29.0
         env:
           TRIVY_OFFLINE_SCAN: true
           TRIVY_DB_REPOSITORY: public.ecr.aws/aquasecurity/trivy-db:2,ghcr.io/aquasecurity/trivy-db:2
@@ -347,7 +347,7 @@ jobs:
           ignore-unfixed: true
           vuln-type: "os,library"
       - name: Upload Trivy scan results to GitHub Security tab
-        uses: github/codeql-action/upload-sarif@v2
+        uses: github/codeql-action/upload-sarif@v3
         with:
           sarif_file: "trivy-results.sarif"
 
@@ -367,7 +367,7 @@ jobs:
         with:
           distribution: "zulu"
           java-version: 17
-      - uses: gradle/actions/setup-gradle@v3
+      - uses: gradle/actions/setup-gradle@v4
       - name: Check out the repo
         uses: acryldata/sane-checkout-action@v3
       - name: Pre-build artifacts for docker image
@@ -404,7 +404,7 @@ jobs:
         with:
           image: ${{ env.DATAHUB_UPGRADE_IMAGE }}:${{ needs.setup.outputs.unique_tag }}
       - name: Run Trivy vulnerability scanner
-        uses: aquasecurity/trivy-action@0.26.0
+        uses: aquasecurity/trivy-action@0.29.0
         env:
           TRIVY_OFFLINE_SCAN: true
           TRIVY_DB_REPOSITORY: public.ecr.aws/aquasecurity/trivy-db:2,ghcr.io/aquasecurity/trivy-db:2
@@ -418,7 +418,7 @@ jobs:
           ignore-unfixed: true
           vuln-type: "os,library"
       - name: Upload Trivy scan results to GitHub Security tab
-        uses: github/codeql-action/upload-sarif@v2
+        uses: github/codeql-action/upload-sarif@v3
         with:
           sarif_file: "trivy-results.sarif"
 
@@ -438,7 +438,7 @@ jobs:
         with:
           distribution: "zulu"
           java-version: 17
-      - uses: gradle/actions/setup-gradle@v3
+      - uses: gradle/actions/setup-gradle@v4
       - name: Check out the repo
         uses: acryldata/sane-checkout-action@v3
       - name: Pre-build artifacts for docker image
@@ -475,7 +475,7 @@ jobs:
         with:
           image: ${{ env.DATAHUB_FRONTEND_IMAGE }}:${{ needs.setup.outputs.unique_tag }}
       - name: Run Trivy vulnerability scanner
-        uses: aquasecurity/trivy-action@0.26.0
+        uses: aquasecurity/trivy-action@0.29.0
         env:
           TRIVY_OFFLINE_SCAN: true
           TRIVY_DB_REPOSITORY: public.ecr.aws/aquasecurity/trivy-db:2,ghcr.io/aquasecurity/trivy-db:2
@@ -489,7 +489,7 @@ jobs:
           ignore-unfixed: true
           vuln-type: "os,library"
       - name: Upload Trivy scan results to GitHub Security tab
-        uses: github/codeql-action/upload-sarif@v2
+        uses: github/codeql-action/upload-sarif@v3
         with:
           sarif_file: "trivy-results.sarif"
 
@@ -536,7 +536,7 @@ jobs:
         with:
           image: ${{ env.DATAHUB_KAFKA_SETUP_IMAGE }}:${{ needs.setup.outputs.unique_tag }}
       - name: Run Trivy vulnerability scanner
-        uses: aquasecurity/trivy-action@0.26.0
+        uses: aquasecurity/trivy-action@0.29.0
         env:
           TRIVY_OFFLINE_SCAN: true
           TRIVY_DB_REPOSITORY: public.ecr.aws/aquasecurity/trivy-db:2,ghcr.io/aquasecurity/trivy-db:2
@@ -550,7 +550,7 @@ jobs:
           ignore-unfixed: true
           vuln-type: "os,library"
       - name: Upload Trivy scan results to GitHub Security tab
-        uses: github/codeql-action/upload-sarif@v2
+        uses: github/codeql-action/upload-sarif@v3
         with:
           sarif_file: "trivy-results.sarif"
 
@@ -597,7 +597,7 @@ jobs:
         with:
           image: ${{ env.DATAHUB_MYSQL_SETUP_IMAGE }}:${{ needs.setup.outputs.unique_tag }}
       - name: Run Trivy vulnerability scanner
-        uses: aquasecurity/trivy-action@0.26.0
+        uses: aquasecurity/trivy-action@0.29.0
         env:
           TRIVY_OFFLINE_SCAN: true
           TRIVY_DB_REPOSITORY: public.ecr.aws/aquasecurity/trivy-db:2,ghcr.io/aquasecurity/trivy-db:2
@@ -611,7 +611,7 @@ jobs:
           ignore-unfixed: true
           vuln-type: "os,library"
       - name: Upload Trivy scan results to GitHub Security tab
-        uses: github/codeql-action/upload-sarif@v2
+        uses: github/codeql-action/upload-sarif@v3
         with:
           sarif_file: "trivy-results.sarif"
 
@@ -658,7 +658,7 @@ jobs:
         with:
           image: ${{ env.DATAHUB_ELASTIC_SETUP_IMAGE }}:${{ needs.setup.outputs.unique_tag }}
       - name: Run Trivy vulnerability scanner
-        uses: aquasecurity/trivy-action@0.26.0
+        uses: aquasecurity/trivy-action@0.29.0
         env:
           TRIVY_OFFLINE_SCAN: true
           TRIVY_DB_REPOSITORY: public.ecr.aws/aquasecurity/trivy-db:2,ghcr.io/aquasecurity/trivy-db:2
@@ -672,7 +672,7 @@ jobs:
           ignore-unfixed: true
           vuln-type: "os,library"
       - name: Upload Trivy scan results to GitHub Security tab
-        uses: github/codeql-action/upload-sarif@v2
+        uses: github/codeql-action/upload-sarif@v3
         with:
           sarif_file: "trivy-results.sarif"
 
@@ -829,7 +829,7 @@ jobs:
         with:
           distribution: "zulu"
           java-version: 17
-      - uses: gradle/actions/setup-gradle@v3
+      - uses: gradle/actions/setup-gradle@v4
       - name: Build codegen
         if: ${{ needs.setup.outputs.ingestion_change == 'true' || needs.setup.outputs.publish == 'true' || needs.setup.outputs.pr-publish =='true' }}
         run: ./gradlew :metadata-ingestion:codegen
@@ -886,7 +886,7 @@ jobs:
         with:
           image: ${{ env.DATAHUB_INGESTION_IMAGE }}:${{ needs.datahub_ingestion_slim_build.outputs.tag }}
       - name: Run Trivy vulnerability scanner Slim Image
-        uses: aquasecurity/trivy-action@0.26.0
+        uses: aquasecurity/trivy-action@0.29.0
         env:
           TRIVY_OFFLINE_SCAN: true
           TRIVY_DB_REPOSITORY: public.ecr.aws/aquasecurity/trivy-db:2,ghcr.io/aquasecurity/trivy-db:2
@@ -901,7 +901,7 @@ jobs:
           vuln-type: "os,library"
           timeout: 15m
       - name: Upload Trivy scan results to GitHub Security tab
-        uses: github/codeql-action/upload-sarif@v2
+        uses: github/codeql-action/upload-sarif@v3
         with:
           sarif_file: "trivy-results.sarif"
 
@@ -930,7 +930,7 @@ jobs:
         with:
           distribution: "zulu"
           java-version: 17
-      - uses: gradle/actions/setup-gradle@v3
+      - uses: gradle/actions/setup-gradle@v4
       - name: Build codegen
         if: ${{ needs.setup.outputs.ingestion_change == 'true' || needs.setup.outputs.publish == 'true' || needs.setup.outputs.pr-publish == 'true' }}
         run: ./gradlew :metadata-ingestion:codegen
@@ -985,7 +985,7 @@ jobs:
         with:
           image: ${{ env.DATAHUB_INGESTION_IMAGE }}:${{ needs.datahub_ingestion_full_build.outputs.tag }}
       - name: Run Trivy vulnerability scanner Full Image
-        uses: aquasecurity/trivy-action@0.26.0
+        uses: aquasecurity/trivy-action@0.29.0
         env:
           TRIVY_OFFLINE_SCAN: true
           TRIVY_DB_REPOSITORY: public.ecr.aws/aquasecurity/trivy-db:2,ghcr.io/aquasecurity/trivy-db:2
@@ -1000,7 +1000,7 @@ jobs:
           vuln-type: "os,library"
           timeout: 15m
       - name: Upload Trivy scan results to GitHub Security tab
-        uses: github/codeql-action/upload-sarif@v2
+        uses: github/codeql-action/upload-sarif@v3
         with:
           sarif_file: "trivy-results.sarif"
 
@@ -1018,8 +1018,8 @@ jobs:
         # python_batch_count is used to split pytests in the smoke-test (batches of actual test functions)
         # cypress_batch_count is used to split the collection of cypress test specs into batches.
         run: |
-          echo "cypress_batch_count=11" >> "$GITHUB_OUTPUT"
-          echo "python_batch_count=5" >> "$GITHUB_OUTPUT"
+          echo "cypress_batch_count=5" >> "$GITHUB_OUTPUT"
+          echo "python_batch_count=3" >> "$GITHUB_OUTPUT"
 
       - id: set-matrix
         # For m batches for python and n batches for cypress, we need a test matrix of python x m + cypress x n.
@@ -1083,7 +1083,7 @@ jobs:
         with:
           distribution: "zulu"
           java-version: 17
-      - uses: gradle/actions/setup-gradle@v3
+      - uses: gradle/actions/setup-gradle@v4
       - name: Login to DockerHub
         uses: docker/login-action@v3
         if: ${{ needs.setup.outputs.docker-login == 'true' }}
diff --git a/.github/workflows/documentation.yml b/.github/workflows/documentation.yml
index e6044badb1b41c..f9d459652086fb 100644
--- a/.github/workflows/documentation.yml
+++ b/.github/workflows/documentation.yml
@@ -5,6 +5,7 @@ on:
     branches:
       - "**"
     paths:
+      - ".github/workflows/documentation.yml"
       - "metadata-ingestion/**"
       - "metadata-models/**"
       - "docs/**"
@@ -13,6 +14,7 @@ on:
     branches:
       - master
     paths:
+      - ".github/workflows/documentation.yml"
       - "metadata-ingestion/**"
       - "metadata-models/**"
       - "docs/**"
@@ -40,7 +42,7 @@ jobs:
         with:
           distribution: "zulu"
           java-version: 17
-      - uses: gradle/actions/setup-gradle@v3
+      - uses: gradle/actions/setup-gradle@v4
       - uses: actions/setup-python@v5
         with:
           python-version: "3.10"
@@ -57,8 +59,12 @@ jobs:
 
       - name: Deploy
         if: github.event_name == 'push' && github.repository == 'datahub-project/datahub'
-        uses: peaceiris/actions-gh-pages@v3
+        uses: peaceiris/actions-gh-pages@v4
         with:
           github_token: ${{ secrets.GITHUB_TOKEN }}
           publish_dir: ./docs-website/build
           cname: datahubproject.io
+          # The gh-pages branch stores the built docs site. We don't need to preserve
+          # the full history of the .html files, since they're generated from our
+          # source files. Doing so significantly reduces the size of the repo's .git dir.
+          force_orphan: true
diff --git a/.github/workflows/gx-plugin.yml b/.github/workflows/gx-plugin.yml
index c28bdbb30eb36d..68d7934d1d6e6f 100644
--- a/.github/workflows/gx-plugin.yml
+++ b/.github/workflows/gx-plugin.yml
@@ -48,7 +48,7 @@ jobs:
         with:
           distribution: "zulu"
           java-version: 17
-      - uses: gradle/actions/setup-gradle@v3
+      - uses: gradle/actions/setup-gradle@v4
       - uses: acryldata/sane-checkout-action@v3
       - uses: actions/setup-python@v5
         with:
diff --git a/.github/workflows/metadata-ingestion.yml b/.github/workflows/metadata-ingestion.yml
index be6026098ce420..ff4531ab0b25c1 100644
--- a/.github/workflows/metadata-ingestion.yml
+++ b/.github/workflows/metadata-ingestion.yml
@@ -54,7 +54,7 @@ jobs:
         with:
           distribution: "zulu"
           java-version: 17
-      - uses: gradle/actions/setup-gradle@v3
+      - uses: gradle/actions/setup-gradle@v4
       - uses: acryldata/sane-checkout-action@v3
       - uses: actions/setup-python@v5
         with:
diff --git a/.github/workflows/metadata-io.yml b/.github/workflows/metadata-io.yml
index 80af03e77eef82..6514b9e7226b60 100644
--- a/.github/workflows/metadata-io.yml
+++ b/.github/workflows/metadata-io.yml
@@ -67,7 +67,7 @@ jobs:
         with:
           distribution: "zulu"
           java-version: 17
-      - uses: gradle/actions/setup-gradle@v3
+      - uses: gradle/actions/setup-gradle@v4
       - name: Gradle build (and test)
         run: |
           ./gradlew :metadata-io:test
diff --git a/.github/workflows/metadata-model.yml b/.github/workflows/metadata-model.yml
index 6f62284afcc172..f9262d89f9ef00 100644
--- a/.github/workflows/metadata-model.yml
+++ b/.github/workflows/metadata-model.yml
@@ -32,7 +32,7 @@ jobs:
         with:
           distribution: "zulu"
           java-version: 17
-      - uses: gradle/actions/setup-gradle@v3
+      - uses: gradle/actions/setup-gradle@v4
       - uses: acryldata/sane-checkout-action@v3
       - uses: actions/setup-python@v5
         with:
diff --git a/.github/workflows/pr-labeler.yml b/.github/workflows/pr-labeler.yml
index de7ad21b3e67bb..55f7fea875388a 100644
--- a/.github/workflows/pr-labeler.yml
+++ b/.github/workflows/pr-labeler.yml
@@ -3,6 +3,10 @@ on:
   pull_request_target:
     types: [opened, reopened]
 
+permissions:
+    contents: read
+    pull-requests: write
+    
 jobs:
   triage:
     permissions:
@@ -10,11 +14,11 @@ jobs:
       pull-requests: write
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/labeler@v4
+      - uses: actions/labeler@v5
         with:
           repo-token: "${{ secrets.GITHUB_TOKEN }}"
           configuration-path: ".github/pr-labeler-config.yml"
-      - uses: actions-ecosystem/action-add-labels@v1.1.0
+      - uses: actions-ecosystem/action-add-labels@v1.1.3
         # only add names of Acryl Data team members here
         if: 
           ${{ 
@@ -25,7 +29,6 @@ jobs:
                 "chriscollins3456",
                 "david-leifker",
                 "shirshanka",
-                "sid-acryl",
                 "swaroopjagadish",
                 "treff7es",
                 "yoonhyejin",
@@ -46,7 +49,9 @@ jobs:
                 "kevinkarchacryl",
                 "sgomezvillamor",
                 "acrylJonny",
-                "chakru-r"
+                "chakru-r",
+                "brock-acryl",
+                "mminichino"
               ]'), 
               github.actor
             ) 
@@ -55,7 +60,7 @@ jobs:
           github_token: ${{ github.token }}
           labels: |
             community-contribution
-      - uses: actions-ecosystem/action-add-labels@v1.1.0
+      - uses: actions-ecosystem/action-add-labels@v1.1.3
         # only add names of champions here. Confirm with DevRel Team
         if: 
           ${{ 
diff --git a/.github/workflows/prefect-plugin.yml b/.github/workflows/prefect-plugin.yml
index 401efa340ae8ca..f1b06f31a05224 100644
--- a/.github/workflows/prefect-plugin.yml
+++ b/.github/workflows/prefect-plugin.yml
@@ -39,7 +39,7 @@ jobs:
         with:
           distribution: "zulu"
           java-version: 17
-      - uses: gradle/actions/setup-gradle@v3
+      - uses: gradle/actions/setup-gradle@v4
       - uses: acryldata/sane-checkout-action@v3
       - uses: actions/setup-python@v5
         with:
diff --git a/.github/workflows/publish-datahub-jars.yml b/.github/workflows/publish-datahub-jars.yml
index 393f9d993e2a2f..fa3554c05d374c 100644
--- a/.github/workflows/publish-datahub-jars.yml
+++ b/.github/workflows/publish-datahub-jars.yml
@@ -57,7 +57,7 @@ jobs:
         with:
           distribution: "zulu"
           java-version: 17
-      - uses: gradle/actions/setup-gradle@v3
+      - uses: gradle/actions/setup-gradle@v4
       - uses: actions/setup-python@v5
         with:
           python-version: "3.10"
@@ -210,7 +210,7 @@ jobs:
         with:
           distribution: "zulu"
           java-version: 17
-      - uses: gradle/actions/setup-gradle@v3
+      - uses: gradle/actions/setup-gradle@v4
       - uses: actions/setup-python@v5
         with:
           python-version: "3.10"
diff --git a/.github/workflows/python-build-pages.yml b/.github/workflows/python-build-pages.yml
index 8971722c374fb7..9e1a2563478039 100644
--- a/.github/workflows/python-build-pages.yml
+++ b/.github/workflows/python-build-pages.yml
@@ -37,7 +37,7 @@ jobs:
         with:
           distribution: "zulu"
           java-version: 17
-      - uses: gradle/actions/setup-gradle@v3
+      - uses: gradle/actions/setup-gradle@v4
       - uses: acryldata/sane-checkout-action@v3
       - uses: actions/setup-python@v5
         with:
diff --git a/.github/workflows/spark-smoke-test.yml b/.github/workflows/spark-smoke-test.yml
index e6a6705a72879c..7a2080b7a9ddf3 100644
--- a/.github/workflows/spark-smoke-test.yml
+++ b/.github/workflows/spark-smoke-test.yml
@@ -35,7 +35,7 @@ jobs:
         with:
           distribution: "zulu"
           java-version: 17
-      - uses: gradle/actions/setup-gradle@v3
+      - uses: gradle/actions/setup-gradle@v4
       - uses: actions/setup-python@v5
         with:
           python-version: "3.10"
diff --git a/.gitignore b/.gitignore
index 43c627f9ed244f..19909b25fefe7f 100644
--- a/.gitignore
+++ b/.gitignore
@@ -85,6 +85,7 @@ metadata-service/plugin/src/test/resources/sample-plugins/**
 smoke-test/rollback-reports
 coverage*.xml
 .vercel
+.envrc
 
 # A long series of binary directories we should ignore
 datahub-frontend/bin/main/
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 3697efa37770e7..103e1680b8d105 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -1,4 +1,4 @@
-# Auto-generated by .github/scripts/generate_pre_commit.py at 2025-01-09 10:08:09 UTC
+# Auto-generated by .github/scripts/generate_pre_commit.py at 2025-01-17 16:43:31 UTC
 # Do not edit this file directly. Run the script to regenerate.
 # Add additional hooks in .github/scripts/pre-commit-override.yaml
 repos:
@@ -442,4 +442,5 @@ repos:
         name: smoke-test cypress Lint Fix
         entry: ./gradlew :smoke-test:cypressLintFix
         language: system
-        files: ^smoke-test/tests/cypress/.*$
+        files: ^smoke-test/tests/cypress/.*\.tsx$
+        pass_filenames: false
diff --git a/build.gradle b/build.gradle
index e4fd70a99e6434..2984812bda13b8 100644
--- a/build.gradle
+++ b/build.gradle
@@ -39,13 +39,13 @@ buildscript {
   ext.springBootVersion = '3.2.9'
   ext.springKafkaVersion = '3.1.6'
   ext.openTelemetryVersion = '1.18.0'
-  ext.neo4jVersion = '5.14.0'
-  ext.neo4jTestVersion = '5.14.0'
-  ext.neo4jApocVersion = '5.14.0'
+  ext.neo4jVersion = '5.20.0'
+  ext.neo4jTestVersion = '5.20.0'
+  ext.neo4jApocVersion = '5.20.0'
   ext.testContainersVersion = '1.17.4'
   ext.elasticsearchVersion = '2.11.1' // ES 7.10, Opensearch 1.x, 2.x
   ext.jacksonVersion = '2.15.3'
-  ext.jettyVersion = '11.0.21'
+  ext.jettyVersion = '12.0.16'
   // see also datahub-frontend/play.gradle
   ext.playVersion = '2.8.22'
   ext.playScalaVersion = '2.13'
@@ -136,7 +136,8 @@ project.ext.externalDependency = [
     'datastaxOssNativeProtocol': 'com.datastax.oss:native-protocol:1.5.1',
     'datastaxOssCore': 'com.datastax.oss:java-driver-core:4.14.1',
     'datastaxOssQueryBuilder': 'com.datastax.oss:java-driver-query-builder:4.14.1',
-    'dgraph4j' : 'io.dgraph:dgraph4j:21.12.0',
+    'dgraph4j' : 'io.dgraph:dgraph4j:24.1.1',
+    'dgraphNetty': 'io.grpc:grpc-netty-shaded:1.69.0',
     'dropwizardMetricsCore': 'io.dropwizard.metrics:metrics-core:4.2.3',
     'dropwizardMetricsJmx': 'io.dropwizard.metrics:metrics-jmx:4.2.3',
     'ebean': 'io.ebean:ebean:' + ebeanVersion,
@@ -176,8 +177,9 @@ project.ext.externalDependency = [
     'jakartaValidation': 'jakarta.validation:jakarta.validation-api:3.1.0-M2',
     'jerseyCore': 'org.glassfish.jersey.core:jersey-client:2.41',
     'jerseyGuava': 'org.glassfish.jersey.bundles.repackaged:jersey-guava:2.25.1',
-    'jettyJaas': "org.eclipse.jetty:jetty-jaas:$jettyVersion",
+    'jettySecurity': "org.eclipse.jetty:jetty-security:$jettyVersion",
     'jettyClient': "org.eclipse.jetty:jetty-client:$jettyVersion",
+    'jettyJmx': "org.eclipse.jetty:jetty-jmx:$jettyVersion",
     'jettison': 'org.codehaus.jettison:jettison:1.5.4',
     'jgrapht': 'org.jgrapht:jgrapht-core:1.5.1',
     'jna': 'net.java.dev.jna:jna:5.12.1',
@@ -193,7 +195,7 @@ project.ext.externalDependency = [
     'junitJupiterEngine': "org.junit.jupiter:junit-jupiter-engine:$junitJupiterVersion",
     // avro-serde includes dependencies for `kafka-avro-serializer` `kafka-schema-registry-client` and `avro`
     'kafkaAvroSerde': "io.confluent:kafka-streams-avro-serde:$kafkaVersion",
-    'kafkaAvroSerializer': 'io.confluent:kafka-avro-serializer:5.1.4',
+    'kafkaAvroSerializer': "io.confluent:kafka-avro-serializer:$kafkaVersion",
     'kafkaClients': "org.apache.kafka:kafka-clients:$kafkaVersion-ccs",
     'snappy': 'org.xerial.snappy:snappy-java:1.1.10.5',
     'logbackClassic': "ch.qos.logback:logback-classic:$logbackClassic",
@@ -380,6 +382,13 @@ configure(subprojects.findAll {! it.name.startsWith('spark-lineage')}) {
     resolutionStrategy.force externalDependency.antlr4Runtime
     resolutionStrategy.force externalDependency.antlr4
     resolutionStrategy.force 'org.apache.mina:mina-core:2.2.4'
+    resolutionStrategy {
+      force "org.eclipse.jetty:jetty-security:${jettyVersion}"
+      force "org.eclipse.jetty:jetty-server:${jettyVersion}"
+      force "org.eclipse.jetty:jetty-ee10-servlet:${jettyVersion}"
+      force "org.eclipse.jetty:jetty-ee10-webapp:${jettyVersion}"
+      force "org.eclipse.jetty:jetty-xml:${jettyVersion}"
+    }
   }
 }
 
@@ -407,7 +416,7 @@ subprojects {
       implementation externalDependency.annotationApi
       constraints {
         implementation("com.google.googlejavaformat:google-java-format:$googleJavaFormatVersion")
-        implementation('io.netty:netty-all:4.1.115.Final')
+        implementation('io.netty:netty-all:4.1.116.Final')
         implementation('org.apache.commons:commons-compress:1.27.1')
         implementation('org.apache.velocity:velocity-engine-core:2.4')
         implementation('org.hibernate:hibernate-validator:6.0.20.Final')
diff --git a/datahub-frontend/app/security/AuthenticationManager.java b/datahub-frontend/app/security/AuthenticationManager.java
index f46dc57c232bd2..8e7d51a0776c23 100644
--- a/datahub-frontend/app/security/AuthenticationManager.java
+++ b/datahub-frontend/app/security/AuthenticationManager.java
@@ -1,68 +1,33 @@
 package security;
 
 import com.google.common.base.Preconditions;
-import java.util.Collections;
 import javax.annotation.Nonnull;
 import javax.naming.AuthenticationException;
-import javax.security.auth.callback.Callback;
-import javax.security.auth.callback.CallbackHandler;
-import javax.security.auth.callback.NameCallback;
-import javax.security.auth.callback.PasswordCallback;
-import javax.security.auth.login.LoginContext;
-import javax.security.auth.login.LoginException;
 import org.apache.commons.lang3.StringUtils;
-import org.eclipse.jetty.jaas.JAASLoginService;
-import org.eclipse.jetty.jaas.PropertyUserStoreManager;
-import play.Logger;
+import org.eclipse.jetty.security.UserPrincipal;
+import org.eclipse.jetty.util.security.Credential;
 
 public class AuthenticationManager {
-
-  private AuthenticationManager(boolean verbose) {}
+  private AuthenticationManager() {} // Prevent instantiation
 
   public static void authenticateJaasUser(@Nonnull String userName, @Nonnull String password)
       throws Exception {
     Preconditions.checkArgument(!StringUtils.isAnyEmpty(userName), "Username cannot be empty");
-    JAASLoginService jaasLoginService = new JAASLoginService("WHZ-Authentication");
-    PropertyUserStoreManager propertyUserStoreManager = new PropertyUserStoreManager();
-    propertyUserStoreManager.start();
-    jaasLoginService.setBeans(Collections.singletonList(propertyUserStoreManager));
-    JAASLoginService.INSTANCE.set(jaasLoginService);
-    try {
-      LoginContext lc =
-          new LoginContext("WHZ-Authentication", new WHZCallbackHandler(userName, password));
-      lc.login();
-    } catch (LoginException le) {
-      AuthenticationException authenticationException =
-          new AuthenticationException(le.getMessage());
-      authenticationException.setRootCause(le);
-      throw authenticationException;
-    }
-  }
 
-  private static class WHZCallbackHandler implements CallbackHandler {
-    private String password;
-    private String username;
-
-    private WHZCallbackHandler(@Nonnull String username, @Nonnull String password) {
-      this.username = username;
-      this.password = password;
-    }
+    try {
+      // Create and configure credentials for authentication
+      UserPrincipal userPrincipal = new UserPrincipal(userName, Credential.getCredential(password));
 
-    @Override
-    public void handle(@Nonnull Callback[] callbacks) {
-      NameCallback nc = null;
-      PasswordCallback pc = null;
-      for (Callback callback : callbacks) {
-        Logger.debug(
-            "The submitted callback is of type: " + callback.getClass() + " : " + callback);
-        if (callback instanceof NameCallback) {
-          nc = (NameCallback) callback;
-          nc.setName(this.username);
-        } else if (callback instanceof PasswordCallback) {
-          pc = (PasswordCallback) callback;
-          pc.setPassword(this.password.toCharArray());
-        }
+      // Verify credentials
+      if (!userPrincipal.authenticate(password)) {
+        throw new AuthenticationException("Invalid credentials for user: " + userName);
       }
+
+    } catch (Exception e) {
+      AuthenticationException authenticationException =
+          new AuthenticationException("Authentication failed");
+      authenticationException.setRootCause(e);
+      throw authenticationException;
     }
   }
 }
diff --git a/datahub-frontend/play.gradle b/datahub-frontend/play.gradle
index d513c3c232d9a0..1a9ffeede56251 100644
--- a/datahub-frontend/play.gradle
+++ b/datahub-frontend/play.gradle
@@ -50,7 +50,7 @@ dependencies {
   implementation externalDependency.springBeans
   implementation externalDependency.springContext
   implementation externalDependency.springBootAutoconfigure
-  implementation externalDependency.jettyJaas
+  implementation externalDependency.jettySecurity
   implementation externalDependency.graphqlJava
   implementation externalDependency.antlr4Runtime
   implementation externalDependency.antlr4
diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/Constants.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/Constants.java
index 69306862a46ef7..aec5352dec1a64 100644
--- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/Constants.java
+++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/Constants.java
@@ -28,6 +28,7 @@ private Constants() {}
   public static final String INCIDENTS_SCHEMA_FILE = "incident.graphql";
   public static final String CONTRACTS_SCHEMA_FILE = "contract.graphql";
   public static final String CONNECTIONS_SCHEMA_FILE = "connection.graphql";
+  public static final String VERSION_SCHEMA_FILE = "versioning.graphql";
   public static final String BROWSE_PATH_DELIMITER = "/";
   public static final String BROWSE_PATH_V2_DELIMITER = "␟";
   public static final String VERSION_STAMP_FIELD_NAME = "versionStamp";
diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java
index b15db80a8487ae..2c5e841322f45f 100644
--- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java
+++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java
@@ -124,6 +124,8 @@
 import com.linkedin.datahub.graphql.generated.TestResult;
 import com.linkedin.datahub.graphql.generated.TypeQualifier;
 import com.linkedin.datahub.graphql.generated.UserUsageCounts;
+import com.linkedin.datahub.graphql.generated.VersionProperties;
+import com.linkedin.datahub.graphql.generated.VersionSet;
 import com.linkedin.datahub.graphql.resolvers.MeResolver;
 import com.linkedin.datahub.graphql.resolvers.assertion.AssertionRunEventResolver;
 import com.linkedin.datahub.graphql.resolvers.assertion.DeleteAssertionResolver;
@@ -324,6 +326,7 @@
 import com.linkedin.datahub.graphql.resolvers.user.ListUsersResolver;
 import com.linkedin.datahub.graphql.resolvers.user.RemoveUserResolver;
 import com.linkedin.datahub.graphql.resolvers.user.UpdateUserStatusResolver;
+import com.linkedin.datahub.graphql.resolvers.versioning.VersionsSearchResolver;
 import com.linkedin.datahub.graphql.resolvers.view.CreateViewResolver;
 import com.linkedin.datahub.graphql.resolvers.view.DeleteViewResolver;
 import com.linkedin.datahub.graphql.resolvers.view.ListGlobalViewsResolver;
@@ -381,6 +384,7 @@
 import com.linkedin.datahub.graphql.types.structuredproperty.StructuredPropertyType;
 import com.linkedin.datahub.graphql.types.tag.TagType;
 import com.linkedin.datahub.graphql.types.test.TestType;
+import com.linkedin.datahub.graphql.types.versioning.VersionSetType;
 import com.linkedin.datahub.graphql.types.view.DataHubViewType;
 import com.linkedin.entity.client.EntityClient;
 import com.linkedin.entity.client.SystemEntityClient;
@@ -537,6 +541,7 @@ public class GmsGraphQLEngine {
   private final IncidentType incidentType;
   private final RestrictedType restrictedType;
   private final DataProcessInstanceType dataProcessInstanceType;
+  private final VersionSetType versionSetType;
 
   private final int graphQLQueryComplexityLimit;
   private final int graphQLQueryDepthLimit;
@@ -658,6 +663,7 @@ public GmsGraphQLEngine(final GmsGraphQLEngineArgs args) {
     this.incidentType = new IncidentType(entityClient);
     this.restrictedType = new RestrictedType(entityClient, restrictedService);
     this.dataProcessInstanceType = new DataProcessInstanceType(entityClient, featureFlags);
+    this.versionSetType = new VersionSetType(entityClient);
 
     this.graphQLQueryComplexityLimit = args.graphQLQueryComplexityLimit;
     this.graphQLQueryDepthLimit = args.graphQLQueryDepthLimit;
@@ -707,6 +713,7 @@ public GmsGraphQLEngine(final GmsGraphQLEngineArgs args) {
                 entityTypeType,
                 formType,
                 incidentType,
+                versionSetType,
                 restrictedType,
                 businessAttributeType,
                 dataProcessInstanceType));
@@ -809,6 +816,8 @@ public void configureRuntimeWiring(final RuntimeWiring.Builder builder) {
     configureConnectionResolvers(builder);
     configureDeprecationResolvers(builder);
     configureMetadataAttributionResolver(builder);
+    configureVersionPropertiesResolvers(builder);
+    configureVersionSetResolvers(builder);
   }
 
   private void configureOrganisationRoleResolvers(RuntimeWiring.Builder builder) {
@@ -863,7 +872,8 @@ public GraphQLEngine.Builder builder() {
         .addSchema(fileBasedSchema(ASSERTIONS_SCHEMA_FILE))
         .addSchema(fileBasedSchema(INCIDENTS_SCHEMA_FILE))
         .addSchema(fileBasedSchema(CONTRACTS_SCHEMA_FILE))
-        .addSchema(fileBasedSchema(COMMON_SCHEMA_FILE));
+        .addSchema(fileBasedSchema(COMMON_SCHEMA_FILE))
+        .addSchema(fileBasedSchema(VERSION_SCHEMA_FILE));
 
     for (GmsGraphQLPlugin plugin : this.graphQLPlugins) {
       List<String> pluginSchemaFiles = plugin.getSchemaFiles();
@@ -1050,6 +1060,7 @@ private void configureQueryResolvers(final RuntimeWiring.Builder builder) {
                 .dataFetcher("form", getResolver(formType))
                 .dataFetcher("view", getResolver(dataHubViewType))
                 .dataFetcher("structuredProperty", getResolver(structuredPropertyType))
+                .dataFetcher("versionSet", getResolver(versionSetType))
                 .dataFetcher("listPolicies", new ListPoliciesResolver(this.entityClient))
                 .dataFetcher("getGrantedPrivileges", new GetGrantedPrivilegesResolver())
                 .dataFetcher("listUsers", new ListUsersResolver(this.entityClient))
@@ -2295,7 +2306,15 @@ private void configureTypeResolvers(final RuntimeWiring.Builder builder) {
         .type(
             "TimeSeriesAspect",
             typeWiring -> typeWiring.typeResolver(new TimeSeriesAspectInterfaceTypeResolver()))
-        .type("ResultsType", typeWiring -> typeWiring.typeResolver(new ResultsTypeResolver()));
+        .type("ResultsType", typeWiring -> typeWiring.typeResolver(new ResultsTypeResolver()))
+        .type(
+            "SupportsVersions",
+            typeWiring ->
+                typeWiring.typeResolver(
+                    new EntityInterfaceTypeResolver(
+                        loadableTypes.stream()
+                            .map(graphType -> (EntityType<?, ?>) graphType)
+                            .collect(Collectors.toList()))));
   }
 
   /** Configures custom type extensions leveraged within our GraphQL schema. */
@@ -3089,16 +3108,6 @@ private void configureDataProcessInstanceResolvers(final RuntimeWiring.Builder b
                               ? dataProcessInstance.getDataPlatformInstance().getUrn()
                               : null;
                         }))
-                .dataFetcher(
-                    "platform",
-                    new LoadableTypeResolver<>(
-                        dataPlatformType,
-                        (env) -> {
-                          final DataProcessInstance dataProcessInstance = env.getSource();
-                          return dataProcessInstance.getPlatform() != null
-                              ? dataProcessInstance.getPlatform().getUrn()
-                              : null;
-                        }))
                 .dataFetcher("parentContainers", new ParentContainersResolver(entityClient))
                 .dataFetcher(
                     "container",
@@ -3322,4 +3331,34 @@ private void configureMetadataAttributionResolver(final RuntimeWiring.Builder bu
                         entityTypes,
                         (env) -> ((MetadataAttribution) env.getSource()).getSource())));
   }
+
+  private void configureVersionPropertiesResolvers(final RuntimeWiring.Builder builder) {
+    builder.type(
+        "VersionProperties",
+        typeWiring ->
+            typeWiring.dataFetcher(
+                "versionSet",
+                new LoadableTypeResolver<>(
+                    versionSetType,
+                    (env) -> {
+                      final VersionProperties versionProperties = env.getSource();
+                      return versionProperties != null
+                          ? versionProperties.getVersionSet().getUrn()
+                          : null;
+                    })));
+  }
+
+  private void configureVersionSetResolvers(final RuntimeWiring.Builder builder) {
+    builder.type(
+        "VersionSet",
+        typeWiring ->
+            typeWiring
+                .dataFetcher(
+                    "latestVersion",
+                    new EntityTypeResolver(
+                        entityTypes, (env) -> ((VersionSet) env.getSource()).getLatestVersion()))
+                .dataFetcher(
+                    "versionsSearch",
+                    new VersionsSearchResolver(this.entityClient, this.viewService)));
+  }
 }
diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/authorization/AuthorizationUtils.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/authorization/AuthorizationUtils.java
index c25d6af75fe76d..29d1c02dacb416 100644
--- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/authorization/AuthorizationUtils.java
+++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/authorization/AuthorizationUtils.java
@@ -232,6 +232,10 @@ public static <T> T restrictEntity(@Nonnull Object entity, Class<T> clazz) {
     try {
       Object[] args =
           allFields.stream()
+              // New versions of graphql.codegen generate serialVersionUID
+              // We need to filter serialVersionUID out because serialVersionUID is
+              // never part of the entity type constructor
+              .filter(field -> !field.getName().contains("serialVersionUID"))
               .map(
                   field -> {
                     // properties are often not required but only because
diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/config/AppConfigResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/config/AppConfigResolver.java
index 3647eb55b2583a..8cdc13a14be87c 100644
--- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/config/AppConfigResolver.java
+++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/config/AppConfigResolver.java
@@ -189,6 +189,7 @@ public CompletableFuture<AppConfig> get(final DataFetchingEnvironment environmen
             .setEditableDatasetNameEnabled(_featureFlags.isEditableDatasetNameEnabled())
             .setShowSeparateSiblings(_featureFlags.isShowSeparateSiblings())
             .setShowManageStructuredProperties(_featureFlags.isShowManageStructuredProperties())
+            .setEntityVersioningEnabled(_featureFlags.isEntityVersioning())
             .build();
 
     appConfig.setFeatureFlags(featureFlagsConfig);
diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/entity/versioning/LinkAssetVersionResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/entity/versioning/LinkAssetVersionResolver.java
index 69e049af1e87b7..f32fd03a384005 100644
--- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/entity/versioning/LinkAssetVersionResolver.java
+++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/entity/versioning/LinkAssetVersionResolver.java
@@ -12,7 +12,9 @@
 import com.linkedin.datahub.graphql.concurrency.GraphQLConcurrencyUtils;
 import com.linkedin.datahub.graphql.exception.AuthorizationException;
 import com.linkedin.datahub.graphql.featureflags.FeatureFlags;
+import com.linkedin.datahub.graphql.generated.EntityType;
 import com.linkedin.datahub.graphql.generated.LinkVersionInput;
+import com.linkedin.datahub.graphql.generated.VersionSet;
 import com.linkedin.metadata.entity.IngestResult;
 import com.linkedin.metadata.entity.versioning.EntityVersioningService;
 import com.linkedin.metadata.entity.versioning.VersionPropertiesInput;
@@ -21,24 +23,22 @@
 import io.datahubproject.metadata.context.OperationContext;
 import java.util.List;
 import java.util.concurrent.CompletableFuture;
+import lombok.RequiredArgsConstructor;
+import lombok.extern.slf4j.Slf4j;
 import org.apache.commons.lang.StringUtils;
 
 /**
  * Currently only supports linking the latest version, but may be modified later to support inserts
  */
-public class LinkAssetVersionResolver implements DataFetcher<CompletableFuture<String>> {
+@Slf4j
+@RequiredArgsConstructor
+public class LinkAssetVersionResolver implements DataFetcher<CompletableFuture<VersionSet>> {
 
   private final EntityVersioningService entityVersioningService;
   private final FeatureFlags featureFlags;
 
-  public LinkAssetVersionResolver(
-      EntityVersioningService entityVersioningService, FeatureFlags featureFlags) {
-    this.entityVersioningService = entityVersioningService;
-    this.featureFlags = featureFlags;
-  }
-
   @Override
-  public CompletableFuture<String> get(DataFetchingEnvironment environment) throws Exception {
+  public CompletableFuture<VersionSet> get(DataFetchingEnvironment environment) throws Exception {
     final QueryContext context = environment.getContext();
     final LinkVersionInput input =
         bindArgument(environment.getArgument("input"), LinkVersionInput.class);
@@ -75,12 +75,22 @@ public CompletableFuture<String> get(DataFetchingEnvironment environment) throws
               entityVersioningService.linkLatestVersion(
                   opContext, versionSetUrn, entityUrn, versionPropertiesInput);
 
-          return linkResults.stream()
-              .filter(
-                  ingestResult -> input.getLinkedEntity().equals(ingestResult.getUrn().toString()))
-              .map(ingestResult -> ingestResult.getUrn().toString())
-              .findAny()
-              .orElse(StringUtils.EMPTY);
+          String successVersionSetUrn =
+              linkResults.stream()
+                  .filter(
+                      ingestResult ->
+                          input.getLinkedEntity().equals(ingestResult.getUrn().toString()))
+                  .map(ingestResult -> ingestResult.getUrn().toString())
+                  .findAny()
+                  .orElse(StringUtils.EMPTY);
+
+          if (StringUtils.isEmpty(successVersionSetUrn)) {
+            return null;
+          }
+          VersionSet versionSet = new VersionSet();
+          versionSet.setUrn(versionSetUrn.toString());
+          versionSet.setType(EntityType.VERSION_SET);
+          return versionSet;
         },
         this.getClass().getSimpleName(),
         "get");
diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/entity/versioning/UnlinkAssetVersionResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/entity/versioning/UnlinkAssetVersionResolver.java
index 3d5027a0d668ac..33ab83a59c6771 100644
--- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/entity/versioning/UnlinkAssetVersionResolver.java
+++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/entity/versioning/UnlinkAssetVersionResolver.java
@@ -12,14 +12,18 @@
 import com.linkedin.datahub.graphql.concurrency.GraphQLConcurrencyUtils;
 import com.linkedin.datahub.graphql.exception.AuthorizationException;
 import com.linkedin.datahub.graphql.featureflags.FeatureFlags;
+import com.linkedin.datahub.graphql.generated.EntityType;
 import com.linkedin.datahub.graphql.generated.UnlinkVersionInput;
+import com.linkedin.datahub.graphql.generated.VersionSet;
+import com.linkedin.metadata.entity.RollbackResult;
 import com.linkedin.metadata.entity.versioning.EntityVersioningService;
 import graphql.schema.DataFetcher;
 import graphql.schema.DataFetchingEnvironment;
 import io.datahubproject.metadata.context.OperationContext;
+import java.util.List;
 import java.util.concurrent.CompletableFuture;
 
-public class UnlinkAssetVersionResolver implements DataFetcher<CompletableFuture<Boolean>> {
+public class UnlinkAssetVersionResolver implements DataFetcher<CompletableFuture<VersionSet>> {
 
   private final EntityVersioningService entityVersioningService;
   private final FeatureFlags featureFlags;
@@ -31,7 +35,7 @@ public UnlinkAssetVersionResolver(
   }
 
   @Override
-  public CompletableFuture<Boolean> get(DataFetchingEnvironment environment) throws Exception {
+  public CompletableFuture<VersionSet> get(DataFetchingEnvironment environment) throws Exception {
     if (!featureFlags.isEntityVersioning()) {
       throw new IllegalAccessError(
           "Entity Versioning is not configured, please enable before attempting to use this feature.");
@@ -58,8 +62,15 @@ public CompletableFuture<Boolean> get(DataFetchingEnvironment environment) throw
     }
     return GraphQLConcurrencyUtils.supplyAsync(
         () -> {
-          entityVersioningService.unlinkVersion(opContext, versionSetUrn, entityUrn);
-          return true;
+          List<RollbackResult> results =
+              entityVersioningService.unlinkVersion(opContext, versionSetUrn, entityUrn);
+          if (results.isEmpty() || results.stream().allMatch(RollbackResult::isNoOp)) {
+            return null;
+          }
+          VersionSet versionSet = new VersionSet();
+          versionSet.setUrn(versionSetUrn.toString());
+          versionSet.setType(EntityType.VERSION_SET);
+          return versionSet;
         },
         this.getClass().getSimpleName(),
         "get");
diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/search/SearchUtils.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/search/SearchUtils.java
index a01b3aaec9c982..f105a72a1273ee 100644
--- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/search/SearchUtils.java
+++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/search/SearchUtils.java
@@ -18,13 +18,18 @@
 
 import com.google.common.collect.ImmutableList;
 import com.linkedin.common.urn.Urn;
+import com.linkedin.common.urn.UrnUtils;
 import com.linkedin.datahub.graphql.QueryContext;
+import com.linkedin.datahub.graphql.concurrency.GraphQLConcurrencyUtils;
 import com.linkedin.datahub.graphql.generated.EntityType;
 import com.linkedin.datahub.graphql.generated.FacetFilterInput;
 import com.linkedin.datahub.graphql.generated.SearchResults;
 import com.linkedin.datahub.graphql.generated.SearchSortInput;
+import com.linkedin.datahub.graphql.resolvers.ResolverUtils;
 import com.linkedin.datahub.graphql.types.common.mappers.SearchFlagsInputMapper;
 import com.linkedin.datahub.graphql.types.entitytype.EntityTypeMapper;
+import com.linkedin.datahub.graphql.types.mappers.UrnSearchResultsMapper;
+import com.linkedin.entity.client.EntityClient;
 import com.linkedin.metadata.query.SearchFlags;
 import com.linkedin.metadata.query.filter.ConjunctiveCriterion;
 import com.linkedin.metadata.query.filter.ConjunctiveCriterionArray;
@@ -33,24 +38,32 @@
 import com.linkedin.metadata.query.filter.Filter;
 import com.linkedin.metadata.query.filter.SortCriterion;
 import com.linkedin.metadata.query.filter.SortOrder;
+import com.linkedin.metadata.search.SearchResult;
 import com.linkedin.metadata.service.ViewService;
 import com.linkedin.view.DataHubViewInfo;
 import io.datahubproject.metadata.context.OperationContext;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
+import java.util.Optional;
 import java.util.Set;
+import java.util.concurrent.CompletableFuture;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
 import lombok.extern.slf4j.Slf4j;
+import org.apache.commons.lang3.StringUtils;
 import org.codehaus.plexus.util.CollectionUtils;
 
 @Slf4j
 public class SearchUtils {
   private SearchUtils() {}
 
+  private static final int DEFAULT_SEARCH_COUNT = 10;
+  private static final int DEFAULT_SCROLL_COUNT = 10;
+  private static final String DEFAULT_SCROLL_KEEP_ALIVE = "5m";
+
   /** Entities that are searched by default in Search Across Entities */
   public static final List<EntityType> SEARCHABLE_ENTITY_TYPES =
       ImmutableList.of(
@@ -348,4 +361,98 @@ public static List<SortCriterion> getSortCriteria(@Nullable final SearchSortInpu
 
     return sortCriteria;
   }
+
+  public static CompletableFuture<SearchResults> searchAcrossEntities(
+      QueryContext inputContext,
+      final EntityClient _entityClient,
+      final ViewService _viewService,
+      List<EntityType> inputEntityTypes,
+      String inputQuery,
+      Filter baseFilter,
+      String viewUrn,
+      List<SortCriterion> sortCriteria,
+      com.linkedin.datahub.graphql.generated.SearchFlags inputSearchFlags,
+      Integer inputCount,
+      Integer inputStart,
+      String className) {
+
+    final List<EntityType> entityTypes =
+        (inputEntityTypes == null || inputEntityTypes.isEmpty())
+            ? SEARCHABLE_ENTITY_TYPES
+            : inputEntityTypes;
+    final List<String> entityNames =
+        entityTypes.stream().map(EntityTypeMapper::getName).collect(Collectors.toList());
+
+    // escape forward slash since it is a reserved character in Elasticsearch, default to * if
+    // blank/empty
+    final String query =
+        StringUtils.isNotBlank(inputQuery) ? ResolverUtils.escapeForwardSlash(inputQuery) : "*";
+
+    final Optional<SearchFlags> searchFlags =
+        Optional.ofNullable(inputSearchFlags)
+            .map((flags) -> SearchFlagsInputMapper.map(inputContext, flags));
+    final OperationContext context =
+        inputContext.getOperationContext().withSearchFlags(searchFlags::orElse);
+
+    final int count = Optional.ofNullable(inputCount).orElse(DEFAULT_SEARCH_COUNT);
+    final int start = Optional.ofNullable(inputStart).orElse(0);
+
+    return GraphQLConcurrencyUtils.supplyAsync(
+        () -> {
+          final OperationContext baseContext = inputContext.getOperationContext();
+          final Optional<DataHubViewInfo> maybeResolvedView =
+              Optional.ofNullable(viewUrn)
+                  .map((urn) -> resolveView(baseContext, _viewService, UrnUtils.getUrn(urn)));
+
+          final List<String> finalEntityNames =
+              maybeResolvedView
+                  .map(
+                      (view) ->
+                          intersectEntityTypes(entityNames, view.getDefinition().getEntityTypes()))
+                  .orElse(entityNames);
+
+          final Filter finalFilters =
+              maybeResolvedView
+                  .map((view) -> combineFilters(baseFilter, view.getDefinition().getFilter()))
+                  .orElse(baseFilter);
+
+          log.debug(
+              "Executing search for multiple entities: entity types {}, query {}, filters: {}, start: {}, count: {}",
+              finalEntityNames,
+              query,
+              finalFilters,
+              start,
+              count);
+
+          try {
+            final SearchResult searchResult =
+                _entityClient.searchAcrossEntities(
+                    context,
+                    finalEntityNames,
+                    query,
+                    finalFilters,
+                    start,
+                    count,
+                    sortCriteria,
+                    null);
+            return UrnSearchResultsMapper.map(inputContext, searchResult);
+          } catch (Exception e) {
+            log.warn(
+                "Failed to execute search for multiple entities: entity types {}, query {}, filters: {}, start: {}, count: {}",
+                finalEntityNames,
+                query,
+                finalFilters,
+                start,
+                count);
+            throw new RuntimeException(
+                "Failed to execute search: "
+                    + String.format(
+                        "entity types %s, query %s, filters: %s, start: %s, count: %s",
+                        finalEntityNames, query, finalFilters, start, count),
+                e);
+          }
+        },
+        className,
+        "searchAcrossEntities");
+  }
 }
diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/versioning/VersionsSearchResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/versioning/VersionsSearchResolver.java
new file mode 100644
index 00000000000000..915e1cf00ebc6b
--- /dev/null
+++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/versioning/VersionsSearchResolver.java
@@ -0,0 +1,87 @@
+package com.linkedin.datahub.graphql.resolvers.versioning;
+
+import static com.linkedin.datahub.graphql.resolvers.ResolverUtils.bindArgument;
+import static com.linkedin.metadata.Constants.*;
+
+import com.linkedin.datahub.graphql.QueryContext;
+import com.linkedin.datahub.graphql.generated.Entity;
+import com.linkedin.datahub.graphql.generated.SearchAcrossEntitiesInput;
+import com.linkedin.datahub.graphql.generated.SearchFlags;
+import com.linkedin.datahub.graphql.generated.SearchResults;
+import com.linkedin.datahub.graphql.resolvers.ResolverUtils;
+import com.linkedin.datahub.graphql.resolvers.search.SearchUtils;
+import com.linkedin.entity.client.EntityClient;
+import com.linkedin.metadata.query.filter.Condition;
+import com.linkedin.metadata.query.filter.ConjunctiveCriterion;
+import com.linkedin.metadata.query.filter.ConjunctiveCriterionArray;
+import com.linkedin.metadata.query.filter.Criterion;
+import com.linkedin.metadata.query.filter.CriterionArray;
+import com.linkedin.metadata.query.filter.Filter;
+import com.linkedin.metadata.query.filter.SortCriterion;
+import com.linkedin.metadata.query.filter.SortOrder;
+import com.linkedin.metadata.service.ViewService;
+import com.linkedin.metadata.utils.CriterionUtils;
+import graphql.schema.DataFetcher;
+import graphql.schema.DataFetchingEnvironment;
+import java.util.List;
+import java.util.Optional;
+import java.util.concurrent.CompletableFuture;
+import java.util.stream.Stream;
+import lombok.RequiredArgsConstructor;
+import lombok.extern.slf4j.Slf4j;
+
+/** Resolver that executes a searchAcrossEntities only on a version set's versioned entities */
+@Slf4j
+@RequiredArgsConstructor
+public class VersionsSearchResolver implements DataFetcher<CompletableFuture<SearchResults>> {
+
+  private static final String VERSION_SET_FIELD_NAME = "versionSet";
+
+  private final EntityClient _entityClient;
+  private final ViewService _viewService;
+
+  @Override
+  public CompletableFuture<SearchResults> get(DataFetchingEnvironment environment) {
+    final Entity entity = environment.getSource();
+    final QueryContext context = environment.getContext();
+    final SearchAcrossEntitiesInput input =
+        bindArgument(environment.getArgument("input"), SearchAcrossEntitiesInput.class);
+
+    final Criterion versionSetFilter =
+        CriterionUtils.buildCriterion(VERSION_SET_FIELD_NAME, Condition.EQUAL, entity.getUrn());
+    final Filter baseFilter =
+        new Filter()
+            .setOr(
+                new ConjunctiveCriterionArray(
+                    new ConjunctiveCriterion().setAnd(new CriterionArray(versionSetFilter))));
+    final Filter inputFilter = ResolverUtils.buildFilter(null, input.getOrFilters());
+
+    final List<SortCriterion> initialSortCriteria =
+        SearchUtils.getSortCriteria(input.getSortInput());
+    final List<SortCriterion> sortCriteria =
+        Stream.concat(
+                initialSortCriteria.stream(),
+                Stream.of(
+                    new SortCriterion()
+                        .setField(VERSION_SORT_ID_FIELD_NAME)
+                        .setOrder(SortOrder.DESCENDING)))
+            .toList();
+
+    SearchFlags searchFlags = Optional.ofNullable(input.getSearchFlags()).orElse(new SearchFlags());
+    searchFlags.setFilterNonLatestVersions(false);
+
+    return SearchUtils.searchAcrossEntities(
+        context,
+        _entityClient,
+        _viewService,
+        input.getTypes(),
+        input.getQuery(),
+        SearchUtils.combineFilters(inputFilter, baseFilter),
+        input.getViewUrn(),
+        sortCriteria,
+        searchFlags,
+        input.getCount(),
+        input.getStart(),
+        this.getClass().getSimpleName());
+  }
+}
diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/common/mappers/SearchFlagsInputMapper.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/common/mappers/SearchFlagsInputMapper.java
index 9f5025ccf303a2..0b3a445175c4c1 100644
--- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/common/mappers/SearchFlagsInputMapper.java
+++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/common/mappers/SearchFlagsInputMapper.java
@@ -69,6 +69,9 @@ public com.linkedin.metadata.query.SearchFlags apply(
       result.setCustomHighlightingFields(
           new StringArray(searchFlags.getCustomHighlightingFields()));
     }
+    if (searchFlags.getFilterNonLatestVersions() != null) {
+      result.setFilterNonLatestVersions(searchFlags.getFilterNonLatestVersions());
+    }
     return result;
   }
 }
diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/common/mappers/UrnToEntityMapper.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/common/mappers/UrnToEntityMapper.java
index eae33e6da2e56d..b815c1b1c1dd9f 100644
--- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/common/mappers/UrnToEntityMapper.java
+++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/common/mappers/UrnToEntityMapper.java
@@ -41,6 +41,7 @@
 import com.linkedin.datahub.graphql.generated.StructuredPropertyEntity;
 import com.linkedin.datahub.graphql.generated.Tag;
 import com.linkedin.datahub.graphql.generated.Test;
+import com.linkedin.datahub.graphql.generated.VersionSet;
 import com.linkedin.datahub.graphql.types.mappers.ModelMapper;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
@@ -231,6 +232,11 @@ public Entity apply(@Nullable QueryContext context, Urn input) {
       ((DataProcessInstance) partialEntity).setUrn(input.toString());
       ((DataProcessInstance) partialEntity).setType(EntityType.DATA_PROCESS_INSTANCE);
     }
+    if (input.getEntityType().equals(VERSION_SET_ENTITY_NAME)) {
+      partialEntity = new VersionSet();
+      ((VersionSet) partialEntity).setUrn(input.toString());
+      ((VersionSet) partialEntity).setType(EntityType.VERSION_SET);
+    }
     return partialEntity;
   }
 }
diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dataprocessinst/mappers/DataProcessInstanceMapper.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dataprocessinst/mappers/DataProcessInstanceMapper.java
index 28c9c8936fdbfb..d721f5a5fb522d 100644
--- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dataprocessinst/mappers/DataProcessInstanceMapper.java
+++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dataprocessinst/mappers/DataProcessInstanceMapper.java
@@ -8,7 +8,6 @@
 import com.linkedin.data.DataMap;
 import com.linkedin.data.template.RecordTemplate;
 import com.linkedin.datahub.graphql.QueryContext;
-import com.linkedin.datahub.graphql.generated.DataPlatform;
 import com.linkedin.datahub.graphql.generated.DataProcessInstance;
 import com.linkedin.datahub.graphql.generated.EntityType;
 import com.linkedin.datahub.graphql.types.common.mappers.AuditStampMapper;
@@ -80,10 +79,6 @@ public DataProcessInstance apply(
           DataPlatformInstance dataPlatformInstance = new DataPlatformInstance(dataMap);
           dataProcessInstance.setDataPlatformInstance(
               DataPlatformInstanceAspectMapper.map(context, dataPlatformInstance));
-          DataPlatform dataPlatform = new DataPlatform();
-          dataPlatform.setUrn(dataPlatformInstance.getPlatform().toString());
-          dataPlatform.setType(EntityType.DATA_PLATFORM);
-          dataProcessInstance.setPlatform(dataPlatform);
         });
     mappingHelper.mapToResult(
         SUB_TYPES_ASPECT_NAME,
diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dataset/DatasetType.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dataset/DatasetType.java
index 6a3f9cb9b21f38..74ef4cf125cd24 100644
--- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dataset/DatasetType.java
+++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dataset/DatasetType.java
@@ -89,7 +89,8 @@ public class DatasetType
           ACCESS_ASPECT_NAME,
           STRUCTURED_PROPERTIES_ASPECT_NAME,
           FORMS_ASPECT_NAME,
-          SUB_TYPES_ASPECT_NAME);
+          SUB_TYPES_ASPECT_NAME,
+          VERSION_PROPERTIES_ASPECT_NAME);
 
   private static final Set<String> FACET_FIELDS = ImmutableSet.of("origin", "platform");
   private static final String ENTITY_NAME = "dataset";
diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dataset/mappers/DatasetMapper.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dataset/mappers/DatasetMapper.java
index e411014c23c89b..aa7033b180e80e 100644
--- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dataset/mappers/DatasetMapper.java
+++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dataset/mappers/DatasetMapper.java
@@ -17,6 +17,7 @@
 import com.linkedin.common.Status;
 import com.linkedin.common.SubTypes;
 import com.linkedin.common.TimeStamp;
+import com.linkedin.common.VersionProperties;
 import com.linkedin.common.urn.Urn;
 import com.linkedin.data.DataMap;
 import com.linkedin.datahub.graphql.QueryContext;
@@ -48,6 +49,7 @@
 import com.linkedin.datahub.graphql.types.rolemetadata.mappers.AccessMapper;
 import com.linkedin.datahub.graphql.types.structuredproperty.StructuredPropertiesMapper;
 import com.linkedin.datahub.graphql.types.tag.mappers.GlobalTagsMapper;
+import com.linkedin.datahub.graphql.types.versioning.VersionPropertiesMapper;
 import com.linkedin.dataset.DatasetDeprecation;
 import com.linkedin.dataset.DatasetProperties;
 import com.linkedin.dataset.EditableDatasetProperties;
@@ -183,6 +185,11 @@ public Dataset apply(
         SUB_TYPES_ASPECT_NAME,
         (dashboard, dataMap) ->
             dashboard.setSubTypes(SubTypesMapper.map(context, new SubTypes(dataMap))));
+    mappingHelper.mapToResult(
+        VERSION_PROPERTIES_ASPECT_NAME,
+        (entity, dataMap) ->
+            entity.setVersionProperties(
+                VersionPropertiesMapper.map(context, new VersionProperties(dataMap))));
 
     if (context != null && !canView(context.getOperationContext(), entityUrn)) {
       return AuthorizationUtils.restrictEntity(mappingHelper.getResult(), Dataset.class);
diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/entitytype/EntityTypeUrnMapper.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/entitytype/EntityTypeUrnMapper.java
index 5b72c2b3c11c5e..334faf753cb8b5 100644
--- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/entitytype/EntityTypeUrnMapper.java
+++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/entitytype/EntityTypeUrnMapper.java
@@ -77,9 +77,6 @@ public class EntityTypeUrnMapper {
           .put(
               Constants.BUSINESS_ATTRIBUTE_ENTITY_NAME,
               "urn:li:entityType:datahub.businessAttribute")
-          .put(
-              Constants.DATA_PROCESS_INSTANCE_ENTITY_NAME,
-              "urn:li:entityType:datahub.dataProcessInstance")
           .build();
 
   private static final Map<String, String> ENTITY_TYPE_URN_TO_NAME =
diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/mappers/MLModelMapper.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/mappers/MLModelMapper.java
index 7102fd4aed9743..11e6b5180f8c1c 100644
--- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/mappers/MLModelMapper.java
+++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/mappers/MLModelMapper.java
@@ -13,6 +13,7 @@
 import com.linkedin.common.InstitutionalMemory;
 import com.linkedin.common.Ownership;
 import com.linkedin.common.Status;
+import com.linkedin.common.VersionProperties;
 import com.linkedin.common.urn.Urn;
 import com.linkedin.data.DataMap;
 import com.linkedin.data.template.RecordTemplate;
@@ -38,6 +39,7 @@
 import com.linkedin.datahub.graphql.types.mappers.ModelMapper;
 import com.linkedin.datahub.graphql.types.structuredproperty.StructuredPropertiesMapper;
 import com.linkedin.datahub.graphql.types.tag.mappers.GlobalTagsMapper;
+import com.linkedin.datahub.graphql.types.versioning.VersionPropertiesMapper;
 import com.linkedin.domain.Domains;
 import com.linkedin.entity.EntityResponse;
 import com.linkedin.entity.EnvelopedAspectMap;
@@ -180,6 +182,11 @@ public MLModel apply(
         FORMS_ASPECT_NAME,
         ((entity, dataMap) ->
             entity.setForms(FormsMapper.map(new Forms(dataMap), entityUrn.toString()))));
+    mappingHelper.mapToResult(
+        VERSION_PROPERTIES_ASPECT_NAME,
+        (entity, dataMap) ->
+            entity.setVersionProperties(
+                VersionPropertiesMapper.map(context, new VersionProperties(dataMap))));
 
     if (context != null && !canView(context.getOperationContext(), entityUrn)) {
       return AuthorizationUtils.restrictEntity(mappingHelper.getResult(), MLModel.class);
diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/versioning/VersionPropertiesMapper.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/versioning/VersionPropertiesMapper.java
new file mode 100644
index 00000000000000..f89ebdc9f2b043
--- /dev/null
+++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/versioning/VersionPropertiesMapper.java
@@ -0,0 +1,53 @@
+package com.linkedin.datahub.graphql.types.versioning;
+
+import com.linkedin.datahub.graphql.QueryContext;
+import com.linkedin.datahub.graphql.generated.EntityType;
+import com.linkedin.datahub.graphql.generated.VersionProperties;
+import com.linkedin.datahub.graphql.generated.VersionSet;
+import com.linkedin.datahub.graphql.types.mappers.MapperUtils;
+import com.linkedin.datahub.graphql.types.mappers.ModelMapper;
+import com.linkedin.datahub.graphql.types.mlmodel.mappers.VersionTagMapper;
+import java.util.stream.Collectors;
+import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
+
+public class VersionPropertiesMapper
+    implements ModelMapper<com.linkedin.common.VersionProperties, VersionProperties> {
+  public static final VersionPropertiesMapper INSTANCE = new VersionPropertiesMapper();
+
+  public static VersionProperties map(
+      @Nullable QueryContext context,
+      @Nonnull final com.linkedin.common.VersionProperties versionProperties) {
+    return INSTANCE.apply(context, versionProperties);
+  }
+
+  @Override
+  public VersionProperties apply(
+      @Nullable QueryContext context, @Nonnull com.linkedin.common.VersionProperties input) {
+    final VersionProperties result = new VersionProperties();
+
+    result.setVersionSet(
+        VersionSet.builder()
+            .setUrn(input.getVersionSet().toString())
+            .setType(EntityType.VERSION_SET)
+            .build());
+
+    result.setVersion(VersionTagMapper.map(context, input.getVersion()));
+    result.setAliases(
+        input.getAliases().stream()
+            .map(alias -> VersionTagMapper.map(context, alias))
+            .collect(Collectors.toList()));
+    result.setComment(input.getComment());
+    result.setIsLatest(Boolean.TRUE.equals(input.isIsLatest()));
+
+    if (input.getMetadataCreatedTimestamp() != null) {
+      result.setCreated(MapperUtils.createResolvedAuditStamp(input.getMetadataCreatedTimestamp()));
+    }
+    if (input.getSourceCreatedTimestamp() != null) {
+      result.setCreatedInSource(
+          MapperUtils.createResolvedAuditStamp(input.getSourceCreatedTimestamp()));
+    }
+
+    return result;
+  }
+}
diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/versioning/VersionSetMapper.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/versioning/VersionSetMapper.java
new file mode 100644
index 00000000000000..3a07115ece5f6e
--- /dev/null
+++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/versioning/VersionSetMapper.java
@@ -0,0 +1,47 @@
+package com.linkedin.datahub.graphql.types.versioning;
+
+import static com.linkedin.metadata.Constants.VERSION_SET_PROPERTIES_ASPECT_NAME;
+
+import com.linkedin.data.DataMap;
+import com.linkedin.datahub.graphql.QueryContext;
+import com.linkedin.datahub.graphql.generated.EntityType;
+import com.linkedin.datahub.graphql.generated.VersionSet;
+import com.linkedin.datahub.graphql.types.common.mappers.UrnToEntityMapper;
+import com.linkedin.datahub.graphql.types.common.mappers.util.MappingHelper;
+import com.linkedin.datahub.graphql.types.mappers.ModelMapper;
+import com.linkedin.entity.EntityResponse;
+import com.linkedin.entity.EnvelopedAspectMap;
+import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
+
+public class VersionSetMapper implements ModelMapper<EntityResponse, VersionSet> {
+
+  public static final VersionSetMapper INSTANCE = new VersionSetMapper();
+
+  public static VersionSet map(
+      @Nullable QueryContext context, @Nonnull final EntityResponse entityResponse) {
+    return INSTANCE.apply(context, entityResponse);
+  }
+
+  @Override
+  public VersionSet apply(@Nullable QueryContext context, @Nonnull EntityResponse entityResponse) {
+    final VersionSet result = new VersionSet();
+    result.setUrn(entityResponse.getUrn().toString());
+    result.setType(EntityType.VERSION_SET);
+
+    EnvelopedAspectMap aspectMap = entityResponse.getAspects();
+    MappingHelper<VersionSet> mappingHelper = new MappingHelper<>(aspectMap, result);
+    mappingHelper.mapToResult(
+        VERSION_SET_PROPERTIES_ASPECT_NAME,
+        (versionSet, dataMap) -> mapVersionSetProperties(context, versionSet, dataMap));
+
+    return result;
+  }
+
+  private void mapVersionSetProperties(
+      @Nullable QueryContext context, @Nonnull VersionSet versionSet, @Nonnull DataMap dataMap) {
+    com.linkedin.versionset.VersionSetProperties versionProperties =
+        new com.linkedin.versionset.VersionSetProperties(dataMap);
+    versionSet.setLatestVersion(UrnToEntityMapper.map(context, versionProperties.getLatest()));
+  }
+}
diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/versioning/VersionSetType.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/versioning/VersionSetType.java
new file mode 100644
index 00000000000000..ed2beff4530949
--- /dev/null
+++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/versioning/VersionSetType.java
@@ -0,0 +1,79 @@
+package com.linkedin.datahub.graphql.types.versioning;
+
+import static com.linkedin.metadata.Constants.*;
+
+import com.google.common.collect.ImmutableSet;
+import com.linkedin.common.urn.Urn;
+import com.linkedin.common.urn.UrnUtils;
+import com.linkedin.datahub.graphql.QueryContext;
+import com.linkedin.datahub.graphql.generated.Entity;
+import com.linkedin.datahub.graphql.generated.EntityType;
+import com.linkedin.datahub.graphql.generated.VersionSet;
+import com.linkedin.entity.EntityResponse;
+import com.linkedin.entity.client.EntityClient;
+import graphql.execution.DataFetcherResult;
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+import javax.annotation.Nonnull;
+import lombok.RequiredArgsConstructor;
+
+@RequiredArgsConstructor
+public class VersionSetType
+    implements com.linkedin.datahub.graphql.types.EntityType<VersionSet, String> {
+
+  public static final Set<String> ASPECTS_TO_FETCH =
+      ImmutableSet.of(VERSION_SET_PROPERTIES_ASPECT_NAME);
+  private final EntityClient _entityClient;
+
+  @Override
+  public EntityType type() {
+    return EntityType.VERSION_SET;
+  }
+
+  @Override
+  public Function<Entity, String> getKeyProvider() {
+    return Entity::getUrn;
+  }
+
+  @Override
+  public Class<VersionSet> objectClass() {
+    return VersionSet.class;
+  }
+
+  @Override
+  public List<DataFetcherResult<VersionSet>> batchLoad(
+      @Nonnull List<String> urns, @Nonnull QueryContext context) throws Exception {
+    final List<Urn> versionSetUrns =
+        urns.stream().map(UrnUtils::getUrn).collect(Collectors.toList());
+
+    try {
+      final Map<Urn, EntityResponse> entities =
+          _entityClient.batchGetV2(
+              context.getOperationContext(),
+              VERSION_SET_ENTITY_NAME,
+              new HashSet<>(versionSetUrns),
+              ASPECTS_TO_FETCH);
+
+      final List<EntityResponse> gmsResults = new ArrayList<>();
+      for (Urn urn : versionSetUrns) {
+        gmsResults.add(entities.getOrDefault(urn, null));
+      }
+      return gmsResults.stream()
+          .map(
+              gmsResult ->
+                  gmsResult == null
+                      ? null
+                      : DataFetcherResult.<VersionSet>newResult()
+                          .data(VersionSetMapper.map(context, gmsResult))
+                          .build())
+          .collect(Collectors.toList());
+    } catch (Exception e) {
+      throw new RuntimeException("Failed to batch load Queries", e);
+    }
+  }
+}
diff --git a/datahub-graphql-core/src/main/resources/app.graphql b/datahub-graphql-core/src/main/resources/app.graphql
index 28688903687235..ca7f89415f6b87 100644
--- a/datahub-graphql-core/src/main/resources/app.graphql
+++ b/datahub-graphql-core/src/main/resources/app.graphql
@@ -531,6 +531,11 @@ type FeatureFlagsConfig {
   If turned on, show the manage structured properties tab in the govern dropdown
   """
   showManageStructuredProperties: Boolean!
+
+  """
+  If turned on, exposes the versioning feature by allowing users to link entities in the UI.
+  """
+  entityVersioningEnabled: Boolean!
 }
 
 """
@@ -573,4 +578,4 @@ type DocPropagationSettings {
   The default doc propagation setting for the platform.
   """
   docColumnPropagation: Boolean
-}
\ No newline at end of file
+}
diff --git a/datahub-graphql-core/src/main/resources/entity.graphql b/datahub-graphql-core/src/main/resources/entity.graphql
index b47be7ae32b2c4..64c79b05745ded 100644
--- a/datahub-graphql-core/src/main/resources/entity.graphql
+++ b/datahub-graphql-core/src/main/resources/entity.graphql
@@ -956,16 +956,6 @@ type Mutation {
     Remove Business Attribute
     """
     removeBusinessAttribute(input: AddBusinessAttributeInput!): Boolean
-
-    """
-    Link the latest versioned entity to a Version Set
-    """
-    linkAssetVersion(input: LinkVersionInput!): String
-
-    """
-    Unlink a versioned entity from a Version Set
-    """
-    unlinkAssetVersion(input: UnlinkVersionInput!): Boolean
 }
 
 """
@@ -1231,6 +1221,11 @@ enum EntityType {
     A Business Attribute
     """
     BUSINESS_ATTRIBUTE
+
+    """
+    A set of versioned entities, representing a single source / logical entity over time
+    """
+    VERSION_SET
 }
 
 """
@@ -12921,56 +12916,6 @@ input ListBusinessAttributesInput {
     query: String
 }
 
-"""
-Input for linking a versioned entity to a Version Set
-"""
-input LinkVersionInput {
-    """
-    The target version set
-    """
-    versionSet: String!
-
-    """
-    The target versioned entity to link
-    """
-    linkedEntity: String!
-
-    """
-    Version Tag label for the version, should be unique within a Version Set
-    """
-    version: String!
-
-    """
-    Optional timestamp from the source system
-    """
-    sourceTimestamp: Long
-
-    """
-    Optional creator from the source system, will be converted to an Urn
-    """
-    sourceCreator: String
-
-    """
-    Optional comment about the version
-    """
-    comment: String
-}
-
-"""
-Input for unlinking a versioned entity from a Version Set
-"""
-input UnlinkVersionInput {
-    """
-    The target version set
-    """
-    versionSet: String
-
-    """
-    The target versioned entity to unlink
-    """
-    unlinkedEntity: String
-}
-
 """
 The result obtained when listing Business Attribute
 """
@@ -13084,11 +13029,6 @@ extend type DataProcessInstance {
     """
     container: Container
 
-    """
-    Standardized platform urn where the data process instance is defined
-    """
-    platform: DataPlatform!
-
     """
     Recursively get the lineage of containers for this entity
     """
diff --git a/datahub-graphql-core/src/main/resources/search.graphql b/datahub-graphql-core/src/main/resources/search.graphql
index 82bfb9ee26fc42..d8f17faa3d11c2 100644
--- a/datahub-graphql-core/src/main/resources/search.graphql
+++ b/datahub-graphql-core/src/main/resources/search.graphql
@@ -172,6 +172,11 @@ input  SearchFlags {
   Whether or not to fetch and request for structured property facets when doing a search
   """
   includeStructuredPropertyFacets: Boolean
+
+  """
+  Determines whether to filter out any non-latest entity version if entity is part of a Version Set, default true
+  """
+  filterNonLatestVersions: Boolean
 }
 
 """
@@ -1497,4 +1502,4 @@ input GroupingCriterion {
   """
   groupingEntityType: EntityType!
 
-}
\ No newline at end of file
+}
diff --git a/datahub-graphql-core/src/main/resources/versioning.graphql b/datahub-graphql-core/src/main/resources/versioning.graphql
new file mode 100644
index 00000000000000..4a63463509c84d
--- /dev/null
+++ b/datahub-graphql-core/src/main/resources/versioning.graphql
@@ -0,0 +1,148 @@
+type VersionSet implements Entity {
+    """
+    The primary key of the VersionSet
+    """
+    urn: String!
+
+    """
+    The standard Entity Type
+    """
+    type: EntityType!
+
+    """
+    Granular API for querying edges extending from this entity
+    """
+    relationships(input: RelationshipsInput!): EntityRelationshipsResult
+
+    """
+    The latest versioned entity linked to in this version set
+    """
+    latestVersion: Entity
+
+    """
+    Executes a search on all versioned entities linked to this version set
+    By default sorts by sortId in descending order
+    """
+    versionsSearch(input: SearchAcrossEntitiesInput!): SearchResults
+}
+
+type VersionProperties {
+    """
+    The linked Version Set entity that ties multiple versioned assets together
+    """
+    versionSet: VersionSet!
+
+    """
+    Label for this versioned asset, should be unique within a version set (not enforced)
+    """
+    version: VersionTag!
+
+    """
+    Additional version identifiers for this versioned asset.
+    """
+    aliases: [VersionTag!]!
+
+    """
+    Comment documenting what this version was created for, changes, or represents
+    """
+    comment: String
+
+    """
+    Whether this version is currently the latest in its verison set
+    """
+    isLatest: Boolean!
+
+    """
+    Timestamp reflecting when the metadata for this version was created in DataHub
+    """
+    created: ResolvedAuditStamp
+
+    """
+    Timestamp reflecting when the metadata for this version was created in DataHub
+    """
+    createdInSource: ResolvedAuditStamp
+}
+
+interface SupportsVersions {
+    """
+    Indicates that this entity is versioned and provides information about the version.
+    """
+    versionProperties: VersionProperties
+}
+
+extend type Dataset implements SupportsVersions {
+    versionProperties: VersionProperties
+}
+
+extend type MLModel implements SupportsVersions {
+    versionProperties: VersionProperties
+}
+
+extend type Query {
+    """
+    Fetch a Version Set by its URN
+    """
+    versionSet(urn: String!): VersionSet
+}
+
+"""
+Input for linking a versioned entity to a Version Set
+"""
+input LinkVersionInput {
+    """
+    The target version set
+    """
+    versionSet: String!
+
+    """
+    The target versioned entity to link
+    """
+    linkedEntity: String!
+
+    """
+    Version Tag label for the version, should be unique within a version set (not enforced)
+    """
+    version: String!
+
+    """
+    Optional timestamp from the source system
+    """
+    sourceTimestamp: Long
+
+    """
+    Optional creator from the source system, will be converted to an Urn
+    """
+    sourceCreator: String
+
+    """
+    Optional comment about the version
+    """
+    comment: String
+}
+
+"""
+Input for unlinking a versioned entity from a Version Set
+"""
+input UnlinkVersionInput {
+    """
+    The target version set
+    """
+    versionSet: String
+
+    """
+    The target versioned entity to unlink
+    """
+    unlinkedEntity: String
+}
+
+extend type Mutation {
+    """
+    Link the latest versioned entity to a Version Set
+    """
+    linkAssetVersion(input: LinkVersionInput!): VersionSet
+
+    """
+    Unlink a versioned entity from a Version Set
+    """
+    unlinkAssetVersion(input: UnlinkVersionInput!): VersionSet
+}
diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/entity/versioning/LinkAssetVersionResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/entity/versioning/LinkAssetVersionResolverTest.java
index 690856263fccc5..c2eb92f4d1cd4c 100644
--- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/entity/versioning/LinkAssetVersionResolverTest.java
+++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/entity/versioning/LinkAssetVersionResolverTest.java
@@ -56,8 +56,7 @@ public void testGetSuccessful() throws Exception {
     Mockito.when(mockEnv.getArgument(Mockito.eq("input"))).thenReturn(input);
     Mockito.when(mockEnv.getContext()).thenReturn(mockContext);
 
-    String result = resolver.get(mockEnv).get();
-    assertEquals(result, TEST_ENTITY_URN);
+    assertEquals(resolver.get(mockEnv).get().getUrn(), TEST_VERSION_SET_URN);
   }
 
   @Test
diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/entity/versioning/UnlinkAssetVersionResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/entity/versioning/UnlinkAssetVersionResolverTest.java
index 0000ad24a04537..e162ce96e627c6 100644
--- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/entity/versioning/UnlinkAssetVersionResolverTest.java
+++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/entity/versioning/UnlinkAssetVersionResolverTest.java
@@ -48,7 +48,7 @@ public void testGetSuccessful() throws Exception {
     Mockito.when(mockEnv.getArgument(Mockito.eq("input"))).thenReturn(input);
     Mockito.when(mockEnv.getContext()).thenReturn(mockContext);
 
-    assertTrue(resolver.get(mockEnv).get());
+    assertEquals(resolver.get(mockEnv).get(), null);
 
     Mockito.verify(mockService)
         .unlinkVersion(
diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/versioning/VersionsSearchResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/versioning/VersionsSearchResolverTest.java
new file mode 100644
index 00000000000000..3554df074df698
--- /dev/null
+++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/versioning/VersionsSearchResolverTest.java
@@ -0,0 +1,294 @@
+package com.linkedin.datahub.graphql.resolvers.versioning;
+
+import static com.linkedin.datahub.graphql.TestUtils.getMockAllowContext;
+import static com.linkedin.metadata.Constants.*;
+import static com.linkedin.metadata.utils.CriterionUtils.*;
+import static org.mockito.ArgumentMatchers.*;
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertThrows;
+
+import com.linkedin.common.AuditStamp;
+import com.linkedin.common.urn.Urn;
+import com.linkedin.common.urn.UrnUtils;
+import com.linkedin.data.template.StringArray;
+import com.linkedin.datahub.graphql.QueryContext;
+import com.linkedin.datahub.graphql.generated.AndFilterInput;
+import com.linkedin.datahub.graphql.generated.EntityType;
+import com.linkedin.datahub.graphql.generated.FacetFilterInput;
+import com.linkedin.datahub.graphql.generated.SearchAcrossEntitiesInput;
+import com.linkedin.datahub.graphql.generated.SearchFlags;
+import com.linkedin.datahub.graphql.generated.SearchResults;
+import com.linkedin.datahub.graphql.generated.SearchSortInput;
+import com.linkedin.datahub.graphql.generated.SortCriterion;
+import com.linkedin.datahub.graphql.generated.SortOrder;
+import com.linkedin.datahub.graphql.generated.VersionSet;
+import com.linkedin.entity.client.EntityClient;
+import com.linkedin.metadata.Constants;
+import com.linkedin.metadata.query.filter.Condition;
+import com.linkedin.metadata.query.filter.ConjunctiveCriterion;
+import com.linkedin.metadata.query.filter.ConjunctiveCriterionArray;
+import com.linkedin.metadata.query.filter.CriterionArray;
+import com.linkedin.metadata.query.filter.Filter;
+import com.linkedin.metadata.search.SearchEntityArray;
+import com.linkedin.metadata.search.SearchResult;
+import com.linkedin.metadata.search.SearchResultMetadata;
+import com.linkedin.metadata.service.ViewService;
+import com.linkedin.metadata.utils.CriterionUtils;
+import com.linkedin.r2.RemoteInvocationException;
+import com.linkedin.view.DataHubViewDefinition;
+import com.linkedin.view.DataHubViewInfo;
+import com.linkedin.view.DataHubViewType;
+import graphql.schema.DataFetchingEnvironment;
+import java.util.List;
+import java.util.concurrent.CompletionException;
+import org.mockito.Mockito;
+import org.testng.annotations.Test;
+
+public class VersionsSearchResolverTest {
+
+  private static final String VERSION_SET_URN = "urn:li:versionSet:(my_version_set,dataset)";
+  private static final Urn TEST_VIEW_URN = UrnUtils.getUrn("urn:li:dataHubView:test");
+  private static final Urn TEST_USER_URN = UrnUtils.getUrn("urn:li:corpuser:test");
+
+  private static final SearchAcrossEntitiesInput BASIC_INPUT =
+      new SearchAcrossEntitiesInput(
+          List.of(EntityType.DATASET), "", 0, 10, null, null, null, null, null);
+
+  private static final SearchAcrossEntitiesInput COMPLEX_INPUT =
+      new SearchAcrossEntitiesInput(
+          List.of(EntityType.CHART, EntityType.DATASET),
+          "query",
+          2,
+          5,
+          null,
+          List.of(
+              AndFilterInput.builder()
+                  .setAnd(
+                      List.of(
+                          FacetFilterInput.builder()
+                              .setField("field1")
+                              .setValues(List.of("1", "2"))
+                              .build(),
+                          FacetFilterInput.builder()
+                              .setField("field2")
+                              .setValues(List.of("a"))
+                              .build()))
+                  .build(),
+              AndFilterInput.builder()
+                  .setAnd(
+                      List.of(
+                          FacetFilterInput.builder()
+                              .setField("field3")
+                              .setValues(List.of("3", "4"))
+                              .build(),
+                          FacetFilterInput.builder()
+                              .setField("field4")
+                              .setValues(List.of("b"))
+                              .build()))
+                  .build()),
+          TEST_VIEW_URN.toString(),
+          SearchFlags.builder().setSkipCache(true).build(),
+          SearchSortInput.builder()
+              .setSortCriteria(
+                  List.of(
+                      SortCriterion.builder()
+                          .setField("sortField1")
+                          .setSortOrder(SortOrder.DESCENDING)
+                          .build(),
+                      SortCriterion.builder()
+                          .setField("sortField2")
+                          .setSortOrder(SortOrder.ASCENDING)
+                          .build()))
+              .build());
+
+  @Test
+  public void testGetSuccessBasic() throws Exception {
+    EntityClient mockEntityClient = initMockEntityClient();
+    ViewService mockViewService = Mockito.mock(ViewService.class);
+    VersionsSearchResolver resolver = new VersionsSearchResolver(mockEntityClient, mockViewService);
+
+    // Execute resolver
+    QueryContext mockContext = getMockAllowContext();
+    DataFetchingEnvironment mockEnv = Mockito.mock(DataFetchingEnvironment.class);
+    Mockito.when(mockEnv.getArgument(Mockito.eq("input"))).thenReturn(BASIC_INPUT);
+    Mockito.when(mockEnv.getContext()).thenReturn(mockContext);
+    VersionSet versionSet = new VersionSet();
+    versionSet.setUrn(VERSION_SET_URN);
+    Mockito.when(mockEnv.getSource()).thenReturn(versionSet);
+
+    SearchResults result = resolver.get(mockEnv).get();
+
+    // Validate the result
+    assertEquals(result.getSearchResults().size(), 0);
+
+    // Validate that we called the search service correctly
+    Mockito.verify(mockEntityClient, Mockito.times(1))
+        .searchAcrossEntities(
+            Mockito.argThat(
+                context ->
+                    !context.getSearchContext().getSearchFlags().isFilterNonLatestVersions()),
+            Mockito.eq(List.of(Constants.DATASET_ENTITY_NAME)),
+            Mockito.eq("*"),
+            Mockito.eq(
+                new Filter()
+                    .setOr(
+                        new ConjunctiveCriterionArray(
+                            new ConjunctiveCriterion()
+                                .setAnd(
+                                    new CriterionArray(
+                                        CriterionUtils.buildCriterion(
+                                            "versionSet", Condition.EQUAL, VERSION_SET_URN)))))),
+            Mockito.eq(0),
+            Mockito.eq(10),
+            Mockito.eq(
+                List.of(
+                    new com.linkedin.metadata.query.filter.SortCriterion()
+                        .setField(VERSION_SORT_ID_FIELD_NAME)
+                        .setOrder(com.linkedin.metadata.query.filter.SortOrder.DESCENDING))),
+            any());
+  }
+
+  @Test
+  public void testGetSuccessComplex() throws Exception {
+    EntityClient mockEntityClient = initMockEntityClient();
+
+    Filter viewFilter =
+        new Filter()
+            .setOr(
+                new ConjunctiveCriterionArray(
+                    new ConjunctiveCriterion()
+                        .setAnd(
+                            new CriterionArray(
+                                List.of(buildCriterion("viewField", Condition.EQUAL, "test"))))));
+    DataHubViewInfo viewInfo =
+        new DataHubViewInfo()
+            .setName("test")
+            .setType(DataHubViewType.GLOBAL)
+            .setCreated(new AuditStamp().setTime(0L).setActor(TEST_USER_URN))
+            .setLastModified(new AuditStamp().setTime(0L).setActor(TEST_USER_URN))
+            .setDefinition(
+                new DataHubViewDefinition()
+                    .setEntityTypes(
+                        new StringArray(
+                            List.of(
+                                Constants.DATASET_ENTITY_NAME, Constants.DASHBOARD_ENTITY_NAME)))
+                    .setFilter(viewFilter));
+    ViewService mockViewService = Mockito.mock(ViewService.class);
+    Mockito.when(mockViewService.getViewInfo(any(), Mockito.eq(TEST_VIEW_URN)))
+        .thenReturn(viewInfo);
+
+    VersionsSearchResolver resolver = new VersionsSearchResolver(mockEntityClient, mockViewService);
+
+    // Execute resolver
+    QueryContext mockContext = getMockAllowContext();
+    DataFetchingEnvironment mockEnv = Mockito.mock(DataFetchingEnvironment.class);
+    Mockito.when(mockEnv.getArgument(Mockito.eq("input"))).thenReturn(COMPLEX_INPUT);
+    Mockito.when(mockEnv.getContext()).thenReturn(mockContext);
+    VersionSet versionSet = new VersionSet();
+    versionSet.setUrn(VERSION_SET_URN);
+    Mockito.when(mockEnv.getSource()).thenReturn(versionSet);
+
+    SearchResults result = resolver.get(mockEnv).get();
+
+    // Validate the result
+    assertEquals(result.getSearchResults().size(), 0);
+
+    // Validate that we called the search service correctly
+    Mockito.verify(mockEntityClient, Mockito.times(1))
+        .searchAcrossEntities(
+            Mockito.argThat(
+                context ->
+                    !context.getSearchContext().getSearchFlags().isFilterNonLatestVersions()
+                        && context.getSearchContext().getSearchFlags().isSkipCache()),
+            Mockito.eq(List.of(Constants.DATASET_ENTITY_NAME)),
+            Mockito.eq("query"),
+            Mockito.eq(
+                new Filter()
+                    .setOr(
+                        new ConjunctiveCriterionArray(
+                            new ConjunctiveCriterion()
+                                .setAnd(
+                                    new CriterionArray(
+                                        CriterionUtils.buildCriterion(
+                                            "field1", Condition.EQUAL, "1", "2"),
+                                        CriterionUtils.buildCriterion(
+                                            "field2", Condition.EQUAL, "a"),
+                                        CriterionUtils.buildCriterion(
+                                            "versionSet", Condition.EQUAL, VERSION_SET_URN),
+                                        CriterionUtils.buildCriterion(
+                                            "viewField", Condition.EQUAL, "test"))),
+                            new ConjunctiveCriterion()
+                                .setAnd(
+                                    new CriterionArray(
+                                        CriterionUtils.buildCriterion(
+                                            "field3", Condition.EQUAL, "3", "4"),
+                                        CriterionUtils.buildCriterion(
+                                            "field4", Condition.EQUAL, "b"),
+                                        CriterionUtils.buildCriterion(
+                                            "versionSet", Condition.EQUAL, VERSION_SET_URN),
+                                        CriterionUtils.buildCriterion(
+                                            "viewField", Condition.EQUAL, "test")))))),
+            Mockito.eq(2),
+            Mockito.eq(5),
+            Mockito.eq(
+                List.of(
+                    new com.linkedin.metadata.query.filter.SortCriterion()
+                        .setField("sortField1")
+                        .setOrder(com.linkedin.metadata.query.filter.SortOrder.DESCENDING),
+                    new com.linkedin.metadata.query.filter.SortCriterion()
+                        .setField("sortField2")
+                        .setOrder(com.linkedin.metadata.query.filter.SortOrder.ASCENDING),
+                    new com.linkedin.metadata.query.filter.SortCriterion()
+                        .setField(VERSION_SORT_ID_FIELD_NAME)
+                        .setOrder(com.linkedin.metadata.query.filter.SortOrder.DESCENDING))),
+            any());
+  }
+
+  @Test
+  public void testThrowsError() throws Exception {
+    EntityClient mockEntityClient = initMockEntityClient();
+    ViewService mockViewService = Mockito.mock(ViewService.class);
+
+    Mockito.when(
+            mockEntityClient.searchAcrossEntities(
+                any(), any(), any(), any(), Mockito.anyInt(), Mockito.anyInt(), any(), any()))
+        .thenThrow(new RemoteInvocationException());
+
+    VersionsSearchResolver resolver = new VersionsSearchResolver(mockEntityClient, mockViewService);
+
+    // Execute resolver
+    QueryContext mockContext = getMockAllowContext();
+    DataFetchingEnvironment mockEnv = Mockito.mock(DataFetchingEnvironment.class);
+    Mockito.when(mockEnv.getArgument(Mockito.eq("input"))).thenReturn(BASIC_INPUT);
+    Mockito.when(mockEnv.getContext()).thenReturn(mockContext);
+    VersionSet versionSet = new VersionSet();
+    versionSet.setUrn(VERSION_SET_URN);
+    Mockito.when(mockEnv.getSource()).thenReturn(versionSet);
+
+    assertThrows(CompletionException.class, () -> resolver.get(mockEnv).join());
+  }
+
+  private EntityClient initMockEntityClient() throws Exception {
+    EntityClient client = Mockito.mock(EntityClient.class);
+
+    Mockito.when(
+            client.searchAcrossEntities(
+                any(),
+                any(),
+                Mockito.anyString(),
+                any(),
+                Mockito.anyInt(),
+                Mockito.anyInt(),
+                any(),
+                Mockito.eq(null)))
+        .thenReturn(
+            new SearchResult()
+                .setEntities(new SearchEntityArray())
+                .setNumEntities(0)
+                .setFrom(0)
+                .setPageSize(0)
+                .setMetadata(new SearchResultMetadata()));
+
+    return client;
+  }
+}
diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/types/dataprocessinst/mappers/DataProcessInstanceMapperTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/types/dataprocessinst/mappers/DataProcessInstanceMapperTest.java
index dc1ce935ad5ecd..cd9d58b54e6b3a 100644
--- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/types/dataprocessinst/mappers/DataProcessInstanceMapperTest.java
+++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/types/dataprocessinst/mappers/DataProcessInstanceMapperTest.java
@@ -80,9 +80,10 @@ public void testMapPlatformInstance() throws Exception {
     DataProcessInstance instance = DataProcessInstanceMapper.map(null, entityResponse);
 
     assertNotNull(instance.getDataPlatformInstance());
-    assertNotNull(instance.getPlatform());
-    assertEquals(instance.getPlatform().getUrn(), TEST_PLATFORM_URN);
-    assertEquals(instance.getPlatform().getType(), EntityType.DATA_PLATFORM);
+    assertNotNull(instance.getDataPlatformInstance().getPlatform());
+    assertEquals(instance.getDataPlatformInstance().getPlatform().getUrn(), TEST_PLATFORM_URN);
+    assertEquals(
+        instance.getDataPlatformInstance().getPlatform().getType(), EntityType.DATA_PLATFORM);
   }
 
   @Test
diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/types/entitytype/EntityTypeMapperTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/types/entitytype/EntityTypeMapperTest.java
new file mode 100644
index 00000000000000..79cc7725b1fc7f
--- /dev/null
+++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/types/entitytype/EntityTypeMapperTest.java
@@ -0,0 +1,20 @@
+package com.linkedin.datahub.graphql.types.entitytype;
+
+import static org.testng.Assert.*;
+
+import com.linkedin.datahub.graphql.generated.EntityType;
+import com.linkedin.metadata.Constants;
+import org.testng.annotations.Test;
+
+public class EntityTypeMapperTest {
+
+  @Test
+  public void testGetType() throws Exception {
+    assertEquals(EntityTypeMapper.getType(Constants.DATASET_ENTITY_NAME), EntityType.DATASET);
+  }
+
+  @Test
+  public void testGetName() throws Exception {
+    assertEquals(EntityTypeMapper.getName(EntityType.DATASET), Constants.DATASET_ENTITY_NAME);
+  }
+}
diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/types/entitytype/EntityTypeUrnMapperTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/types/entitytype/EntityTypeUrnMapperTest.java
new file mode 100644
index 00000000000000..ed16226d0685ee
--- /dev/null
+++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/types/entitytype/EntityTypeUrnMapperTest.java
@@ -0,0 +1,30 @@
+package com.linkedin.datahub.graphql.types.entitytype;
+
+import static org.testng.Assert.*;
+
+import com.linkedin.datahub.graphql.generated.EntityType;
+import com.linkedin.metadata.Constants;
+import org.testng.annotations.Test;
+
+public class EntityTypeUrnMapperTest {
+
+  @Test
+  public void testGetName() throws Exception {
+    assertEquals(
+        EntityTypeUrnMapper.getName("urn:li:entityType:datahub.dataset"),
+        Constants.DATASET_ENTITY_NAME);
+  }
+
+  @Test
+  public void testGetEntityType() throws Exception {
+    assertEquals(
+        EntityTypeUrnMapper.getEntityType("urn:li:entityType:datahub.dataset"), EntityType.DATASET);
+  }
+
+  @Test
+  public void testGetEntityTypeUrn() throws Exception {
+    assertEquals(
+        EntityTypeUrnMapper.getEntityTypeUrn(Constants.DATASET_ENTITY_NAME),
+        "urn:li:entityType:datahub.dataset");
+  }
+}
diff --git a/datahub-web-react/package.json b/datahub-web-react/package.json
index a608698d7602c4..0ff68de2481ed0 100644
--- a/datahub-web-react/package.json
+++ b/datahub-web-react/package.json
@@ -152,7 +152,7 @@
         "prettier": "^2.8.8",
         "source-map-explorer": "^2.5.2",
         "storybook": "^8.1.11",
-        "vite": "^4.5.5",
+        "vite": "^4.5.6",
         "vite-plugin-babel-macros": "^1.0.6",
         "vite-plugin-static-copy": "^0.17.0",
         "vite-plugin-svgr": "^4.1.0",
diff --git a/datahub-web-react/src/Mocks.tsx b/datahub-web-react/src/Mocks.tsx
index 2da9e733eb4072..063b784920e234 100644
--- a/datahub-web-react/src/Mocks.tsx
+++ b/datahub-web-react/src/Mocks.tsx
@@ -645,6 +645,7 @@ export const dataset3 = {
     structuredProperties: null,
     forms: null,
     activeIncidents: null,
+    versionProperties: null,
 } as Dataset;
 
 export const dataset3WithSchema = {
diff --git a/datahub-web-react/src/app/context/CustomUserContext.tsx b/datahub-web-react/src/app/context/CustomUserContext.tsx
new file mode 100644
index 00000000000000..016bbe29684ea5
--- /dev/null
+++ b/datahub-web-react/src/app/context/CustomUserContext.tsx
@@ -0,0 +1,7 @@
+/**
+ * Custom User Context State - This is a custom user context state and can be overriden in specific fork of DataHub.
+ * The below type can be customized with specific object properties as well if needed.
+ */
+export type CustomUserContextState = Record<string, any>;
+
+export const DEFAULT_CUSTOM_STATE: CustomUserContextState = {};
diff --git a/datahub-web-react/src/app/context/userContext.tsx b/datahub-web-react/src/app/context/userContext.tsx
index c9b8adafd9722f..a728e01ddc29ae 100644
--- a/datahub-web-react/src/app/context/userContext.tsx
+++ b/datahub-web-react/src/app/context/userContext.tsx
@@ -1,5 +1,6 @@
 import React from 'react';
 import { CorpUser, PlatformPrivileges } from '../../types.generated';
+import { CustomUserContextState, DEFAULT_CUSTOM_STATE } from './CustomUserContext';
 
 /**
  * Local State is persisted to local storage.
@@ -22,6 +23,7 @@ export type State = {
         loadedPersonalDefaultViewUrn: boolean;
         hasSetDefaultView: boolean;
     };
+    customState?: CustomUserContextState;
 };
 
 /**
@@ -51,6 +53,7 @@ export const DEFAULT_STATE: State = {
         loadedPersonalDefaultViewUrn: false,
         hasSetDefaultView: false,
     },
+    customState: DEFAULT_CUSTOM_STATE,
 };
 
 export const DEFAULT_CONTEXT = {
diff --git a/datahub-web-react/src/app/entity/dataProcessInstance/DataProcessInstanceEntity.tsx b/datahub-web-react/src/app/entity/dataProcessInstance/DataProcessInstanceEntity.tsx
index 4834a026ad94a3..9bb9bd745d1ee6 100644
--- a/datahub-web-react/src/app/entity/dataProcessInstance/DataProcessInstanceEntity.tsx
+++ b/datahub-web-react/src/app/entity/dataProcessInstance/DataProcessInstanceEntity.tsx
@@ -161,6 +161,7 @@ export class DataProcessInstanceEntity implements Entity<DataProcessInstance> {
         return {
             name,
             externalUrl,
+            platform: processInstance?.dataPlatformInstance?.platform,
         };
     };
 
@@ -174,9 +175,10 @@ export class DataProcessInstanceEntity implements Entity<DataProcessInstance> {
                 subType={data.subTypes?.typeNames?.[0]}
                 description=""
                 platformName={
-                    data?.platform?.properties?.displayName || capitalizeFirstLetterOnly(data?.platform?.name)
+                    data?.dataPlatformInstance?.platform?.properties?.displayName ||
+                    capitalizeFirstLetterOnly(data?.dataPlatformInstance?.platform?.name)
                 }
-                platformLogo={data.platform.properties?.logoUrl}
+                platformLogo={data?.dataPlatformInstance?.platform?.properties?.logoUrl}
                 owners={null}
                 globalTags={null}
                 // domain={data.domain?.domain}
@@ -201,9 +203,10 @@ export class DataProcessInstanceEntity implements Entity<DataProcessInstance> {
                 subType={data.subTypes?.typeNames?.[0]}
                 description=""
                 platformName={
-                    data?.platform?.properties?.displayName || capitalizeFirstLetterOnly(data?.platform?.name)
+                    data?.dataPlatformInstance?.platform?.properties?.displayName ||
+                    capitalizeFirstLetterOnly(data?.dataPlatformInstance?.platform?.name)
                 }
-                platformLogo={data.platform.properties?.logoUrl}
+                platformLogo={data.dataPlatformInstance?.platform?.properties?.logoUrl}
                 platformInstanceId={data.dataPlatformInstance?.instanceId}
                 owners={null}
                 globalTags={null}
@@ -231,8 +234,8 @@ export class DataProcessInstanceEntity implements Entity<DataProcessInstance> {
             name: this.displayName(entity),
             type: EntityType.DataProcessInstance,
             subtype: entity?.subTypes?.typeNames?.[0],
-            icon: entity?.platform?.properties?.logoUrl || undefined,
-            platform: entity?.platform,
+            icon: entity?.dataPlatformInstance?.platform?.properties?.logoUrl || undefined,
+            platform: entity?.dataPlatformInstance?.platform,
             container: entity?.container,
             //            health: entity?.health || undefined,
         };
diff --git a/datahub-web-react/src/app/entity/domain/DataProductsTab/DataProductsTab.tsx b/datahub-web-react/src/app/entity/domain/DataProductsTab/DataProductsTab.tsx
index 15cc99127f3500..39f89979dd95a0 100644
--- a/datahub-web-react/src/app/entity/domain/DataProductsTab/DataProductsTab.tsx
+++ b/datahub-web-react/src/app/entity/domain/DataProductsTab/DataProductsTab.tsx
@@ -68,6 +68,7 @@ export default function DataProductsTab() {
                 searchFlags: { skipCache: true },
             },
         },
+        fetchPolicy: 'no-cache',
     });
     const totalResults = data?.searchAcrossEntities?.total || 0;
     const searchResults = data?.searchAcrossEntities?.searchResults?.map((r) => r.entity) || [];
diff --git a/datahub-web-react/src/app/entity/shared/containers/profile/header/EntityName.tsx b/datahub-web-react/src/app/entity/shared/containers/profile/header/EntityName.tsx
index 8976629d9ef0b1..549724bd1945d9 100644
--- a/datahub-web-react/src/app/entity/shared/containers/profile/header/EntityName.tsx
+++ b/datahub-web-react/src/app/entity/shared/containers/profile/header/EntityName.tsx
@@ -48,9 +48,9 @@ function EntityName(props: Props) {
             setIsEditing(false);
             return;
         }
-        setUpdatedName(name);
         updateName({ variables: { input: { name, urn } } })
             .then(() => {
+                setUpdatedName(name);
                 setIsEditing(false);
                 message.success({ content: 'Name Updated', duration: 2 });
                 refetch();
diff --git a/datahub-web-react/src/app/entity/shared/containers/profile/sidebar/Ownership/EditOwnersModal.tsx b/datahub-web-react/src/app/entity/shared/containers/profile/sidebar/Ownership/EditOwnersModal.tsx
index 62b967e8f7b30d..e57666471df1a6 100644
--- a/datahub-web-react/src/app/entity/shared/containers/profile/sidebar/Ownership/EditOwnersModal.tsx
+++ b/datahub-web-react/src/app/entity/shared/containers/profile/sidebar/Ownership/EditOwnersModal.tsx
@@ -78,10 +78,26 @@ export const EditOwnersModal = ({
     const renderSearchResult = (entity: Entity) => {
         const avatarUrl =
             (entity.type === EntityType.CorpUser && (entity as CorpUser).editableProperties?.pictureLink) || undefined;
+        const corpUserDepartmentName =
+            (entity.type === EntityType.CorpUser && (entity as CorpUser).properties?.departmentName) || '';
+        const corpUserId = (entity.type === EntityType.CorpUser && (entity as CorpUser).username) || '';
+        const corpUserTitle = (entity.type === EntityType.CorpUser && (entity as CorpUser).properties?.title) || '';
         const displayName = entityRegistry.getDisplayName(entity.type, entity);
+
         return (
-            <Select.Option value={entity.urn} key={entity.urn}>
-                <OwnerLabel name={displayName} avatarUrl={avatarUrl} type={entity.type} />
+            <Select.Option
+                key={entity.urn}
+                value={entity.urn}
+                label={<OwnerLabel name={displayName} avatarUrl={avatarUrl} type={entity.type} />}
+            >
+                <OwnerLabel
+                    name={displayName}
+                    avatarUrl={avatarUrl}
+                    type={entity.type}
+                    corpUserId={corpUserId}
+                    corpUserTitle={corpUserTitle}
+                    corpUserDepartmentName={corpUserDepartmentName}
+                />
             </Select.Option>
         );
     };
@@ -381,6 +397,7 @@ export const EditOwnersModal = ({
                                 value: owner.value.ownerUrn,
                                 label: owner.label,
                             }))}
+                            optionLabelProp="label"
                         >
                             {ownerSearchOptions}
                         </SelectInput>
diff --git a/datahub-web-react/src/app/ingest/source/builder/NameSourceStep.tsx b/datahub-web-react/src/app/ingest/source/builder/NameSourceStep.tsx
index 898fbd6a6d9268..68e6c8d3436fb9 100644
--- a/datahub-web-react/src/app/ingest/source/builder/NameSourceStep.tsx
+++ b/datahub-web-react/src/app/ingest/source/builder/NameSourceStep.tsx
@@ -200,7 +200,7 @@ export const NameSourceStep = ({ state, updateState, prev, submit }: StepProps)
                             <Input
                                 data-testid="cli-version-input"
                                 className="cli-version-input"
-                                placeholder="(e.g. 0.12.0)"
+                                placeholder="(e.g. 0.15.0)"
                                 value={state.config?.version || ''}
                                 onChange={(event) => setVersion(event.target.value)}
                                 onBlur={(event) => handleBlur(event, setVersion)}
diff --git a/datahub-web-react/src/app/shared/OwnerLabel.tsx b/datahub-web-react/src/app/shared/OwnerLabel.tsx
index de3c03dea2ba4a..fb670aa56d7881 100644
--- a/datahub-web-react/src/app/shared/OwnerLabel.tsx
+++ b/datahub-web-react/src/app/shared/OwnerLabel.tsx
@@ -20,14 +20,22 @@ type Props = {
     name: string;
     avatarUrl: string | undefined;
     type: EntityType;
+    corpUserId?: string;
+    corpUserTitle?: string;
+    corpUserDepartmentName?: string;
 };
 
-export const OwnerLabel = ({ name, avatarUrl, type }: Props) => {
+export const OwnerLabel = ({ name, avatarUrl, type, corpUserId, corpUserTitle, corpUserDepartmentName }: Props) => {
+    const subHeader = [corpUserId, corpUserTitle, corpUserDepartmentName].filter(Boolean).join(' - ');
+
     return (
         <OwnerContainerWrapper>
             <OwnerContentWrapper>
                 <CustomAvatar size={24} name={name} photoUrl={avatarUrl} isGroup={type === EntityType.CorpGroup} />
-                <div>{name}</div>
+                <div>
+                    <div>{name}</div>
+                    {subHeader && <div style={{ color: 'gray' }}>{subHeader}</div>}
+                </div>
             </OwnerContentWrapper>
         </OwnerContainerWrapper>
     );
diff --git a/datahub-web-react/src/appConfigContext.tsx b/datahub-web-react/src/appConfigContext.tsx
index d7fef85db4b625..8ac18d0142b4e9 100644
--- a/datahub-web-react/src/appConfigContext.tsx
+++ b/datahub-web-react/src/appConfigContext.tsx
@@ -57,6 +57,7 @@ export const DEFAULT_APP_CONFIG = {
         editableDatasetNameEnabled: false,
         showSeparateSiblings: false,
         showManageStructuredProperties: false,
+        entityVersioningEnabled: false,
     },
 };
 
diff --git a/datahub-web-react/src/graphql/app.graphql b/datahub-web-react/src/graphql/app.graphql
index 0d1999f82f77cd..c1fe50d7620a3c 100644
--- a/datahub-web-react/src/graphql/app.graphql
+++ b/datahub-web-react/src/graphql/app.graphql
@@ -72,6 +72,7 @@ query appConfig {
             editableDatasetNameEnabled
             showSeparateSiblings
             showManageStructuredProperties
+            entityVersioningEnabled
         }
     }
 }
diff --git a/datahub-web-react/src/graphql/dataProcessInstance.graphql b/datahub-web-react/src/graphql/dataProcessInstance.graphql
index 8f55ca4903d527..442f8db0a933b2 100644
--- a/datahub-web-react/src/graphql/dataProcessInstance.graphql
+++ b/datahub-web-react/src/graphql/dataProcessInstance.graphql
@@ -67,9 +67,6 @@ fragment processInstanceRelationshipResults on EntityRelationshipsResult {
 fragment dataProcessInstanceFields on DataProcessInstance {
     urn
     type
-    platform {
-        ...platformFields
-    }
     parentContainers {
         ...parentContainersFields
     }
@@ -125,9 +122,6 @@ query getDataProcessInstance($urn: String!) {
     dataProcessInstance(urn: $urn) {
         urn
         type
-        platform {
-            ...platformFields
-        }
         parentContainers {
             ...parentContainersFields
         }
diff --git a/datahub-web-react/src/graphql/dataset.graphql b/datahub-web-react/src/graphql/dataset.graphql
index fcca919f614235..8bbeb304aae2cc 100644
--- a/datahub-web-react/src/graphql/dataset.graphql
+++ b/datahub-web-react/src/graphql/dataset.graphql
@@ -172,6 +172,7 @@ fragment nonSiblingDatasetFields on Dataset {
     forms {
         ...formsFields
     }
+    ...entityProfileVersionProperties
 }
 
 query getRecentQueries($urn: String!) {
diff --git a/datahub-web-react/src/graphql/mlModel.graphql b/datahub-web-react/src/graphql/mlModel.graphql
index 2192888caef701..ad97c7c6f530a1 100644
--- a/datahub-web-react/src/graphql/mlModel.graphql
+++ b/datahub-web-react/src/graphql/mlModel.graphql
@@ -34,5 +34,6 @@ query getMLModel($urn: String!) {
         forms {
             ...formsFields
         }
+        ...entityProfileVersionProperties
     }
 }
diff --git a/datahub-web-react/src/graphql/preview.graphql b/datahub-web-react/src/graphql/preview.graphql
index 1bee614dd7adbe..8000f59f2bf258 100644
--- a/datahub-web-react/src/graphql/preview.graphql
+++ b/datahub-web-react/src/graphql/preview.graphql
@@ -346,4 +346,9 @@ fragment entityPreview on Entity {
     ... on Container {
         ...entityContainer
     }
+    ... on SupportsVersions {
+        versionProperties {
+            ...versionProperties
+        }
+    }
 }
diff --git a/datahub-web-react/src/graphql/search.graphql b/datahub-web-react/src/graphql/search.graphql
index 72e7d347187828..9edd6754022866 100644
--- a/datahub-web-react/src/graphql/search.graphql
+++ b/datahub-web-react/src/graphql/search.graphql
@@ -433,6 +433,8 @@ fragment searchResultsWithoutSchemaField on Entity {
             lastName
             fullName
             email
+            departmentName
+            title
         }
         info {
             active
@@ -442,6 +444,8 @@ fragment searchResultsWithoutSchemaField on Entity {
             lastName
             fullName
             email
+            departmentName
+            title
         }
         editableProperties {
             displayName
@@ -945,6 +949,11 @@ fragment searchResultsWithoutSchemaField on Entity {
     ... on StructuredPropertyEntity {
         ...structuredPropertyFields
     }
+    ... on SupportsVersions {
+        versionProperties {
+            ...versionProperties
+        }
+    }
 }
 
 fragment searchResultFields on Entity {
diff --git a/datahub-web-react/src/graphql/versioning.graphql b/datahub-web-react/src/graphql/versioning.graphql
new file mode 100644
index 00000000000000..e9b6b82494b6ed
--- /dev/null
+++ b/datahub-web-react/src/graphql/versioning.graphql
@@ -0,0 +1,89 @@
+fragment versionProperties on VersionProperties {
+    versionSet {
+        urn
+        type
+    }
+    isLatest
+    version {
+        versionTag
+    }
+    aliases {
+        versionTag
+    }
+    comment
+    created {
+        time
+        actor {
+            urn
+            ...entityDisplayNameFields
+            editableProperties {
+                displayName
+                pictureLink
+            }
+        }
+    }
+    createdInSource {
+        time
+        actor {
+            urn
+            ...entityDisplayNameFields
+            editableProperties {
+                displayName
+                pictureLink
+            }
+        }
+    }
+}
+
+fragment versionsSearchResults on SearchResults {
+    count
+    total
+    searchResults {
+        entity {
+            urn
+            type
+            ... on SupportsVersions {
+                versionProperties {
+                    ...versionProperties
+                }
+            }
+        }
+    }
+}
+
+fragment entityProfileVersionProperties on SupportsVersions {
+    versionProperties {
+        ...versionProperties
+        versionSet {
+            urn
+            type
+            versionsSearch(input: { query: "*", count: 5, searchFlags: { skipCache: true } }) {
+                ...versionsSearchResults
+            }
+        }
+    }
+}
+
+query searchAcrossVersions($versionSetUrn: String!, $input: SearchAcrossEntitiesInput!) {
+    versionSet(urn: $versionSetUrn) {
+        versionsSearch(input: $input) {
+            ...versionsSearchResults
+        }
+    }
+}
+
+mutation linkAssetVersion($input: LinkVersionInput!) {
+    linkAssetVersion(input: $input) {
+        versionsSearch(input: { query: "*", count: 5, searchFlags: { skipCache: true } }) {
+            ...versionsSearchResults
+        }
+    }
+}
+
+mutation unlinkAssetVersion($input: UnlinkVersionInput!) {
+    unlinkAssetVersion(input: $input) {
+        versionsSearch(input: { query: "*", count: 5, searchFlags: { skipCache: true } }) {
+            ...versionsSearchResults
+        }
+    }
+}
diff --git a/datahub-web-react/yarn.lock b/datahub-web-react/yarn.lock
index f16e8aa506e2cc..6e946e584f18a5 100644
--- a/datahub-web-react/yarn.lock
+++ b/datahub-web-react/yarn.lock
@@ -12066,10 +12066,10 @@ vite-plugin-svgr@^4.1.0:
     "@svgr/core" "^8.1.0"
     "@svgr/plugin-jsx" "^8.1.0"
 
-"vite@^3.0.0 || ^4.0.0 || ^5.0.0-0", "vite@^3.1.0 || ^4.0.0 || ^5.0.0-0", vite@^4.5.5:
-  version "4.5.5"
-  resolved "https://registry.yarnpkg.com/vite/-/vite-4.5.5.tgz#639b9feca5c0a3bfe3c60cb630ef28bf219d742e"
-  integrity sha512-ifW3Lb2sMdX+WU91s3R0FyQlAyLxOzCSCP37ujw0+r5POeHPwe6udWVIElKQq8gk3t7b8rkmvqC6IHBpCff4GQ==
+"vite@^3.0.0 || ^4.0.0 || ^5.0.0-0", "vite@^3.1.0 || ^4.0.0 || ^5.0.0-0", vite@^4.5.6:
+  version "4.5.6"
+  resolved "https://registry.yarnpkg.com/vite/-/vite-4.5.6.tgz#48bbd97fe06e8241df2e625b31c581707e10b57d"
+  integrity sha512-ElBNuVvJKslxcfY2gMmae5IjaKGqCYGicCNZ+8R56sAznobeE3pI9ctzI17cBS/6OJh5YuQNMSN4BP4dRjugBg==
   dependencies:
     esbuild "^0.18.10"
     postcss "^8.4.27"
diff --git a/docker/README.md b/docker/README.md
index ad847dc70cf3c0..b756ae3eaed752 100644
--- a/docker/README.md
+++ b/docker/README.md
@@ -16,7 +16,7 @@ area.
 ## Quickstart
 
 The easiest way to bring up and test DataHub is using DataHub [Docker](https://www.docker.com) images 
-which are continuously deployed to [Docker Hub](https://hub.docker.com/u/linkedin) with every commit to repository.
+which are continuously deployed to [Docker Hub](https://hub.docker.com/u/acryldata) with every commit to repository.
 
 You can easily download and run all these images and their dependencies with our
 [quick start guide](../docs/quickstart.md).
diff --git a/docker/datahub-gms/Dockerfile b/docker/datahub-gms/Dockerfile
index 52cc507f9268d1..5462d4f70002c1 100644
--- a/docker/datahub-gms/Dockerfile
+++ b/docker/datahub-gms/Dockerfile
@@ -26,7 +26,6 @@ RUN go install github.com/jwilder/dockerize@$DOCKERIZE_VERSION
 FROM alpine:3.21 AS base
 
 ENV JMX_VERSION=0.18.0
-ENV JETTY_VERSION=11.0.21
 
 # Re-declaring args from above to make them available in this stage (will inherit default values)
 ARG ALPINE_REPO_URL
@@ -42,9 +41,6 @@ RUN apk --no-cache --update-cache --available upgrade \
     && apk --no-cache add curl bash coreutils gcompat sqlite libc6-compat snappy \
     && apk --no-cache add openjdk17-jre-headless --repository=${ALPINE_REPO_URL}/edge/community \
     && apk --no-cache add jattach --repository ${ALPINE_REPO_URL}/edge/community/ \
-    && curl -sS ${MAVEN_CENTRAL_REPO_URL}/org/eclipse/jetty/jetty-runner/${JETTY_VERSION}/jetty-runner-${JETTY_VERSION}.jar --output jetty-runner.jar \
-    && curl -sS ${MAVEN_CENTRAL_REPO_URL}/org/eclipse/jetty/jetty-jmx/${JETTY_VERSION}/jetty-jmx-${JETTY_VERSION}.jar --output jetty-jmx.jar \
-    && curl -sS ${MAVEN_CENTRAL_REPO_URL}/org/eclipse/jetty/jetty-util/${JETTY_VERSION}/jetty-util-${JETTY_VERSION}.jar --output jetty-util.jar \
     && wget --no-verbose ${GITHUB_REPO_URL}/open-telemetry/opentelemetry-java-instrumentation/releases/download/v1.24.0/opentelemetry-javaagent.jar \
     && wget --no-verbose ${MAVEN_CENTRAL_REPO_URL}/io/prometheus/jmx/jmx_prometheus_javaagent/${JMX_VERSION}/jmx_prometheus_javaagent-${JMX_VERSION}.jar -O jmx_prometheus_javaagent.jar \
     && cp /usr/lib/jvm/java-17-openjdk/jre/lib/security/cacerts /tmp/kafka.client.truststore.jks
@@ -56,8 +52,6 @@ FROM base AS prod-install
 COPY war.war /datahub/datahub-gms/bin/war.war
 COPY metadata-models/src/main/resources/entity-registry.yml /datahub/datahub-gms/resources/entity-registry.yml
 COPY docker/datahub-gms/start.sh /datahub/datahub-gms/scripts/start.sh
-COPY docker/datahub-gms/jetty.xml /datahub/datahub-gms/scripts/jetty.xml
-COPY docker/datahub-gms/jetty-jmx.xml /datahub/datahub-gms/scripts/jetty-jmx.xml
 COPY docker/monitoring/client-prometheus-config.yaml /datahub/datahub-gms/scripts/prometheus-config.yaml
 RUN chmod +x /datahub/datahub-gms/scripts/start.sh
 
@@ -70,7 +64,7 @@ FROM ${APP_ENV}-install AS final
 RUN mkdir -p /etc/datahub/plugins/auth/resources
 
 RUN addgroup -S datahub && adduser -S datahub -G datahub
-RUN chown -R datahub:datahub /etc/datahub
+RUN chown -R datahub:datahub /etc/datahub /datahub
 USER datahub
 
 ENV JMX_OPTS=""
diff --git a/docker/datahub-gms/jetty-jmx.xml b/docker/datahub-gms/jetty-jmx.xml
deleted file mode 100644
index 5aadbb66a70ed1..00000000000000
--- a/docker/datahub-gms/jetty-jmx.xml
+++ /dev/null
@@ -1,31 +0,0 @@
-<?xml version="1.0"?>
-<!DOCTYPE Configure PUBLIC "-//Jetty//Configure//EN" "https://www.eclipse.org/jetty/configure_10_0.dtd">
-
-<Configure id="Server" class="org.eclipse.jetty.server.Server">
-
-  <!-- =========================================================== -->
-  <!-- Get the platform MBeanServer                                -->
-  <!-- =========================================================== -->
-  <Call id="MBeanServer" class="java.lang.management.ManagementFactory"
-    name="getPlatformMBeanServer" />
-
-  <!-- =========================================================== -->
-  <!-- Initialize the Jetty MBeanContainer                         -->
-  <!-- =========================================================== -->
-  <Call name="addBean">
-    <Arg>
-      <New id="MBeanContainer" class="org.eclipse.jetty.jmx.MBeanContainer">
-        <Arg>
-          <Ref refid="MBeanServer" />
-        </Arg>
-        <Call name="beanAdded">
-          <Arg/>
-          <Arg>
-            <Get name="ILoggerFactory" class="org.slf4j.LoggerFactory"/>
-          </Arg>
-        </Call>
-      </New>
-    </Arg>
-  </Call>
-</Configure>
-
diff --git a/docker/datahub-gms/jetty.xml b/docker/datahub-gms/jetty.xml
deleted file mode 100644
index 3f04635d9498ca..00000000000000
--- a/docker/datahub-gms/jetty.xml
+++ /dev/null
@@ -1,57 +0,0 @@
-<?xml version="1.0"?><!DOCTYPE Configure PUBLIC "-//Jetty//Configure//EN" "http://www.eclipse.org/jetty/configure_9_3.dtd">
-
-<!-- =============================================================== -->
-<!-- Documentation of this file format can be found at:              -->
-<!-- https://www.eclipse.org/jetty/documentation/current/       -->
-<!--                                                                 -->
-<!-- Additional configuration files are available in $JETTY_HOME/etc -->
-<!-- and can be mixed in. See start.ini file for the default         -->
-<!-- configuration files.                                            -->
-<!--                                                                 -->
-<!-- For a description of the configuration mechanism, see the       -->
-<!-- output of:                                                      -->
-<!--   java -jar start.jar -?                                        -->
-<!-- =============================================================== -->
-
-<!-- =============================================================== -->
-<!-- Configure a Jetty Server instance with an ID "Server"           -->
-<!-- Other configuration files may also configure the "Server"       -->
-<!-- ID, in which case they are adding configuration to the same     -->
-<!-- instance.  If other configuration have a different ID, they     -->
-<!-- will create and configure another instance of Jetty.            -->
-<!-- Consult the javadoc of o.e.j.server.Server for all              -->
-<!-- configuration that may be set here.                             -->
-<!-- =============================================================== -->
-<Configure id="Server" class="org.eclipse.jetty.server.Server">
-
-    <!-- =========================================================== -->
-    <!-- Change requestHeaderSize 8k to 16k                          -->
-    <!-- Some large cookies will return 431.                         -->
-    <!-- <h1>Bad Message 431</h1>                                    -->
-    <!-- <pre>reason: Request Header Fields Too Large</pre>          -->
-    <!-- =========================================================== -->
-    <New id="httpConfig" class="org.eclipse.jetty.server.HttpConfiguration">
-      <Set name="requestHeaderSize"><Property name="jetty.httpConfig.requestHeaderSize" deprecated="jetty.request.header.size" default="32768" /></Set>
-    </New>
-
-    <Call name="addConnector">
-        <Arg>
-            <New id="httpConnector" class="org.eclipse.jetty.server.ServerConnector">
-                <Arg name="server"><Ref refid="Server" /></Arg>
-                <Arg name="acceptors" type="int"><Property name="jetty.http.acceptors" deprecated="http.acceptors" default="-1"/></Arg>
-                <Arg name="selectors" type="int"><Property name="jetty.http.selectors" deprecated="http.selectors" default="-1"/></Arg>
-                <Arg name="factories">
-                    <Array type="org.eclipse.jetty.server.ConnectionFactory">
-                        <Item>
-                            <New class="org.eclipse.jetty.server.HttpConnectionFactory">
-                                <Arg name="config"><Ref refid="httpConfig" /></Arg>
-                            </New>
-                        </Item>
-                    </Array>
-                </Arg>
-                <Set name="port"><Property name="jetty.http.port" deprecated="jetty.port" default="8080" /></Set>
-            </New>
-        </Arg>
-    </Call>
-
-</Configure>
diff --git a/docker/datahub-gms/start.sh b/docker/datahub-gms/start.sh
index c91580eed83cb9..cac36920022749 100755
--- a/docker/datahub-gms/start.sh
+++ b/docker/datahub-gms/start.sh
@@ -62,13 +62,8 @@ COMMON="
     java $JAVA_OPTS $JMX_OPTS \
     $OTEL_AGENT \
     $PROMETHEUS_AGENT \
-    -jar /jetty-runner.jar \
-    --stats unsecure \
-    --jar jetty-util.jar \
-    --jar jetty-jmx.jar \
-    --config /datahub/datahub-gms/scripts/jetty.xml \
-    --config /datahub/datahub-gms/scripts/jetty-jmx.xml \
-    /datahub/datahub-gms/bin/war.war"
+    -Dstats=unsecure \
+    -jar /datahub/datahub-gms/bin/war.war"
 
 if [[ $SKIP_ELASTICSEARCH_CHECK != true ]]; then
   exec dockerize \
diff --git a/docker/datahub-ingestion-base/smoke.Dockerfile b/docker/datahub-ingestion-base/smoke.Dockerfile
index 34654faaad729d..81a6bd0e20cacc 100644
--- a/docker/datahub-ingestion-base/smoke.Dockerfile
+++ b/docker/datahub-ingestion-base/smoke.Dockerfile
@@ -20,9 +20,9 @@ RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y \
 COPY . /datahub-src
 ARG RELEASE_VERSION
 RUN cd /datahub-src && \
-    sed -i.bak "s/__version__ = \"1\!0.0.0.dev0\"/__version__ = \"$(echo $RELEASE_VERSION|sed s/-/+/)\"/" metadata-ingestion/src/datahub/__init__.py && \
-    sed -i.bak "s/__version__ = \"1\!0.0.0.dev0\"/__version__ = \"$(echo $RELEASE_VERSION|sed s/-/+/)\"/" metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/__init__.py && \
-    cat metadata-ingestion/src/datahub/__init__.py && \
+    sed -i.bak "s/__version__ = .*$/__version__ = \"$(echo $RELEASE_VERSION|sed s/-/+/)\"/" metadata-ingestion/src/datahub/_version.py && \
+    sed -i.bak "s/__version__ = .*$/__version__ = \"$(echo $RELEASE_VERSION|sed s/-/+/)\"/" metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/_version.py && \
+    cat metadata-ingestion/src/datahub/_version.py && \
     ./gradlew :metadata-ingestion:codegen && \
     pip install file:metadata-ingestion-modules/airflow-plugin#egg=acryl-datahub-airflow-plugin file:metadata-ingestion#egg=acryl-datahub
 
diff --git a/docker/datahub-ingestion/Dockerfile b/docker/datahub-ingestion/Dockerfile
index ee0333e1cb1d1f..a9fd3a6662d1bb 100644
--- a/docker/datahub-ingestion/Dockerfile
+++ b/docker/datahub-ingestion/Dockerfile
@@ -26,10 +26,10 @@ COPY --chown=datahub ./metadata-ingestion-modules/airflow-plugin /metadata-inges
 
 ARG RELEASE_VERSION
 WORKDIR /metadata-ingestion
-RUN sed -i.bak "s/__version__ = \"1\!0.0.0.dev0\"/__version__ = \"$(echo $RELEASE_VERSION|sed s/-/+/)\"/" src/datahub/__init__.py && \
-    sed -i.bak "s/__version__ = \"1\!0.0.0.dev0\"/__version__ = \"$(echo $RELEASE_VERSION|sed s/-/+/)\"/" airflow-plugin/src/datahub_airflow_plugin/__init__.py && \
-    cat src/datahub/__init__.py | grep __version__ && \
-    cat airflow-plugin/src/datahub_airflow_plugin/__init__.py | grep __version__
+RUN sed -i.bak "s/__version__ = .*$/__version__ = \"$(echo $RELEASE_VERSION|sed s/-/+/)\"/" src/datahub/_version.py && \
+    sed -i.bak "s/__version__ = .*$/__version__ = \"$(echo $RELEASE_VERSION|sed s/-/+/)\"/" airflow-plugin/src/datahub_airflow_plugin/_version.py && \
+    cat src/datahub/_version.py | grep __version__ && \
+    cat airflow-plugin/src/datahub_airflow_plugin/_version.py | grep __version__
 
 FROM base AS slim-install
 
diff --git a/docker/datahub-ingestion/Dockerfile-slim-only b/docker/datahub-ingestion/Dockerfile-slim-only
index 6ade262f2feded..80abff204df9f2 100644
--- a/docker/datahub-ingestion/Dockerfile-slim-only
+++ b/docker/datahub-ingestion/Dockerfile-slim-only
@@ -15,8 +15,8 @@ COPY --chown=datahub ./metadata-ingestion /metadata-ingestion
 
 ARG RELEASE_VERSION
 WORKDIR /metadata-ingestion
-RUN sed -i.bak "s/__version__ = \"1\!0.0.0.dev0\"/__version__ = \"$(echo $RELEASE_VERSION|sed s/-/+/)\"/" src/datahub/__init__.py && \
-    cat src/datahub/__init__.py
+RUN sed -i.bak "s/__version__ = .*$/__version__ = \"$(echo $RELEASE_VERSION|sed s/-/+/)\"/" src/datahub/_version.py && \
+    cat src/datahub/_version.py
 
 FROM base as slim-install
 
diff --git a/docker/docker-compose.dev.yml b/docker/docker-compose.dev.yml
index c68a4c1f5a8fcf..0cbb2aee903c84 100644
--- a/docker/docker-compose.dev.yml
+++ b/docker/docker-compose.dev.yml
@@ -49,8 +49,6 @@ services:
     - KAFKA_CONSUMER_STOP_ON_DESERIALIZATION_ERROR=${KAFKA_CONSUMER_STOP_ON_DESERIALIZATION_ERROR:-true}
     volumes:
     - ./datahub-gms/start.sh:/datahub/datahub-gms/scripts/start.sh
-    - ./datahub-gms/jetty.xml:/datahub/datahub-gms/scripts/jetty.xml
-    - ./datahub-gms/jetty-jmx.xml:/datahub/datahub-gms/scripts/jetty-jmx.xml
     - ./monitoring/client-prometheus-config.yaml:/datahub/datahub-gms/scripts/prometheus-config.yaml
     - ../metadata-models/src/main/resources/:/datahub/datahub-gms/resources
     - ../metadata-service/war/build/libs/:/datahub/datahub-gms/bin
diff --git a/docker/profiles/docker-compose.frontend.yml b/docker/profiles/docker-compose.frontend.yml
index c6b15a7016670d..b278cd41b0b231 100644
--- a/docker/profiles/docker-compose.frontend.yml
+++ b/docker/profiles/docker-compose.frontend.yml
@@ -26,6 +26,7 @@ x-datahub-frontend-service-dev: &datahub-frontend-service-dev
     DATAHUB_ANALYTICS_ENABLED: ${DATAHUB_ANALYTICS_ENABLED:-true}
   volumes:
     - ../../datahub-frontend/build/stage/main:/datahub-frontend
+    - ./monitoring/client-prometheus-config.yaml:/datahub-frontend/client-prometheus-config.yaml
 
 services:
   frontend-quickstart:
diff --git a/docker/profiles/docker-compose.gms.yml b/docker/profiles/docker-compose.gms.yml
index 2147d6b5a0247f..d4ea7dde9f8481 100644
--- a/docker/profiles/docker-compose.gms.yml
+++ b/docker/profiles/docker-compose.gms.yml
@@ -130,15 +130,13 @@ x-datahub-gms-service-dev: &datahub-gms-service-dev
     <<: [*datahub-dev-telemetry-env, *datahub-gms-env]
     SKIP_ELASTICSEARCH_CHECK: false
     JAVA_TOOL_OPTIONS: '-agentlib:jdwp=transport=dt_socket,server=y,suspend=n,address=*:5001'
-    BOOTSTRAP_SYSTEM_UPDATE_WAIT_FOR_SYSTEM_UPDATE: false
+    BOOTSTRAP_SYSTEM_UPDATE_WAIT_FOR_SYSTEM_UPDATE: ${BOOTSTRAP_SYSTEM_UPDATE_WAIT_FOR_SYSTEM_UPDATE:-false}
     SEARCH_SERVICE_ENABLE_CACHE: false
     LINEAGE_SEARCH_CACHE_ENABLED: false
     SHOW_BROWSE_V2: true
     ENTITY_VERSIONING_ENABLED: ${ENTITY_VERSIONING_ENABLED:-true}
   volumes:
     - ./datahub-gms/start.sh:/datahub/datahub-gms/scripts/start.sh
-    - ./datahub-gms/jetty.xml:/datahub/datahub-gms/scripts/jetty.xml
-    - ./datahub-gms/jetty-jmx.xml:/datahub/datahub-gms/scripts/jetty-jmx.xml
     - ./monitoring/client-prometheus-config.yaml:/datahub/datahub-gms/scripts/prometheus-config.yaml
     - ../../metadata-models/src/main/resources/:/datahub/datahub-gms/resources
     - ../../metadata-service/war/build/libs/:/datahub/datahub-gms/bin
diff --git a/docs-website/docusaurus.config.js b/docs-website/docusaurus.config.js
index fb2f418e37c044..686c28fa3d1b93 100644
--- a/docs-website/docusaurus.config.js
+++ b/docs-website/docusaurus.config.js
@@ -77,7 +77,7 @@ module.exports = {
     announcementBar: {
           id: "announcement-3",
           content:
-            '<div style="display: flex; justify-content: center; align-items: center;width: 100%;"><!--img src="/img/acryl-logo-white-mark.svg" / --><!--div style="font-size: .8rem; font-weight: 600; background-color: white; color: #111; padding: 0px 8px; border-radius: 4px; margin-right:12px;">NEW</div--><p>Watch Metadata & AI Summit sessions on-demand.</p><a href="https://www.youtube.com/@DataHubProject/videos" target="_blank" class="button">Watch Now<span> →</span></a></div>',
+            '<div style="display: flex; justify-content: center; align-items: center;width: 100%;"><!--img src="/img/acryl-logo-white-mark.svg" / --><!--div style="font-size: .8rem; font-weight: 600; background-color: white; color: #111; padding: 0px 8px; border-radius: 4px; margin-right:12px;">NEW</div--><p>Learn about DataHub 1.0 launching at our 5th birthday party!</p><a href="https://lu.ma/0j5jcocn" target="_blank" class="button">Register<span> →</span></a></div>',
           backgroundColor: "#111",
           textColor: "#ffffff",
           isCloseable: false,
diff --git a/docs-website/sidebars.js b/docs-website/sidebars.js
index bcb06affedcff0..fbd35b60aedba9 100644
--- a/docs-website/sidebars.js
+++ b/docs-website/sidebars.js
@@ -368,6 +368,7 @@ module.exports = {
         },
         {
           "DataHub Cloud Release History": [
+            "docs/managed-datahub/release-notes/v_0_3_8",
             "docs/managed-datahub/release-notes/v_0_3_7",
             "docs/managed-datahub/release-notes/v_0_3_6",
             "docs/managed-datahub/release-notes/v_0_3_5",
diff --git a/docs-website/src/components/SolutionsDropdown/SolutionsDropdownContent/solutionsDropdownContent.js b/docs-website/src/components/SolutionsDropdown/SolutionsDropdownContent/solutionsDropdownContent.js
index abede0f11735d4..ad7278a438cf81 100644
--- a/docs-website/src/components/SolutionsDropdown/SolutionsDropdownContent/solutionsDropdownContent.js
+++ b/docs-website/src/components/SolutionsDropdown/SolutionsDropdownContent/solutionsDropdownContent.js
@@ -24,7 +24,7 @@ const solutionsDropdownContent = {
             title: "DataHub Core",
             description: "Get started with the Open Source platform.",
             iconImage: "/img/solutions/icon-dropdown-core.png",
-            href: "/",
+            href: "/docs/quickstart",
         },
         {
             title: "Cloud vs Core",
diff --git a/docs-website/src/theme/NavbarItem/DocsVersionDropdownNavbarItem.js b/docs-website/src/theme/NavbarItem/DocsVersionDropdownNavbarItem.js
index ed083e4e0c3987..c9e5098a016c2d 100644
--- a/docs-website/src/theme/NavbarItem/DocsVersionDropdownNavbarItem.js
+++ b/docs-website/src/theme/NavbarItem/DocsVersionDropdownNavbarItem.js
@@ -42,6 +42,22 @@ export default function DocsVersionDropdownNavbarItem({
       type: 'html',
       value: '<div class="dropdown__link"><b>Archived versions</b></div>',
     },
+    {
+      value: `
+         <a class="dropdown__link" href="https://docs-website-eue2qafvn-acryldata.vercel.app//docs/features">0.14.0
+         <svg width="12" height="12" aria-hidden="true" viewBox="0 0 24 24"><path fill="currentColor" d="M21 13v10h-21v-19h12v2h-10v15h17v-8h2zm3-12h-10.988l4.035 4-6.977 7.07 2.828 2.828 6.977-7.07 4.125 4.172v-11z"></path></svg>
+         </a>
+      `,
+      type: "html",
+    },
+    {
+      value: `
+         <a class="dropdown__link" href="https://docs-website-psat3nzgi-acryldata.vercel.app/docs/features">0.13.1
+         <svg width="12" height="12" aria-hidden="true" viewBox="0 0 24 24"><path fill="currentColor" d="M21 13v10h-21v-19h12v2h-10v15h17v-8h2zm3-12h-10.988l4.035 4-6.977 7.07 2.828 2.828 6.977-7.07 4.125 4.172v-11z"></path></svg>
+         </a>
+      `,
+      type: "html",
+    },
     {
       value: `
          <a class="dropdown__link" href="https://docs-website-lzxh86531-acryldata.vercel.app/docs/features">0.13.0
diff --git a/docs-website/versions.json b/docs-website/versions.json
index 5288c42437c779..ab242fa47dce5b 100644
--- a/docs-website/versions.json
+++ b/docs-website/versions.json
@@ -1,4 +1,4 @@
 [
-  "0.14.0",
-  "0.13.1"
+  "0.15.0",
+  "0.14.1"
 ]
diff --git a/docs/actions/README.md b/docs/actions/README.md
index 4fa44eec588bed..a906c3c78f119e 100644
--- a/docs/actions/README.md
+++ b/docs/actions/README.md
@@ -189,11 +189,10 @@ Two event types are currently supported. Read more about them below.
 
 ## Supported Event Sources
 
-Currently, the only event source that is officially supported is `kafka`, which polls for events
-via a Kafka Consumer. 
+Currently, the following event sources are supported:
 
 - [Kafka Event Source](sources/kafka-event-source.md)
-
+- [DataHub Cloud Event Source](sources/datahub-cloud-event-source.md)
 
 ## Supported Actions
 
diff --git a/docs/actions/sources/datahub-cloud-event-source.md b/docs/actions/sources/datahub-cloud-event-source.md
new file mode 100644
index 00000000000000..656fe4a3a6329f
--- /dev/null
+++ b/docs/actions/sources/datahub-cloud-event-source.md
@@ -0,0 +1,113 @@
+# DataHub Cloud Event Source
+
+## Prerequisites
+
+### Compatibility
+
+The **DataHub Cloud Event Source** is only compatible with versions of DataHub Cloud above  `v0.3.7`.
+
+### Privileges
+
+By default, users do not have access to the Events API of DataHub Cloud. In order to access the API, the user or service account
+associated with the access token used to configure this events source _must_ have the `Get Platform Events` platform privilege, which
+can be granted using an [Access Policy](https://datahubproject.io/docs/authorization/access-policies-guide/).
+
+## Overview
+
+The DataHub Cloud Event Source allows you to use DataHub Actions with an instance of DataHub Cloud hosted by [Acryl](https://acryl.io).
+
+Under the hood, the DataHub Cloud Event Source communicates with DataHub Cloud to extract change events in realtime.
+The state of progress is automatically saved to DataHub Cloud after messages are processed, allowing you to seamlessly pause and restart the consumer, using the provided `name` to uniquely identify the consumer state.
+
+On initial startup of a new consumer id, the DataHub event source will automatically begin the _latest_ events by default. Afterwards, the message stream processed offsets will be continually saved. However, the source can also optionally be configured to "look back" in time
+by a certain number of days on initial bootstrap using the `lookback_days` parameter. To reset all previously saved offsets for a consumer,
+you can set `reset_offsets` to `True`.
+
+### Processing Guarantees
+
+This event source implements an "ack" function which is invoked if and only if an event is successfully processed
+by the Actions framework, meaning that the event made it through the Transformers and into the Action without
+any errors. Under the hood, the "ack" method synchronously commits DataHub Cloud Consumer Offsets on behalf of the Action. This means that by default, the framework provides *at-least once* processing semantics. That is, in the unusual case that a failure occurs when attempting to commit offsets back to Kafka, that event may be replayed on restart of the Action.
+
+If you've configured your Action pipeline `failure_mode` to be `CONTINUE` (the default), then events which
+fail to be processed will simply be logged to a `failed_events.log` file for further investigation (dead letter queue). The DataHub Cloud Event Source will continue to make progress against the underlying topics and continue to commit offsets even in the case of failed messages.
+
+If you've configured your Action pipeline `failure_mode` to be `THROW`, then events which fail to be processed result in an Action Pipeline error. This in turn terminates the pipeline before committing offsets back to DataHub Cloud. Thus the message will not be marked as "processed" by the Action consumer.
+
+## Supported Events
+
+The DataHub Cloud Event Source produces
+
+- [Entity Change Event V1](../../managed-datahub/datahub-api/entity-events-api.md)
+
+Note that the DataHub Cloud Event Source does _not_ yet support the full [Metadata Change Log V1](../events/metadata-change-log-event.md) event stream.
+
+## Configure the Event Source
+
+Use the following config(s) to get started with the DataHub Cloud Event Source.
+
+### Quickstart
+
+To start listening for new events from now, you can use the following recipe:
+
+```yml
+name: "unique-action-name"
+datahub:
+  server: "https://<your-organization>.acryl.io"
+  token: "<your-datahub-cloud-token>"
+source:
+  type: "datahub-cloud"
+action:
+  # action configs
+```
+
+Note that the `datahub` configuration block is **required** to connect to your DataHub Cloud instance.
+
+### Advanced Configurations
+
+To reset the offsets for the action pipeline and start consuming events from 7 days ago, you can use the following recipe:
+
+```yml
+name: "unique-action-name"
+datahub:
+  server: "https://<your-organization>.acryl.io"
+  token: "<your-datahub-cloud-token>"
+source:
+  type: "datahub-cloud"
+  config:
+    lookback_days: 7                           # Look back 7 days for events
+    reset_offsets: true                        # Ignore stored offsets and start fresh
+    kill_after_idle_timeout: true              # Enable shutdown after idle period
+    idle_timeout_duration_seconds: 60          # Idle timeout set to 60 seconds
+    event_processing_time_max_duration_seconds: 45  # Max processing time of 45 seconds per batch
+action:
+  # action configs
+```
+
+Note that the `datahub` configuration block is **required** to connect to your DataHub Cloud instance.
+
+<details>
+  <summary>View All Configuration Options</summary>
+
+| Field                                 | Required | Default                       | Description                                                                               |
+  | ------------------------------------- | :------: | :---------------------------: | ----------------------------------------------------------------------------------------- |
+| `topic`                               |    ❌    | `PlatformEvent_v1`            | The name of the topic from which events will be consumed. Do not change this unless you know what you're doing!                                |
+| `lookback_days`                       |    ❌    | None                           | Optional number of days to look back when polling for events.                             |
+| `reset_offsets`                       |    ❌    | `False`                       | When set to `True`, the consumer will ignore any stored offsets and start fresh.          |
+| `kill_after_idle_timeout`             |    ❌    | `False`                       | If `True`, stops the consumer after being idle for the specified timeout duration.        |
+| `idle_timeout_duration_seconds`       |    ❌    | `30`                          | Duration in seconds after which, if no events are received, the consumer is considered idle. |
+| `event_processing_time_max_duration_seconds` | ❌  | `30`                          | Maximum allowed time in seconds for processing events before timing out.                  |
+</details>
+
+
+## FAQ
+
+1. Is there a way to always start processing from the end of the topics on Actions start?
+
+Yes, simply set `reset_offsets` to True for a single run of the action. Remember to disable this for subsequent runs if you don't want to miss any events!
+
+2. What happens if I have multiple actions with the same pipeline `name` running? Can I scale out horizontally?
+
+Today, there is undefined behavior deploying multiple actions with the same name using the DataHub Cloud Events Source.
+All events must be processed by a single running action
+
diff --git a/docs/api/datahub-apis.md b/docs/api/datahub-apis.md
index c46aacde3a0cb5..62136406e6ff66 100644
--- a/docs/api/datahub-apis.md
+++ b/docs/api/datahub-apis.md
@@ -12,6 +12,10 @@ DataHub has several APIs to manipulate metadata on the platform. Here's the list
 In general, **Python and Java SDKs** are our most recommended tools for extending and customizing the behavior of your DataHub instance.
 We don't recommend using the **OpenAPI** directly, as it's more complex and less user-friendly than the other APIs.
 
+:::warning
+About async usage of APIs - DataHub's asynchronous APIs perform only basic schema validation when receiving MCP requests, similar to direct production to MCP Kafka topics. While requests must conform to the MCP schema to be accepted, actual processing happens later in the pipeline. Any processing failures that occur after the initial acceptance are captured in the Failed MCP topic, but these failures are not immediately surfaced to the API caller since they happen asynchronously.
+:::
+
 ## Python and Java SDK
 
 We offer an SDK for both Python and Java that provide full functionality when it comes to CRUD operations and any complex functionality you may want to build into DataHub. We recommend using the SDKs for most use cases. Here are the examples of how to use the SDKs:
diff --git a/docs/automations/bigquery-metadata-sync.md b/docs/automations/bigquery-metadata-sync.md
index 3d896f57d341a4..78bdbdd453e9f7 100644
--- a/docs/automations/bigquery-metadata-sync.md
+++ b/docs/automations/bigquery-metadata-sync.md
@@ -67,7 +67,7 @@ Ensure your service account has the following permissions:
     | Propagation Type | DataHub Entity | BigQuery Entity | Note |
     | -------- | ------- | ------- | ------- |
     | Table Tags as Labels | [Table Tag](https://datahubproject.io/docs/tags/) | [BigQuery Label](https://cloud.google.com/bigquery/docs/labels-intro) | - |
-    | Column Glossary Terms as Policy Tags | [Glossary Term on Table Column](https://datahubproject.io/docs/0.14.0/glossary/business-glossary/) | [Policy Tag](https://cloud.google.com/bigquery/docs/best-practices-policy-tags) | <ul><li>Assigned Policy tags are created under DataHub taxonomy.</li></ul><ul><li>Only the latest assigned glossary term set as policy tag. BigQuery only supports one assigned policy tag.</li></ul> <ul><li>Policy Tags are not synced to DataHub as glossary term from BigQuery.</li></ul>
+    | Column Glossary Terms as Policy Tags | [Glossary Term on Table Column](https://datahubproject.io/docs/glossary/business-glossary/) | [Policy Tag](https://cloud.google.com/bigquery/docs/best-practices-policy-tags) | <ul><li>Assigned Policy tags are created under DataHub taxonomy.</li></ul><ul><li>Only the latest assigned glossary term set as policy tag. BigQuery only supports one assigned policy tag.</li></ul> <ul><li>Policy Tags are not synced to DataHub as glossary term from BigQuery.</li></ul>
     | Table Descriptions | [Table Description](https://datahubproject.io/docs/api/tutorials/descriptions/) | Table Description | - |
     | Column Descriptions | [Column Description](https://datahubproject.io/docs/api/tutorials/descriptions/) | Column Description | - |
 
diff --git a/docs/cli.md b/docs/cli.md
index 1c38077d0d12ef..10f330a3ecfc6e 100644
--- a/docs/cli.md
+++ b/docs/cli.md
@@ -57,24 +57,30 @@ Options:
   --help                          Show this message and exit.
 
 Commands:
-  actions      <disabled due to missing dependencies>
-  check        Helper commands for checking various aspects of DataHub.
-  dataproduct  A group of commands to interact with the DataProduct entity in DataHub.
-  delete       Delete metadata from datahub using a single urn or a combination of filters
-  docker       Helper commands for setting up and interacting with a local DataHub instance using Docker.
-  exists       A group of commands to check existence of entities in DataHub.
-  get          A group of commands to get metadata from DataHub.
-  group        A group of commands to interact with the Group entity in DataHub.
-  ingest       Ingest metadata into DataHub.
-  init         Configure which datahub instance to connect to
-  lite         A group of commands to work with a DataHub Lite instance
-  migrate      Helper commands for migrating metadata within DataHub.
-  put          A group of commands to put metadata in DataHub.
-  state        Managed state stored in DataHub by stateful ingestion.
-  telemetry    Toggle telemetry.
-  timeline     Get timeline for an entity based on certain categories
-  user         A group of commands to interact with the User entity in DataHub.
-  version      Print version number and exit.
+  actions       <disabled due to missing dependencies>
+  assertions    A group of commands to interact with the Assertion entity in DataHub.
+  check         Helper commands for checking various aspects of DataHub.
+  container     A group of commands to interact with containers in DataHub.
+  datacontract  A group of commands to interact with the DataContract entity in DataHub.
+  dataproduct   A group of commands to interact with the DataProduct entity in DataHub.
+  dataset       A group of commands to interact with the Dataset entity in DataHub.
+  delete        Delete metadata from DataHub.
+  docker        Helper commands for setting up and interacting with a local DataHub instance using Docker.
+  exists        A group of commands to check existence of entities in DataHub.
+  forms         A group of commands to interact with forms in DataHub.
+  get           A group of commands to get metadata from DataHub.
+  group         A group of commands to interact with the Group entity in DataHub.
+  ingest        Ingest metadata into DataHub.
+  init          Configure which datahub instance to connect to
+  lite          A group of commands to work with a DataHub Lite instance
+  migrate       Helper commands for migrating metadata within DataHub.
+  properties    A group of commands to interact with structured properties in DataHub.
+  put           A group of commands to put metadata in DataHub.
+  state         Managed state stored in DataHub by stateful ingestion.
+  telemetry     Toggle telemetry.
+  timeline      Get timeline for an entity based on certain categories
+  user          A group of commands to interact with the User entity in DataHub.
+  version       Print version number and exit.
 ```
 
 The following top-level commands listed below are here mainly to give the reader a high-level picture of what are the kinds of things you can accomplish with the cli.
@@ -115,17 +121,17 @@ datahub ingest -c ./examples/recipes/example_to_datahub_rest.dhub.yaml --dry-run
 datahub ingest -c ./examples/recipes/example_to_datahub_rest.dhub.yaml -n
 ```
 
-#### ingest --list-source-runs
+#### ingest list-source-runs
 
-The `--list-source-runs` option of the `ingest` command lists the previous runs, displaying their run ID, source name, 
+The `list-source-runs` option of the `ingest` command lists the previous runs, displaying their run ID, source name, 
 start time, status, and source URN. This command allows you to filter results using the --urn option for URN-based 
 filtering or the --source option to filter by source name (partial or complete matches are supported).
 
 ```shell
 # List all ingestion runs
-datahub ingest --list-source-runs
+datahub ingest list-source-runs
 # Filter runs by a source name containing "demo"
-datahub ingest --list-source-runs --source "demo"
+datahub ingest list-source-runs --source "demo"
 ```
 
 #### ingest --preview
@@ -274,6 +280,18 @@ DATAHUB_TELEMETRY_TIMEOUT=10
 DATAHUB_DEBUG=false
 ```
 
+### container
+
+A group of commands to interact with containers in DataHub.
+
+e.g. You can use this to apply a tag to all datasets recursively in this container.
+```shell
+datahub container tag --container-urn "urn:li:container:0e9e46bd6d5cf645f33d5a8f0254bc2d" --tag-urn "urn:li:tag:tag1"
+datahub container domain --container-urn "urn:li:container:3f2effd1fbe154a4d60b597263a41e41" --domain-urn  "urn:li:domain:ajsajo-b832-4ab3-8881-7ed5e991a44c"
+datahub container owner --container-urn "urn:li:container:3f2effd1fbe154a4d60b597263a41e41" --owner-urn  "urn:li:corpGroup:eng@example.com"
+datahub container term --container-urn "urn:li:container:3f2effd1fbe154a4d60b597263a41e41" --term-urn  "urn:li:term:PII"
+```
+
 ### check
 
 The datahub package is composed of different plugins that allow you to connect to different metadata sources and ingest metadata from them.
diff --git a/docs/how/updating-datahub.md b/docs/how/updating-datahub.md
index eb5a792216d981..945a00088f1f74 100644
--- a/docs/how/updating-datahub.md
+++ b/docs/how/updating-datahub.md
@@ -1,8 +1,3 @@
-# Known Issues
-
-- Async APIs - DataHub's asynchronous APIs perform only basic schema validation when receiving MCP requests, similar to direct production to MCP Kafka topics. While requests must conform to the MCP schema to be accepted, actual processing happens later in the pipeline. Any processing failures that occur after the initial acceptance are captured in the Failed MCP topic, but these failures are not immediately surfaced to the API caller since they happen asynchronously.
-
-
 # Updating DataHub
 
 <!--
@@ -22,37 +17,28 @@
 This file documents any backwards-incompatible changes in DataHub and assists people when migrating to a new version.
 
 ## Next
-- #12191 - Configs `include_view_lineage` and `include_view_column_lineage` are removed from snowflake ingestion source. View and External Table DDL lineage will always be ingested when definitions are available.
-- #12181 - Configs `include_view_lineage`, `include_view_column_lineage` and `lineage_parse_view_ddl` are removed from bigquery ingestion source. View and Snapshot lineage will always be ingested when definitions are available.
-- #11560 - The PowerBI ingestion source configuration option include_workspace_name_in_dataset_urn determines whether the workspace name is included in the PowerBI dataset's URN.<br/> PowerBI allows to have identical name of semantic model and their tables across the workspace, It will overwrite the semantic model in-case of multi-workspace ingestion.<br/>
-  Entity urn with `include_workspace_name_in_dataset_urn: false`
 
-  ```
-   urn:li:dataset:(urn:li:dataPlatform:powerbi,[<PlatformInstance>.]<SemanticModelName>.<TableName>,<ENV>)
-  ```
+### Breaking Changes
 
-  Entity urn with `include_workspace_name_in_dataset_urn: true`
+- #12408: The `platform` field in the DataPlatformInstance GraphQL type is removed. Clients need to retrieve the platform via the optional `dataPlatformInstance` field.
 
-  ```
-   urn:li:dataset:(urn:li:dataPlatform:powerbi,[<PlatformInstance>.].<WorkspaceName>.<SemanticModelName>.<TableName>,<ENV>)
-  ```
+### Potential Downtime
 
-  The config `include_workspace_name_in_dataset_urn` is default to `false` for backward compatibility, However, we recommend enabling this flag after performing the necessary cleanup.
-  If stateful ingestion is enabled, running ingestion with the latest CLI version will handle the cleanup automatically. Otherwise, we recommend soft deleting all powerbi data via the DataHub CLI:
-  `datahub delete --platform powerbi --soft` and then re-ingest with the latest CLI version, ensuring the `include_workspace_name_in_dataset_urn` configuration is set to true.
+### Deprecations
 
-- #11701: The Fivetran `sources_to_database` field is deprecated in favor of setting directly within `sources_to_platform_instance.<key>.database`.
-- #11742: For PowerBi ingestion, `use_powerbi_email` is now enabled by default when extracting ownership information.
-- #12056: The DataHub Airflow plugin no longer supports Airflow 2.1 and Airflow 2.2.
-- #12056: The DataHub Airflow plugin now defaults to the v2 plugin implementation.
-- #12077: `Kafka` source no longer ingests schemas from schema registry as separate entities by default, set `ingest_schemas_as_entities` to `true` to ingest them
-- OpenAPI Update: PIT Keep Alive parameter added to scroll. NOTE: This parameter requires the `pointInTimeCreationEnabled` feature flag to be enabled and the `elasticSearch.implementation` configuration to be `elasticsearch`. This feature is not supported for OpenSearch at this time and the parameter will not be respected without both of these set.
+### Other Notable Changes
+
+## 0.15.0
+
+- OpenAPI Update: PIT Keep Alive parameter added to scroll endpoints. NOTE: This parameter requires the `pointInTimeCreationEnabled` feature flag to be enabled and the `elasticSearch.implementation` configuration to be `elasticsearch`. This feature is not supported for OpenSearch at this time and the parameter will not be respected without both of these set.
 - OpenAPI Update 2: Previously there was an incorrectly marked parameter named `sort` on the generic list entities endpoint for v3. This parameter is deprecated and only supports a single string value while the documentation indicates it supports a list of strings. This documentation error has been fixed and the correct field, `sortCriteria`, is now documented which supports a list of strings. 
-- #12223: For dbt Cloud ingestion, the "View in dbt" link will point at the "Explore" page in the dbt Cloud UI. You can revert to the old behavior of linking to the dbt Cloud IDE by setting `external_url_mode: ide".
-- #12236: Data flow and data job entities may additionally produce container aspect that will require a corresponding upgrade of server. Otherwise server can reject the aspect.
 
 ### Breaking Changes
 
+- #12223: For dbt Cloud ingestion, the "View in dbt" link will point at the "Explore" page in the dbt Cloud UI. You can revert to the old behavior of linking to the dbt Cloud IDE by setting `external_url_mode: ide".
+- #12191 - Configs `include_view_lineage` and `include_view_column_lineage` are removed from snowflake ingestion source. View and External Table DDL lineage will always be ingested when definitions are available.
+- #12181 - Configs `include_view_lineage`, `include_view_column_lineage` and `lineage_parse_view_ddl` are removed from bigquery ingestion source. View and Snapshot lineage will always be ingested when definitions are available.
+- #12077: `Kafka` source no longer ingests schemas from schema registry as separate entities by default, set `ingest_schemas_as_entities` to `true` to ingest them
 - #11486 - Criterion's `value` parameter has been previously deprecated. Use of `value` instead of `values` is no longer supported and will be completely removed on the next major version.
 - #11484 - Metadata service authentication enabled by default
 - #11484 - Rest API authorization enabled by default
@@ -78,8 +64,30 @@ This file documents any backwards-incompatible changes in DataHub and assists pe
 
 ### Deprecations
 
+- #12056: The DataHub Airflow plugin no longer supports Airflow 2.1 and Airflow 2.2.
+- #11701: The Fivetran `sources_to_database` field is deprecated in favor of setting directly within `sources_to_platform_instance.<key>.database`.
+- #11560 - The PowerBI ingestion source configuration option include_workspace_name_in_dataset_urn determines whether the workspace name is included in the PowerBI dataset's URN.<br/> PowerBI allows to have identical name of semantic model and their tables across the workspace, It will overwrite the semantic model in-case of multi-workspace ingestion.<br/>
+  Entity urn with `include_workspace_name_in_dataset_urn: false`
+
+  ```
+   urn:li:dataset:(urn:li:dataPlatform:powerbi,[<PlatformInstance>.]<SemanticModelName>.<TableName>,<ENV>)
+  ```
+
+  Entity urn with `include_workspace_name_in_dataset_urn: true`
+
+  ```
+   urn:li:dataset:(urn:li:dataPlatform:powerbi,[<PlatformInstance>.].<WorkspaceName>.<SemanticModelName>.<TableName>,<ENV>)
+  ```
+
+  The config `include_workspace_name_in_dataset_urn` is default to `false` for backward compatibility, However, we recommend enabling this flag after performing the necessary cleanup.
+  If stateful ingestion is enabled, running ingestion with the latest CLI version will handle the cleanup automatically. Otherwise, we recommend soft deleting all powerbi data via the DataHub CLI:
+  `datahub delete --platform powerbi --soft` and then re-ingest with the latest CLI version, ensuring the `include_workspace_name_in_dataset_urn` configuration is set to true.
+
 ### Other Notable Changes
 
+- #12236: Data flow and data job entities may additionally produce container aspect that will require a corresponding upgrade of server. Otherwise server can reject the aspect.
+- #12056: The DataHub Airflow plugin now defaults to the v2 plugin implementation.
+- #11742: For PowerBi ingestion, `use_powerbi_email` is now enabled by default when extracting ownership information.
 - #11549 - Manage Operations Privilege is extended from throttle control to all system management and operations APIs.
 
 ## 0.14.1
diff --git a/docs/managed-datahub/datahub-api/entity-events-api.md b/docs/managed-datahub/datahub-api/entity-events-api.md
index e59f1650c7d766..377f2fd01e813b 100644
--- a/docs/managed-datahub/datahub-api/entity-events-api.md
+++ b/docs/managed-datahub/datahub-api/entity-events-api.md
@@ -15,6 +15,7 @@ The Events API allows you to integrate changes happening on the DataHub Metadata
 ### Supported Integrations
 
 * [AWS EventBridge](docs/managed-datahub/operator-guide/setting-up-events-api-on-aws-eventbridge.md)
+* [DataHub Cloud Event Source](docs/actions/sources/datahub-cloud-event-source.md)
 
 ### Use Cases
 
diff --git a/docs/managed-datahub/release-notes/v_0_3_8.md b/docs/managed-datahub/release-notes/v_0_3_8.md
new file mode 100644
index 00000000000000..75bba525688d89
--- /dev/null
+++ b/docs/managed-datahub/release-notes/v_0_3_8.md
@@ -0,0 +1,81 @@
+# v0.3.8
+---
+
+Release Availability Date
+---
+21-Jan-2025
+
+Recommended CLI/SDK
+---
+0.15.0.1
+
+## Known Issues
+
+- Async APIs - DataHub's asynchronous APIs perform only basic schema validation when receiving MCP requests, similar to direct production to MCP Kafka topics. While requests must conform to the MCP schema to be accepted, actual processing happens later in the pipeline. Any processing failures that occur after the initial acceptance are captured in the Failed MCP topic, but these failures are not immediately surfaced to the API caller since they happen asynchronously.
+
+## Release Changelog
+---
+
+### v0.3.8
+
+- All changes in https://github.com/datahub-project/datahub/releases/tag/v0.14.1
+    - Note Breaking Changes: https://datahubproject.io/docs/how/updating-datahub/#0141
+
+- Breaking Changes
+    - Soft-deleted and SCIM suspended users are no longer able to access the UI/APIs with a pre-existing valid token
+
+- Bug Fixes
+  - Ingestion page sorts across all sources rather than just the ones loaded on the page.
+  - Data Products Preview cards now render correctly in scenarios where there are several fully filled out products in a list
+  - Snowflake tag propagation looks for case-insensitive table match on tables with special characters
+  - Lineage patch properly encodes special characters
+  - Fixed multiple bugs around lineage visualization's handling of graph cycles
+  - Fixed lineage filter node search on schema field level lineage
+  - Search csv export properly includes data products
+  - Properly show lineage to an upstream dbt source that has no further upstreams
+  - Soft-deleted & suspended users with valid tokens are rejected
+  - Fixes scroll behavior on quality tab with small screens
+  - Add backwards compatibility to the UI for old policy filters in the UI
+  - Fixes nesting logic in Properties tab for structured properties that prefix matched each other's qualified names
+  - Show error when editing help link if there is an error
+  - Fix UI bug on schema table when clicking hides part of table
+  - Fix rendering schema field references as structured property values to show field name, not urn
+  - Partially fix a bug where dbt sync back would generate invalid YAML. This won't happen anymore, but we might miss some changes because of it
+  - Stability and error message improvements for AI docs generation
+  - Various quality and performance improvements to AI classification
+  - Fix searching for nested columns in the Columns tab of Dataset profile page
+  - Remote Executor: make cron schedules in ingestions/assertions scheduler POSIX-compliant.
+  - Remote Executor: limit the number of automated restarts of aborted tasks.
+  - Remote Executor: always update task status synchronously.
+
+
+- Product
+    - [BETA] Introducing the new Navigation Bar / App redesign! Starting with this release, you can try out a new look and feel to your DataHub app. To enable this BETA feature, reach out to your Acryl representative.
+    - [BETA] Add ability to create tag propagation automations to downstream and sibling assets via the UI in the Automations tab
+    - Support the Notes feature for all entity types with profile pages
+    - Test notifications button is now easier to discover.
+    - UX polishes to the Assertions Assertions tab
+    - UX polishes to the Column sidebar
+    - UX polishes to column and table deprecation, including deprecation visible in lineage graph
+    - Add the ability to do AND for filters of a given type. For example, we now allow you to filter for assets that have tag1 AND tag2
+    - Combine sibling entities on home page insights
+    - Hardened structured properties feature by supporting broken references and adding better validation to the ID/qualified name
+    - Hardened forms feature by supporting invalid / unresolvable references
+    - Recurrent metadata sharing: Any assets shared between DataHub assets will be re-shared automatically once per day. 
+  
+
+- Platform
+    - Optional config to reduce Slack app permissions
+    - Support batch processing for MCP consumption to increase throughput of the Kafka consumer significantly
+    - Resumability for DataHub Upgrade jobs
+    - Async actions for Metadata Tests corner case where some actions could be missed has been fixed
+    - Structured Properties now have a stricter validation for no spaces in the qualified name
+    - Urn Validation logic has been improved to reduce delayed error conditions
+    - Database transaction and locking updates to fix conditions where retryable persistence exceptions would cause ingestion batches to fail
+    - Fix for Metadata Tests to support primitive PDL Typeref fields like timestamps
+    - Entity Versioning initial feature support, disabled by default and can be turned on via feature flag
+    - OpenAPI fixes for some faulty parameters
+    - Improved memory performance of the search ranking system ingestion source
+    - Remote Executor: Added registration/discovery mechanism.
+    - Remote Executor: Added weight-based task scheduling to allow variable concurrency for memory-greedy tasks.
+
diff --git a/gradle/versioning/versioning.gradle b/gradle/versioning/versioning.gradle
index 39a8a3faf80117..2ffb085869f7b2 100644
--- a/gradle/versioning/versioning.gradle
+++ b/gradle/versioning/versioning.gradle
@@ -21,7 +21,7 @@ Produces the following variables and supports token replacement
 import org.apache.tools.ant.filters.ReplaceTokens
 
 def detailedVersionString = "0.0.0-unknown-SNAPSHOT"
-def cliMajorVersion = "0.12.0" // base default cli major version
+def cliMajorVersion = "0.15.0" // base default cli major version
 def snapshotVersion = false
 if (project.hasProperty("releaseVersion")) {
   version = releaseVersion
diff --git a/li-utils/src/main/javaPegasus/com/linkedin/common/urn/DataPlatformInstanceUrn.java b/li-utils/src/main/javaPegasus/com/linkedin/common/urn/DataPlatformInstanceUrn.java
new file mode 100644
index 00000000000000..dfce6dc1e51085
--- /dev/null
+++ b/li-utils/src/main/javaPegasus/com/linkedin/common/urn/DataPlatformInstanceUrn.java
@@ -0,0 +1,79 @@
+package com.linkedin.common.urn;
+
+import com.linkedin.data.template.Custom;
+import com.linkedin.data.template.DirectCoercer;
+import com.linkedin.data.template.TemplateOutputCastException;
+import java.net.URISyntaxException;
+
+public final class DataPlatformInstanceUrn extends Urn {
+
+  public static final String ENTITY_TYPE = "dataPlatformInstance";
+
+  private final DataPlatformUrn _platform;
+  private final String _instanceId;
+
+  public DataPlatformInstanceUrn(DataPlatformUrn platform, String instanceId) {
+    super(ENTITY_TYPE, TupleKey.create(platform, instanceId));
+    this._platform = platform;
+    this._instanceId = instanceId;
+  }
+
+  public DataPlatformUrn getPlatformEntity() {
+    return _platform;
+  }
+
+  public String getInstance() {
+    return _instanceId;
+  }
+
+  public static DataPlatformInstanceUrn createFromString(String rawUrn) throws URISyntaxException {
+    return createFromUrn(Urn.createFromString(rawUrn));
+  }
+
+  public static DataPlatformInstanceUrn createFromUrn(Urn urn) throws URISyntaxException {
+    if (!"li".equals(urn.getNamespace())) {
+      throw new URISyntaxException(urn.toString(), "Urn namespace type should be 'li'.");
+    } else if (!ENTITY_TYPE.equals(urn.getEntityType())) {
+      throw new URISyntaxException(
+          urn.toString(), "Urn entity type should be 'dataPlatformInstance'.");
+    } else {
+      TupleKey key = urn.getEntityKey();
+      if (key.size() != 2) {
+        throw new URISyntaxException(urn.toString(), "Invalid number of keys.");
+      } else {
+        try {
+          return new DataPlatformInstanceUrn(
+              (DataPlatformUrn) key.getAs(0, DataPlatformUrn.class),
+              (String) key.getAs(1, String.class));
+        } catch (Exception e) {
+          throw new URISyntaxException(urn.toString(), "Invalid URN Parameter: '" + e.getMessage());
+        }
+      }
+    }
+  }
+
+  public static DataPlatformInstanceUrn deserialize(String rawUrn) throws URISyntaxException {
+    return createFromString(rawUrn);
+  }
+
+  static {
+    Custom.initializeCustomClass(DataPlatformUrn.class);
+    Custom.initializeCustomClass(DataPlatformInstanceUrn.class);
+    Custom.registerCoercer(
+        new DirectCoercer<DataPlatformInstanceUrn>() {
+          public Object coerceInput(DataPlatformInstanceUrn object) throws ClassCastException {
+            return object.toString();
+          }
+
+          public DataPlatformInstanceUrn coerceOutput(Object object)
+              throws TemplateOutputCastException {
+            try {
+              return DataPlatformInstanceUrn.createFromString((String) object);
+            } catch (URISyntaxException e) {
+              throw new TemplateOutputCastException("Invalid URN syntax: " + e.getMessage(), e);
+            }
+          }
+        },
+        DataPlatformInstanceUrn.class);
+  }
+}
diff --git a/li-utils/src/main/pegasus/com/linkedin/common/DataPlatformInstanceUrn.pdl b/li-utils/src/main/pegasus/com/linkedin/common/DataPlatformInstanceUrn.pdl
new file mode 100644
index 00000000000000..168e0ee7611d31
--- /dev/null
+++ b/li-utils/src/main/pegasus/com/linkedin/common/DataPlatformInstanceUrn.pdl
@@ -0,0 +1,27 @@
+namespace com.linkedin.common
+
+/**
+ * Standardized dataset identifier.
+ */
+@java.class = "com.linkedin.common.urn.DataPlatformInstanceUrn"
+@validate.`com.linkedin.common.validator.TypedUrnValidator` = {
+  "accessible" : true,
+  "owningTeam" : "urn:li:internalTeam:datahub",
+  "entityType" : "dataPlatformInstance",
+  "constructable" : true,
+  "namespace" : "li",
+  "name" : "DataPlatformInstance",
+  "doc" : "Standardized data platform instance identifier.",
+  "owners" : [ "urn:li:corpuser:fbar", "urn:li:corpuser:bfoo" ],
+  "fields" : [ {
+    "type" : "com.linkedin.common.urn.DataPlatformUrn",
+    "name" : "platform",
+    "doc" : "Standardized platform urn."
+  }, {
+    "name" : "instance",
+    "doc" : "Instance of the data platform (e.g. db instance)",
+    "type" : "string",
+  } ],
+  "maxLength" : 100
+}
+typeref DataPlatformInstanceUrn = string
diff --git a/metadata-ingestion-modules/airflow-plugin/.gitignore b/metadata-ingestion-modules/airflow-plugin/.gitignore
index d0108e8361a060..6801b785ea1e4b 100644
--- a/metadata-ingestion-modules/airflow-plugin/.gitignore
+++ b/metadata-ingestion-modules/airflow-plugin/.gitignore
@@ -1,5 +1,4 @@
 .envrc
-src/datahub_airflow_plugin/__init__.py.bak
 .vscode/
 output
 pvenv36/
diff --git a/metadata-ingestion-modules/airflow-plugin/build.gradle b/metadata-ingestion-modules/airflow-plugin/build.gradle
index 68a35c0dfc417b..95b4ee3118f03e 100644
--- a/metadata-ingestion-modules/airflow-plugin/build.gradle
+++ b/metadata-ingestion-modules/airflow-plugin/build.gradle
@@ -73,18 +73,15 @@ task lint(type: Exec, dependsOn: installDev) {
   commandLine 'bash', '-c',
     "find ${venv_name}/lib -path *airflow/_vendor/connexion/spec.py -exec sed -i.bak -e '169,169s/  # type: List\\[str\\]//g' {} \\; && " +
     "source ${venv_name}/bin/activate && set -x && " +
-    "black --check --diff src/ tests/ && " +
-    "isort --check --diff src/ tests/ && " +
-    "flake8 --count --statistics src/ tests/ && " +
+    "ruff check src/ tests/ && " +
+    "ruff format --check src/ tests/ && " +
     "mypy --show-traceback --show-error-codes src/ tests/"
 }
 task lintFix(type: Exec, dependsOn: installDev) {
   commandLine 'bash', '-c',
     "source ${venv_name}/bin/activate && set -x && " +
-    "black src/ tests/ && " +
-    "isort src/ tests/ && " +
-    "flake8 src/ tests/ && " +
-    "mypy src/ tests/ "
+    "ruff check --fix src/ tests/ && " +
+    "ruff format src/ tests/ "
 }
 
 // HACK: Some of the Airflow constraint files conflict with packages that we install (e.g. black).
@@ -121,5 +118,8 @@ clean {
   delete venv_name
   delete 'build'
   delete 'dist'
+  delete '.ruff_cache'
+  delete '.mypy_cache'
+  delete '.pytest_cache'
 }
 clean.dependsOn cleanPythonCache
diff --git a/metadata-ingestion-modules/airflow-plugin/pyproject.toml b/metadata-ingestion-modules/airflow-plugin/pyproject.toml
index 648040c1951db8..d1e1d0ad479442 100644
--- a/metadata-ingestion-modules/airflow-plugin/pyproject.toml
+++ b/metadata-ingestion-modules/airflow-plugin/pyproject.toml
@@ -2,19 +2,58 @@
 build-backend = "setuptools.build_meta"
 requires = ["setuptools>=54.0.0", "wheel", "pip>=21.0.0"]
 
-[tool.black]
-extend-exclude = '''
-# A regex preceded with ^/ will apply only to files and directories
-# in the root of the project.
-^/tmp
-'''
-include = '\.pyi?$'
+[tool.ruff]
+line-length = 88
+target-version = "py38"
+exclude = [
+    ".git",
+    "venv",
+    ".tox",
+    "__pycache__",
+]
 
-[tool.isort]
-indent = '    '
-known_future_library = ['__future__', 'datahub.utilities._markupsafe_compat', 'datahub_provider._airflow_compat']
-profile = 'black'
-sections = 'FUTURE,STDLIB,THIRDPARTY,FIRSTPARTY,LOCALFOLDER'
+[tool.ruff.format]
+quote-style = "double"
+indent-style = "space"
+skip-magic-trailing-comma = false
+line-ending = "auto"
 
-[tool.pyright]
-extraPaths = ['tests']
\ No newline at end of file
+[tool.ruff.lint.isort]
+combine-as-imports = true
+known-first-party = ["datahub"]
+extra-standard-library = ["__future__", "datahub.utilities._markupsafe_compat", "datahub.sql_parsing._sqlglot_patch"]
+section-order = ["future", "standard-library", "third-party", "first-party", "local-folder"]
+force-sort-within-sections = false
+force-wrap-aliases = false
+split-on-trailing-comma = false
+order-by-type = true
+relative-imports-order = "closest-to-furthest"
+force-single-line = false
+single-line-exclusions = ["typing"]
+length-sort = false
+from-first = false
+required-imports = []
+classes = ["typing"]
+
+[tool.ruff.lint]
+extend-select = [
+    "B",    # flake8-bugbear
+    "C90",  # mccabe complexity
+    "E",    # pycodestyle errors
+    "F",    # pyflakes
+    "G010", # logging.warn -> logging.warning
+    "I",    # isort
+    "TID",  # flake8-tidy-imports
+]
+ignore = [
+    "E501",  # Line length violations (handled by formatter)
+]
+
+[tool.ruff.lint.mccabe]
+max-complexity = 15
+
+[tool.ruff.lint.flake8-tidy-imports]
+ban-relative-imports = "all"
+
+[tool.ruff.lint.per-file-ignores]
+"__init__.py" = ["F401"]
\ No newline at end of file
diff --git a/metadata-ingestion-modules/airflow-plugin/scripts/release.sh b/metadata-ingestion-modules/airflow-plugin/scripts/release.sh
index 8f23f72082c2ca..994c1ae145ce55 100755
--- a/metadata-ingestion-modules/airflow-plugin/scripts/release.sh
+++ b/metadata-ingestion-modules/airflow-plugin/scripts/release.sh
@@ -1,26 +1,31 @@
 #!/bin/bash
+# Auto-generated by python-build/generate_release_scripts.py. Do not edit manually.
+
 set -euxo pipefail
 
+ROOT=../..
+MODULE=datahub_airflow_plugin
+
 if [[ ! ${RELEASE_SKIP_TEST:-} ]] && [[ ! ${RELEASE_SKIP_INSTALL:-} ]]; then
-	../../gradlew build  # also runs tests
+    ${ROOT}/gradlew build  # also runs tests
 elif [[ ! ${RELEASE_SKIP_INSTALL:-} ]]; then
-	../../gradlew install
+    ${ROOT}/gradlew install
 fi
 
-MODULE=datahub_airflow_plugin
-
 # Check packaging constraint.
 python -c 'import setuptools; where="./src"; assert setuptools.find_packages(where) == setuptools.find_namespace_packages(where), "you seem to be missing or have extra __init__.py files"'
-if [[ ${RELEASE_VERSION:-} ]]; then
-    # Replace version with RELEASE_VERSION env variable
-    sed -i.bak "s/__version__ = \"1\!0.0.0.dev0\"/__version__ = \"$(echo $RELEASE_VERSION|sed s/-/+/)\"/" src/${MODULE}/__init__.py
-else
-    vim src/${MODULE}/__init__.py
+
+# Update the release version.
+if [[ ! ${RELEASE_VERSION:-} ]]; then
+    echo "RELEASE_VERSION is not set"
+    exit 1
 fi
+sed -i.bak "s/__version__ = .*$/__version__ = \"$(echo $RELEASE_VERSION|sed s/-/+/)\"/" src/${MODULE}/_version.py
 
+# Build and upload the release.
 rm -rf build dist || true
 python -m build
 if [[ ! ${RELEASE_SKIP_UPLOAD:-} ]]; then
     python -m twine upload 'dist/*'
 fi
-mv src/${MODULE}/__init__.py.bak src/${MODULE}/__init__.py
+mv src/${MODULE}/_version.py.bak src/${MODULE}/_version.py
diff --git a/metadata-ingestion-modules/airflow-plugin/setup.cfg b/metadata-ingestion-modules/airflow-plugin/setup.cfg
index c25256c5751b8d..abb9040ab3535a 100644
--- a/metadata-ingestion-modules/airflow-plugin/setup.cfg
+++ b/metadata-ingestion-modules/airflow-plugin/setup.cfg
@@ -1,24 +1,3 @@
-[flake8]
-max-complexity = 15
-ignore =
-    # Ignore: line length issues, since black's formatter will take care of them.
-    E501,
-    # Ignore: 1 blank line required before class docstring.
-    D203,
-    # See https://stackoverflow.com/a/57074416.
-    W503,
-    # See https://github.com/psf/black/issues/315.
-    E203
-exclude =
-    .git,
-    venv,
-    .tox,
-    __pycache__
-per-file-ignores =
-    # imported but unused
-    __init__.py: F401
-ban-relative-imports = true
-
 [mypy]
 plugins =
     sqlmypy,
diff --git a/metadata-ingestion-modules/airflow-plugin/setup.py b/metadata-ingestion-modules/airflow-plugin/setup.py
index d07063dbffc5c4..d03ed824c9a261 100644
--- a/metadata-ingestion-modules/airflow-plugin/setup.py
+++ b/metadata-ingestion-modules/airflow-plugin/setup.py
@@ -5,7 +5,7 @@
 import setuptools
 
 package_metadata: dict = {}
-with open("./src/datahub_airflow_plugin/__init__.py") as fp:
+with open("./src/datahub_airflow_plugin/_version.py") as fp:
     exec(fp.read(), package_metadata)
 
 
@@ -71,11 +71,8 @@ def get_long_description():
 dev_requirements = {
     *base_requirements,
     *mypy_stubs,
-    "black==22.12.0",
     "coverage>=5.1",
-    "flake8>=3.8.3",
-    "flake8-tidy-imports>=4.3.0",
-    "isort>=5.7.0",
+    "ruff==0.9.2",
     "mypy==1.10.1",
     # pydantic 1.8.2 is incompatible with mypy 0.910.
     # See https://github.com/samuelcolvin/pydantic/pull/3175#issuecomment-995382910.
diff --git a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/__init__.py b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/__init__.py
index e4040e3a17dfdc..7743c8ab2bab1a 100644
--- a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/__init__.py
+++ b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/__init__.py
@@ -1,23 +1,12 @@
-# Published at https://pypi.org/project/acryl-datahub/.
-__package_name__ = "acryl-datahub-airflow-plugin"
-__version__ = "1!0.0.0.dev0"
+from datahub_airflow_plugin._version import __package_name__, __version__
 
 
-def is_dev_mode() -> bool:
-    return __version__.endswith("dev0")
-
-
-def nice_version_name() -> str:
-    if is_dev_mode():
-        return "unavailable (installed in develop mode)"
-    return __version__
-
-
-def get_provider_info():
+def get_provider_info() -> dict:
+    # Register our hooks with Airflow.
     return {
         "package-name": f"{__package_name__}",
         "name": f"{__package_name__}",
-        "description": "Datahub metadata collector plugin",
+        "description": "DataHub metadata collector plugin",
         "connection-types": [
             {
                 "hook-class-name": "datahub_airflow_plugin.hooks.datahub.DatahubRestHook",
diff --git a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/_config.py b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/_config.py
index c4964712cf9f7d..6d6ba601556788 100644
--- a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/_config.py
+++ b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/_config.py
@@ -1,11 +1,12 @@
 from enum import Enum
 from typing import TYPE_CHECKING, Optional
 
-import datahub.emitter.mce_builder as builder
 from airflow.configuration import conf
-from datahub.configuration.common import AllowDenyPattern, ConfigModel
 from pydantic.fields import Field
 
+import datahub.emitter.mce_builder as builder
+from datahub.configuration.common import AllowDenyPattern, ConfigModel
+
 if TYPE_CHECKING:
     from datahub_airflow_plugin.hooks.datahub import DatahubGenericHook
 
diff --git a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/_datahub_ol_adapter.py b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/_datahub_ol_adapter.py
index 69de61aced0a59..72cdcd8813252a 100644
--- a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/_datahub_ol_adapter.py
+++ b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/_datahub_ol_adapter.py
@@ -1,8 +1,9 @@
 import logging
 
-import datahub.emitter.mce_builder as builder
 from openlineage.client.run import Dataset as OpenLineageDataset
 
+import datahub.emitter.mce_builder as builder
+
 logger = logging.getLogger(__name__)
 
 
diff --git a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/_extractors.py b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/_extractors.py
index 28d5775f61f542..5904ce1e9e978c 100644
--- a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/_extractors.py
+++ b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/_extractors.py
@@ -3,17 +3,11 @@
 import unittest.mock
 from typing import TYPE_CHECKING, Optional
 
-import datahub.emitter.mce_builder as builder
-from datahub.ingestion.source.sql.sqlalchemy_uri_mapper import (
-    get_platform_from_sqlalchemy_uri,
+from openlineage.airflow.extractors import (
+    BaseExtractor,
+    ExtractorManager as OLExtractorManager,
+    TaskMetadata,
 )
-from datahub.sql_parsing.sqlglot_lineage import (
-    SqlParsingResult,
-    create_lineage_sql_parsed_result,
-)
-from openlineage.airflow.extractors import BaseExtractor
-from openlineage.airflow.extractors import ExtractorManager as OLExtractorManager
-from openlineage.airflow.extractors import TaskMetadata
 from openlineage.airflow.extractors.snowflake_extractor import SnowflakeExtractor
 from openlineage.airflow.extractors.sql_extractor import SqlExtractor
 from openlineage.airflow.utils import get_operator_class, try_import_from_string
@@ -23,11 +17,20 @@
     SqlJobFacet,
 )
 
+import datahub.emitter.mce_builder as builder
+from datahub.ingestion.source.sql.sqlalchemy_uri_mapper import (
+    get_platform_from_sqlalchemy_uri,
+)
+from datahub.sql_parsing.sqlglot_lineage import (
+    SqlParsingResult,
+    create_lineage_sql_parsed_result,
+)
 from datahub_airflow_plugin._airflow_shims import Operator
 from datahub_airflow_plugin._datahub_ol_adapter import OL_SCHEME_TWEAKS
 
 if TYPE_CHECKING:
     from airflow.models import DagRun, TaskInstance
+
     from datahub.ingestion.graph.client import DataHubGraph
 
 logger = logging.getLogger(__name__)
@@ -60,9 +63,9 @@ def __init__(self):
 
         self.task_to_extractor.extractors["AthenaOperator"] = AthenaOperatorExtractor
 
-        self.task_to_extractor.extractors[
-            "BigQueryInsertJobOperator"
-        ] = BigQueryInsertJobOperatorExtractor
+        self.task_to_extractor.extractors["BigQueryInsertJobOperator"] = (
+            BigQueryInsertJobOperatorExtractor
+        )
 
         self._graph: Optional["DataHubGraph"] = None
 
diff --git a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/_version.py b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/_version.py
new file mode 100644
index 00000000000000..efda3f6bf31247
--- /dev/null
+++ b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/_version.py
@@ -0,0 +1,3 @@
+# Published at https://pypi.org/project/acryl-datahub-airflow-plugin/.
+__package_name__ = "acryl-datahub-airflow-plugin"
+__version__ = "1!0.0.0.dev0"
diff --git a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/client/airflow_generator.py b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/client/airflow_generator.py
index e9f93c0c1eab0a..c1ccdaeb0a1fbd 100644
--- a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/client/airflow_generator.py
+++ b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/client/airflow_generator.py
@@ -2,6 +2,7 @@
 from typing import TYPE_CHECKING, Dict, List, Optional, Set, Tuple, Union, cast
 
 from airflow.configuration import conf
+
 from datahub.api.entities.datajob import DataFlow, DataJob
 from datahub.api.entities.dataprocess.dataprocess_instance import (
     DataProcessInstance,
@@ -11,7 +12,6 @@
 from datahub.metadata.schema_classes import DataProcessTypeClass
 from datahub.utilities.urns.data_flow_urn import DataFlowUrn
 from datahub.utilities.urns.data_job_urn import DataJobUrn
-
 from datahub_airflow_plugin._airflow_compat import AIRFLOW_PATCHED
 from datahub_airflow_plugin._config import DatahubLineageConfig, DatajobUrl
 
diff --git a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/datahub_listener.py b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/datahub_listener.py
index 640991a90a1d28..b2ca61e3de3bf5 100644
--- a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/datahub_listener.py
+++ b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/datahub_listener.py
@@ -8,9 +8,13 @@
 from typing import TYPE_CHECKING, Callable, Dict, List, Optional, TypeVar, cast
 
 import airflow
-import datahub.emitter.mce_builder as builder
 from airflow.models import Variable
 from airflow.models.serialized_dag import SerializedDagModel
+from openlineage.airflow.listener import TaskHolder
+from openlineage.airflow.utils import redact_with_exclusions
+from openlineage.client.serde import Serde
+
+import datahub.emitter.mce_builder as builder
 from datahub.api.entities.datajob import DataJob
 from datahub.api.entities.dataprocess.dataprocess_instance import InstanceRunResult
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
@@ -30,10 +34,6 @@
 )
 from datahub.sql_parsing.sqlglot_lineage import SqlParsingResult
 from datahub.telemetry import telemetry
-from openlineage.airflow.listener import TaskHolder
-from openlineage.airflow.utils import redact_with_exclusions
-from openlineage.client.serde import Serde
-
 from datahub_airflow_plugin._airflow_shims import (
     HAS_AIRFLOW_DAG_LISTENER_API,
     HAS_AIRFLOW_DATASET_LISTENER_API,
@@ -286,9 +286,9 @@ def _extract_lineage(
         if sql_parsing_result:
             if error := sql_parsing_result.debug_info.error:
                 logger.info(f"SQL parsing error: {error}", exc_info=error)
-                datajob.properties[
-                    "datahub_sql_parser_error"
-                ] = f"{type(error).__name__}: {error}"
+                datajob.properties["datahub_sql_parser_error"] = (
+                    f"{type(error).__name__}: {error}"
+                )
             if not sql_parsing_result.debug_info.table_error:
                 input_urns.extend(sql_parsing_result.in_tables)
                 output_urns.extend(sql_parsing_result.out_tables)
diff --git a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/datahub_plugin.py b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/datahub_plugin.py
index 137cf97f69280a..2aeaaee4a542d4 100644
--- a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/datahub_plugin.py
+++ b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/datahub_plugin.py
@@ -4,20 +4,20 @@
 
 from airflow.plugins_manager import AirflowPlugin
 
-from datahub_airflow_plugin import __package_name__
 from datahub_airflow_plugin._airflow_compat import AIRFLOW_PATCHED
 from datahub_airflow_plugin._airflow_shims import (
     HAS_AIRFLOW_LISTENER_API,
     NEEDS_AIRFLOW_LISTENER_MODULE,
 )
+from datahub_airflow_plugin._version import __package_name__
 
 assert AIRFLOW_PATCHED
 logger = logging.getLogger(__name__)
 
 
-_USE_AIRFLOW_LISTENER_INTERFACE = HAS_AIRFLOW_LISTENER_API and not os.getenv(
+_USE_AIRFLOW_LISTENER_INTERFACE = HAS_AIRFLOW_LISTENER_API and os.getenv(
     "DATAHUB_AIRFLOW_PLUGIN_USE_V1_PLUGIN", "false"
-).lower() in ("true", "1")
+).lower() not in ("true", "1")
 
 if _USE_AIRFLOW_LISTENER_INTERFACE:
     try:
@@ -32,7 +32,7 @@
 
 
 with contextlib.suppress(Exception):
-    if not os.getenv("DATAHUB_AIRFLOW_PLUGIN_SKIP_FORK_PATCH", "false").lower() in (
+    if os.getenv("DATAHUB_AIRFLOW_PLUGIN_SKIP_FORK_PATCH", "false").lower() not in (
         "true",
         "1",
     ):
diff --git a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/datahub_plugin_v22.py b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/datahub_plugin_v22.py
index db47f37bed562e..99b0a40fd3c13e 100644
--- a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/datahub_plugin_v22.py
+++ b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/datahub_plugin_v22.py
@@ -7,9 +7,9 @@
 from airflow.lineage import PIPELINE_OUTLETS
 from airflow.models.baseoperator import BaseOperator
 from airflow.utils.module_loading import import_string
+
 from datahub.api.entities.dataprocess.dataprocess_instance import InstanceRunResult
 from datahub.telemetry import telemetry
-
 from datahub_airflow_plugin._airflow_shims import (
     MappedOperator,
     get_task_inlets,
@@ -44,11 +44,9 @@ def get_task_inlets_advanced(task: BaseOperator, context: Any) -> Iterable[Any]:
 
     if task_inlets and isinstance(task_inlets, list):
         inlets = []
-        task_ids = (
-            {o for o in task_inlets if isinstance(o, str)}
-            .union(op.task_id for op in task_inlets if isinstance(op, BaseOperator))
-            .intersection(task.get_flat_relative_ids(upstream=True))
-        )
+        task_ids = {o for o in task_inlets if isinstance(o, str)}.union(
+            op.task_id for op in task_inlets if isinstance(op, BaseOperator)
+        ).intersection(task.get_flat_relative_ids(upstream=True))
 
         from airflow.lineage import AUTO
         from cattr import structure
diff --git a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/entities.py b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/entities.py
index 30b35ac6d6198b..f3fd17259c9f63 100644
--- a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/entities.py
+++ b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/entities.py
@@ -2,6 +2,7 @@
 from typing import List, Optional
 
 import attr
+
 import datahub.emitter.mce_builder as builder
 from datahub.utilities.urns.data_job_urn import DataJobUrn
 from datahub.utilities.urns.dataset_urn import DatasetUrn
diff --git a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/example_dags/generic_recipe_sample_dag.py b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/example_dags/generic_recipe_sample_dag.py
index ff8dba457066fd..ac620852c6f288 100644
--- a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/example_dags/generic_recipe_sample_dag.py
+++ b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/example_dags/generic_recipe_sample_dag.py
@@ -9,6 +9,7 @@
 from airflow import DAG
 from airflow.operators.python import PythonOperator
 from airflow.utils.dates import days_ago
+
 from datahub.configuration.config_loader import load_config_file
 from datahub.ingestion.run.pipeline import Pipeline
 
diff --git a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/example_dags/graph_usage_sample_dag.py b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/example_dags/graph_usage_sample_dag.py
index d72ba67c23cd72..7951d6f7fd21ef 100644
--- a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/example_dags/graph_usage_sample_dag.py
+++ b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/example_dags/graph_usage_sample_dag.py
@@ -4,8 +4,8 @@
 
 import pendulum
 from airflow.decorators import dag, task
-from datahub.ingestion.graph.client import DataHubGraph, RemovedStatusFilter
 
+from datahub.ingestion.graph.client import DataHubGraph, RemovedStatusFilter
 from datahub_airflow_plugin.hooks.datahub import DatahubRestHook
 
 
diff --git a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/example_dags/lineage_emission_dag.py b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/example_dags/lineage_emission_dag.py
index 0d7cdb6b6e90a5..24e89211dd3c5b 100644
--- a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/example_dags/lineage_emission_dag.py
+++ b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/example_dags/lineage_emission_dag.py
@@ -2,13 +2,14 @@
 
 This example demonstrates how to emit lineage to DataHub within an Airflow DAG.
 """
+
 from datetime import timedelta
 
-import datahub.emitter.mce_builder as builder
 from airflow import DAG
 from airflow.operators.bash import BashOperator
 from airflow.utils.dates import days_ago
 
+import datahub.emitter.mce_builder as builder
 from datahub_airflow_plugin.operators.datahub import DatahubEmitterOperator
 
 default_args = {
diff --git a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/hooks/datahub.py b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/hooks/datahub.py
index 5f4d787fb893d3..26c5026c075bd7 100644
--- a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/hooks/datahub.py
+++ b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/hooks/datahub.py
@@ -2,6 +2,7 @@
 
 from airflow.exceptions import AirflowException
 from airflow.hooks.base import BaseHook
+
 from datahub.emitter.generic_emitter import Emitter
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.metadata.com.linkedin.pegasus2avro.mxe import (
@@ -11,6 +12,7 @@
 
 if TYPE_CHECKING:
     from airflow.models.connection import Connection
+
     from datahub.emitter.kafka_emitter import DatahubKafkaEmitter
     from datahub.emitter.rest_emitter import DataHubRestEmitter
     from datahub.emitter.synchronized_file_emitter import SynchronizedFileEmitter
diff --git a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/lineage/_lineage_core.py b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/lineage/_lineage_core.py
index 638458b0efd6ab..db50c48dfaf08a 100644
--- a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/lineage/_lineage_core.py
+++ b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/lineage/_lineage_core.py
@@ -2,7 +2,6 @@
 from typing import TYPE_CHECKING, Dict, List
 
 from datahub.api.entities.dataprocess.dataprocess_instance import InstanceRunResult
-
 from datahub_airflow_plugin._config import DatahubLineageConfig
 from datahub_airflow_plugin.client.airflow_generator import AirflowGenerator
 from datahub_airflow_plugin.entities import (
diff --git a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/operators/datahub.py b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/operators/datahub.py
index 41d855512aa066..817db6b7480c4b 100644
--- a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/operators/datahub.py
+++ b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/operators/datahub.py
@@ -3,9 +3,9 @@
 from airflow.models import BaseOperator
 from airflow.utils.decorators import apply_defaults
 from avrogen.dict_wrapper import DictWrapper
+
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.metadata.com.linkedin.pegasus2avro.mxe import MetadataChangeEvent
-
 from datahub_airflow_plugin.hooks.datahub import (
     DatahubGenericHook,
     DatahubKafkaHook,
diff --git a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/operators/datahub_assertion_operator.py b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/operators/datahub_assertion_operator.py
index 6f93c09a9e2872..3a440b0ec14e07 100644
--- a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/operators/datahub_assertion_operator.py
+++ b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/operators/datahub_assertion_operator.py
@@ -2,11 +2,11 @@
 from typing import Any, List, Optional, Sequence, Union
 
 from airflow.models import BaseOperator
+
 from datahub.api.circuit_breaker import (
     AssertionCircuitBreaker,
     AssertionCircuitBreakerConfig,
 )
-
 from datahub_airflow_plugin.hooks.datahub import DatahubRestHook
 
 
diff --git a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/operators/datahub_assertion_sensor.py b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/operators/datahub_assertion_sensor.py
index 16e5d1cbe8b1f4..6a446ba1f3b55e 100644
--- a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/operators/datahub_assertion_sensor.py
+++ b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/operators/datahub_assertion_sensor.py
@@ -2,11 +2,11 @@
 from typing import Any, List, Optional, Sequence, Union
 
 from airflow.sensors.base import BaseSensorOperator
+
 from datahub.api.circuit_breaker import (
     AssertionCircuitBreaker,
     AssertionCircuitBreakerConfig,
 )
-
 from datahub_airflow_plugin.hooks.datahub import DatahubRestHook
 
 
diff --git a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/operators/datahub_operation_operator.py b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/operators/datahub_operation_operator.py
index 94e105309537b6..eb5fe8168bccf8 100644
--- a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/operators/datahub_operation_operator.py
+++ b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/operators/datahub_operation_operator.py
@@ -2,11 +2,11 @@
 from typing import Any, List, Optional, Sequence, Union
 
 from airflow.sensors.base import BaseSensorOperator
+
 from datahub.api.circuit_breaker import (
     OperationCircuitBreaker,
     OperationCircuitBreakerConfig,
 )
-
 from datahub_airflow_plugin.hooks.datahub import DatahubRestHook
 
 
diff --git a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/operators/datahub_operation_sensor.py b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/operators/datahub_operation_sensor.py
index 434c60754064d0..89e20e46a0074a 100644
--- a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/operators/datahub_operation_sensor.py
+++ b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/operators/datahub_operation_sensor.py
@@ -2,11 +2,11 @@
 from typing import Any, List, Optional, Sequence, Union
 
 from airflow.sensors.base import BaseSensorOperator
+
 from datahub.api.circuit_breaker import (
     OperationCircuitBreaker,
     OperationCircuitBreakerConfig,
 )
-
 from datahub_airflow_plugin.hooks.datahub import DatahubRestHook
 
 
diff --git a/metadata-ingestion-modules/airflow-plugin/tests/integration/dags/datahub_emitter_operator_jinja_template_dag.py b/metadata-ingestion-modules/airflow-plugin/tests/integration/dags/datahub_emitter_operator_jinja_template_dag.py
index c1b4aa4d7b94f4..04845e601d674d 100644
--- a/metadata-ingestion-modules/airflow-plugin/tests/integration/dags/datahub_emitter_operator_jinja_template_dag.py
+++ b/metadata-ingestion-modules/airflow-plugin/tests/integration/dags/datahub_emitter_operator_jinja_template_dag.py
@@ -1,6 +1,7 @@
 from datetime import datetime, timedelta
 
 from airflow import DAG
+
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.metadata.com.linkedin.pegasus2avro.mxe import MetadataChangeEvent
 from datahub.metadata.schema_classes import (
@@ -9,7 +10,6 @@
     DatasetPropertiesClass,
     DatasetSnapshotClass,
 )
-
 from datahub_airflow_plugin.operators.datahub import DatahubEmitterOperator
 
 default_args = {
diff --git a/metadata-ingestion-modules/airflow-plugin/tests/integration/test_plugin.py b/metadata-ingestion-modules/airflow-plugin/tests/integration/test_plugin.py
index 3b2c9140e4632f..2744c26021cde3 100644
--- a/metadata-ingestion-modules/airflow-plugin/tests/integration/test_plugin.py
+++ b/metadata-ingestion-modules/airflow-plugin/tests/integration/test_plugin.py
@@ -17,9 +17,9 @@
 import requests
 import tenacity
 from airflow.models.connection import Connection
+
 from datahub.ingestion.sink.file import write_metadata_file
 from datahub.testing.compare_metadata_json import assert_metadata_files_equal
-
 from datahub_airflow_plugin._airflow_shims import (
     AIRFLOW_VERSION,
     HAS_AIRFLOW_DAG_LISTENER_API,
@@ -273,13 +273,21 @@ def _run_airflow(
             subprocess.check_call(
                 [
                     # fmt: off
-                    "airflow", "users", "create",
-                    "--username", "airflow",
-                    "--password", "airflow",
-                    "--firstname", "admin",
-                    "--lastname", "admin",
-                    "--role", "Admin",
-                    "--email", "airflow@example.com",
+                    "airflow",
+                    "users",
+                    "create",
+                    "--username",
+                    "airflow",
+                    "--password",
+                    "airflow",
+                    "--firstname",
+                    "admin",
+                    "--lastname",
+                    "admin",
+                    "--role",
+                    "Admin",
+                    "--email",
+                    "airflow@example.com",
                     # fmt: on
                 ],
                 env=environment,
diff --git a/metadata-ingestion-modules/airflow-plugin/tests/unit/test_airflow.py b/metadata-ingestion-modules/airflow-plugin/tests/unit/test_airflow.py
index 79620f81a437b0..4219c5fb9cefb3 100644
--- a/metadata-ingestion-modules/airflow-plugin/tests/unit/test_airflow.py
+++ b/metadata-ingestion-modules/airflow-plugin/tests/unit/test_airflow.py
@@ -8,12 +8,12 @@
 
 import airflow.configuration
 import airflow.version
-import datahub.emitter.mce_builder as builder
 import packaging.version
 import pytest
 from airflow.lineage import apply_lineage, prepare_lineage
 from airflow.models import DAG, Connection, DagBag, DagRun, TaskInstance
 
+import datahub.emitter.mce_builder as builder
 from datahub_airflow_plugin import get_provider_info
 from datahub_airflow_plugin._airflow_shims import (
     AIRFLOW_PATCHED,
@@ -242,9 +242,7 @@ def test_lineage_backend(mock_emit, inlets, outlets, capture_executions):
         },
     ), mock.patch("airflow.models.BaseOperator.xcom_pull"), mock.patch(
         "airflow.models.BaseOperator.xcom_push"
-    ), patch_airflow_connection(
-        datahub_rest_connection_config
-    ):
+    ), patch_airflow_connection(datahub_rest_connection_config):
         func = mock.Mock()
         func.__name__ = "foo"
 
@@ -275,7 +273,10 @@ def test_lineage_backend(mock_emit, inlets, outlets, capture_executions):
         if AIRFLOW_VERSION < packaging.version.parse("2.2.0"):
             ti = TaskInstance(task=op2, execution_date=DEFAULT_DATE)
             # Ignoring type here because DagRun state is just a sring at Airflow 1
-            dag_run = DagRun(state="success", run_id=f"scheduled_{DEFAULT_DATE.isoformat()}")  # type: ignore
+            dag_run = DagRun(
+                state="success",  # type: ignore[arg-type]
+                run_id=f"scheduled_{DEFAULT_DATE.isoformat()}",
+            )
         else:
             from airflow.utils.state import DagRunState
 
diff --git a/metadata-ingestion-modules/airflow-plugin/tests/unit/test_packaging.py b/metadata-ingestion-modules/airflow-plugin/tests/unit/test_packaging.py
index a9c8b7ec65fa3c..a822527582c2cd 100644
--- a/metadata-ingestion-modules/airflow-plugin/tests/unit/test_packaging.py
+++ b/metadata-ingestion-modules/airflow-plugin/tests/unit/test_packaging.py
@@ -1,6 +1,6 @@
 import setuptools
-from datahub.testing.check_imports import ensure_no_indirect_model_imports
 
+from datahub.testing.check_imports import ensure_no_indirect_model_imports
 from tests.utils import PytestConfig
 
 
diff --git a/metadata-ingestion-modules/dagster-plugin/.gitignore b/metadata-ingestion-modules/dagster-plugin/.gitignore
index 4ff42af3e16cff..6801b785ea1e4b 100644
--- a/metadata-ingestion-modules/dagster-plugin/.gitignore
+++ b/metadata-ingestion-modules/dagster-plugin/.gitignore
@@ -1,5 +1,4 @@
 .envrc
-src/datahub_dagster_plugin/__init__.py.bak
 .vscode/
 output
 pvenv36/
diff --git a/metadata-ingestion-modules/dagster-plugin/build.gradle b/metadata-ingestion-modules/dagster-plugin/build.gradle
index 0d57bb5bfdff70..7dd7036e276151 100644
--- a/metadata-ingestion-modules/dagster-plugin/build.gradle
+++ b/metadata-ingestion-modules/dagster-plugin/build.gradle
@@ -54,18 +54,15 @@ task installDev(type: Exec, dependsOn: [install]) {
 task lint(type: Exec, dependsOn: installDev) {
   commandLine 'bash', '-c',
     "source ${venv_name}/bin/activate && set -x && " +
-    "black --check --diff src/ tests/ examples/ && " +
-    "isort --check --diff src/ tests/ examples/ && " +
-    "flake8 --count --statistics src/ tests/ examples/ && " +
+    "ruff check src/ tests/ examples/ && " +
+    "ruff format --check src/ tests/ && " +
     "mypy --show-traceback --show-error-codes src/ tests/ examples/"
 }
 task lintFix(type: Exec, dependsOn: installDev) {
   commandLine 'bash', '-x', '-c',
     "source ${venv_name}/bin/activate && " +
-    "black src/ tests/ examples/ && " +
-    "isort src/ tests/ examples/ && " +
-    "flake8 src/ tests/ examples/ && " +
-    "mypy src/ tests/ examples/"
+    "ruff check --fix src/ tests/ examples/ && " +
+    "ruff format src/ tests/ examples/ "
 }
 
 task installDevTest(type: Exec, dependsOn: [installDev]) {
@@ -107,5 +104,8 @@ clean {
   delete venv_name
   delete 'build'
   delete 'dist'
+  delete '.ruff_cache'
+  delete '.mypy_cache'
+  delete '.pytest_cache'
 }
 clean.dependsOn cleanPythonCache
diff --git a/metadata-ingestion-modules/dagster-plugin/examples/advanced_ops_jobs.py b/metadata-ingestion-modules/dagster-plugin/examples/advanced_ops_jobs.py
index 7b7616b1ec11de..75cab237b05a3e 100644
--- a/metadata-ingestion-modules/dagster-plugin/examples/advanced_ops_jobs.py
+++ b/metadata-ingestion-modules/dagster-plugin/examples/advanced_ops_jobs.py
@@ -9,9 +9,9 @@
     job,
     op,
 )
+
 from datahub.ingestion.graph.client import DatahubClientConfig, DataHubGraph
 from datahub.utilities.urns.dataset_urn import DatasetUrn
-
 from datahub_dagster_plugin.client.dagster_generator import (
     DagsterGenerator,
     DatasetLineage,
diff --git a/metadata-ingestion-modules/dagster-plugin/examples/assets_job.py b/metadata-ingestion-modules/dagster-plugin/examples/assets_job.py
index 1ed3f2f915061b..9b26b502d770f2 100644
--- a/metadata-ingestion-modules/dagster-plugin/examples/assets_job.py
+++ b/metadata-ingestion-modules/dagster-plugin/examples/assets_job.py
@@ -7,9 +7,9 @@
     define_asset_job,
     multi_asset,
 )
+
 from datahub.ingestion.graph.config import DatahubClientConfig
 from datahub.utilities.urns.dataset_urn import DatasetUrn
-
 from datahub_dagster_plugin.sensors.datahub_sensors import (
     DatahubDagsterSourceConfig,
     make_datahub_sensor,
diff --git a/metadata-ingestion-modules/dagster-plugin/examples/basic_setup.py b/metadata-ingestion-modules/dagster-plugin/examples/basic_setup.py
index 300cf9df022c66..2eeff225697261 100644
--- a/metadata-ingestion-modules/dagster-plugin/examples/basic_setup.py
+++ b/metadata-ingestion-modules/dagster-plugin/examples/basic_setup.py
@@ -1,6 +1,6 @@
 from dagster import Definitions
-from datahub.ingestion.graph.client import DatahubClientConfig
 
+from datahub.ingestion.graph.client import DatahubClientConfig
 from datahub_dagster_plugin.sensors.datahub_sensors import (
     DatahubDagsterSourceConfig,
     make_datahub_sensor,
diff --git a/metadata-ingestion-modules/dagster-plugin/examples/ops_job.py b/metadata-ingestion-modules/dagster-plugin/examples/ops_job.py
index a17fc89e6922df..aa2902ee5c708b 100644
--- a/metadata-ingestion-modules/dagster-plugin/examples/ops_job.py
+++ b/metadata-ingestion-modules/dagster-plugin/examples/ops_job.py
@@ -1,7 +1,7 @@
 from dagster import Definitions, In, Out, PythonObjectDagsterType, job, op
+
 from datahub.ingestion.graph.config import DatahubClientConfig
 from datahub.utilities.urns.dataset_urn import DatasetUrn
-
 from datahub_dagster_plugin.sensors.datahub_sensors import (
     DatahubDagsterSourceConfig,
     make_datahub_sensor,
diff --git a/metadata-ingestion-modules/dagster-plugin/pyproject.toml b/metadata-ingestion-modules/dagster-plugin/pyproject.toml
index fba81486b9f677..d1e1d0ad479442 100644
--- a/metadata-ingestion-modules/dagster-plugin/pyproject.toml
+++ b/metadata-ingestion-modules/dagster-plugin/pyproject.toml
@@ -2,18 +2,58 @@
 build-backend = "setuptools.build_meta"
 requires = ["setuptools>=54.0.0", "wheel", "pip>=21.0.0"]
 
-[tool.black]
-extend-exclude = '''
-# A regex preceded with ^/ will apply only to files and directories
-# in the root of the project.
-^/tmp
-'''
-include = '\.pyi?$'
+[tool.ruff]
+line-length = 88
+target-version = "py38"
+exclude = [
+    ".git",
+    "venv",
+    ".tox",
+    "__pycache__",
+]
 
-[tool.isort]
-indent = '    '
-profile = 'black'
-sections = 'FUTURE,STDLIB,THIRDPARTY,FIRSTPARTY,LOCALFOLDER'
+[tool.ruff.format]
+quote-style = "double"
+indent-style = "space"
+skip-magic-trailing-comma = false
+line-ending = "auto"
 
-[tool.pyright]
-extraPaths = ['tests']
\ No newline at end of file
+[tool.ruff.lint.isort]
+combine-as-imports = true
+known-first-party = ["datahub"]
+extra-standard-library = ["__future__", "datahub.utilities._markupsafe_compat", "datahub.sql_parsing._sqlglot_patch"]
+section-order = ["future", "standard-library", "third-party", "first-party", "local-folder"]
+force-sort-within-sections = false
+force-wrap-aliases = false
+split-on-trailing-comma = false
+order-by-type = true
+relative-imports-order = "closest-to-furthest"
+force-single-line = false
+single-line-exclusions = ["typing"]
+length-sort = false
+from-first = false
+required-imports = []
+classes = ["typing"]
+
+[tool.ruff.lint]
+extend-select = [
+    "B",    # flake8-bugbear
+    "C90",  # mccabe complexity
+    "E",    # pycodestyle errors
+    "F",    # pyflakes
+    "G010", # logging.warn -> logging.warning
+    "I",    # isort
+    "TID",  # flake8-tidy-imports
+]
+ignore = [
+    "E501",  # Line length violations (handled by formatter)
+]
+
+[tool.ruff.lint.mccabe]
+max-complexity = 15
+
+[tool.ruff.lint.flake8-tidy-imports]
+ban-relative-imports = "all"
+
+[tool.ruff.lint.per-file-ignores]
+"__init__.py" = ["F401"]
\ No newline at end of file
diff --git a/metadata-ingestion-modules/dagster-plugin/scripts/release.sh b/metadata-ingestion-modules/dagster-plugin/scripts/release.sh
index 10cb816d9ffc04..ffd5201574891c 100755
--- a/metadata-ingestion-modules/dagster-plugin/scripts/release.sh
+++ b/metadata-ingestion-modules/dagster-plugin/scripts/release.sh
@@ -1,26 +1,31 @@
 #!/bin/bash
+# Auto-generated by python-build/generate_release_scripts.py. Do not edit manually.
+
 set -euxo pipefail
 
+ROOT=../..
+MODULE=datahub_dagster_plugin
+
 if [[ ! ${RELEASE_SKIP_TEST:-} ]] && [[ ! ${RELEASE_SKIP_INSTALL:-} ]]; then
-	../../gradlew build  # also runs tests
+    ${ROOT}/gradlew build  # also runs tests
 elif [[ ! ${RELEASE_SKIP_INSTALL:-} ]]; then
-	../../gradlew install
+    ${ROOT}/gradlew install
 fi
 
-MODULE=datahub_dagster_plugin
-
 # Check packaging constraint.
 python -c 'import setuptools; where="./src"; assert setuptools.find_packages(where) == setuptools.find_namespace_packages(where), "you seem to be missing or have extra __init__.py files"'
-if [[ ${RELEASE_VERSION:-} ]]; then
-    # Replace version with RELEASE_VERSION env variable
-    sed -i.bak "s/__version__ = \"1\!0.0.0.dev0\"/__version__ = \"$(echo $RELEASE_VERSION|sed s/-/+/)\"/" src/${MODULE}/__init__.py
-else
-    vim src/${MODULE}/__init__.py
+
+# Update the release version.
+if [[ ! ${RELEASE_VERSION:-} ]]; then
+    echo "RELEASE_VERSION is not set"
+    exit 1
 fi
+sed -i.bak "s/__version__ = .*$/__version__ = \"$(echo $RELEASE_VERSION|sed s/-/+/)\"/" src/${MODULE}/_version.py
 
+# Build and upload the release.
 rm -rf build dist || true
 python -m build
 if [[ ! ${RELEASE_SKIP_UPLOAD:-} ]]; then
     python -m twine upload 'dist/*'
 fi
-mv src/${MODULE}/__init__.py.bak src/${MODULE}/__init__.py
+mv src/${MODULE}/_version.py.bak src/${MODULE}/_version.py
diff --git a/metadata-ingestion-modules/dagster-plugin/setup.cfg b/metadata-ingestion-modules/dagster-plugin/setup.cfg
index 20a903914332aa..89b28ae45f9648 100644
--- a/metadata-ingestion-modules/dagster-plugin/setup.cfg
+++ b/metadata-ingestion-modules/dagster-plugin/setup.cfg
@@ -1,24 +1,3 @@
-[flake8]
-max-complexity = 15
-ignore =
-    # Ignore: line length issues, since black's formatter will take care of them.
-    E501,
-    # Ignore: 1 blank line required before class docstring.
-    D203,
-    # See https://stackoverflow.com/a/57074416.
-    W503,
-    # See https://github.com/psf/black/issues/315.
-    E203
-exclude =
-    .git,
-    venv,
-    .tox,
-    __pycache__
-per-file-ignores =
-    # imported but unused
-    __init__.py: F401
-ban-relative-imports = true
-
 [mypy]
 plugins =
     pydantic.mypy
diff --git a/metadata-ingestion-modules/dagster-plugin/setup.py b/metadata-ingestion-modules/dagster-plugin/setup.py
index 22c15497bd8070..6e2e013f719f51 100644
--- a/metadata-ingestion-modules/dagster-plugin/setup.py
+++ b/metadata-ingestion-modules/dagster-plugin/setup.py
@@ -4,7 +4,7 @@
 import setuptools
 
 package_metadata: dict = {}
-with open("./src/datahub_dagster_plugin/__init__.py") as fp:
+with open("./src/datahub_dagster_plugin/_version.py") as fp:
     exec(fp.read(), package_metadata)
 
 
@@ -51,12 +51,8 @@ def get_long_description():
     "dagster-aws >= 0.11.0",
     "dagster-snowflake >= 0.11.0",
     "dagster-snowflake-pandas >= 0.11.0",
-    "black==22.12.0",
     "coverage>=5.1",
-    "flake8>=6.0.0",
-    "flake8-tidy-imports>=4.3.0",
-    "flake8-bugbear==23.3.12",
-    "isort>=5.7.0",
+    "ruff==0.9.2",
     "mypy>=1.4.0",
     # pydantic 1.8.2 is incompatible with mypy 0.910.
     # See https://github.com/samuelcolvin/pydantic/pull/3175#issuecomment-995382910.
diff --git a/metadata-ingestion-modules/dagster-plugin/src/datahub_dagster_plugin/__init__.py b/metadata-ingestion-modules/dagster-plugin/src/datahub_dagster_plugin/__init__.py
index 1c7d60666a085f..20baf254135355 100644
--- a/metadata-ingestion-modules/dagster-plugin/src/datahub_dagster_plugin/__init__.py
+++ b/metadata-ingestion-modules/dagster-plugin/src/datahub_dagster_plugin/__init__.py
@@ -1,21 +1 @@
-# Published at https://pypi.org/project/acryl-datahub/.
-__package_name__ = "acryl-datahub-dagster-plugin"
-__version__ = "1!0.0.0.dev0"
-
-
-def is_dev_mode() -> bool:
-    return __version__.endswith("dev0")
-
-
-def nice_version_name() -> str:
-    if is_dev_mode():
-        return "unavailable (installed in develop mode)"
-    return __version__
-
-
-def get_provider_info():
-    return {
-        "package-name": f"{__package_name__}",
-        "name": f"{__package_name__}",
-        "description": "Datahub metadata collector plugin",
-    }
+from datahub_dagster_plugin._version import __package_name__, __version__
diff --git a/metadata-ingestion-modules/dagster-plugin/src/datahub_dagster_plugin/_version.py b/metadata-ingestion-modules/dagster-plugin/src/datahub_dagster_plugin/_version.py
new file mode 100644
index 00000000000000..e287b6bf32f5dc
--- /dev/null
+++ b/metadata-ingestion-modules/dagster-plugin/src/datahub_dagster_plugin/_version.py
@@ -0,0 +1,3 @@
+# Published at https://pypi.org/project/acryl-datahub-dagster-plugin/.
+__package_name__ = "acryl-datahub-dagster-plugin"
+__version__ = "1!0.0.0.dev0"
diff --git a/metadata-ingestion-modules/dagster-plugin/src/datahub_dagster_plugin/client/dagster_generator.py b/metadata-ingestion-modules/dagster-plugin/src/datahub_dagster_plugin/client/dagster_generator.py
index a87f490f2d947e..033d3967145017 100644
--- a/metadata-ingestion-modules/dagster-plugin/src/datahub_dagster_plugin/client/dagster_generator.py
+++ b/metadata-ingestion-modules/dagster-plugin/src/datahub_dagster_plugin/client/dagster_generator.py
@@ -13,6 +13,7 @@
     TableSchemaMetadataValue,
 )
 from dagster._core.execution.stats import RunStepKeyStatsSnapshot, StepEventStatus
+
 from datahub.sql_parsing.sqlglot_utils import get_query_fingerprint
 
 try:
@@ -23,6 +24,7 @@
 
 from dagster._core.snap.node import OpDefSnap
 from dagster._core.storage.dagster_run import DagsterRun, DagsterRunStatsSnapshot
+
 from datahub.api.entities.datajob import DataFlow, DataJob
 from datahub.api.entities.dataprocess.dataprocess_instance import (
     DataProcessInstance,
@@ -505,7 +507,7 @@ def generate_datajob(
         job_property_bag: Dict[str, str] = {}
         if input_datasets:
             self.logger.info(
-                f"Input datasets for {op_def_snap.name} are { list(input_datasets.get(op_def_snap.name, []))}"
+                f"Input datasets for {op_def_snap.name} are {list(input_datasets.get(op_def_snap.name, []))}"
             )
             inlets.update(input_datasets.get(op_def_snap.name, []))
 
@@ -513,7 +515,7 @@ def generate_datajob(
 
         if output_datasets:
             self.logger.info(
-                f"Output datasets for {op_def_snap.name} are { list(output_datasets.get(op_def_snap.name, []))}"
+                f"Output datasets for {op_def_snap.name} are {list(output_datasets.get(op_def_snap.name, []))}"
             )
             datajob.outlets = list(output_datasets.get(op_def_snap.name, []))
 
@@ -604,7 +606,7 @@ def emit_job_run(
         if run.status not in status_result_map:
             raise Exception(
                 f"Job run status should be either complete, failed or cancelled and it was "
-                f"{run.status }"
+                f"{run.status}"
             )
 
         if run_stats.start_time is not None:
@@ -671,7 +673,7 @@ def emit_op_run(
         if run_step_stats.status not in status_result_map:
             raise Exception(
                 f"Step run status should be either complete, failed or cancelled and it was "
-                f"{run_step_stats.status }"
+                f"{run_step_stats.status}"
             )
 
         if run_step_stats.start_time is not None:
diff --git a/metadata-ingestion-modules/dagster-plugin/src/datahub_dagster_plugin/sensors/datahub_sensors.py b/metadata-ingestion-modules/dagster-plugin/src/datahub_dagster_plugin/sensors/datahub_sensors.py
index bccdb4ac7922a5..5f049d55c16a12 100644
--- a/metadata-ingestion-modules/dagster-plugin/src/datahub_dagster_plugin/sensors/datahub_sensors.py
+++ b/metadata-ingestion-modules/dagster-plugin/src/datahub_dagster_plugin/sensors/datahub_sensors.py
@@ -35,7 +35,9 @@
 try:
     from dagster._core.definitions.sensor_definition import SensorReturnTypesUnion
 except ImportError:
-    from dagster._core.definitions.sensor_definition import RawSensorEvaluationFunctionReturn as SensorReturnTypesUnion  # type: ignore
+    from dagster._core.definitions.sensor_definition import (  # type: ignore
+        RawSensorEvaluationFunctionReturn as SensorReturnTypesUnion,
+    )
 
 from dagster._core.definitions.target import ExecutableDefinition
 from dagster._core.definitions.unresolved_asset_job_definition import (
@@ -43,6 +45,7 @@
 )
 from dagster._core.events import DagsterEventType, HandledOutputData, LoadedInputData
 from dagster._core.execution.stats import RunStepKeyStatsSnapshot
+
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.ingestion.graph.client import DatahubClientConfig, DataHubGraph
 from datahub.metadata.schema_classes import SubTypesClass
@@ -52,7 +55,6 @@
 )
 from datahub.utilities.urns.dataset_urn import DatasetUrn
 from datahub.utilities.urns.error import InvalidUrnError
-
 from datahub_dagster_plugin.client.dagster_generator import (
     DATAHUB_ASSET_GROUP_NAME_CACHE,
     Constant,
@@ -260,7 +262,6 @@ def get_dagster_environment(
             and context.dagster_run.job_code_origin.repository_origin
             and context.dagster_run.job_code_origin.repository_origin.code_pointer
         ):
-
             code_pointer = (
                 context.dagster_run.job_code_origin.repository_origin.code_pointer
             )
diff --git a/metadata-ingestion-modules/dagster-plugin/tests/unit/test_dagster.py b/metadata-ingestion-modules/dagster-plugin/tests/unit/test_dagster.py
index c951b959f85d43..9a69822984bb80 100644
--- a/metadata-ingestion-modules/dagster-plugin/tests/unit/test_dagster.py
+++ b/metadata-ingestion-modules/dagster-plugin/tests/unit/test_dagster.py
@@ -22,11 +22,11 @@
     RepositoryDefinition,
 )
 from dagster._core.definitions.resource_definition import ResourceDefinition
-from datahub.emitter.mcp import MetadataChangeProposalWrapper
-from datahub.ingestion.graph.client import DatahubClientConfig
 from freezegun import freeze_time
 from utils.utils import PytestConfig, check_golden_file
 
+from datahub.emitter.mcp import MetadataChangeProposalWrapper
+from datahub.ingestion.graph.client import DatahubClientConfig
 from datahub_dagster_plugin.client.dagster_generator import DatahubDagsterSourceConfig
 from datahub_dagster_plugin.sensors.datahub_sensors import (
     DatahubSensors,
diff --git a/metadata-ingestion-modules/gx-plugin/.gitignore b/metadata-ingestion-modules/gx-plugin/.gitignore
index 8c01744589e35e..6801b785ea1e4b 100644
--- a/metadata-ingestion-modules/gx-plugin/.gitignore
+++ b/metadata-ingestion-modules/gx-plugin/.gitignore
@@ -1,5 +1,4 @@
 .envrc
-src/datahub_gx_plugin/__init__.py.bak
 .vscode/
 output
 pvenv36/
diff --git a/metadata-ingestion-modules/gx-plugin/build.gradle b/metadata-ingestion-modules/gx-plugin/build.gradle
index 70bc3e0360a3a4..57a1ed0b2169d3 100644
--- a/metadata-ingestion-modules/gx-plugin/build.gradle
+++ b/metadata-ingestion-modules/gx-plugin/build.gradle
@@ -25,7 +25,7 @@ task environmentSetup(type: Exec, dependsOn: checkPythonVersion) {
   outputs.file(sentinel_file)
   commandLine 'bash', '-c',
       "${python_executable} -m venv ${venv_name} && " +
-      "${venv_name}/bin/python -m pip install --upgrade uv && " +
+      "${venv_name}/bin/pip install --upgrade uv && " +
       "touch ${sentinel_file}"
 }
 
@@ -54,18 +54,15 @@ task installDev(type: Exec, dependsOn: [install]) {
 task lint(type: Exec, dependsOn: installDev) {
   commandLine 'bash', '-c',
     "source ${venv_name}/bin/activate && set -x && " +
-    "black --check --diff src/ tests/ && " +
-    "isort --check --diff src/ tests/ && " +
-    "flake8 --count --statistics src/ tests/ && " +
+    "ruff check src/ tests/ && " +
+    "ruff format --check src/ tests/ && " +
     "mypy --show-traceback --show-error-codes src/ tests/"
 }
 task lintFix(type: Exec, dependsOn: installDev) {
   commandLine 'bash', '-x', '-c',
     "source ${venv_name}/bin/activate && " +
-    "black src/ tests/ && " +
-    "isort src/ tests/ && " +
-    "flake8 src/ tests/ && " +
-    "mypy src/ tests/"
+    "ruff check --fix src/ tests/ && " +
+    "ruff format src/ tests/ "
 }
 
 task installDevTest(type: Exec, dependsOn: [installDev]) {
@@ -107,5 +104,8 @@ clean {
   delete venv_name
   delete 'build'
   delete 'dist'
+  delete '.ruff_cache'
+  delete '.mypy_cache'
+  delete '.pytest_cache'
 }
 clean.dependsOn cleanPythonCache
diff --git a/metadata-ingestion-modules/gx-plugin/pyproject.toml b/metadata-ingestion-modules/gx-plugin/pyproject.toml
index fba81486b9f677..d1e1d0ad479442 100644
--- a/metadata-ingestion-modules/gx-plugin/pyproject.toml
+++ b/metadata-ingestion-modules/gx-plugin/pyproject.toml
@@ -2,18 +2,58 @@
 build-backend = "setuptools.build_meta"
 requires = ["setuptools>=54.0.0", "wheel", "pip>=21.0.0"]
 
-[tool.black]
-extend-exclude = '''
-# A regex preceded with ^/ will apply only to files and directories
-# in the root of the project.
-^/tmp
-'''
-include = '\.pyi?$'
+[tool.ruff]
+line-length = 88
+target-version = "py38"
+exclude = [
+    ".git",
+    "venv",
+    ".tox",
+    "__pycache__",
+]
 
-[tool.isort]
-indent = '    '
-profile = 'black'
-sections = 'FUTURE,STDLIB,THIRDPARTY,FIRSTPARTY,LOCALFOLDER'
+[tool.ruff.format]
+quote-style = "double"
+indent-style = "space"
+skip-magic-trailing-comma = false
+line-ending = "auto"
 
-[tool.pyright]
-extraPaths = ['tests']
\ No newline at end of file
+[tool.ruff.lint.isort]
+combine-as-imports = true
+known-first-party = ["datahub"]
+extra-standard-library = ["__future__", "datahub.utilities._markupsafe_compat", "datahub.sql_parsing._sqlglot_patch"]
+section-order = ["future", "standard-library", "third-party", "first-party", "local-folder"]
+force-sort-within-sections = false
+force-wrap-aliases = false
+split-on-trailing-comma = false
+order-by-type = true
+relative-imports-order = "closest-to-furthest"
+force-single-line = false
+single-line-exclusions = ["typing"]
+length-sort = false
+from-first = false
+required-imports = []
+classes = ["typing"]
+
+[tool.ruff.lint]
+extend-select = [
+    "B",    # flake8-bugbear
+    "C90",  # mccabe complexity
+    "E",    # pycodestyle errors
+    "F",    # pyflakes
+    "G010", # logging.warn -> logging.warning
+    "I",    # isort
+    "TID",  # flake8-tidy-imports
+]
+ignore = [
+    "E501",  # Line length violations (handled by formatter)
+]
+
+[tool.ruff.lint.mccabe]
+max-complexity = 15
+
+[tool.ruff.lint.flake8-tidy-imports]
+ban-relative-imports = "all"
+
+[tool.ruff.lint.per-file-ignores]
+"__init__.py" = ["F401"]
\ No newline at end of file
diff --git a/metadata-ingestion-modules/gx-plugin/scripts/release.sh b/metadata-ingestion-modules/gx-plugin/scripts/release.sh
index 058add495821cb..06605f03a78aa3 100755
--- a/metadata-ingestion-modules/gx-plugin/scripts/release.sh
+++ b/metadata-ingestion-modules/gx-plugin/scripts/release.sh
@@ -1,26 +1,31 @@
 #!/bin/bash
+# Auto-generated by python-build/generate_release_scripts.py. Do not edit manually.
+
 set -euxo pipefail
 
+ROOT=../..
+MODULE=datahub_gx_plugin
+
 if [[ ! ${RELEASE_SKIP_TEST:-} ]] && [[ ! ${RELEASE_SKIP_INSTALL:-} ]]; then
-	../../gradlew build  # also runs tests
+    ${ROOT}/gradlew build  # also runs tests
 elif [[ ! ${RELEASE_SKIP_INSTALL:-} ]]; then
-	../../gradlew install
+    ${ROOT}/gradlew install
 fi
 
-MODULE=datahub_gx_plugin
-
 # Check packaging constraint.
 python -c 'import setuptools; where="./src"; assert setuptools.find_packages(where) == setuptools.find_namespace_packages(where), "you seem to be missing or have extra __init__.py files"'
-if [[ ${RELEASE_VERSION:-} ]]; then
-    # Replace version with RELEASE_VERSION env variable
-    sed -i.bak "s/__version__ = \"1\!0.0.0.dev0\"/__version__ = \"$(echo $RELEASE_VERSION|sed s/-/+/)\"/" src/${MODULE}/__init__.py
-else
-    vim src/${MODULE}/__init__.py
+
+# Update the release version.
+if [[ ! ${RELEASE_VERSION:-} ]]; then
+    echo "RELEASE_VERSION is not set"
+    exit 1
 fi
+sed -i.bak "s/__version__ = .*$/__version__ = \"$(echo $RELEASE_VERSION|sed s/-/+/)\"/" src/${MODULE}/_version.py
 
+# Build and upload the release.
 rm -rf build dist || true
 python -m build
 if [[ ! ${RELEASE_SKIP_UPLOAD:-} ]]; then
     python -m twine upload 'dist/*'
 fi
-mv src/${MODULE}/__init__.py.bak src/${MODULE}/__init__.py
+mv src/${MODULE}/_version.py.bak src/${MODULE}/_version.py
diff --git a/metadata-ingestion-modules/gx-plugin/setup.cfg b/metadata-ingestion-modules/gx-plugin/setup.cfg
index bbdd85f0fdc4ed..de78d3597e74d2 100644
--- a/metadata-ingestion-modules/gx-plugin/setup.cfg
+++ b/metadata-ingestion-modules/gx-plugin/setup.cfg
@@ -1,24 +1,3 @@
-[flake8]
-max-complexity = 15
-ignore =
-    # Ignore: line length issues, since black's formatter will take care of them.
-    E501,
-    # Ignore: 1 blank line required before class docstring.
-    D203,
-    # See https://stackoverflow.com/a/57074416.
-    W503,
-    # See https://github.com/psf/black/issues/315.
-    E203
-exclude =
-    .git,
-    venv,
-    .tox,
-    __pycache__
-per-file-ignores =
-    # imported but unused
-    __init__.py: F401
-ban-relative-imports = true
-
 [mypy]
 plugins =
     pydantic.mypy
diff --git a/metadata-ingestion-modules/gx-plugin/setup.py b/metadata-ingestion-modules/gx-plugin/setup.py
index 40afc81a98f9c8..43495673a7ff12 100644
--- a/metadata-ingestion-modules/gx-plugin/setup.py
+++ b/metadata-ingestion-modules/gx-plugin/setup.py
@@ -4,7 +4,7 @@
 import setuptools
 
 package_metadata: dict = {}
-with open("./src/datahub_gx_plugin/__init__.py") as fp:
+with open("./src/datahub_gx_plugin/_version.py") as fp:
     exec(fp.read(), package_metadata)
 
 
@@ -58,12 +58,8 @@ def get_long_description():
 base_dev_requirements = {
     *base_requirements,
     *mypy_stubs,
-    "black==22.12.0",
     "coverage>=5.1",
-    "flake8>=6.0.0",
-    "flake8-tidy-imports>=4.3.0",
-    "flake8-bugbear==23.3.12",
-    "isort>=5.7.0",
+    "ruff==0.9.1",
     "mypy>=1.4.0",
     # pydantic 1.8.2 is incompatible with mypy 0.910.
     # See https://github.com/samuelcolvin/pydantic/pull/3175#issuecomment-995382910.
diff --git a/metadata-ingestion-modules/gx-plugin/src/datahub_gx_plugin/__init__.py b/metadata-ingestion-modules/gx-plugin/src/datahub_gx_plugin/__init__.py
index a7689be82a5d99..b3f8638c28088d 100644
--- a/metadata-ingestion-modules/gx-plugin/src/datahub_gx_plugin/__init__.py
+++ b/metadata-ingestion-modules/gx-plugin/src/datahub_gx_plugin/__init__.py
@@ -1,21 +1 @@
-# Published at https://pypi.org/project/acryl-datahub/.
-__package_name__ = "acryl-datahub-gx-plugin"
-__version__ = "1!0.0.0.dev0"
-
-
-def is_dev_mode() -> bool:
-    return __version__.endswith("dev0")
-
-
-def nice_version_name() -> str:
-    if is_dev_mode():
-        return "unavailable (installed in develop mode)"
-    return __version__
-
-
-def get_provider_info():
-    return {
-        "package-name": f"{__package_name__}",
-        "name": f"{__package_name__}",
-        "description": "Datahub metadata collector plugin",
-    }
+from datahub_gx_plugin._version import __package_name__, __version__
diff --git a/metadata-ingestion-modules/gx-plugin/src/datahub_gx_plugin/_version.py b/metadata-ingestion-modules/gx-plugin/src/datahub_gx_plugin/_version.py
new file mode 100644
index 00000000000000..02dbb70a9b4833
--- /dev/null
+++ b/metadata-ingestion-modules/gx-plugin/src/datahub_gx_plugin/_version.py
@@ -0,0 +1,3 @@
+# Published at https://pypi.org/project/acryl-datahub-gx-plugin/.
+__package_name__ = "acryl-datahub-gx-plugin"
+__version__ = "1!0.0.0.dev0"
diff --git a/metadata-ingestion-modules/gx-plugin/src/datahub_gx_plugin/action.py b/metadata-ingestion-modules/gx-plugin/src/datahub_gx_plugin/action.py
index 2d89d26997d1f3..1070d4d3d5d66d 100644
--- a/metadata-ingestion-modules/gx-plugin/src/datahub_gx_plugin/action.py
+++ b/metadata-ingestion-modules/gx-plugin/src/datahub_gx_plugin/action.py
@@ -3,12 +3,37 @@
 import sys
 import time
 from dataclasses import dataclass
+from datahub.utilities._markupsafe_compat import MARKUPSAFE_PATCHED
 from datetime import timezone
 from decimal import Decimal
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
 
-import datahub.emitter.mce_builder as builder
 import packaging.version
+from great_expectations.checkpoint.actions import ValidationAction
+from great_expectations.core.batch import Batch
+from great_expectations.core.batch_spec import (
+    RuntimeDataBatchSpec,
+    RuntimeQueryBatchSpec,
+    SqlAlchemyDatasourceBatchSpec,
+)
+from great_expectations.core.expectation_validation_result import (
+    ExpectationSuiteValidationResult,
+)
+from great_expectations.data_asset.data_asset import DataAsset
+from great_expectations.data_context import AbstractDataContext
+from great_expectations.data_context.types.resource_identifiers import (
+    ExpectationSuiteIdentifier,
+    ValidationResultIdentifier,
+)
+from great_expectations.execution_engine import PandasExecutionEngine
+from great_expectations.execution_engine.sqlalchemy_execution_engine import (
+    SqlAlchemyExecutionEngine,
+)
+from great_expectations.validator.validator import Validator
+from sqlalchemy.engine.base import Connection, Engine
+from sqlalchemy.engine.url import make_url
+
+import datahub.emitter.mce_builder as builder
 from datahub.cli.env_utils import get_boolean_env_variable
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.emitter.rest_emitter import DatahubRestEmitter
@@ -35,31 +60,7 @@
 from datahub.metadata.com.linkedin.pegasus2avro.common import DataPlatformInstance
 from datahub.metadata.schema_classes import PartitionSpecClass, PartitionTypeClass
 from datahub.sql_parsing.sqlglot_lineage import create_lineage_sql_parsed_result
-from datahub.utilities._markupsafe_compat import MARKUPSAFE_PATCHED
 from datahub.utilities.urns.dataset_urn import DatasetUrn
-from great_expectations.checkpoint.actions import ValidationAction
-from great_expectations.core.batch import Batch
-from great_expectations.core.batch_spec import (
-    RuntimeDataBatchSpec,
-    RuntimeQueryBatchSpec,
-    SqlAlchemyDatasourceBatchSpec,
-)
-from great_expectations.core.expectation_validation_result import (
-    ExpectationSuiteValidationResult,
-)
-from great_expectations.data_asset.data_asset import DataAsset
-from great_expectations.data_context import AbstractDataContext
-from great_expectations.data_context.types.resource_identifiers import (
-    ExpectationSuiteIdentifier,
-    ValidationResultIdentifier,
-)
-from great_expectations.execution_engine import PandasExecutionEngine
-from great_expectations.execution_engine.sqlalchemy_execution_engine import (
-    SqlAlchemyExecutionEngine,
-)
-from great_expectations.validator.validator import Validator
-from sqlalchemy.engine.base import Connection, Engine
-from sqlalchemy.engine.url import make_url
 
 # TODO: move this and version check used in tests to some common module
 try:
@@ -107,7 +108,6 @@ def __init__(
         convert_urns_to_lowercase: bool = False,
         name: str = "DataHubValidationAction",
     ):
-
         if has_name_positional_arg:
             if len(args) >= 1 and isinstance(args[0], str):
                 name = args[0]
@@ -163,9 +163,7 @@ def _run(
             if isinstance(
                 validation_result_suite_identifier, ValidationResultIdentifier
             ):
-                expectation_suite_name = (
-                    validation_result_suite_identifier.expectation_suite_identifier.expectation_suite_name
-                )
+                expectation_suite_name = validation_result_suite_identifier.expectation_suite_identifier.expectation_suite_name
                 run_id = validation_result_suite_identifier.run_id
                 batch_identifier = validation_result_suite_identifier.batch_identifier
 
diff --git a/metadata-ingestion-modules/gx-plugin/tests/integration/test_great_expectations.py b/metadata-ingestion-modules/gx-plugin/tests/integration/test_great_expectations.py
index b03681dc780584..9919c1db25508c 100644
--- a/metadata-ingestion-modules/gx-plugin/tests/integration/test_great_expectations.py
+++ b/metadata-ingestion-modules/gx-plugin/tests/integration/test_great_expectations.py
@@ -5,12 +5,13 @@
 
 import packaging.version
 import pytest
+from freezegun import freeze_time
+from great_expectations.data_context import FileDataContext
+
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.ingestion.sink.file import write_metadata_file
 from datahub.testing.compare_metadata_json import assert_metadata_files_equal
 from datahub.testing.docker_utils import wait_for_port
-from freezegun import freeze_time
-from great_expectations.data_context import FileDataContext
 
 try:
     from great_expectations import __version__ as GX_VERSION  # type: ignore
diff --git a/metadata-ingestion-modules/gx-plugin/tests/unit/test_great_expectations_action.py b/metadata-ingestion-modules/gx-plugin/tests/unit/test_great_expectations_action.py
index fcad17fa4b71dd..3eb3e5af6e16bd 100644
--- a/metadata-ingestion-modules/gx-plugin/tests/unit/test_great_expectations_action.py
+++ b/metadata-ingestion-modules/gx-plugin/tests/unit/test_great_expectations_action.py
@@ -4,22 +4,6 @@
 
 import pandas as pd
 import pytest
-from datahub.emitter.mcp import MetadataChangeProposalWrapper
-from datahub.metadata.schema_classes import (
-    AssertionInfoClass,
-    AssertionResultClass,
-    AssertionResultTypeClass,
-    AssertionRunEventClass,
-    AssertionRunStatusClass,
-    AssertionStdParameterClass,
-    AssertionStdParametersClass,
-    AssertionTypeClass,
-    BatchSpecClass,
-    DataPlatformInstanceClass,
-    DatasetAssertionInfoClass,
-    DatasetAssertionScopeClass,
-    PartitionSpecClass,
-)
 from great_expectations.core.batch import Batch, BatchDefinition, BatchRequest
 from great_expectations.core.batch_spec import (
     RuntimeDataBatchSpec,
@@ -46,6 +30,22 @@
 )
 from great_expectations.validator.validator import Validator
 
+from datahub.emitter.mcp import MetadataChangeProposalWrapper
+from datahub.metadata.schema_classes import (
+    AssertionInfoClass,
+    AssertionResultClass,
+    AssertionResultTypeClass,
+    AssertionRunEventClass,
+    AssertionRunStatusClass,
+    AssertionStdParameterClass,
+    AssertionStdParametersClass,
+    AssertionTypeClass,
+    BatchSpecClass,
+    DataPlatformInstanceClass,
+    DatasetAssertionInfoClass,
+    DatasetAssertionScopeClass,
+    PartitionSpecClass,
+)
 from datahub_gx_plugin.action import DataHubValidationAction
 
 logger = logging.getLogger(__name__)
diff --git a/metadata-ingestion-modules/prefect-plugin/.gitignore b/metadata-ingestion-modules/prefect-plugin/.gitignore
index 1d2916d00eabde..6801b785ea1e4b 100644
--- a/metadata-ingestion-modules/prefect-plugin/.gitignore
+++ b/metadata-ingestion-modules/prefect-plugin/.gitignore
@@ -1,5 +1,4 @@
 .envrc
-src/prefect_datahub/__init__.py.bak
 .vscode/
 output
 pvenv36/
diff --git a/metadata-ingestion-modules/prefect-plugin/build.gradle b/metadata-ingestion-modules/prefect-plugin/build.gradle
index a2b240cc8bf708..d13c9fe3c9abe7 100644
--- a/metadata-ingestion-modules/prefect-plugin/build.gradle
+++ b/metadata-ingestion-modules/prefect-plugin/build.gradle
@@ -54,18 +54,15 @@ task installDev(type: Exec, dependsOn: [install]) {
 task lint(type: Exec, dependsOn: installDev) {
   commandLine 'bash', '-c',
     "source ${venv_name}/bin/activate && set -x && " +
-    "black --check --diff src/ tests/ && " +
-    "isort --check --diff src/ tests/ && " +
-    "flake8 --count --statistics src/ tests/ && " +
+    "ruff check src/ tests/ && " +
+    "ruff format --check src/ tests/ && " +
     "mypy --show-traceback --show-error-codes src/ tests/"
 }
 task lintFix(type: Exec, dependsOn: installDev) {
   commandLine 'bash', '-x', '-c',
     "source ${venv_name}/bin/activate && " +
-    "black src/ tests/ && " +
-    "isort src/ tests/ && " +
-    "flake8 src/ tests/ && " +
-    "mypy src/ tests/ "
+    "ruff check --fix src/ tests/ && " +
+    "ruff format src/ tests/ "
 }
 
 task installDevTest(type: Exec, dependsOn: [installDev]) {
@@ -113,5 +110,8 @@ clean {
   delete venv_name
   delete 'build'
   delete 'dist'
+  delete '.ruff_cache'
+  delete '.mypy_cache'
+  delete '.pytest_cache'
 }
 clean.dependsOn cleanPythonCache
diff --git a/metadata-ingestion-modules/prefect-plugin/pyproject.toml b/metadata-ingestion-modules/prefect-plugin/pyproject.toml
index fba81486b9f677..d1e1d0ad479442 100644
--- a/metadata-ingestion-modules/prefect-plugin/pyproject.toml
+++ b/metadata-ingestion-modules/prefect-plugin/pyproject.toml
@@ -2,18 +2,58 @@
 build-backend = "setuptools.build_meta"
 requires = ["setuptools>=54.0.0", "wheel", "pip>=21.0.0"]
 
-[tool.black]
-extend-exclude = '''
-# A regex preceded with ^/ will apply only to files and directories
-# in the root of the project.
-^/tmp
-'''
-include = '\.pyi?$'
+[tool.ruff]
+line-length = 88
+target-version = "py38"
+exclude = [
+    ".git",
+    "venv",
+    ".tox",
+    "__pycache__",
+]
 
-[tool.isort]
-indent = '    '
-profile = 'black'
-sections = 'FUTURE,STDLIB,THIRDPARTY,FIRSTPARTY,LOCALFOLDER'
+[tool.ruff.format]
+quote-style = "double"
+indent-style = "space"
+skip-magic-trailing-comma = false
+line-ending = "auto"
 
-[tool.pyright]
-extraPaths = ['tests']
\ No newline at end of file
+[tool.ruff.lint.isort]
+combine-as-imports = true
+known-first-party = ["datahub"]
+extra-standard-library = ["__future__", "datahub.utilities._markupsafe_compat", "datahub.sql_parsing._sqlglot_patch"]
+section-order = ["future", "standard-library", "third-party", "first-party", "local-folder"]
+force-sort-within-sections = false
+force-wrap-aliases = false
+split-on-trailing-comma = false
+order-by-type = true
+relative-imports-order = "closest-to-furthest"
+force-single-line = false
+single-line-exclusions = ["typing"]
+length-sort = false
+from-first = false
+required-imports = []
+classes = ["typing"]
+
+[tool.ruff.lint]
+extend-select = [
+    "B",    # flake8-bugbear
+    "C90",  # mccabe complexity
+    "E",    # pycodestyle errors
+    "F",    # pyflakes
+    "G010", # logging.warn -> logging.warning
+    "I",    # isort
+    "TID",  # flake8-tidy-imports
+]
+ignore = [
+    "E501",  # Line length violations (handled by formatter)
+]
+
+[tool.ruff.lint.mccabe]
+max-complexity = 15
+
+[tool.ruff.lint.flake8-tidy-imports]
+ban-relative-imports = "all"
+
+[tool.ruff.lint.per-file-ignores]
+"__init__.py" = ["F401"]
\ No newline at end of file
diff --git a/metadata-ingestion-modules/prefect-plugin/scripts/release.sh b/metadata-ingestion-modules/prefect-plugin/scripts/release.sh
index f398db98b60290..b3b99d61c904a3 100755
--- a/metadata-ingestion-modules/prefect-plugin/scripts/release.sh
+++ b/metadata-ingestion-modules/prefect-plugin/scripts/release.sh
@@ -1,26 +1,31 @@
 #!/bin/bash
+# Auto-generated by python-build/generate_release_scripts.py. Do not edit manually.
+
 set -euxo pipefail
 
+ROOT=../..
+MODULE=prefect_datahub
+
 if [[ ! ${RELEASE_SKIP_TEST:-} ]] && [[ ! ${RELEASE_SKIP_INSTALL:-} ]]; then
-	../../gradlew build  # also runs tests
+    ${ROOT}/gradlew build  # also runs tests
 elif [[ ! ${RELEASE_SKIP_INSTALL:-} ]]; then
-	../../gradlew install
+    ${ROOT}/gradlew install
 fi
 
-MODULE=prefect_datahub
-
 # Check packaging constraint.
 python -c 'import setuptools; where="./src"; assert setuptools.find_packages(where) == setuptools.find_namespace_packages(where), "you seem to be missing or have extra __init__.py files"'
-if [[ ${RELEASE_VERSION:-} ]]; then
-    # Replace version with RELEASE_VERSION env variable
-    sed -i.bak "s/__version__ = \"1\!0.0.0.dev0\"/__version__ = \"$(echo $RELEASE_VERSION|sed s/-/+/)\"/" src/${MODULE}/__init__.py
-else
-    vim src/${MODULE}/__init__.py
+
+# Update the release version.
+if [[ ! ${RELEASE_VERSION:-} ]]; then
+    echo "RELEASE_VERSION is not set"
+    exit 1
 fi
+sed -i.bak "s/__version__ = .*$/__version__ = \"$(echo $RELEASE_VERSION|sed s/-/+/)\"/" src/${MODULE}/_version.py
 
+# Build and upload the release.
 rm -rf build dist || true
 python -m build
 if [[ ! ${RELEASE_SKIP_UPLOAD:-} ]]; then
     python -m twine upload 'dist/*'
 fi
-mv src/${MODULE}/__init__.py.bak src/${MODULE}/__init__.py
\ No newline at end of file
+mv src/${MODULE}/_version.py.bak src/${MODULE}/_version.py
diff --git a/metadata-ingestion-modules/prefect-plugin/setup.cfg b/metadata-ingestion-modules/prefect-plugin/setup.cfg
index c59a99fa8aec0a..828597e6a812b0 100644
--- a/metadata-ingestion-modules/prefect-plugin/setup.cfg
+++ b/metadata-ingestion-modules/prefect-plugin/setup.cfg
@@ -1,24 +1,3 @@
-[flake8]
-max-complexity = 15
-ignore =
-    # Ignore: line length issues, since black's formatter will take care of them.
-    E501,
-    # Ignore: 1 blank line required before class docstring.
-    D203,
-    # See https://stackoverflow.com/a/57074416.
-    W503,
-    # See https://github.com/psf/black/issues/315.
-    E203
-exclude =
-    .git,
-    venv,
-    .tox,
-    __pycache__
-per-file-ignores =
-    # imported but unused
-    __init__.py: F401
-ban-relative-imports = true
-
 [mypy]
 plugins =
     sqlmypy,
diff --git a/metadata-ingestion-modules/prefect-plugin/setup.py b/metadata-ingestion-modules/prefect-plugin/setup.py
index 70b0e958195645..87feb810b8e5a6 100644
--- a/metadata-ingestion-modules/prefect-plugin/setup.py
+++ b/metadata-ingestion-modules/prefect-plugin/setup.py
@@ -4,7 +4,7 @@
 import setuptools
 
 package_metadata: dict = {}
-with open("./src/prefect_datahub/__init__.py") as fp:
+with open("./src/prefect_datahub/_version.py") as fp:
     exec(fp.read(), package_metadata)
 
 
@@ -30,9 +30,7 @@ def get_long_description():
     # Temporary pinning to 2.0.0 until we can upgrade to 3.0.0
     "prefect >= 2.0.0,<3.0.0",
     *rest_common,
-    # Ignoring the dependency below because it causes issues with the vercel built wheel install
-    # f"acryl-datahub[datahub-rest]{_self_pin}",
-    "acryl-datahub[datahub-rest]",
+    f"acryl-datahub[datahub-rest]{_self_pin}",
 }
 
 
@@ -57,11 +55,8 @@ def get_long_description():
 dev_requirements = {
     *base_requirements,
     *mypy_stubs,
-    "black==22.12.0",
     "coverage>=5.1",
-    "flake8>=3.8.3",
-    "flake8-tidy-imports>=4.3.0",
-    "isort>=5.7.0",
+    "ruff==0.9.2",
     "mypy>=1.4.0",
     # pydantic 1.8.2 is incompatible with mypy 0.910.
     # See https://github.com/samuelcolvin/pydantic/pull/3175#issuecomment-995382910.
diff --git a/metadata-ingestion-modules/prefect-plugin/src/prefect_datahub/__init__.py b/metadata-ingestion-modules/prefect-plugin/src/prefect_datahub/__init__.py
index 8cc65f9010613d..f38863a1f31e28 100644
--- a/metadata-ingestion-modules/prefect-plugin/src/prefect_datahub/__init__.py
+++ b/metadata-ingestion-modules/prefect-plugin/src/prefect_datahub/__init__.py
@@ -1,21 +1 @@
-# Published at https://pypi.org/project/acryl-datahub/.
-__package_name__ = "prefect-datahub"
-__version__ = "1!0.0.0.dev0"
-
-
-def is_dev_mode() -> bool:
-    return __version__.endswith("dev0")
-
-
-def nice_version_name() -> str:
-    if is_dev_mode():
-        return "unavailable (installed in develop mode)"
-    return __version__
-
-
-def get_provider_info():
-    return {
-        "package-name": f"{__package_name__}",
-        "name": f"{__package_name__}",
-        "description": "Datahub prefect block to capture executions and send to Datahub",
-    }
+from prefect_datahub._version import __package_name__, __version__
diff --git a/metadata-ingestion-modules/prefect-plugin/src/prefect_datahub/_version.py b/metadata-ingestion-modules/prefect-plugin/src/prefect_datahub/_version.py
new file mode 100644
index 00000000000000..8c45e7d04b3679
--- /dev/null
+++ b/metadata-ingestion-modules/prefect-plugin/src/prefect_datahub/_version.py
@@ -0,0 +1,3 @@
+# Published at https://pypi.org/project/prefect-datahub/.
+__package_name__ = "prefect-datahub"
+__version__ = "1!0.0.0.dev0"
diff --git a/metadata-ingestion-modules/prefect-plugin/src/prefect_datahub/datahub_emitter.py b/metadata-ingestion-modules/prefect-plugin/src/prefect_datahub/datahub_emitter.py
index 8617381cf16139..190a249a912d1a 100644
--- a/metadata-ingestion-modules/prefect-plugin/src/prefect_datahub/datahub_emitter.py
+++ b/metadata-ingestion-modules/prefect-plugin/src/prefect_datahub/datahub_emitter.py
@@ -5,6 +5,15 @@
 from typing import Any, Dict, List, Optional, cast
 from uuid import UUID
 
+from prefect import get_run_logger
+from prefect.blocks.core import Block
+from prefect.client import cloud, orchestration
+from prefect.client.schemas import FlowRun, TaskRun, Workspace
+from prefect.client.schemas.objects import Flow
+from prefect.context import FlowRunContext, TaskRunContext
+from prefect.settings import PREFECT_API_URL
+from pydantic.v1 import SecretStr
+
 import datahub.emitter.mce_builder as builder
 from datahub.api.entities.datajob import DataFlow, DataJob
 from datahub.api.entities.dataprocess.dataprocess_instance import (
@@ -17,15 +26,6 @@
 from datahub.utilities.urns.data_flow_urn import DataFlowUrn
 from datahub.utilities.urns.data_job_urn import DataJobUrn
 from datahub.utilities.urns.dataset_urn import DatasetUrn
-from prefect import get_run_logger
-from prefect.blocks.core import Block
-from prefect.client import cloud, orchestration
-from prefect.client.schemas import FlowRun, TaskRun, Workspace
-from prefect.client.schemas.objects import Flow
-from prefect.context import FlowRunContext, TaskRunContext
-from prefect.settings import PREFECT_API_URL
-from pydantic.v1 import SecretStr
-
 from prefect_datahub.entities import _Entity
 
 ORCHESTRATOR = "prefect"
@@ -351,9 +351,9 @@ def _emit_tasks(
 
             for prefect_future in flow_run_ctx.task_run_futures:
                 if prefect_future.task_run is not None:
-                    task_run_key_map[
-                        str(prefect_future.task_run.id)
-                    ] = prefect_future.task_run.task_key
+                    task_run_key_map[str(prefect_future.task_run.id)] = (
+                        prefect_future.task_run.task_key
+                    )
 
             for node in graph_json:
                 datajob_urn = DataJobUrn.create_from_ids(
diff --git a/metadata-ingestion-modules/prefect-plugin/src/prefect_datahub/entities.py b/metadata-ingestion-modules/prefect-plugin/src/prefect_datahub/entities.py
index e2711d0925d977..c8f858b8e59c19 100644
--- a/metadata-ingestion-modules/prefect-plugin/src/prefect_datahub/entities.py
+++ b/metadata-ingestion-modules/prefect-plugin/src/prefect_datahub/entities.py
@@ -2,6 +2,7 @@
 from typing import Optional
 
 import attr
+
 import datahub.emitter.mce_builder as builder
 from datahub.utilities.urns.urn import guess_entity_type
 
diff --git a/metadata-ingestion-modules/prefect-plugin/tests/unit/test_datahub_emitter.py b/metadata-ingestion-modules/prefect-plugin/tests/unit/test_datahub_emitter.py
index ba50cddc986b66..71555f76e04f71 100644
--- a/metadata-ingestion-modules/prefect-plugin/tests/unit/test_datahub_emitter.py
+++ b/metadata-ingestion-modules/prefect-plugin/tests/unit/test_datahub_emitter.py
@@ -6,14 +6,14 @@
 from uuid import UUID
 
 import pytest
-from datahub.api.entities.datajob import DataJob
-from datahub.utilities.urns.dataset_urn import DatasetUrn
 from prefect.client.schemas import FlowRun, TaskRun, Workspace
 from prefect.futures import PrefectFuture
 from prefect.server.schemas.core import Flow
 from prefect.task_runners import SequentialTaskRunner
 from requests.models import Response
 
+from datahub.api.entities.datajob import DataJob
+from datahub.utilities.urns.dataset_urn import DatasetUrn
 from prefect_datahub.datahub_emitter import DatahubEmitter
 from prefect_datahub.entities import Dataset, _Entity
 
diff --git a/metadata-ingestion/.gitignore b/metadata-ingestion/.gitignore
index acc15c45988698..c7a781ded1d688 100644
--- a/metadata-ingestion/.gitignore
+++ b/metadata-ingestion/.gitignore
@@ -1,5 +1,4 @@
 .envrc
-src/datahub/__init__.py.bak
 .vscode/
 output
 pvenv36/
diff --git a/metadata-ingestion/LICENSE b/metadata-ingestion/LICENSE
new file mode 100644
index 00000000000000..5e8a14cef59b4c
--- /dev/null
+++ b/metadata-ingestion/LICENSE
@@ -0,0 +1,202 @@
+
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright 2015 LinkedIn Corp. All rights reserved.
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.
\ No newline at end of file
diff --git a/metadata-ingestion/build.gradle b/metadata-ingestion/build.gradle
index 6629a0b9ef022e..be9d69a2f0e4b6 100644
--- a/metadata-ingestion/build.gradle
+++ b/metadata-ingestion/build.gradle
@@ -23,8 +23,10 @@ task environmentSetup(type: Exec, dependsOn: checkPythonVersion) {
   inputs.file file('setup.py')
   outputs.file(sentinel_file)
   commandLine 'bash', '-c',
-    "${python_executable} -m venv ${venv_name} && set -x && " +
-    "${venv_name}/bin/python -m pip install --upgrade uv && " +
+    "if [ ! -d ${venv_name} ] || [ ! -f ${venv_name}/bin/python ]; then ${python_executable} -m venv ${venv_name}; fi && " +
+    "set -x && " +
+    // If we already have uv available, use it to upgrade uv. Otherwise, install it with pip.
+    "if [ ! -f ${venv_name}/bin/uv ]; then ${venv_name}/bin/python -m pip install --upgrade uv; else ${venv_name}/bin/python -m uv pip install --upgrade uv; fi && " +
     "touch ${sentinel_file}"
 }
 
@@ -108,16 +110,16 @@ task modelDocUpload(type: Exec, dependsOn: [modelDocGen]) {
 task lint(type: Exec, dependsOn: installDev) {
   commandLine 'bash', '-c',
     "source ${venv_name}/bin/activate && set -x && " +
-    "black --check --diff src/ tests/ examples/ && " +
-    "isort --check --diff src/ tests/ examples/ && " +
-    "flake8 --count --statistics src/ tests/ examples/ && " +
+    "ruff check src/ tests/ examples/ && " +
+    "ruff format --check src/ tests/ examples/ && " +
     "mypy --show-traceback --show-error-codes src/ tests/ examples/"
 }
+
 task lintFix(type: Exec, dependsOn: installDev) {
   commandLine 'bash', '-c',
     "source ${venv_name}/bin/activate && set -x && " +
-    "black src/ tests/ examples/ && " +
-    "isort src/ tests/ examples/"
+    "ruff check --fix src/ tests/ examples/ && " +
+    "ruff format src/ tests/ examples/ "
 }
 
 def pytest_default_env = "PYTHONDEVMODE=1"
@@ -214,6 +216,7 @@ clean {
   delete 'src/datahub/metadata'
   delete '../docs/generated'
   delete 'generated'
+  delete '.ruff_cache'
   delete '.mypy_cache'
   delete '.pytest_cache'
   delete '.preflight_sentinel'
diff --git a/metadata-ingestion/developing.md b/metadata-ingestion/developing.md
index 5e4bc2d4d71158..005b0427a5e6ad 100644
--- a/metadata-ingestion/developing.md
+++ b/metadata-ingestion/developing.md
@@ -177,13 +177,11 @@ The architecture of this metadata ingestion framework is heavily inspired by [Ap
 
 ## Code style
 
-We use black, isort, flake8, and mypy to ensure consistent code style and quality.
+We use ruff, and mypy to ensure consistent code style and quality.
 
 ```shell
 # Assumes: pip install -e '.[dev]' and venv is activated
-black src/ tests/
-isort src/ tests/
-flake8 src/ tests/
+ruff check src/ tests/
 mypy src/ tests/
 ```
 
diff --git a/metadata-ingestion/docs/sources/redshift/redshift_recipe.yml b/metadata-ingestion/docs/sources/redshift/redshift_recipe.yml
index a561405d3de47e..e6d2dc3de551d6 100644
--- a/metadata-ingestion/docs/sources/redshift/redshift_recipe.yml
+++ b/metadata-ingestion/docs/sources/redshift/redshift_recipe.yml
@@ -1,3 +1,4 @@
+source:
   type: redshift
   config:
     # Coordinates
diff --git a/metadata-ingestion/docs/sources/tableau/tableau_pre.md b/metadata-ingestion/docs/sources/tableau/tableau_pre.md
index 65ff08367fdc8f..a3ac85818a51aa 100644
--- a/metadata-ingestion/docs/sources/tableau/tableau_pre.md
+++ b/metadata-ingestion/docs/sources/tableau/tableau_pre.md
@@ -16,9 +16,12 @@ Either way, the user/token must have the **Site Administrator Explorer** site ro
 
 :::info
 
-We need the `Site Administrator Explorer` site role in order to get complete metadata from Tableau.
+We need the **Site Administrator Explorer** site role in order to get complete metadata from Tableau.
 
-With any lower role, the Tableau Metadata API returns missing/partial metadata. This particularly affects data source fields and definitions, which impacts our ability to extract columns and generate column lineage. As such, other site roles like `Viewer` are insufficient with the current Tableau Metadata API.
+With any lower role, the Tableau Metadata API returns missing/partial metadata.
+This particularly affects data source fields and definitions, which impacts our ability to extract most columns and generate column lineage. Some table-level lineage is also impacted.
+
+Other site roles, including Site Administrator Creator and Viewer, are insufficient due to these limitations in the current Tableau Metadata API.
 
 :::
 
diff --git a/metadata-ingestion/examples/library/run_assertion.py b/metadata-ingestion/examples/library/run_assertion.py
index 414e5f46cc7f91..e7c717837eed3c 100644
--- a/metadata-ingestion/examples/library/run_assertion.py
+++ b/metadata-ingestion/examples/library/run_assertion.py
@@ -16,5 +16,5 @@
 assertion_result = graph.run_assertion(urn=assertion_urn, save_result=True)
 
 log.info(
-    f'Assertion result (SUCCESS / FAILURE / ERROR): {assertion_result.get("type")}'
+    f"Assertion result (SUCCESS / FAILURE / ERROR): {assertion_result.get('type')}"
 )
diff --git a/metadata-ingestion/pyproject.toml b/metadata-ingestion/pyproject.toml
index 94e06fd53a70ec..1d434eb8c3a94f 100644
--- a/metadata-ingestion/pyproject.toml
+++ b/metadata-ingestion/pyproject.toml
@@ -2,30 +2,53 @@
 build-backend = "setuptools.build_meta"
 requires = ["setuptools>=63.0.0", "wheel"]
 
-[tool.black]
-extend-exclude = '''
-# A regex preceded with ^/ will apply only to files and directories
-# in the root of the project.
-^/tmp
-'''
-include = '\.pyi?$'
-target-version = ['py37', 'py38', 'py39', 'py310']
+[tool.ruff.lint.isort]
+section-order = ["future", "patch", "standard-library", "third-party", "first-party", "local-folder"]
+sections = { "patch" = ["datahub.utilities._markupsafe_compat", "datahub.sql_parsing._sqlglot_patch"] }
+combine-as-imports = true
 
-[tool.isort]
-combine_as_imports = true
-indent = '    '
-known_future_library = ['__future__', 'datahub.utilities._markupsafe_compat', 'datahub.sql_parsing._sqlglot_patch']
-profile = 'black'
-sections = 'FUTURE,STDLIB,THIRDPARTY,FIRSTPARTY,LOCALFOLDER'
-skip_glob = 'src/datahub/metadata'
+[tool.ruff]
+target-version = "py38"
+# Same as Black.
+line-length = 88
+# Exclude directories matching these patterns.
+exclude = [
+    ".git",
+    "src/datahub/metadata",
+    "venv",
+    ".tox",
+    "__pycache__",
+]
 
-[tool.pyright]
-extraPaths = ['tests']
+[tool.ruff.lint]
+extend-select = [
+    "B",    # flake8-bugbear
+    "C90",  # mccabe complexity
+    "E",    # pycodestyle errors
+    "F",    # pyflakes
+    "G010", # logging.warn -> logging.warning
+    "I",    # isort
+    "TID",  # flake8-tidy-imports
+]
+extend-ignore = [
+    "E501",  # Handled by formatter
+    "E111",  # Handled by formatter
+    "E114",  # Handled by formatter
+    "E117",  # Handled by formatter
+    "E203",  # Ignore whitespace before ':' (matches Black)
+    "B019",  # Allow usages of functools.lru_cache
+    "B008",  # Allow function call in argument defaults
+    # TODO: Enable these later
+    "B006",  # Mutable args
+    "B017",  # Do not assert blind exception
+    "B904",  # Checks for raise statements in exception handlers that lack a from clause
+]
 
-[tool.vulture]
-exclude = ["src/datahub/metadata/"]
-ignore_decorators = ["@click.*", "@validator", "@root_validator", "@pydantic.validator", "@pydantic.root_validator", "@pytest.fixture"]
-ignore_names = ["*Source", "*Sink", "*Report"]
-# min_confidence = 80
-paths = ["src"]
-sort_by_size = true
+[tool.ruff.lint.mccabe]
+max-complexity = 20
+
+[tool.ruff.lint.flake8-tidy-imports]
+ban-relative-imports = "all"
+
+[tool.ruff.lint.per-file-ignores]
+"__init__.py" = ["F401"]
diff --git a/metadata-ingestion/scripts/avro_codegen.py b/metadata-ingestion/scripts/avro_codegen.py
index 0fe79a2c6a8e47..7e75cba9833810 100644
--- a/metadata-ingestion/scripts/avro_codegen.py
+++ b/metadata-ingestion/scripts/avro_codegen.py
@@ -277,7 +277,7 @@ def annotate_aspects(aspects: List[dict], schema_class_file: Path) -> None:
 """
 
     for aspect in aspects:
-        className = f'{aspect["name"]}Class'
+        className = f"{aspect['name']}Class"
         aspectName = aspect["Aspect"]["name"]
         class_def_original = f"class {className}(DictWrapper):"
 
@@ -299,9 +299,9 @@ def annotate_aspects(aspects: List[dict], schema_class_file: Path) -> None:
         schema_classes_lines[empty_line] = "\n"
         schema_classes_lines[empty_line] += f"\n    ASPECT_NAME = '{aspectName}'"
         if "type" in aspect["Aspect"]:
-            schema_classes_lines[
-                empty_line
-            ] += f"\n    ASPECT_TYPE = '{aspect['Aspect']['type']}'"
+            schema_classes_lines[empty_line] += (
+                f"\n    ASPECT_TYPE = '{aspect['Aspect']['type']}'"
+            )
 
         aspect_info = {
             k: v for k, v in aspect["Aspect"].items() if k not in {"name", "type"}
@@ -315,7 +315,7 @@ def annotate_aspects(aspects: List[dict], schema_class_file: Path) -> None:
     schema_classes_lines.append(
         f"""
 ASPECT_CLASSES: List[Type[_Aspect]] = [
-    {f',{newline}    '.join(f"{aspect['name']}Class" for aspect in aspects)}
+    {f",{newline}    ".join(f"{aspect['name']}Class" for aspect in aspects)}
 ]
 
 ASPECT_NAME_MAP: Dict[str, Type[_Aspect]] = {{
@@ -326,11 +326,11 @@ def annotate_aspects(aspects: List[dict], schema_class_file: Path) -> None:
 from typing_extensions import TypedDict
 
 class AspectBag(TypedDict, total=False):
-    {f'{newline}    '.join(f"{aspect['Aspect']['name']}: {aspect['name']}Class" for aspect in aspects)}
+    {f"{newline}    ".join(f"{aspect['Aspect']['name']}: {aspect['name']}Class" for aspect in aspects)}
 
 
 KEY_ASPECTS: Dict[str, Type[_Aspect]] = {{
-    {f',{newline}    '.join(f"'{aspect['Aspect']['keyForEntity']}': {aspect['name']}Class" for aspect in aspects if aspect['Aspect'].get('keyForEntity'))}
+    {f",{newline}    ".join(f"'{aspect['Aspect']['keyForEntity']}': {aspect['name']}Class" for aspect in aspects if aspect["Aspect"].get("keyForEntity"))}
 }}
 """
     )
@@ -546,6 +546,7 @@ def generate_urn_class(entity_type: str, key_aspect: dict) -> str:
         assert field_name(fields[0]) == "guid"
         assert fields[0]["type"] == ["null", "string"]
         fields[0]["type"] = "string"
+    arg_count = len(fields)
 
     field_urn_type_classes = {}
     for field in fields:
@@ -561,6 +562,12 @@ def generate_urn_class(entity_type: str, key_aspect: dict) -> str:
                 field_urn_type_class = "Urn"
 
         field_urn_type_classes[field_name(field)] = field_urn_type_class
+    if arg_count == 1:
+        field = fields[0]
+
+        if field_urn_type_classes[field_name(field)] is None:
+            # All single-arg urn types should accept themselves.
+            field_urn_type_classes[field_name(field)] = class_name
 
     _init_arg_parts: List[str] = []
     for field in fields:
@@ -579,7 +586,6 @@ def generate_urn_class(entity_type: str, key_aspect: dict) -> str:
 
     super_init_args = ", ".join(field_name(field) for field in fields)
 
-    arg_count = len(fields)
     parse_ids_mapping = ", ".join(
         f"{field_name(field)}=entity_ids[{i}]" for i, field in enumerate(fields)
     )
@@ -601,8 +607,26 @@ def generate_urn_class(entity_type: str, key_aspect: dict) -> str:
 
         # Generalized mechanism for validating embedded urns.
         field_urn_type_class = field_urn_type_classes[field_name(field)]
-        if field_urn_type_class:
-            init_validation += f"{field_name(field)} = str({field_name(field)})\n"
+        if field_urn_type_class and field_urn_type_class == class_name:
+            # First, we need to extract the main piece from the urn type.
+            init_validation += (
+                f"if isinstance({field_name(field)}, {field_urn_type_class}):\n"
+                f"    {field_name(field)} = {field_name(field)}.{field_name(field)}\n"
+            )
+
+            # If it's still an urn type, that's a problem.
+            init_validation += (
+                f"elif isinstance({field_name(field)}, Urn):\n"
+                f"    raise InvalidUrnError(f'Expecting a {field_urn_type_class} but got {{{field_name(field)}}}')\n"
+            )
+
+            # Then, we do character validation as normal.
+            init_validation += (
+                f"if UrnEncoder.contains_reserved_char({field_name(field)}):\n"
+                f"    raise InvalidUrnError(f'{class_name} {field_name(field)} contains reserved characters')\n"
+            )
+        elif field_urn_type_class:
+            init_validation += f"{field_name(field)} = str({field_name(field)})  # convert urn type to str\n"
             init_validation += (
                 f"assert {field_urn_type_class}.from_string({field_name(field)})\n"
             )
@@ -611,17 +635,29 @@ def generate_urn_class(entity_type: str, key_aspect: dict) -> str:
                 f"if UrnEncoder.contains_reserved_char({field_name(field)}):\n"
                 f"    raise InvalidUrnError(f'{class_name} {field_name(field)} contains reserved characters')\n"
             )
+        # TODO add ALL_ENV_TYPES validation
 
+        # Field coercion logic.
         if field_name(field) == "env":
             init_coercion += "env = env.upper()\n"
-        # TODO add ALL_ENV_TYPES validation
-        elif entity_type == "dataPlatform" and field_name(field) == "platform_name":
-            init_coercion += 'if platform_name.startswith("urn:li:dataPlatform:"):\n'
-            init_coercion += "    platform_name = DataPlatformUrn.from_string(platform_name).platform_name\n"
-
-        if field_name(field) == "platform":
-            init_coercion += "platform = platform.urn() if isinstance(platform, DataPlatformUrn) else DataPlatformUrn(platform).urn()\n"
-        elif field_urn_type_class is None:
+        elif field_name(field) == "platform":
+            # For platform names in particular, we also qualify them when they don't have the prefix.
+            # We can rely on the DataPlatformUrn constructor to do this prefixing.
+            init_coercion += "platform = DataPlatformUrn(platform).urn()\n"
+        elif field_urn_type_class is not None:
+            # For urn types, we need to parse them into urn types where appropriate.
+            # Otherwise, we just need to encode special characters.
+            init_coercion += (
+                f"if isinstance({field_name(field)}, str):\n"
+                f"    if {field_name(field)}.startswith('urn:li:'):\n"
+                f"        try:\n"
+                f"            {field_name(field)} = {field_urn_type_class}.from_string({field_name(field)})\n"
+                f"        except InvalidUrnError:\n"
+                f"            raise InvalidUrnError(f'Expecting a {field_urn_type_class} but got {{{field_name(field)}}}')\n"
+                f"    else:\n"
+                f"        {field_name(field)} = UrnEncoder.encode_string({field_name(field)})\n"
+            )
+        else:
             # For all non-urns, run the value through the UrnEncoder.
             init_coercion += (
                 f"{field_name(field)} = UrnEncoder.encode_string({field_name(field)})\n"
@@ -642,10 +678,10 @@ class {class_name}(_SpecificUrn):
     def __init__(self, {init_args}, *, _allow_coercion: bool = True) -> None:
         if _allow_coercion:
             # Field coercion logic (if any is required).
-{textwrap.indent(init_coercion.strip(), prefix=" "*4*3)}
+{textwrap.indent(init_coercion.strip(), prefix=" " * 4 * 3)}
 
         # Validation logic.
-{textwrap.indent(init_validation.strip(), prefix=" "*4*2)}
+{textwrap.indent(init_validation.strip(), prefix=" " * 4 * 2)}
 
         super().__init__(self.ENTITY_TYPE, [{super_init_args}])
 
@@ -729,7 +765,7 @@ def generate(
             and aspect["Aspect"]["keyForEntity"] != entity.name
         ):
             raise ValueError(
-                f'Entity key {entity.keyAspect} is used by {aspect["Aspect"]["keyForEntity"]} and {entity.name}'
+                f"Entity key {entity.keyAspect} is used by {aspect['Aspect']['keyForEntity']} and {entity.name}"
             )
 
         # Also require that the aspect list is deduplicated.
diff --git a/metadata-ingestion/scripts/release.sh b/metadata-ingestion/scripts/release.sh
index a18dd6f934b431..9b4e62aef581ae 100755
--- a/metadata-ingestion/scripts/release.sh
+++ b/metadata-ingestion/scripts/release.sh
@@ -1,26 +1,31 @@
 #!/bin/bash
+# Auto-generated by python-build/generate_release_scripts.py. Do not edit manually.
+
 set -euxo pipefail
 
+ROOT=..
+MODULE=datahub
+
 if [[ ! ${RELEASE_SKIP_TEST:-} ]] && [[ ! ${RELEASE_SKIP_INSTALL:-} ]]; then
-	../gradlew build  # also runs tests
+    ${ROOT}/gradlew build  # also runs tests
 elif [[ ! ${RELEASE_SKIP_INSTALL:-} ]]; then
-	../gradlew install
+    ${ROOT}/gradlew install
 fi
 
-MODULE=datahub
-
 # Check packaging constraint.
 python -c 'import setuptools; where="./src"; assert setuptools.find_packages(where) == setuptools.find_namespace_packages(where), "you seem to be missing or have extra __init__.py files"'
-if [[ ${RELEASE_VERSION:-} ]]; then
-    # Replace version with RELEASE_VERSION env variable
-    sed -i.bak "s/__version__ = \"1\!0.0.0.dev0\"/__version__ = \"$(echo $RELEASE_VERSION|sed s/-/+/)\"/" src/${MODULE}/__init__.py
-else
-    vim src/${MODULE}/__init__.py
+
+# Update the release version.
+if [[ ! ${RELEASE_VERSION:-} ]]; then
+    echo "RELEASE_VERSION is not set"
+    exit 1
 fi
+sed -i.bak "s/__version__ = .*$/__version__ = \"$(echo $RELEASE_VERSION|sed s/-/+/)\"/" src/${MODULE}/_version.py
 
+# Build and upload the release.
 rm -rf build dist || true
 python -m build
 if [[ ! ${RELEASE_SKIP_UPLOAD:-} ]]; then
     python -m twine upload 'dist/*'
 fi
-mv src/${MODULE}/__init__.py.bak src/${MODULE}/__init__.py
+mv src/${MODULE}/_version.py.bak src/${MODULE}/_version.py
diff --git a/metadata-ingestion/setup.cfg b/metadata-ingestion/setup.cfg
index 057779bc87c622..b7cf43b80b149e 100644
--- a/metadata-ingestion/setup.cfg
+++ b/metadata-ingestion/setup.cfg
@@ -1,39 +1,3 @@
-[flake8]
-max-complexity = 20
-ignore =
-    # Ignore: line length issues, since black's formatter will take care of them.
-    E501,
-    # Ignore compound statements, since they're used for ellipsis by black
-    # See https://github.com/psf/black/issues/3887
-    E704,
-    # Ignore: 1 blank line required before class docstring.
-    D203,
-    # See https://stackoverflow.com/a/57074416.
-    W503,
-    # See https://github.com/psf/black/issues/315.
-    E203,
-    # Allow usages of functools.lru_cache.
-    B019,
-    # This rule flags the use of function calls in argument defaults.
-    # There's some good reasons to do this, so we're ok with it.
-    B008,
-    # TODO: However, we should enable B006 to catch issues with mutable args.
-    B006,
-    # TODO: Enable B007 - unused loop control variable.
-    B007
-    # TODO: Enable B902 - require self/cls naming.
-    # TODO: Enable B904 - use raise from in except clauses.
-exclude =
-    .git,
-    src/datahub/metadata,
-    venv,
-    .tox,
-    __pycache__
-per-file-ignores =
-    # imported but unused
-    __init__.py: F401, I250
-ban-relative-imports = true
-
 [mypy]
 plugins =
     ./tests/test_helpers/sqlalchemy_mypy_plugin.py,
diff --git a/metadata-ingestion/setup.py b/metadata-ingestion/setup.py
index d5dbb98d3cb17b..b3175989309849 100644
--- a/metadata-ingestion/setup.py
+++ b/metadata-ingestion/setup.py
@@ -3,7 +3,7 @@
 import setuptools
 
 package_metadata: dict = {}
-with open("./src/datahub/__init__.py") as fp:
+with open("./src/datahub/_version.py") as fp:
     exec(fp.read(), package_metadata)
 
 _version: str = package_metadata["__version__"]
@@ -312,7 +312,10 @@
 
 powerbi_report_server = {"requests", "requests_ntlm"}
 
-slack = {"slack-sdk==3.18.1"}
+slack = {
+    "slack-sdk==3.18.1",
+    "tenacity>=8.0.1",
+}
 
 databricks = {
     # 0.1.11 appears to have authentication issues with azure databricks
@@ -505,12 +508,10 @@
     "starburst-trino-usage": sql_common | usage_common | trino,
     "nifi": {"requests", "packaging", "requests-gssapi"},
     "powerbi": (
-        (
-            microsoft_common
-            | {"lark[regex]==1.1.4", "sqlparse", "more-itertools"}
-            | sqlglot_lib
-            | threading_timeout_common
-        )
+        microsoft_common
+        | {"lark[regex]==1.1.4", "sqlparse", "more-itertools"}
+        | sqlglot_lib
+        | threading_timeout_common
     ),
     "powerbi-report-server": powerbi_report_server,
     "vertica": sql_common | {"vertica-sqlalchemy-dialect[vertica-python]==0.0.8.2"},
@@ -592,11 +593,7 @@
 lint_requirements = {
     # This is pinned only to avoid spurious errors in CI.
     # We should make an effort to keep it up to date.
-    "black==23.3.0",
-    "flake8>=6.0.0",
-    "flake8-tidy-imports>=4.3.0",
-    "flake8-bugbear==23.3.12",
-    "isort>=5.7.0",
+    "ruff==0.9.2",
     "mypy==1.10.1",
 }
 
diff --git a/metadata-ingestion/src/datahub/__init__.py b/metadata-ingestion/src/datahub/__init__.py
index b254deb7fa30e5..8b8ef52d27bb92 100644
--- a/metadata-ingestion/src/datahub/__init__.py
+++ b/metadata-ingestion/src/datahub/__init__.py
@@ -1,25 +1 @@
-import sys
-import warnings
-
-# Published at https://pypi.org/project/acryl-datahub/.
-__package_name__ = "acryl-datahub"
-__version__ = "1!0.0.0.dev0"
-
-
-def is_dev_mode() -> bool:
-    return __version__.endswith("dev0")
-
-
-def nice_version_name() -> str:
-    if is_dev_mode():
-        return "unavailable (installed in develop mode)"
-    return __version__
-
-
-if sys.version_info < (3, 8):
-    warnings.warn(
-        "DataHub requires Python 3.8 or newer. "
-        "Please upgrade your Python version to continue using DataHub.",
-        FutureWarning,
-        stacklevel=2,
-    )
+from datahub._version import __package_name__, __version__
diff --git a/metadata-ingestion/src/datahub/_version.py b/metadata-ingestion/src/datahub/_version.py
new file mode 100644
index 00000000000000..a34748ac942a17
--- /dev/null
+++ b/metadata-ingestion/src/datahub/_version.py
@@ -0,0 +1,13 @@
+# Published at https://pypi.org/project/acryl-datahub/.
+__package_name__ = "acryl-datahub"
+__version__ = "1!0.0.0.dev0"
+
+
+def is_dev_mode() -> bool:
+    return __version__.endswith("dev0")
+
+
+def nice_version_name() -> str:
+    if is_dev_mode():
+        return "unavailable (installed in develop mode)"
+    return __version__
diff --git a/metadata-ingestion/src/datahub/api/entities/assertion/assertion_operator.py b/metadata-ingestion/src/datahub/api/entities/assertion/assertion_operator.py
index 8704ed13cb6c30..a05386798495de 100644
--- a/metadata-ingestion/src/datahub/api/entities/assertion/assertion_operator.py
+++ b/metadata-ingestion/src/datahub/api/entities/assertion/assertion_operator.py
@@ -20,15 +20,13 @@ class Operator(Protocol):
 
     operator: str
 
-    def id(self) -> str:
-        ...
+    def id(self) -> str: ...
 
-    def generate_parameters(self) -> AssertionStdParametersClass:
-        ...
+    def generate_parameters(self) -> AssertionStdParametersClass: ...
 
 
 def _generate_assertion_std_parameter(
-    value: Union[str, int, float, list]
+    value: Union[str, int, float, list],
 ) -> AssertionStdParameterClass:
     if isinstance(value, str):
         return AssertionStdParameterClass(
diff --git a/metadata-ingestion/src/datahub/api/entities/corpgroup/corpgroup.py b/metadata-ingestion/src/datahub/api/entities/corpgroup/corpgroup.py
index a898e35bb810ec..bf58d2fbbda913 100644
--- a/metadata-ingestion/src/datahub/api/entities/corpgroup/corpgroup.py
+++ b/metadata-ingestion/src/datahub/api/entities/corpgroup/corpgroup.py
@@ -114,7 +114,7 @@ def generate_mcp(
                 )
                 urns_created.add(m.urn)
             else:
-                logger.warn(
+                logger.warning(
                     f"Suppressing emission of member {m.urn} before we already emitted metadata for it"
                 )
 
diff --git a/metadata-ingestion/src/datahub/api/entities/datacontract/assertion_operator.py b/metadata-ingestion/src/datahub/api/entities/datacontract/assertion_operator.py
index dc0c97d1c74e56..145a6097d7336c 100644
--- a/metadata-ingestion/src/datahub/api/entities/datacontract/assertion_operator.py
+++ b/metadata-ingestion/src/datahub/api/entities/datacontract/assertion_operator.py
@@ -19,15 +19,13 @@ class Operator(Protocol):
 
     operator: str
 
-    def id(self) -> str:
-        ...
+    def id(self) -> str: ...
 
-    def generate_parameters(self) -> AssertionStdParametersClass:
-        ...
+    def generate_parameters(self) -> AssertionStdParametersClass: ...
 
 
 def _generate_assertion_std_parameter(
-    value: Union[str, int, float]
+    value: Union[str, int, float],
 ) -> AssertionStdParameterClass:
     if isinstance(value, str):
         return AssertionStdParameterClass(
diff --git a/metadata-ingestion/src/datahub/api/entities/dataproduct/dataproduct.py b/metadata-ingestion/src/datahub/api/entities/dataproduct/dataproduct.py
index 2097922c151366..d2035d560716ae 100644
--- a/metadata-ingestion/src/datahub/api/entities/dataproduct/dataproduct.py
+++ b/metadata-ingestion/src/datahub/api/entities/dataproduct/dataproduct.py
@@ -321,9 +321,9 @@ def from_yaml(
 
     @classmethod
     def from_datahub(cls, graph: DataHubGraph, id: str) -> DataProduct:
-        data_product_properties: Optional[
-            DataProductPropertiesClass
-        ] = graph.get_aspect(id, DataProductPropertiesClass)
+        data_product_properties: Optional[DataProductPropertiesClass] = (
+            graph.get_aspect(id, DataProductPropertiesClass)
+        )
         domains: Optional[DomainsClass] = graph.get_aspect(id, DomainsClass)
         assert domains, "Data Product must have an associated domain. Found none."
         owners: Optional[OwnershipClass] = graph.get_aspect(id, OwnershipClass)
@@ -438,7 +438,7 @@ def _patch_ownership(
             for replace_index, replace_value in patches_replace.items():
                 list_to_manipulate[replace_index] = replace_value
 
-            for drop_index, drop_value in patches_drop.items():
+            for drop_value in patches_drop.values():
                 list_to_manipulate.remove(drop_value)
 
             for add_value in patches_add:
diff --git a/metadata-ingestion/src/datahub/api/entities/dataset/dataset.py b/metadata-ingestion/src/datahub/api/entities/dataset/dataset.py
index 315f2249d2e5cd..bf824a11a77b5d 100644
--- a/metadata-ingestion/src/datahub/api/entities/dataset/dataset.py
+++ b/metadata-ingestion/src/datahub/api/entities/dataset/dataset.py
@@ -266,7 +266,8 @@ def generate_mcp(
             if self.schema_metadata.fields:
                 for field in self.schema_metadata.fields:
                     field_urn = field.urn or make_schema_field_urn(
-                        self.urn, field.id  # type: ignore[arg-type]
+                        self.urn,  # type: ignore[arg-type]
+                        field.id,  # type: ignore[arg-type]
                     )
                     assert field_urn.startswith("urn:li:schemaField:")
 
diff --git a/metadata-ingestion/src/datahub/api/entities/structuredproperties/structuredproperties.py b/metadata-ingestion/src/datahub/api/entities/structuredproperties/structuredproperties.py
index 179dbdb231c912..b0b434751ad2cc 100644
--- a/metadata-ingestion/src/datahub/api/entities/structuredproperties/structuredproperties.py
+++ b/metadata-ingestion/src/datahub/api/entities/structuredproperties/structuredproperties.py
@@ -118,9 +118,9 @@ def fqn(self) -> str:
         id = StructuredPropertyUrn.from_string(self.urn).id
         if self.qualified_name is not None:
             # ensure that qualified name and ID match
-            assert (
-                self.qualified_name == id
-            ), "ID in the urn and the qualified_name must match"
+            assert self.qualified_name == id, (
+                "ID in the urn and the qualified_name must match"
+            )
         return id
 
     @validator("urn", pre=True, always=True)
@@ -184,9 +184,9 @@ def create(file: str, graph: DataHubGraph) -> None:
 
     @classmethod
     def from_datahub(cls, graph: DataHubGraph, urn: str) -> "StructuredProperties":
-        structured_property: Optional[
-            StructuredPropertyDefinitionClass
-        ] = graph.get_aspect(urn, StructuredPropertyDefinitionClass)
+        structured_property: Optional[StructuredPropertyDefinitionClass] = (
+            graph.get_aspect(urn, StructuredPropertyDefinitionClass)
+        )
         if structured_property is None:
             raise Exception(
                 "StructuredPropertyDefinition aspect is None. Unable to create structured property."
diff --git a/metadata-ingestion/src/datahub/cli/check_cli.py b/metadata-ingestion/src/datahub/cli/check_cli.py
index fbe07b64f0e154..6b3124fc37393a 100644
--- a/metadata-ingestion/src/datahub/cli/check_cli.py
+++ b/metadata-ingestion/src/datahub/cli/check_cli.py
@@ -9,7 +9,7 @@
 
 import click
 
-from datahub import __package_name__
+from datahub._version import __package_name__
 from datahub.cli.json_file import check_mce_file
 from datahub.configuration import config_loader
 from datahub.configuration.common import AllowDenyPattern
diff --git a/metadata-ingestion/src/datahub/cli/cli_utils.py b/metadata-ingestion/src/datahub/cli/cli_utils.py
index f6b5ba6176c59d..26f4117e151f9e 100644
--- a/metadata-ingestion/src/datahub/cli/cli_utils.py
+++ b/metadata-ingestion/src/datahub/cli/cli_utils.py
@@ -9,7 +9,7 @@
 import requests
 from requests.sessions import Session
 
-import datahub
+import datahub._version as datahub_version
 from datahub.cli import config_utils
 from datahub.emitter.aspect import ASPECT_MAP, TIMESERIES_ASPECT_MAP
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
@@ -412,7 +412,7 @@ def generate_access_token(
 def ensure_has_system_metadata(
     event: Union[
         MetadataChangeProposal, MetadataChangeProposalWrapper, MetadataChangeEvent
-    ]
+    ],
 ) -> None:
     if event.systemMetadata is None:
         event.systemMetadata = SystemMetadataClass()
@@ -422,5 +422,5 @@ def ensure_has_system_metadata(
     if metadata.properties is None:
         metadata.properties = {}
     props = metadata.properties
-    props["clientId"] = datahub.__package_name__
-    props["clientVersion"] = datahub.__version__
+    props["clientId"] = datahub_version.__package_name__
+    props["clientVersion"] = datahub_version.__version__
diff --git a/metadata-ingestion/src/datahub/cli/container_cli.py b/metadata-ingestion/src/datahub/cli/container_cli.py
new file mode 100644
index 00000000000000..5f1dc35edd3ee1
--- /dev/null
+++ b/metadata-ingestion/src/datahub/cli/container_cli.py
@@ -0,0 +1,108 @@
+import logging
+from typing import Any, List
+
+import click
+import progressbar
+
+from datahub.emitter.mcp import MetadataChangeProposalWrapper
+from datahub.ingestion.graph.client import get_default_graph
+from datahub.metadata.schema_classes import (
+    DomainsClass,
+    GlossaryTermAssociationClass,
+    OwnerClass,
+    OwnershipTypeClass,
+    TagAssociationClass,
+)
+from datahub.specific.dataset import DatasetPatchBuilder
+
+logger = logging.getLogger(__name__)
+
+
+@click.group()
+def container() -> None:
+    """A group of commands to interact with containers in DataHub."""
+    pass
+
+
+def apply_association_to_container(
+    container_urn: str,
+    association_urn: str,
+    association_type: str,
+) -> None:
+    """
+    Common function to add either tags, terms, domains, or owners to child datasets (for now).
+
+    Args:
+        container_urn: The URN of the container
+        association_urn: The URN of the tag, term, or user to apply
+        association_type: One of 'tag', 'term', 'domain' or 'owner'
+    """
+    urns: List[str] = []
+    graph = get_default_graph()
+    logger.info(f"Using {graph}")
+    urns.extend(
+        graph.get_urns_by_filter(
+            container=container_urn, batch_size=1000, entity_types=["dataset"]
+        )
+    )
+
+    all_patches: List[Any] = []
+    for urn in urns:
+        builder = DatasetPatchBuilder(urn)
+        patches: List[Any] = []
+        if association_type == "tag":
+            patches = builder.add_tag(TagAssociationClass(association_urn)).build()
+        elif association_type == "term":
+            patches = builder.add_term(
+                GlossaryTermAssociationClass(association_urn)
+            ).build()
+        elif association_type == "owner":
+            patches = builder.add_owner(
+                OwnerClass(
+                    owner=association_urn,
+                    type=OwnershipTypeClass.TECHNICAL_OWNER,
+                )
+            ).build()
+        elif association_type == "domain":
+            patches = [
+                MetadataChangeProposalWrapper(
+                    entityUrn=urn,
+                    aspect=DomainsClass(domains=[association_urn]),
+                )
+            ]
+        all_patches.extend(patches)
+    mcps_iter = progressbar.progressbar(all_patches, redirect_stdout=True)
+    for mcp in mcps_iter:
+        graph.emit(mcp)
+
+
+@container.command()
+@click.option("--container-urn", required=True, type=str)
+@click.option("--tag-urn", required=True, type=str)
+def tag(container_urn: str, tag_urn: str) -> None:
+    """Add patch to add a tag to all datasets in a container"""
+    apply_association_to_container(container_urn, tag_urn, "tag")
+
+
+@container.command()
+@click.option("--container-urn", required=True, type=str)
+@click.option("--term-urn", required=True, type=str)
+def term(container_urn: str, term_urn: str) -> None:
+    """Add patch to add a term to all datasets in a container"""
+    apply_association_to_container(container_urn, term_urn, "term")
+
+
+@container.command()
+@click.option("--container-urn", required=True, type=str)
+@click.option("--owner-urn", required=True, type=str)
+def owner(container_urn: str, owner_urn: str) -> None:
+    """Add patch to add a owner to all datasets in a container"""
+    apply_association_to_container(container_urn, owner_urn, "owner")
+
+
+@container.command()
+@click.option("--container-urn", required=True, type=str)
+@click.option("--domain-urn", required=True, type=str)
+def domain(container_urn: str, domain_urn: str) -> None:
+    """Add patch to add a domain to all datasets in a container"""
+    apply_association_to_container(container_urn, domain_urn, "domain")
diff --git a/metadata-ingestion/src/datahub/cli/delete_cli.py b/metadata-ingestion/src/datahub/cli/delete_cli.py
index 8501cf71f0d544..f77156b4bd7544 100644
--- a/metadata-ingestion/src/datahub/cli/delete_cli.py
+++ b/metadata-ingestion/src/datahub/cli/delete_cli.py
@@ -265,6 +265,11 @@ def undo_by_filter(
     type=str,
     help="Urn of the entity to delete, for single entity deletion",
 )
+@click.option(
+    "--urn-file",
+    required=False,
+    help="Path of file with urns (one per line) to be deleted",
+)
 @click.option(
     "-a",
     "--aspect",
@@ -353,6 +358,7 @@ def undo_by_filter(
 @telemetry.with_telemetry()
 def by_filter(
     urn: Optional[str],
+    urn_file: Optional[str],
     aspect: Optional[str],
     force: bool,
     soft: bool,
@@ -373,6 +379,7 @@ def by_filter(
     # Validate the cli arguments.
     _validate_user_urn_and_filters(
         urn=urn,
+        urn_file=urn_file,
         entity_type=entity_type,
         platform=platform,
         env=env,
@@ -429,6 +436,12 @@ def by_filter(
                         batch_size=batch_size,
                     )
                 )
+    elif urn_file:
+        with open(urn_file, "r") as r:
+            urns = []
+            for line in r.readlines():
+                urn = line.strip().strip('"')
+                urns.append(urn)
     else:
         urns = list(
             graph.get_urns_by_filter(
@@ -537,6 +550,7 @@ def process_urn(urn):
 
 def _validate_user_urn_and_filters(
     urn: Optional[str],
+    urn_file: Optional[str],
     entity_type: Optional[str],
     platform: Optional[str],
     env: Optional[str],
@@ -549,9 +563,9 @@ def _validate_user_urn_and_filters(
             raise click.UsageError(
                 "You cannot provide both an urn and a filter rule (entity-type / platform / env / query)."
             )
-    elif not urn and not (entity_type or platform or env or query):
+    elif not urn and not urn_file and not (entity_type or platform or env or query):
         raise click.UsageError(
-            "You must provide either an urn or at least one filter (entity-type / platform / env / query) in order to delete entities."
+            "You must provide either an urn or urn_file or at least one filter (entity-type / platform / env / query) in order to delete entities."
         )
     elif query:
         logger.warning(
diff --git a/metadata-ingestion/src/datahub/cli/docker_cli.py b/metadata-ingestion/src/datahub/cli/docker_cli.py
index 86bcd7eff1cbfc..b744ac573aed6e 100644
--- a/metadata-ingestion/src/datahub/cli/docker_cli.py
+++ b/metadata-ingestion/src/datahub/cli/docker_cli.py
@@ -296,9 +296,9 @@ def _restore(
     restore_indices: Optional[bool],
     primary_restore_file: Optional[str],
 ) -> int:
-    assert (
-        restore_primary or restore_indices
-    ), "Either restore_primary or restore_indices must be set"
+    assert restore_primary or restore_indices, (
+        "Either restore_primary or restore_indices must be set"
+    )
     msg = "datahub> "
     if restore_primary:
         msg += f"Will restore primary database from {primary_restore_file}. "
@@ -314,9 +314,9 @@ def _restore(
         assert primary_restore_file
         resolved_restore_file = os.path.expanduser(primary_restore_file)
         logger.info(f"Restoring primary db from backup at {resolved_restore_file}")
-        assert os.path.exists(
-            resolved_restore_file
-        ), f"File {resolved_restore_file} does not exist"
+        assert os.path.exists(resolved_restore_file), (
+            f"File {resolved_restore_file} does not exist"
+        )
         with open(resolved_restore_file) as fp:
             result = subprocess.run(
                 [
diff --git a/metadata-ingestion/src/datahub/cli/ingest_cli.py b/metadata-ingestion/src/datahub/cli/ingest_cli.py
index fcab07a1c2aaf6..e2a2f35a36631b 100644
--- a/metadata-ingestion/src/datahub/cli/ingest_cli.py
+++ b/metadata-ingestion/src/datahub/cli/ingest_cli.py
@@ -12,7 +12,7 @@
 from click_default_group import DefaultGroup
 from tabulate import tabulate
 
-import datahub as datahub_package
+from datahub._version import nice_version_name
 from datahub.cli import cli_utils
 from datahub.cli.config_utils import CONDENSED_DATAHUB_CONFIG_PATH
 from datahub.configuration.common import ConfigModel, GraphError
@@ -147,7 +147,7 @@ def run_pipeline_to_completion(pipeline: Pipeline) -> int:
                 return ret
 
     # main function begins
-    logger.info("DataHub CLI version: %s", datahub_package.nice_version_name())
+    logger.info("DataHub CLI version: %s", nice_version_name())
 
     pipeline_config = load_config_file(
         config,
@@ -507,15 +507,11 @@ def list_source_runs(page_offset: int, page_size: int, urn: str, source: str) ->
         click.echo("No response received from the server.")
         return
 
-    # when urn or source filter does not match, exit gracefully
-    if (
-        not isinstance(data.get("data"), dict)
-        or "listIngestionSources" not in data["data"]
-    ):
-        click.echo("No matching ingestion sources found. Please check your filters.")
-        return
+    # a lot of responses can be null if there's errors in the run
+    ingestion_sources = (
+        data.get("data", {}).get("listIngestionSources", {}).get("ingestionSources", [])
+    )
 
-    ingestion_sources = data["data"]["listIngestionSources"]["ingestionSources"]
     if not ingestion_sources:
         click.echo("No ingestion sources or executions found.")
         return
@@ -526,18 +522,32 @@ def list_source_runs(page_offset: int, page_size: int, urn: str, source: str) ->
         name = ingestion_source.get("name", "N/A")
 
         executions = ingestion_source.get("executions", {}).get("executionRequests", [])
+
         for execution in executions:
+            if execution is None:
+                continue
+
             execution_id = execution.get("id", "N/A")
-            start_time = execution.get("result", {}).get("startTimeMs", "N/A")
-            start_time = (
-                datetime.fromtimestamp(start_time / 1000).strftime("%Y-%m-%d %H:%M:%S")
-                if start_time != "N/A"
-                else "N/A"
-            )
-            status = execution.get("result", {}).get("status", "N/A")
+            result = execution.get("result") or {}
+            status = result.get("status", "N/A")
+
+            try:
+                start_time = (
+                    datetime.fromtimestamp(
+                        result.get("startTimeMs", 0) / 1000
+                    ).strftime("%Y-%m-%d %H:%M:%S")
+                    if status != "DUPLICATE" and result.get("startTimeMs") is not None
+                    else "N/A"
+                )
+            except (TypeError, ValueError):
+                start_time = "N/A"
 
             rows.append([execution_id, name, start_time, status, urn])
 
+    if not rows:
+        click.echo("No execution data found.")
+        return
+
     click.echo(
         tabulate(
             rows,
diff --git a/metadata-ingestion/src/datahub/cli/lite_cli.py b/metadata-ingestion/src/datahub/cli/lite_cli.py
index 957ee16245dd81..90bbb353deab18 100644
--- a/metadata-ingestion/src/datahub/cli/lite_cli.py
+++ b/metadata-ingestion/src/datahub/cli/lite_cli.py
@@ -176,7 +176,7 @@ def get(
             )
         )
     end_time = time.time()
-    logger.debug(f"Time taken: {int((end_time - start_time)*1000.0)} millis")
+    logger.debug(f"Time taken: {int((end_time - start_time) * 1000.0)} millis")
 
 
 @lite.command()
@@ -228,7 +228,7 @@ def ls(path: Optional[str]) -> None:
     try:
         browseables = lite.ls(path)
         end_time = time.time()
-        logger.debug(f"Time taken: {int((end_time - start_time)*1000.0)} millis")
+        logger.debug(f"Time taken: {int((end_time - start_time) * 1000.0)} millis")
         auto_complete: List[AutoComplete] = [
             b.auto_complete for b in browseables if b.auto_complete is not None
         ]
diff --git a/metadata-ingestion/src/datahub/cli/migrate.py b/metadata-ingestion/src/datahub/cli/migrate.py
index 1bf1211674f596..3bd1b6fc4dc124 100644
--- a/metadata-ingestion/src/datahub/cli/migrate.py
+++ b/metadata-ingestion/src/datahub/cli/migrate.py
@@ -426,9 +426,9 @@ def batch_get_ids(
             entities_yielded += 1
             log.debug(f"yielding {x}")
             yield x
-        assert (
-            entities_yielded == num_entities
-        ), "Did not delete all entities, try running this command again!"
+        assert entities_yielded == num_entities, (
+            "Did not delete all entities, try running this command again!"
+        )
     else:
         log.error(f"Failed to execute batch get with {str(response.content)}")
         response.raise_for_status()
diff --git a/metadata-ingestion/src/datahub/cli/specific/assertions_cli.py b/metadata-ingestion/src/datahub/cli/specific/assertions_cli.py
index dad724bfe11157..c0d93af90ada00 100644
--- a/metadata-ingestion/src/datahub/cli/specific/assertions_cli.py
+++ b/metadata-ingestion/src/datahub/cli/specific/assertions_cli.py
@@ -136,9 +136,9 @@ def extras_list_to_dict(extras: List[str]) -> Dict[str, str]:
     extra_properties: Dict[str, str] = dict()
     for x in extras:
         parts = x.split("=")
-        assert (
-            len(parts) == 2
-        ), f"Invalid value for extras {x}, should be in format key=value"
+        assert len(parts) == 2, (
+            f"Invalid value for extras {x}, should be in format key=value"
+        )
         extra_properties[parts[0]] = parts[1]
     return extra_properties
 
diff --git a/metadata-ingestion/src/datahub/cli/timeline_cli.py b/metadata-ingestion/src/datahub/cli/timeline_cli.py
index 37089e6f051f0d..174ce63e84ef4c 100644
--- a/metadata-ingestion/src/datahub/cli/timeline_cli.py
+++ b/metadata-ingestion/src/datahub/cli/timeline_cli.py
@@ -50,7 +50,7 @@ def pretty_id(id: Optional[str]) -> str:
     if id.startswith("urn:li:dataset"):
         dataset_key = dataset_urn_to_key(id)
         if dataset_key:
-            return f"{click.style('dataset', fg='cyan')}:{click.style(dataset_key.platform[len('urn:li:dataPlatform:'):], fg='white')}:{click.style(dataset_key.name, fg='white')}"
+            return f"{click.style('dataset', fg='cyan')}:{click.style(dataset_key.platform[len('urn:li:dataPlatform:') :], fg='white')}:{click.style(dataset_key.name, fg='white')}"
     # failed to prettify, return original
     return id
 
diff --git a/metadata-ingestion/src/datahub/configuration/common.py b/metadata-ingestion/src/datahub/configuration/common.py
index 08817d9d5fdb93..8052de1b0669c4 100644
--- a/metadata-ingestion/src/datahub/configuration/common.py
+++ b/metadata-ingestion/src/datahub/configuration/common.py
@@ -200,8 +200,7 @@ class IgnorableError(MetaError):
 
 @runtime_checkable
 class ExceptionWithProps(Protocol):
-    def get_telemetry_props(self) -> Dict[str, Any]:
-        ...
+    def get_telemetry_props(self) -> Dict[str, Any]: ...
 
 
 def should_show_stack_trace(exc: Exception) -> bool:
diff --git a/metadata-ingestion/src/datahub/configuration/config_loader.py b/metadata-ingestion/src/datahub/configuration/config_loader.py
index c64e81e7c714c9..16105f69d584de 100644
--- a/metadata-ingestion/src/datahub/configuration/config_loader.py
+++ b/metadata-ingestion/src/datahub/configuration/config_loader.py
@@ -19,64 +19,87 @@
 Environ = Mapping[str, str]
 
 
-def _resolve_element(element: str, environ: Environ) -> str:
-    if re.search(r"(\$\{).+(\})", element):
-        return expand(element, nounset=True, environ=environ)
-    elif element.startswith("$"):
-        try:
-            return expand(element, nounset=True, environ=environ)
-        except UnboundVariable:
-            return element
-    else:
-        return element
-
-
-def _resolve_list(ele_list: list, environ: Environ) -> list:
-    new_v: list = []
-    for ele in ele_list:
-        if isinstance(ele, str):
-            new_v.append(_resolve_element(ele, environ=environ))
-        elif isinstance(ele, list):
-            new_v.append(_resolve_list(ele, environ=environ))
-        elif isinstance(ele, dict):
-            new_v.append(resolve_env_variables(ele, environ=environ))
-        else:
-            new_v.append(ele)
-    return new_v
-
-
 def resolve_env_variables(config: dict, environ: Environ) -> dict:
-    new_dict: Dict[Any, Any] = {}
-    for k, v in config.items():
-        if isinstance(v, dict):
-            new_dict[k] = resolve_env_variables(v, environ=environ)
-        elif isinstance(v, list):
-            new_dict[k] = _resolve_list(v, environ=environ)
-        elif isinstance(v, str):
-            new_dict[k] = _resolve_element(v, environ=environ)
-        else:
-            new_dict[k] = v
-    return new_dict
+    # TODO: This is kept around for backwards compatibility.
+    return EnvResolver(environ).resolve(config)
 
 
 def list_referenced_env_variables(config: dict) -> Set[str]:
-    # This is a bit of a hack, but expandvars does a bunch of escaping
-    # and other logic that we don't want to duplicate here.
+    # TODO: This is kept around for backwards compatibility.
+    return EnvResolver(environ=os.environ).list_referenced_variables(config)
+
+
+class EnvResolver:
+    def __init__(self, environ: Environ, strict_env_syntax: bool = False):
+        self.environ = environ
+        self.strict_env_syntax = strict_env_syntax
 
-    vars = set()
+    def resolve(self, config: dict) -> dict:
+        return self._resolve_dict(config)
 
-    def mock_get_env(key: str, default: Optional[str] = None) -> str:
-        vars.add(key)
-        if default is not None:
-            return default
-        return "mocked_value"
+    @classmethod
+    def list_referenced_variables(
+        cls,
+        config: dict,
+        strict_env_syntax: bool = False,
+    ) -> Set[str]:
+        # This is a bit of a hack, but expandvars does a bunch of escaping
+        # and other logic that we don't want to duplicate here.
 
-    mock = unittest.mock.MagicMock()
-    mock.get.side_effect = mock_get_env
+        vars = set()
 
-    resolve_env_variables(config, environ=mock)
+        def mock_get_env(key: str, default: Optional[str] = None) -> str:
+            vars.add(key)
+            if default is not None:
+                return default
+            return "mocked_value"
+
+        mock = unittest.mock.MagicMock()
+        mock.get.side_effect = mock_get_env
+
+        resolver = EnvResolver(environ=mock, strict_env_syntax=strict_env_syntax)
+        resolver._resolve_dict(config)
+
+        return vars
+
+    def _resolve_element(self, element: str) -> str:
+        if re.search(r"(\$\{).+(\})", element):
+            return expand(element, nounset=True, environ=self.environ)
+        elif not self.strict_env_syntax and element.startswith("$"):
+            try:
+                return expand(element, nounset=True, environ=self.environ)
+            except UnboundVariable:
+                # TODO: This fallback is kept around for backwards compatibility, but
+                # doesn't make a ton of sense from first principles.
+                return element
+        else:
+            return element
 
-    return vars
+    def _resolve_list(self, ele_list: list) -> list:
+        new_v: list = []
+        for ele in ele_list:
+            if isinstance(ele, str):
+                new_v.append(self._resolve_element(ele))
+            elif isinstance(ele, list):
+                new_v.append(self._resolve_list(ele))
+            elif isinstance(ele, dict):
+                new_v.append(self._resolve_dict(ele))
+            else:
+                new_v.append(ele)
+        return new_v
+
+    def _resolve_dict(self, config: dict) -> dict:
+        new_dict: Dict[Any, Any] = {}
+        for k, v in config.items():
+            if isinstance(v, dict):
+                new_dict[k] = self._resolve_dict(v)
+            elif isinstance(v, list):
+                new_dict[k] = self._resolve_list(v)
+            elif isinstance(v, str):
+                new_dict[k] = self._resolve_element(v)
+            else:
+                new_dict[k] = v
+        return new_dict
 
 
 WRITE_TO_FILE_DIRECTIVE_PREFIX = "__DATAHUB_TO_FILE_"
@@ -159,7 +182,7 @@ def load_config_file(
 
     config = raw_config.copy()
     if resolve_env_vars:
-        config = resolve_env_variables(config, environ=os.environ)
+        config = EnvResolver(environ=os.environ).resolve(config)
     if process_directives:
         config = _process_directives(config)
 
diff --git a/metadata-ingestion/src/datahub/configuration/git.py b/metadata-ingestion/src/datahub/configuration/git.py
index e7e9bfd43adca5..7e68e9f80da4ff 100644
--- a/metadata-ingestion/src/datahub/configuration/git.py
+++ b/metadata-ingestion/src/datahub/configuration/git.py
@@ -121,9 +121,9 @@ def infer_repo_ssh_locator(
 
         repo: str = values["repo"]
         if repo.startswith(_GITHUB_PREFIX):
-            return f"git@github.com:{repo[len(_GITHUB_PREFIX):]}.git"
+            return f"git@github.com:{repo[len(_GITHUB_PREFIX) :]}.git"
         elif repo.startswith(_GITLAB_PREFIX):
-            return f"git@gitlab.com:{repo[len(_GITLAB_PREFIX):]}.git"
+            return f"git@gitlab.com:{repo[len(_GITLAB_PREFIX) :]}.git"
         else:
             raise ValueError(
                 "Unable to infer repo_ssh_locator from repo. Please set repo_ssh_locator manually."
diff --git a/metadata-ingestion/src/datahub/configuration/time_window_config.py b/metadata-ingestion/src/datahub/configuration/time_window_config.py
index b3cc0316091173..5fabcf904d3219 100644
--- a/metadata-ingestion/src/datahub/configuration/time_window_config.py
+++ b/metadata-ingestion/src/datahub/configuration/time_window_config.py
@@ -47,7 +47,10 @@ class BaseTimeWindowConfig(ConfigModel):
         default_factory=lambda: datetime.now(tz=timezone.utc),
         description="Latest date of lineage/usage to consider. Default: Current time in UTC",
     )
-    start_time: datetime = Field(default=None, description="Earliest date of lineage/usage to consider. Default: Last full day in UTC (or hour, depending on `bucket_duration`). You can also specify relative time with respect to end_time such as '-7 days' Or '-7d'.")  # type: ignore
+    start_time: datetime = Field(
+        default=None,
+        description="Earliest date of lineage/usage to consider. Default: Last full day in UTC (or hour, depending on `bucket_duration`). You can also specify relative time with respect to end_time such as '-7 days' Or '-7d'.",
+    )  # type: ignore
 
     @pydantic.validator("start_time", pre=True, always=True)
     def default_start_time(
@@ -63,12 +66,14 @@ def default_start_time(
             # This is where start_time str is resolved to datetime
             try:
                 delta = parse_relative_timespan(v)
-                assert delta < timedelta(
-                    0
-                ), "Relative start time should start with minus sign (-) e.g. '-2 days'."
+                assert delta < timedelta(0), (
+                    "Relative start time should start with minus sign (-) e.g. '-2 days'."
+                )
                 assert abs(delta) >= get_bucket_duration_delta(
                     values["bucket_duration"]
-                ), "Relative start time should be in terms of configured bucket duration. e.g '-2 days' or '-2 hours'."
+                ), (
+                    "Relative start time should be in terms of configured bucket duration. e.g '-2 days' or '-2 hours'."
+                )
 
                 # The end_time's default value is not yet populated, in which case
                 # we can just manually generate it here.
diff --git a/metadata-ingestion/src/datahub/emitter/enum_helpers.py b/metadata-ingestion/src/datahub/emitter/enum_helpers.py
index 89949ab3717ff3..89916ffd899c3f 100644
--- a/metadata-ingestion/src/datahub/emitter/enum_helpers.py
+++ b/metadata-ingestion/src/datahub/emitter/enum_helpers.py
@@ -1,11 +1,13 @@
 from typing import List, Type
 
+from typing_extensions import LiteralString
 
-def get_enum_options(_class: Type[object]) -> List[str]:
+
+def get_enum_options(class_: Type[object]) -> List[LiteralString]:
     """Get the valid values for an enum in the datahub.metadata.schema_classes module."""
 
     return [
         value
-        for name, value in vars(_class).items()
+        for name, value in vars(class_).items()
         if not callable(value) and not name.startswith("_")
     ]
diff --git a/metadata-ingestion/src/datahub/emitter/mce_builder.py b/metadata-ingestion/src/datahub/emitter/mce_builder.py
index f095fffbaea6b4..9fa060266a7ab8 100644
--- a/metadata-ingestion/src/datahub/emitter/mce_builder.py
+++ b/metadata-ingestion/src/datahub/emitter/mce_builder.py
@@ -88,13 +88,11 @@ def get_sys_time() -> int:
 
 
 @overload
-def make_ts_millis(ts: None) -> None:
-    ...
+def make_ts_millis(ts: None) -> None: ...
 
 
 @overload
-def make_ts_millis(ts: datetime) -> int:
-    ...
+def make_ts_millis(ts: datetime) -> int: ...
 
 
 def make_ts_millis(ts: Optional[datetime]) -> Optional[int]:
@@ -105,13 +103,11 @@ def make_ts_millis(ts: Optional[datetime]) -> Optional[int]:
 
 
 @overload
-def parse_ts_millis(ts: float) -> datetime:
-    ...
+def parse_ts_millis(ts: float) -> datetime: ...
 
 
 @overload
-def parse_ts_millis(ts: None) -> None:
-    ...
+def parse_ts_millis(ts: None) -> None: ...
 
 
 def parse_ts_millis(ts: Optional[float]) -> Optional[datetime]:
@@ -444,6 +440,10 @@ def can_add_aspect_to_snapshot(
 
 
 def can_add_aspect(mce: MetadataChangeEventClass, AspectType: Type[Aspect]) -> bool:
+    # TODO: This is specific to snapshot types. We have a more general method
+    # in `entity_supports_aspect`, which should be used instead. This method
+    # should be deprecated, and all usages should be replaced.
+
     SnapshotType = type(mce.proposedSnapshot)
 
     return can_add_aspect_to_snapshot(SnapshotType, AspectType)
diff --git a/metadata-ingestion/src/datahub/emitter/mcp_builder.py b/metadata-ingestion/src/datahub/emitter/mcp_builder.py
index 581f903d0eef0d..7466504987e05d 100644
--- a/metadata-ingestion/src/datahub/emitter/mcp_builder.py
+++ b/metadata-ingestion/src/datahub/emitter/mcp_builder.py
@@ -90,6 +90,25 @@ def property_dict(self) -> Dict[str, str]:
     def as_urn(self) -> str:
         return make_container_urn(guid=self.guid())
 
+    def parent_key(self) -> Optional["ContainerKey"]:
+        # Find the immediate base class of self.
+        # This is a bit of a hack, but it works.
+        base_classes = self.__class__.__bases__
+        if len(base_classes) != 1:
+            # TODO: Raise a more specific error.
+            raise ValueError(
+                f"Unable to determine parent key for {self.__class__}: {self}"
+            )
+        base_class = base_classes[0]
+        if base_class is DatahubKey or base_class is ContainerKey:
+            return None
+
+        # We need to use `__dict__` instead of `pydantic.BaseModel.dict()`
+        # in order to include "excluded" fields e.g. `backcompat_env_as_instance`.
+        # Tricky: this only works because DatahubKey is a BaseModel and hence
+        # allows extra fields.
+        return base_class(**self.__dict__)
+
 
 # DEPRECATION: Keeping the `PlatformKey` name around for backwards compatibility.
 PlatformKey = ContainerKey
diff --git a/metadata-ingestion/src/datahub/emitter/mcp_patch_builder.py b/metadata-ingestion/src/datahub/emitter/mcp_patch_builder.py
index 17026a4114c128..e51c37d96e90f0 100644
--- a/metadata-ingestion/src/datahub/emitter/mcp_patch_builder.py
+++ b/metadata-ingestion/src/datahub/emitter/mcp_patch_builder.py
@@ -33,8 +33,7 @@
 
 @runtime_checkable
 class SupportsToObj(Protocol):
-    def to_obj(self) -> Any:
-        ...
+    def to_obj(self) -> Any: ...
 
 
 def _recursive_to_obj(obj: Any) -> Any:
diff --git a/metadata-ingestion/src/datahub/emitter/rest_emitter.py b/metadata-ingestion/src/datahub/emitter/rest_emitter.py
index 7271f784bf881e..4e7a152204da89 100644
--- a/metadata-ingestion/src/datahub/emitter/rest_emitter.py
+++ b/metadata-ingestion/src/datahub/emitter/rest_emitter.py
@@ -22,7 +22,7 @@
 from requests.adapters import HTTPAdapter, Retry
 from requests.exceptions import HTTPError, RequestException
 
-from datahub import nice_version_name
+from datahub._version import nice_version_name
 from datahub.cli import config_utils
 from datahub.cli.cli_utils import ensure_has_system_metadata, fixup_gms_url, get_or_else
 from datahub.cli.env_utils import get_boolean_env_variable
diff --git a/metadata-ingestion/src/datahub/entrypoints.py b/metadata-ingestion/src/datahub/entrypoints.py
index 85968f050a3716..73d35381d5df29 100644
--- a/metadata-ingestion/src/datahub/entrypoints.py
+++ b/metadata-ingestion/src/datahub/entrypoints.py
@@ -6,7 +6,7 @@
 
 import click
 
-import datahub as datahub_package
+import datahub._version as datahub_version
 from datahub.cli.check_cli import check
 from datahub.cli.cli_utils import (
     fixup_gms_url,
@@ -14,6 +14,7 @@
     make_shim_command,
 )
 from datahub.cli.config_utils import DATAHUB_CONFIG_PATH, write_gms_config
+from datahub.cli.container_cli import container
 from datahub.cli.delete_cli import delete
 from datahub.cli.docker_cli import docker
 from datahub.cli.env_utils import get_boolean_env_variable
@@ -45,6 +46,12 @@
 
 MAX_CONTENT_WIDTH = 120
 
+if sys.version_info >= (3, 12):
+    click.secho(
+        "Python versions above 3.11 are not tested with. Please use Python 3.11.",
+        fg="red",
+    )
+
 
 @click.group(
     context_settings=dict(
@@ -67,8 +74,8 @@
     help="Write debug-level logs to a file.",
 )
 @click.version_option(
-    version=datahub_package.nice_version_name(),
-    prog_name=datahub_package.__package_name__,
+    version=datahub_version.nice_version_name(),
+    prog_name=datahub_version.__package_name__,
 )
 def datahub(
     debug: bool,
@@ -105,7 +112,7 @@ def datahub(
 def version(include_server: bool = False) -> None:
     """Print version number and exit."""
 
-    click.echo(f"DataHub CLI version: {datahub_package.nice_version_name()}")
+    click.echo(f"DataHub CLI version: {datahub_version.nice_version_name()}")
     click.echo(f"Models: {model_version_name()}")
     click.echo(f"Python version: {sys.version}")
     if include_server:
@@ -174,6 +181,7 @@ def init(use_password: bool = False) -> None:
 datahub.add_command(forms)
 datahub.add_command(datacontract)
 datahub.add_command(assertions)
+datahub.add_command(container)
 
 try:
     from datahub.cli.lite_cli import lite
@@ -215,7 +223,7 @@ def main(**kwargs):
             logger.exception(f"Command failed: {exc}")
 
         logger.debug(
-            f"DataHub CLI version: {datahub_package.__version__} at {datahub_package.__file__}"
+            f"DataHub CLI version: {datahub_version.__version__} at {__file__}"
         )
         logger.debug(
             f"Python version: {sys.version} at {sys.executable} on {platform.platform()}"
diff --git a/metadata-ingestion/src/datahub/ingestion/api/decorators.py b/metadata-ingestion/src/datahub/ingestion/api/decorators.py
index d32c0b85ceef4c..b521dc5e9e7f5d 100644
--- a/metadata-ingestion/src/datahub/ingestion/api/decorators.py
+++ b/metadata-ingestion/src/datahub/ingestion/api/decorators.py
@@ -25,6 +25,8 @@ def wrapper(cls: Type) -> Type:
             # add the create method only if it has not been overridden from the base Source.create method
             cls.create = classmethod(default_create)
 
+            # TODO: Once we're on Python 3.10, we should call abc.update_abstractmethods here.
+
         return cls
 
     return wrapper
diff --git a/metadata-ingestion/src/datahub/ingestion/api/incremental_lineage_helper.py b/metadata-ingestion/src/datahub/ingestion/api/incremental_lineage_helper.py
index 78a091f1ffe689..92ee158661d3d4 100644
--- a/metadata-ingestion/src/datahub/ingestion/api/incremental_lineage_helper.py
+++ b/metadata-ingestion/src/datahub/ingestion/api/incremental_lineage_helper.py
@@ -55,15 +55,9 @@ def convert_chart_info_to_patch(
         aspect.externalUrl
     ).set_type(aspect.type).set_title(aspect.title).set_access(
         aspect.access
-    ).set_last_modified(
-        aspect.lastModified
-    ).set_last_refreshed(
+    ).set_last_modified(aspect.lastModified).set_last_refreshed(
         aspect.lastRefreshed
-    ).set_description(
-        aspect.description
-    ).add_inputs(
-        aspect.inputs
-    )
+    ).set_description(aspect.description).add_inputs(aspect.inputs)
 
     values = patch_builder.build()
     if values:
diff --git a/metadata-ingestion/src/datahub/ingestion/api/registry.py b/metadata-ingestion/src/datahub/ingestion/api/registry.py
index 6b50e0b8a8a5d4..91ee98865e78e4 100644
--- a/metadata-ingestion/src/datahub/ingestion/api/registry.py
+++ b/metadata-ingestion/src/datahub/ingestion/api/registry.py
@@ -17,7 +17,7 @@
 
 import typing_inspect
 
-from datahub import __package_name__
+from datahub._version import __package_name__
 from datahub.configuration.common import ConfigurationError
 
 if sys.version_info < (3, 10):
@@ -173,8 +173,10 @@ def get(self, key: str) -> Type[T]:
 
         tp = self._ensure_not_lazy(key)
         if isinstance(tp, ModuleNotFoundError):
+            # TODO: Once we're on Python 3.11 (with PEP 678), we can use .add_note()
+            # to enrich the error instead of wrapping it.
             raise ConfigurationError(
-                f"{key} is disabled; try running: pip install '{__package_name__}[{key}]'"
+                f"{key} is disabled due to a missing dependency: {tp.name}; try running `pip install '{__package_name__}[{key}]'`"
             ) from tp
         elif isinstance(tp, Exception):
             raise ConfigurationError(
diff --git a/metadata-ingestion/src/datahub/ingestion/api/report.py b/metadata-ingestion/src/datahub/ingestion/api/report.py
index 32810189acd00b..8cfca5782bee40 100644
--- a/metadata-ingestion/src/datahub/ingestion/api/report.py
+++ b/metadata-ingestion/src/datahub/ingestion/api/report.py
@@ -21,8 +21,7 @@
 
 @runtime_checkable
 class SupportsAsObj(Protocol):
-    def as_obj(self) -> dict:
-        ...
+    def as_obj(self) -> dict: ...
 
 
 @dataclass
diff --git a/metadata-ingestion/src/datahub/ingestion/api/sink.py b/metadata-ingestion/src/datahub/ingestion/api/sink.py
index 655e6bb22fa8d1..44148c751df828 100644
--- a/metadata-ingestion/src/datahub/ingestion/api/sink.py
+++ b/metadata-ingestion/src/datahub/ingestion/api/sink.py
@@ -110,6 +110,10 @@ def __init__(self, ctx: PipelineContext, config: SinkConfig):
         self.__post_init__()
 
     def __post_init__(self) -> None:
+        """Hook called after the sink's main initialization is complete.
+
+        Sink subclasses can override this method to customize initialization.
+        """
         pass
 
     @classmethod
@@ -117,9 +121,17 @@ def create(cls, config_dict: dict, ctx: PipelineContext) -> "Self":
         return cls(ctx, cls.get_config_class().parse_obj(config_dict))
 
     def handle_work_unit_start(self, workunit: WorkUnit) -> None:
+        """Called at the start of each new workunit.
+
+        This method is deprecated and will be removed in a future release.
+        """
         pass
 
     def handle_work_unit_end(self, workunit: WorkUnit) -> None:
+        """Called at the end of each workunit.
+
+        This method is deprecated and will be removed in a future release.
+        """
         pass
 
     @abstractmethod
diff --git a/metadata-ingestion/src/datahub/ingestion/api/source.py b/metadata-ingestion/src/datahub/ingestion/api/source.py
index b04ffdb3258934..e401d6e976b6d0 100644
--- a/metadata-ingestion/src/datahub/ingestion/api/source.py
+++ b/metadata-ingestion/src/datahub/ingestion/api/source.py
@@ -404,8 +404,11 @@ def create(cls, config_dict: dict, ctx: PipelineContext) -> Self:
         # Technically, this method should be abstract. However, the @config_class
         # decorator automatically generates a create method at runtime if one is
         # not defined. Python still treats the class as abstract because it thinks
-        # the create method is missing. To avoid the class becoming abstract, we
-        # can't make this method abstract.
+        # the create method is missing.
+        #
+        # Once we're on Python 3.10, we can use the abc.update_abstractmethods(cls)
+        # method in the config_class decorator. That would allow us to make this
+        # method abstract.
         raise NotImplementedError('sources must implement "create"')
 
     def get_workunit_processors(self) -> List[Optional[MetadataWorkUnitProcessor]]:
diff --git a/metadata-ingestion/src/datahub/ingestion/api/source_helpers.py b/metadata-ingestion/src/datahub/ingestion/api/source_helpers.py
index f3e5b1db6a1c85..08af39cd24982a 100644
--- a/metadata-ingestion/src/datahub/ingestion/api/source_helpers.py
+++ b/metadata-ingestion/src/datahub/ingestion/api/source_helpers.py
@@ -48,7 +48,7 @@
 
 
 def auto_workunit(
-    stream: Iterable[Union[MetadataChangeEventClass, MetadataChangeProposalWrapper]]
+    stream: Iterable[Union[MetadataChangeEventClass, MetadataChangeProposalWrapper]],
 ) -> Iterable[MetadataWorkUnit]:
     """Convert a stream of MCEs and MCPs to a stream of :class:`MetadataWorkUnit`s."""
 
diff --git a/metadata-ingestion/src/datahub/ingestion/extractor/json_schema_util.py b/metadata-ingestion/src/datahub/ingestion/extractor/json_schema_util.py
index 88d1fcc52e2196..1c440642e06d8b 100644
--- a/metadata-ingestion/src/datahub/ingestion/extractor/json_schema_util.py
+++ b/metadata-ingestion/src/datahub/ingestion/extractor/json_schema_util.py
@@ -131,9 +131,9 @@ def get_recursive(self, schema: Dict) -> Optional[str]:
             for i, schema_type in enumerate(p.schema_types):
                 if schema_type == schema_str:
                     # return the corresponding type for the schema that's a match
-                    assert (
-                        len(p.type) > i
-                    ), f"p.type({len(p.type)})) and p.schema_types({len(p.schema_types)}) should have the same length"
+                    assert len(p.type) > i, (
+                        f"p.type({len(p.type)})) and p.schema_types({len(p.schema_types)}) should have the same length"
+                    )
                     return p.type[i]
         return None
 
diff --git a/metadata-ingestion/src/datahub/ingestion/extractor/schema_util.py b/metadata-ingestion/src/datahub/ingestion/extractor/schema_util.py
index d5af4f7a2389c0..dbb851c74e7e34 100644
--- a/metadata-ingestion/src/datahub/ingestion/extractor/schema_util.py
+++ b/metadata-ingestion/src/datahub/ingestion/extractor/schema_util.py
@@ -263,15 +263,13 @@ def _get_type_annotation(schema: SchemaOrField) -> str:
     @overload
     def _get_underlying_type_if_option_as_union(
         schema: SchemaOrField, default: SchemaOrField
-    ) -> SchemaOrField:
-        ...
+    ) -> SchemaOrField: ...
 
     @staticmethod
     @overload
     def _get_underlying_type_if_option_as_union(
         schema: SchemaOrField, default: Optional[SchemaOrField] = None
-    ) -> Optional[SchemaOrField]:
-        ...
+    ) -> Optional[SchemaOrField]: ...
 
     @staticmethod
     def _get_underlying_type_if_option_as_union(
@@ -386,7 +384,7 @@ def emit(self) -> Iterable[SchemaField]:
 
                 if "deprecated" in merged_props:
                     description = (
-                        f"<span style=\"color:red\">DEPRECATED: {merged_props['deprecated']}</span>\n"
+                        f'<span style="color:red">DEPRECATED: {merged_props["deprecated"]}</span>\n'
                         + description
                         if description
                         else ""
diff --git a/metadata-ingestion/src/datahub/ingestion/fs/s3_fs.py b/metadata-ingestion/src/datahub/ingestion/fs/s3_fs.py
index 9c34c4f83b0a93..beec42724529e6 100644
--- a/metadata-ingestion/src/datahub/ingestion/fs/s3_fs.py
+++ b/metadata-ingestion/src/datahub/ingestion/fs/s3_fs.py
@@ -17,9 +17,9 @@ def parse_s3_path(path: str) -> "S3Path":
 
 def assert_ok_status(s3_response):
     is_ok = s3_response["ResponseMetadata"]["HTTPStatusCode"] == 200
-    assert (
-        is_ok
-    ), f"Failed to fetch S3 object, error message: {s3_response['Error']['Message']}"
+    assert is_ok, (
+        f"Failed to fetch S3 object, error message: {s3_response['Error']['Message']}"
+    )
 
 
 @dataclass
diff --git a/metadata-ingestion/src/datahub/ingestion/glossary/datahub_classifier.py b/metadata-ingestion/src/datahub/ingestion/glossary/datahub_classifier.py
index 50268768d0ce9f..ba03083854e785 100644
--- a/metadata-ingestion/src/datahub/ingestion/glossary/datahub_classifier.py
+++ b/metadata-ingestion/src/datahub/ingestion/glossary/datahub_classifier.py
@@ -148,9 +148,9 @@ def input_config_selectively_overrides_default_config(cls, info_types_config):
                 weight,
             ) in custom_infotype_config.Prediction_Factors_and_Weights.dict().items():
                 if weight > 0:
-                    assert (
-                        getattr(custom_infotype_config, factor) is not None
-                    ), f"Missing Configuration for Prediction Factor {factor} for Custom Info Type {custom_infotype}"
+                    assert getattr(custom_infotype_config, factor) is not None, (
+                        f"Missing Configuration for Prediction Factor {factor} for Custom Info Type {custom_infotype}"
+                    )
 
             # Custom infotype supports only regex based prediction for column values
             if custom_infotype_config.Prediction_Factors_and_Weights.Values > 0:
@@ -158,7 +158,9 @@ def input_config_selectively_overrides_default_config(cls, info_types_config):
                 assert (
                     custom_infotype_config.Values.prediction_type
                     == ValuePredictionType.REGEX
-                ), f"Invalid Prediction Type for Values for Custom Info Type {custom_infotype}. Only `regex` is supported."
+                ), (
+                    f"Invalid Prediction Type for Values for Custom Info Type {custom_infotype}. Only `regex` is supported."
+                )
 
         return info_types_config
 
diff --git a/metadata-ingestion/src/datahub/ingestion/graph/client.py b/metadata-ingestion/src/datahub/ingestion/graph/client.py
index 8c5f894a072d93..48a008536ed1ed 100644
--- a/metadata-ingestion/src/datahub/ingestion/graph/client.py
+++ b/metadata-ingestion/src/datahub/ingestion/graph/client.py
@@ -519,9 +519,9 @@ def get_aspects_for_entity(
         :return: Optionally, a map of aspect_name to aspect_value as a dictionary if present, aspect_value will be set to None if that aspect was not found. Returns None on HTTP status 404.
         :raises HttpError: if the HTTP response is not a 200
         """
-        assert len(aspects) == len(
-            aspect_types
-        ), f"number of aspects requested ({len(aspects)}) should be the same as number of aspect types provided ({len(aspect_types)})"
+        assert len(aspects) == len(aspect_types), (
+            f"number of aspects requested ({len(aspects)}) should be the same as number of aspect types provided ({len(aspect_types)})"
+        )
 
         # TODO: generate aspects list from type classes
         response_json = self.get_entity_raw(entity_urn, aspects)
@@ -1576,9 +1576,7 @@ def run_assertion(
                     ... assertionResult
                 }
             }
-        """ % (
-            self._assertion_result_shared()
-        )
+        """ % (self._assertion_result_shared())
 
         variables = {
             "assertionUrn": urn,
diff --git a/metadata-ingestion/src/datahub/ingestion/reporting/datahub_ingestion_run_summary_provider.py b/metadata-ingestion/src/datahub/ingestion/reporting/datahub_ingestion_run_summary_provider.py
index c143a8b49f4b7c..fc790535cfe03d 100644
--- a/metadata-ingestion/src/datahub/ingestion/reporting/datahub_ingestion_run_summary_provider.py
+++ b/metadata-ingestion/src/datahub/ingestion/reporting/datahub_ingestion_run_summary_provider.py
@@ -3,7 +3,7 @@
 import time
 from typing import Any, Dict, Optional
 
-from datahub import nice_version_name
+from datahub._version import nice_version_name
 from datahub.configuration.common import (
     ConfigModel,
     DynamicTypedConfig,
diff --git a/metadata-ingestion/src/datahub/ingestion/run/connection.py b/metadata-ingestion/src/datahub/ingestion/run/connection.py
index 54b0ab9f22c65e..d42a1ba8767f96 100644
--- a/metadata-ingestion/src/datahub/ingestion/run/connection.py
+++ b/metadata-ingestion/src/datahub/ingestion/run/connection.py
@@ -1,6 +1,6 @@
 import logging
 
-from datahub import __version__
+from datahub._version import __version__
 from datahub.ingestion.api.source import TestableSource, TestConnectionReport
 from datahub.ingestion.source.source_registry import source_registry
 
diff --git a/metadata-ingestion/src/datahub/ingestion/run/pipeline.py b/metadata-ingestion/src/datahub/ingestion/run/pipeline.py
index ee1c1608cd48c6..120cf6a79bc02c 100644
--- a/metadata-ingestion/src/datahub/ingestion/run/pipeline.py
+++ b/metadata-ingestion/src/datahub/ingestion/run/pipeline.py
@@ -15,7 +15,7 @@
 import humanfriendly
 import psutil
 
-import datahub
+from datahub._version import nice_version_name
 from datahub.configuration.common import (
     ConfigModel,
     IgnorableError,
@@ -76,8 +76,9 @@ def on_failure(
         failure_metadata: dict,
     ) -> None:
         logger.error(
-            f"{self.name} failed to write record with workunit {record_envelope.metadata['workunit_id']}"
-            f" with {failure_exception} and info {failure_metadata}"
+            f"{self.name} failed to write record with workunit {record_envelope.metadata['workunit_id']}",
+            extra={"failure_metadata": failure_metadata},
+            exc_info=failure_exception,
         )
 
 
@@ -108,9 +109,9 @@ def on_failure(
                         mcp.systemMetadata.properties = {}
                     if "workunit_id" not in mcp.systemMetadata.properties:
                         # update the workunit id
-                        mcp.systemMetadata.properties[
-                            "workunit_id"
-                        ] = record_envelope.metadata["workunit_id"]
+                        mcp.systemMetadata.properties["workunit_id"] = (
+                            record_envelope.metadata["workunit_id"]
+                        )
                 record_envelope.record = mcp
         self.file_sink.write_record_async(record_envelope, self.logging_callback)
 
@@ -143,8 +144,8 @@ def _add_init_error_context(step: str) -> Iterator[None]:
 
 @dataclass
 class CliReport(Report):
-    cli_version: str = datahub.nice_version_name()
-    cli_entry_location: str = datahub.__file__
+    cli_version: str = nice_version_name()
+    cli_entry_location: str = __file__
     models_version: str = model_version_name()
     py_version: str = sys.version
     py_exec_path: str = sys.executable
@@ -700,7 +701,7 @@ def pretty_print_summary(
             num_failures_sink = len(self.sink.get_report().failures)
             click.secho(
                 message_template.format(
-                    status=f"with at least {num_failures_source+num_failures_sink} failures"
+                    status=f"with at least {num_failures_source + num_failures_sink} failures"
                 ),
                 fg=self._get_text_color(
                     running=currently_running, failures=True, warnings=False
@@ -718,7 +719,7 @@ def pretty_print_summary(
             num_warn_global = len(global_warnings)
             click.secho(
                 message_template.format(
-                    status=f"with at least {num_warn_source+num_warn_sink+num_warn_global} warnings"
+                    status=f"with at least {num_warn_source + num_warn_sink + num_warn_global} warnings"
                 ),
                 fg=self._get_text_color(
                     running=currently_running, failures=False, warnings=True
diff --git a/metadata-ingestion/src/datahub/ingestion/run/pipeline_config.py b/metadata-ingestion/src/datahub/ingestion/run/pipeline_config.py
index 7a4e7ec52a8e96..53e31aa2ea96e1 100644
--- a/metadata-ingestion/src/datahub/ingestion/run/pipeline_config.py
+++ b/metadata-ingestion/src/datahub/ingestion/run/pipeline_config.py
@@ -92,9 +92,9 @@ class PipelineConfig(ConfigModel):
     pipeline_name: Optional[str] = None
     failure_log: FailureLoggingConfig = FailureLoggingConfig()
 
-    _raw_dict: Optional[
-        dict
-    ] = None  # the raw dict that was parsed to construct this config
+    _raw_dict: Optional[dict] = (
+        None  # the raw dict that was parsed to construct this config
+    )
 
     @validator("run_id", pre=True, always=True)
     def run_id_should_be_semantic(
diff --git a/metadata-ingestion/src/datahub/ingestion/source/abs/datalake_profiler_config.py b/metadata-ingestion/src/datahub/ingestion/source/abs/datalake_profiler_config.py
index 9f6d13a08b182e..d12ff7415faefc 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/abs/datalake_profiler_config.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/abs/datalake_profiler_config.py
@@ -85,8 +85,8 @@ def ensure_field_level_settings_are_normalized(
                 if field_level_metric.startswith("include_field_"):
                     values.setdefault(field_level_metric, False)
 
-            assert (
-                max_num_fields_to_profile is None
-            ), f"{max_num_fields_to_profile_key} should be set to None"
+            assert max_num_fields_to_profile is None, (
+                f"{max_num_fields_to_profile_key} should be set to None"
+            )
 
         return values
diff --git a/metadata-ingestion/src/datahub/ingestion/source/abs/source.py b/metadata-ingestion/src/datahub/ingestion/source/abs/source.py
index ad2bc36cf558b5..586e7a3af3bcd1 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/abs/source.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/abs/source.py
@@ -508,7 +508,12 @@ def abs_browser(
                         ):
                             abs_path = self.create_abs_path(obj.name)
                             logger.debug(f"Sampling file: {abs_path}")
-                            yield abs_path, obj.name, obj.last_modified, obj.size,
+                            yield (
+                                abs_path,
+                                obj.name,
+                                obj.last_modified,
+                                obj.size,
+                            )
                 except Exception as e:
                     # This odd check if being done because boto does not have a proper exception to catch
                     # The exception that appears in stacktrace cannot actually be caught without a lot more work
@@ -552,9 +557,12 @@ def local_browser(
         if os.path.isfile(prefix):
             logger.debug(f"Scanning single local file: {prefix}")
             file_name = prefix
-            yield prefix, file_name, datetime.utcfromtimestamp(
-                os.path.getmtime(prefix)
-            ), os.path.getsize(prefix)
+            yield (
+                prefix,
+                file_name,
+                datetime.utcfromtimestamp(os.path.getmtime(prefix)),
+                os.path.getsize(prefix),
+            )
         else:
             logger.debug(f"Scanning files under local folder: {prefix}")
             for root, dirs, files in os.walk(prefix):
@@ -565,9 +573,12 @@ def local_browser(
                     full_path = PurePath(
                         os.path.normpath(os.path.join(root, file))
                     ).as_posix()
-                    yield full_path, file, datetime.utcfromtimestamp(
-                        os.path.getmtime(full_path)
-                    ), os.path.getsize(full_path)
+                    yield (
+                        full_path,
+                        file,
+                        datetime.utcfromtimestamp(os.path.getmtime(full_path)),
+                        os.path.getsize(full_path),
+                    )
 
     def get_workunits_internal(self) -> Iterable[MetadataWorkUnit]:
         self.container_WU_creator = ContainerWUCreator(
@@ -613,7 +624,7 @@ def get_workunits_internal(self) -> Iterable[MetadataWorkUnit]:
                                 table_data.table_path
                             ].timestamp = table_data.timestamp
 
-                for guid, table_data in table_dict.items():
+                for _, table_data in table_dict.items():
                     yield from self.ingest_table(table_data, path_spec)
 
     def get_workunit_processors(self) -> List[Optional[MetadataWorkUnitProcessor]]:
diff --git a/metadata-ingestion/src/datahub/ingestion/source/aws/glue.py b/metadata-ingestion/src/datahub/ingestion/source/aws/glue.py
index 30e81643837375..2ace71b6ff6c1d 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/aws/glue.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/aws/glue.py
@@ -521,7 +521,7 @@ def process_dataflow_node(
         # otherwise, a node represents a transformation
         else:
             node_urn = mce_builder.make_data_job_urn_with_flow(
-                flow_urn, job_id=f'{node["NodeType"]}-{node["Id"]}'
+                flow_urn, job_id=f"{node['NodeType']}-{node['Id']}"
             )
 
         return {
@@ -679,7 +679,7 @@ def get_datajob_wu(self, node: Dict[str, Any], job_name: str) -> MetadataWorkUni
             )
         )
 
-        return MetadataWorkUnit(id=f'{job_name}-{node["Id"]}', mce=mce)
+        return MetadataWorkUnit(id=f"{job_name}-{node['Id']}", mce=mce)
 
     def get_all_databases(self) -> Iterable[Mapping[str, Any]]:
         logger.debug("Getting all databases")
@@ -738,11 +738,17 @@ def get_all_databases_and_tables(
         self,
     ) -> Tuple[List[Mapping[str, Any]], List[Dict]]:
         all_databases = [*self.get_all_databases()]
-        all_tables = [
-            tables
-            for database in all_databases
-            for tables in self.get_tables_from_database(database)
-        ]
+        all_tables = []
+        for database in all_databases:
+            try:
+                for tables in self.get_tables_from_database(database):
+                    all_tables.append(tables)
+            except Exception as e:
+                self.report.failure(
+                    message="Failed to get tables from database",
+                    context=database["Name"],
+                    exc=e,
+                )
         return all_databases, all_tables
 
     def get_lineage_if_enabled(
@@ -750,13 +756,13 @@ def get_lineage_if_enabled(
     ) -> Optional[MetadataWorkUnit]:
         if self.source_config.emit_s3_lineage:
             # extract dataset properties aspect
-            dataset_properties: Optional[
-                DatasetPropertiesClass
-            ] = mce_builder.get_aspect_if_available(mce, DatasetPropertiesClass)
+            dataset_properties: Optional[DatasetPropertiesClass] = (
+                mce_builder.get_aspect_if_available(mce, DatasetPropertiesClass)
+            )
             # extract dataset schema aspect
-            schema_metadata: Optional[
-                SchemaMetadataClass
-            ] = mce_builder.get_aspect_if_available(mce, SchemaMetadataClass)
+            schema_metadata: Optional[SchemaMetadataClass] = (
+                mce_builder.get_aspect_if_available(mce, SchemaMetadataClass)
+            )
 
             if dataset_properties and "Location" in dataset_properties.customProperties:
                 location = dataset_properties.customProperties["Location"]
@@ -765,9 +771,9 @@ def get_lineage_if_enabled(
                         location, self.source_config.env
                     )
                     assert self.ctx.graph
-                    schema_metadata_for_s3: Optional[
-                        SchemaMetadataClass
-                    ] = self.ctx.graph.get_schema_metadata(s3_dataset_urn)
+                    schema_metadata_for_s3: Optional[SchemaMetadataClass] = (
+                        self.ctx.graph.get_schema_metadata(s3_dataset_urn)
+                    )
 
                     if self.source_config.glue_s3_lineage_direction == "upstream":
                         fine_grained_lineages = None
diff --git a/metadata-ingestion/src/datahub/ingestion/source/aws/s3_boto_utils.py b/metadata-ingestion/src/datahub/ingestion/source/aws/s3_boto_utils.py
index 682cddca8f4154..87a6f8a5baf2e3 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/aws/s3_boto_utils.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/aws/s3_boto_utils.py
@@ -40,7 +40,7 @@ def get_s3_tags(
                 ]
             )
         except s3.meta.client.exceptions.ClientError:
-            logger.warn(f"No tags found for bucket={bucket_name}")
+            logger.warning(f"No tags found for bucket={bucket_name}")
 
     if use_s3_object_tags and key_name is not None:
         s3_client = aws_config.get_s3_client()
@@ -53,7 +53,7 @@ def get_s3_tags(
         else:
             # Unlike bucket tags, if an object does not have tags, it will just return an empty array
             # as opposed to an exception.
-            logger.warn(f"No tags found for bucket={bucket_name} key={key_name}")
+            logger.warning(f"No tags found for bucket={bucket_name} key={key_name}")
     if len(tags_to_add) == 0:
         return None
     if ctx.graph is not None:
@@ -65,7 +65,7 @@ def get_s3_tags(
         if current_tags:
             tags_to_add.extend([current_tag.tag for current_tag in current_tags.tags])
     else:
-        logger.warn("Could not connect to DatahubApi. No current tags to maintain")
+        logger.warning("Could not connect to DatahubApi. No current tags to maintain")
     # Remove duplicate tags
     tags_to_add = sorted(list(set(tags_to_add)))
     new_tags = GlobalTagsClass(
diff --git a/metadata-ingestion/src/datahub/ingestion/source/aws/s3_util.py b/metadata-ingestion/src/datahub/ingestion/source/aws/s3_util.py
index 360f18aa448f27..878b8dd1bb9a51 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/aws/s3_util.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/aws/s3_util.py
@@ -1,11 +1,6 @@
 import logging
 import os
-from collections import defaultdict
-from typing import TYPE_CHECKING, Dict, Iterable, List, Optional
-
-if TYPE_CHECKING:
-    from mypy_boto3_s3.service_resource import ObjectSummary
-
+from typing import Optional
 
 S3_PREFIXES = ["s3://", "s3n://", "s3a://"]
 
@@ -73,21 +68,3 @@ def get_key_prefix(s3_uri: str) -> str:
             f"Not an S3 URI. Must start with one of the following prefixes: {str(S3_PREFIXES)}"
         )
     return strip_s3_prefix(s3_uri).split("/", maxsplit=1)[1]
-
-
-def group_s3_objects_by_dirname(
-    s3_objects: Iterable["ObjectSummary"],
-) -> Dict[str, List["ObjectSummary"]]:
-    """
-    Groups S3 objects by their directory name.
-
-    If a s3_object in the root directory (i.e., s3://bucket/file.txt), it is grouped under '/'.
-    """
-    grouped_s3_objs = defaultdict(list)
-    for obj in s3_objects:
-        if "/" in obj.key:
-            dirname = obj.key.rsplit("/", 1)[0]
-        else:
-            dirname = "/"
-        grouped_s3_objs[dirname].append(obj)
-    return grouped_s3_objs
diff --git a/metadata-ingestion/src/datahub/ingestion/source/aws/sagemaker_processors/feature_groups.py b/metadata-ingestion/src/datahub/ingestion/source/aws/sagemaker_processors/feature_groups.py
index c4561b9d9e676a..d46d1c099383fe 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/aws/sagemaker_processors/feature_groups.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/aws/sagemaker_processors/feature_groups.py
@@ -257,7 +257,7 @@ def get_feature_wu(
             mce = MetadataChangeEvent(proposedSnapshot=feature_snapshot)
 
         return MetadataWorkUnit(
-            id=f'{feature_group_details["FeatureGroupName"]}-{feature["FeatureName"]}',
+            id=f"{feature_group_details['FeatureGroupName']}-{feature['FeatureName']}",
             mce=mce,
         )
 
diff --git a/metadata-ingestion/src/datahub/ingestion/source/aws/sagemaker_processors/models.py b/metadata-ingestion/src/datahub/ingestion/source/aws/sagemaker_processors/models.py
index 0f433aaecf2d96..f1374117af775f 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/aws/sagemaker_processors/models.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/aws/sagemaker_processors/models.py
@@ -212,7 +212,7 @@ def get_endpoint_wu(
         mce = MetadataChangeEvent(proposedSnapshot=endpoint_snapshot)
 
         return MetadataWorkUnit(
-            id=f'{endpoint_details["EndpointName"]}',
+            id=f"{endpoint_details['EndpointName']}",
             mce=mce,
         )
 
@@ -503,7 +503,7 @@ def get_model_wu(
         mce = MetadataChangeEvent(proposedSnapshot=model_snapshot)
 
         return MetadataWorkUnit(
-            id=f'{model_details["ModelName"]}',
+            id=f"{model_details['ModelName']}",
             mce=mce,
         )
 
diff --git a/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery.py b/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery.py
index 508b4bbaa277dc..ceb010a7f0675f 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery.py
@@ -132,9 +132,9 @@ def __init__(self, ctx: PipelineContext, config: BigQueryV2Config):
         self.filters = BigQueryFilter(self.config, self.report)
         self.identifiers = BigQueryIdentifierBuilder(self.config, self.report)
 
-        redundant_lineage_run_skip_handler: Optional[
-            RedundantLineageRunSkipHandler
-        ] = None
+        redundant_lineage_run_skip_handler: Optional[RedundantLineageRunSkipHandler] = (
+            None
+        )
         if self.config.enable_stateful_lineage_ingestion:
             redundant_lineage_run_skip_handler = RedundantLineageRunSkipHandler(
                 source=self,
diff --git a/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery_audit.py b/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery_audit.py
index 42f82704c81b99..d35c5265878c03 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery_audit.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery_audit.py
@@ -37,9 +37,9 @@ class BigqueryTableIdentifier:
 
     # Note: this regex may get overwritten by the sharded_table_pattern config.
     # The class-level constant, however, will not be overwritten.
-    _BIGQUERY_DEFAULT_SHARDED_TABLE_REGEX: ClassVar[
-        str
-    ] = _BIGQUERY_DEFAULT_SHARDED_TABLE_REGEX
+    _BIGQUERY_DEFAULT_SHARDED_TABLE_REGEX: ClassVar[str] = (
+        _BIGQUERY_DEFAULT_SHARDED_TABLE_REGEX
+    )
     _BIGQUERY_WILDCARD_REGEX: ClassVar[str] = "((_(\\d+)?)\\*$)|\\*$"
     _BQ_SHARDED_TABLE_SUFFIX: str = "_yyyymmdd"
 
diff --git a/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery_config.py b/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery_config.py
index afbe919df4dcae..57bfa2e3090d31 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery_config.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery_config.py
@@ -137,9 +137,9 @@ class BigQueryCredential(ConfigModel):
     @root_validator(skip_on_failure=True)
     def validate_config(cls, values: Dict[str, Any]) -> Dict[str, Any]:
         if values.get("client_x509_cert_url") is None:
-            values[
-                "client_x509_cert_url"
-            ] = f'https://www.googleapis.com/robot/v1/metadata/x509/{values["client_email"]}'
+            values["client_x509_cert_url"] = (
+                f"https://www.googleapis.com/robot/v1/metadata/x509/{values['client_email']}"
+            )
         return values
 
     def create_credential_temp_file(self) -> str:
@@ -611,9 +611,9 @@ def validate_bigquery_audit_metadata_datasets(
         cls, v: Optional[List[str]], values: Dict
     ) -> Optional[List[str]]:
         if values.get("use_exported_bigquery_audit_metadata"):
-            assert (
-                v and len(v) > 0
-            ), "`bigquery_audit_metadata_datasets` should be set if using `use_exported_bigquery_audit_metadata: True`."
+            assert v and len(v) > 0, (
+                "`bigquery_audit_metadata_datasets` should be set if using `use_exported_bigquery_audit_metadata: True`."
+            )
 
         return v
 
diff --git a/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery_platform_resource_helper.py b/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery_platform_resource_helper.py
index 9da2aceb19220a..7dc0e4195d5dc9 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery_platform_resource_helper.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery_platform_resource_helper.py
@@ -87,9 +87,9 @@ def get_platform_resource(
             key=platform_resource_key, graph_client=self.graph
         )
         if platform_resource:
-            self.platform_resource_cache[
-                platform_resource_key.primary_key
-            ] = platform_resource
+            self.platform_resource_cache[platform_resource_key.primary_key] = (
+                platform_resource
+            )
             return platform_resource
         return None
 
@@ -115,7 +115,11 @@ def generate_label_platform_resource(
                 and platform_resource.resource_info.value
             ):
                 try:
-                    existing_info: Optional[BigQueryLabelInfo] = platform_resource.resource_info.value.as_pydantic_object(BigQueryLabelInfo)  # type: ignore
+                    existing_info: Optional[BigQueryLabelInfo] = (
+                        platform_resource.resource_info.value.as_pydantic_object(  # type: ignore
+                            BigQueryLabelInfo
+                        )
+                    )
                 except ValidationError as e:
                     logger.error(
                         f"Error converting existing value to BigQueryLabelInfo: {e}. Creating new one. Maybe this is because of a non backward compatible schema change."
diff --git a/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery_schema_gen.py b/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery_schema_gen.py
index 56e930dfb811f1..1cace31a5e4c49 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery_schema_gen.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery_schema_gen.py
@@ -2,7 +2,6 @@
 import re
 from base64 import b32decode
 from collections import defaultdict
-from itertools import groupby
 from typing import Dict, Iterable, List, Optional, Set, Type, Union, cast
 
 from google.cloud.bigquery.table import TableListItem
@@ -101,6 +100,7 @@
 from datahub.metadata.urns import TagUrn
 from datahub.sql_parsing.schema_resolver import SchemaResolver
 from datahub.utilities.file_backed_collections import FileBackedDict
+from datahub.utilities.groupby import groupby_unsorted
 from datahub.utilities.hive_schema_to_avro import (
     HiveColumnToAvroConverter,
     get_schema_fields_for_hive_column,
@@ -311,8 +311,10 @@ def gen_dataset_containers(
                         platform_resource: PlatformResource = self.platform_resource_helper.generate_label_platform_resource(
                             label, tag_urn, managed_by_datahub=False
                         )
-                        label_info: BigQueryLabelInfo = platform_resource.resource_info.value.as_pydantic_object(  # type: ignore
-                            BigQueryLabelInfo
+                        label_info: BigQueryLabelInfo = (
+                            platform_resource.resource_info.value.as_pydantic_object(  # type: ignore
+                                BigQueryLabelInfo
+                            )
                         )
                         tag_urn = TagUrn.from_string(label_info.datahub_urn)
 
@@ -728,7 +730,7 @@ def gen_foreign_keys(
         foreign_keys: List[BigqueryTableConstraint] = list(
             filter(lambda x: x.type == "FOREIGN KEY", table.constraints)
         )
-        for key, group in groupby(
+        for key, group in groupby_unsorted(
             foreign_keys,
             lambda x: f"{x.referenced_project_id}.{x.referenced_dataset}.{x.referenced_table_name}",
         ):
@@ -820,8 +822,10 @@ def gen_table_dataset_workunits(
                         platform_resource: PlatformResource = self.platform_resource_helper.generate_label_platform_resource(
                             label, tag_urn, managed_by_datahub=False
                         )
-                        label_info: BigQueryLabelInfo = platform_resource.resource_info.value.as_pydantic_object(  # type: ignore
-                            BigQueryLabelInfo
+                        label_info: BigQueryLabelInfo = (
+                            platform_resource.resource_info.value.as_pydantic_object(  # type: ignore
+                                BigQueryLabelInfo
+                            )
                         )
                         tag_urn = TagUrn.from_string(label_info.datahub_urn)
 
@@ -860,8 +864,10 @@ def gen_view_dataset_workunits(
                         platform_resource: PlatformResource = self.platform_resource_helper.generate_label_platform_resource(
                             label, tag_urn, managed_by_datahub=False
                         )
-                        label_info: BigQueryLabelInfo = platform_resource.resource_info.value.as_pydantic_object(  # type: ignore
-                            BigQueryLabelInfo
+                        label_info: BigQueryLabelInfo = (
+                            platform_resource.resource_info.value.as_pydantic_object(  # type: ignore
+                                BigQueryLabelInfo
+                            )
                         )
                         tag_urn = TagUrn.from_string(label_info.datahub_urn)
 
@@ -1203,9 +1209,9 @@ def get_tables_for_dataset(
                     report=self.report,
                 )
 
-        self.report.metadata_extraction_sec[
-            f"{project_id}.{dataset.name}"
-        ] = timer.elapsed_seconds(digits=2)
+        self.report.metadata_extraction_sec[f"{project_id}.{dataset.name}"] = (
+            timer.elapsed_seconds(digits=2)
+        )
 
     def get_core_table_details(
         self, dataset_name: str, project_id: str, temp_table_dataset_prefix: str
diff --git a/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/lineage.py b/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/lineage.py
index 433282a21fdb66..da82c6a06f0395 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/lineage.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/lineage.py
@@ -697,7 +697,7 @@ def _create_lineage_map(
                         if parsed_queries[-1]:
                             query = f"""create table `{destination_table.get_sanitized_table_ref().table_identifier.get_table_name()}` AS
                             (
-                                {parsed_queries[-1].sql(dialect='bigquery')}
+                                {parsed_queries[-1].sql(dialect="bigquery")}
                             )"""
                         else:
                             query = e.query
@@ -809,11 +809,11 @@ def get_upstream_tables(
                             upstream_lineage, temp_table_upstream
                         )
 
-                        upstreams[
-                            ref_temp_table_upstream
-                        ] = _merge_lineage_edge_columns(
-                            upstreams.get(ref_temp_table_upstream),
-                            collapsed_lineage,
+                        upstreams[ref_temp_table_upstream] = (
+                            _merge_lineage_edge_columns(
+                                upstreams.get(ref_temp_table_upstream),
+                                collapsed_lineage,
+                            )
                         )
             else:
                 upstreams[upstream_table_ref] = _merge_lineage_edge_columns(
@@ -1004,9 +1004,9 @@ def get_lineage_for_external_table(
                 dataset_urn
             )
             for gcs_dataset_urn in gcs_urns:
-                schema_metadata_for_gcs: Optional[
-                    SchemaMetadataClass
-                ] = graph.get_schema_metadata(gcs_dataset_urn)
+                schema_metadata_for_gcs: Optional[SchemaMetadataClass] = (
+                    graph.get_schema_metadata(gcs_dataset_urn)
+                )
                 if schema_metadata and schema_metadata_for_gcs:
                     fine_grained_lineage = self.get_fine_grained_lineages_with_gcs(
                         dataset_urn,
diff --git a/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/queries.py b/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/queries.py
index 2ac40a48de4cc7..8a558d7736a389 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/queries.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/queries.py
@@ -387,9 +387,7 @@ class BigqueryQuery:
     OR
     protoPayload.metadata.tableDataRead.reason = "JOB"
 )
-""".strip(
-    "\t \n"
-)
+""".strip("\t \n")
 
 
 def bigquery_audit_metadata_query_template_lineage(
diff --git a/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/queries_extractor.py b/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/queries_extractor.py
index 08c9beaa73c53b..0f9471219c6590 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/queries_extractor.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/queries_extractor.py
@@ -271,9 +271,9 @@ def get_workunits_internal(
             # Preprocessing stage that deduplicates the queries using query hash per usage bucket
             # Note: FileBackedDict is an ordered dictionary, so the order of execution of
             # queries is inherently maintained
-            queries_deduped: FileBackedDict[
-                Dict[int, ObservedQuery]
-            ] = self.deduplicate_queries(queries)
+            queries_deduped: FileBackedDict[Dict[int, ObservedQuery]] = (
+                self.deduplicate_queries(queries)
+            )
             self.report.num_unique_queries = len(queries_deduped)
             logger.info(f"Found {self.report.num_unique_queries} unique queries")
 
diff --git a/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/usage.py b/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/usage.py
index f2f6cc731858d1..c2b849e58fc6dc 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/usage.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/usage.py
@@ -763,9 +763,9 @@ def _create_operational_custom_properties(
                     )
 
                 if event.query_event.default_dataset:
-                    custom_properties[
-                        "defaultDatabase"
-                    ] = event.query_event.default_dataset
+                    custom_properties["defaultDatabase"] = (
+                        event.query_event.default_dataset
+                    )
             if event.read_event:
                 if event.read_event.readReason:
                     custom_properties["readReason"] = event.read_event.readReason
diff --git a/metadata-ingestion/src/datahub/ingestion/source/cassandra/cassandra.py b/metadata-ingestion/src/datahub/ingestion/source/cassandra/cassandra.py
index dcdccc08ce0483..062c64d45767fc 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/cassandra/cassandra.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/cassandra/cassandra.py
@@ -91,7 +91,6 @@ class KeyspaceKey(ContainerKey):
     supported=True,
 )
 class CassandraSource(StatefulIngestionSourceBase):
-
     """
     This plugin extracts the following:
 
diff --git a/metadata-ingestion/src/datahub/ingestion/source/cassandra/cassandra_utils.py b/metadata-ingestion/src/datahub/ingestion/source/cassandra/cassandra_utils.py
index 75a0ba0c617734..b467ca0aca6be4 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/cassandra/cassandra_utils.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/cassandra/cassandra_utils.py
@@ -107,10 +107,10 @@ class CassandraToSchemaFieldConverter:
 
     @staticmethod
     def get_column_type(cassandra_column_type: str) -> SchemaFieldDataType:
-        type_class: Optional[
-            Type
-        ] = CassandraToSchemaFieldConverter._field_type_to_schema_field_type.get(
-            cassandra_column_type
+        type_class: Optional[Type] = (
+            CassandraToSchemaFieldConverter._field_type_to_schema_field_type.get(
+                cassandra_column_type
+            )
         )
         if type_class is None:
             logger.warning(
diff --git a/metadata-ingestion/src/datahub/ingestion/source/confluent_schema_registry.py b/metadata-ingestion/src/datahub/ingestion/source/confluent_schema_registry.py
index 2b75d0dca53cb7..5ba4dd13fb2ac9 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/confluent_schema_registry.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/confluent_schema_registry.py
@@ -293,9 +293,9 @@ def _get_schema_and_fields(
     def _load_json_schema_with_resolved_references(
         self, schema: Schema, name: str, subject: str
     ) -> dict:
-        imported_json_schemas: List[
-            JsonSchemaWrapper
-        ] = self.get_schemas_from_confluent_ref_json(schema, name=name, subject=subject)
+        imported_json_schemas: List[JsonSchemaWrapper] = (
+            self.get_schemas_from_confluent_ref_json(schema, name=name, subject=subject)
+        )
         schema_dict = json.loads(schema.schema_str)
         reference_map = {}
         for imported_schema in imported_json_schemas:
@@ -332,9 +332,9 @@ def _get_schema_fields(
             )
 
         elif schema.schema_type == "PROTOBUF":
-            imported_schemas: List[
-                ProtobufSchema
-            ] = self.get_schemas_from_confluent_ref_protobuf(schema)
+            imported_schemas: List[ProtobufSchema] = (
+                self.get_schemas_from_confluent_ref_protobuf(schema)
+            )
             base_name: str = topic.replace(".", "_")
             fields = protobuf_util.protobuf_schema_to_mce_fields(
                 ProtobufSchema(
diff --git a/metadata-ingestion/src/datahub/ingestion/source/csv_enricher.py b/metadata-ingestion/src/datahub/ingestion/source/csv_enricher.py
index 42e025073b534e..8ebb7b9ef7fbdf 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/csv_enricher.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/csv_enricher.py
@@ -371,11 +371,11 @@ def get_resource_workunits(
         domain: Optional[str],
         description: Optional[str],
     ) -> Iterable[MetadataWorkUnit]:
-        maybe_terms_wu: Optional[
-            MetadataWorkUnit
-        ] = self.get_resource_glossary_terms_work_unit(
-            entity_urn=entity_urn,
-            term_associations=term_associations,
+        maybe_terms_wu: Optional[MetadataWorkUnit] = (
+            self.get_resource_glossary_terms_work_unit(
+                entity_urn=entity_urn,
+                term_associations=term_associations,
+            )
         )
         if maybe_terms_wu:
             self.report.num_glossary_term_workunits_produced += 1
@@ -389,31 +389,31 @@ def get_resource_workunits(
             self.report.num_tag_workunits_produced += 1
             yield maybe_tags_wu
 
-        maybe_owners_wu: Optional[
-            MetadataWorkUnit
-        ] = self.get_resource_owners_work_unit(
-            entity_urn=entity_urn,
-            owners=owners,
+        maybe_owners_wu: Optional[MetadataWorkUnit] = (
+            self.get_resource_owners_work_unit(
+                entity_urn=entity_urn,
+                owners=owners,
+            )
         )
         if maybe_owners_wu:
             self.report.num_owners_workunits_produced += 1
             yield maybe_owners_wu
 
-        maybe_domain_wu: Optional[
-            MetadataWorkUnit
-        ] = self.get_resource_domain_work_unit(
-            entity_urn=entity_urn,
-            domain=domain,
+        maybe_domain_wu: Optional[MetadataWorkUnit] = (
+            self.get_resource_domain_work_unit(
+                entity_urn=entity_urn,
+                domain=domain,
+            )
         )
         if maybe_domain_wu:
             self.report.num_domain_workunits_produced += 1
             yield maybe_domain_wu
 
-        maybe_description_wu: Optional[
-            MetadataWorkUnit
-        ] = self.get_resource_description_work_unit(
-            entity_urn=entity_urn,
-            description=description,
+        maybe_description_wu: Optional[MetadataWorkUnit] = (
+            self.get_resource_description_work_unit(
+                entity_urn=entity_urn,
+                description=description,
+            )
         )
         if maybe_description_wu:
             self.report.num_description_workunits_produced += 1
@@ -426,9 +426,9 @@ def process_sub_resource_row(
         needs_write: bool,
     ) -> Tuple[EditableSchemaMetadataClass, bool]:
         field_path: str = sub_resource_row.field_path
-        term_associations: List[
-            GlossaryTermAssociationClass
-        ] = sub_resource_row.term_associations
+        term_associations: List[GlossaryTermAssociationClass] = (
+            sub_resource_row.term_associations
+        )
         tag_associations: List[TagAssociationClass] = sub_resource_row.tag_associations
         description: Optional[str] = sub_resource_row.description
         has_terms: bool = len(term_associations) > 0
@@ -517,9 +517,9 @@ def get_sub_resource_work_units(self) -> Iterable[MetadataWorkUnit]:
             # Boolean field to tell whether we need to write an MCPW.
             needs_write = False
 
-            current_editable_schema_metadata: Optional[
-                EditableSchemaMetadataClass
-            ] = None
+            current_editable_schema_metadata: Optional[EditableSchemaMetadataClass] = (
+                None
+            )
             if self.ctx.graph and not self.should_overwrite:
                 # Fetch the current editable schema metadata
                 current_editable_schema_metadata = self.ctx.graph.get_aspect(
@@ -655,9 +655,9 @@ def get_workunits_internal(self) -> Iterable[MetadataWorkUnit]:
             entity_urn = row["resource"]
             entity_type = Urn.from_string(row["resource"]).get_type()
 
-            term_associations: List[
-                GlossaryTermAssociationClass
-            ] = self.maybe_extract_glossary_terms(row)
+            term_associations: List[GlossaryTermAssociationClass] = (
+                self.maybe_extract_glossary_terms(row)
+            )
             tag_associations: List[TagAssociationClass] = self.maybe_extract_tags(row)
             owners: List[OwnerClass] = self.maybe_extract_owners(row, is_resource_row)
 
diff --git a/metadata-ingestion/src/datahub/ingestion/source/datahub/config.py b/metadata-ingestion/src/datahub/ingestion/source/datahub/config.py
index 8622e221940317..37ccce8f8b9657 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/datahub/config.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/datahub/config.py
@@ -25,6 +25,10 @@
     "globalSettingsKey",
     "globalSettingsInfo",
     "testResults",
+    "dataHubExecutionRequestKey",
+    "dataHubExecutionRequestInput",
+    "dataHubExecutionRequestSignal",
+    "dataHubExecutionRequestResult",
 }
 
 
diff --git a/metadata-ingestion/src/datahub/ingestion/source/datahub/datahub_database_reader.py b/metadata-ingestion/src/datahub/ingestion/source/datahub/datahub_database_reader.py
index ee105f4862caba..51a25829d21dba 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/datahub/datahub_database_reader.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/datahub/datahub_database_reader.py
@@ -152,7 +152,9 @@ def execute_server_cursor(
     ) -> Iterable[Dict[str, Any]]:
         with self.engine.connect() as conn:
             if self.engine.dialect.name in ["postgresql", "mysql", "mariadb"]:
-                with conn.begin():  # Transaction required for PostgreSQL server-side cursor
+                with (
+                    conn.begin()
+                ):  # Transaction required for PostgreSQL server-side cursor
                     # Note that stream_results=True is mainly supported by PostgreSQL and MySQL-based dialects.
                     # https://docs.sqlalchemy.org/en/14/core/connections.html#sqlalchemy.engine.Connection.execution_options.params.stream_results
                     conn = conn.execution_options(
@@ -222,7 +224,7 @@ def _parse_row(
             )
         except Exception as e:
             logger.warning(
-                f'Failed to parse metadata for {row["urn"]}: {e}', exc_info=True
+                f"Failed to parse metadata for {row['urn']}: {e}", exc_info=True
             )
             self.report.num_database_parse_errors += 1
             self.report.database_parse_errors.setdefault(
diff --git a/metadata-ingestion/src/datahub/ingestion/source/dbt/dbt_cloud.py b/metadata-ingestion/src/datahub/ingestion/source/dbt/dbt_cloud.py
index 5042f6d69b261a..41b59a9c8b892c 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/dbt/dbt_cloud.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/dbt/dbt_cloud.py
@@ -194,20 +194,20 @@ def infer_metadata_endpoint(access_url: str) -> Optional[str]:
 
 _DBT_FIELDS_BY_TYPE = {
     "models": f"""
-    { _DBT_GRAPHQL_COMMON_FIELDS }
-    { _DBT_GRAPHQL_NODE_COMMON_FIELDS }
-    { _DBT_GRAPHQL_MODEL_SEED_SNAPSHOT_FIELDS }
+    {_DBT_GRAPHQL_COMMON_FIELDS}
+    {_DBT_GRAPHQL_NODE_COMMON_FIELDS}
+    {_DBT_GRAPHQL_MODEL_SEED_SNAPSHOT_FIELDS}
     dependsOn
     materializedType
 """,
     "seeds": f"""
-    { _DBT_GRAPHQL_COMMON_FIELDS }
-    { _DBT_GRAPHQL_NODE_COMMON_FIELDS }
-    { _DBT_GRAPHQL_MODEL_SEED_SNAPSHOT_FIELDS }
+    {_DBT_GRAPHQL_COMMON_FIELDS}
+    {_DBT_GRAPHQL_NODE_COMMON_FIELDS}
+    {_DBT_GRAPHQL_MODEL_SEED_SNAPSHOT_FIELDS}
 """,
     "sources": f"""
-    { _DBT_GRAPHQL_COMMON_FIELDS }
-    { _DBT_GRAPHQL_NODE_COMMON_FIELDS }
+    {_DBT_GRAPHQL_COMMON_FIELDS}
+    {_DBT_GRAPHQL_NODE_COMMON_FIELDS}
     identifier
     sourceName
     sourceDescription
@@ -218,9 +218,9 @@ def infer_metadata_endpoint(access_url: str) -> Optional[str]:
     loader
 """,
     "snapshots": f"""
-    { _DBT_GRAPHQL_COMMON_FIELDS }
-    { _DBT_GRAPHQL_NODE_COMMON_FIELDS }
-    { _DBT_GRAPHQL_MODEL_SEED_SNAPSHOT_FIELDS }
+    {_DBT_GRAPHQL_COMMON_FIELDS}
+    {_DBT_GRAPHQL_NODE_COMMON_FIELDS}
+    {_DBT_GRAPHQL_MODEL_SEED_SNAPSHOT_FIELDS}
     parentsSources {{
       uniqueId
     }}
@@ -229,7 +229,7 @@ def infer_metadata_endpoint(access_url: str) -> Optional[str]:
     }}
 """,
     "tests": f"""
-    { _DBT_GRAPHQL_COMMON_FIELDS }
+    {_DBT_GRAPHQL_COMMON_FIELDS}
     state
     columnName
     status
@@ -315,7 +315,7 @@ def _send_graphql_query(
             res = response.json()
             if "errors" in res:
                 raise ValueError(
-                    f'Unable to fetch metadata from dbt Cloud: {res["errors"]}'
+                    f"Unable to fetch metadata from dbt Cloud: {res['errors']}"
                 )
             data = res["data"]
         except JSONDecodeError as e:
diff --git a/metadata-ingestion/src/datahub/ingestion/source/dbt/dbt_common.py b/metadata-ingestion/src/datahub/ingestion/source/dbt/dbt_common.py
index 499e7e1231d050..99fa09b8b20b9b 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/dbt/dbt_common.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/dbt/dbt_common.py
@@ -1,4 +1,3 @@
-import itertools
 import logging
 import re
 from abc import abstractmethod
@@ -111,6 +110,7 @@
     parse_statements_and_pick,
     try_format_query,
 )
+from datahub.utilities.groupby import groupby_unsorted
 from datahub.utilities.lossy_collections import LossyList
 from datahub.utilities.mapping import Constants, OperationProcessor
 from datahub.utilities.time import datetime_to_ts_millis
@@ -506,16 +506,18 @@ class DBTNode:
     materialization: Optional[str]  # table, view, ephemeral, incremental, snapshot
     # see https://docs.getdbt.com/reference/artifacts/manifest-json
     catalog_type: Optional[str]
-    missing_from_catalog: bool  # indicates if the node was missing from the catalog.json
+    missing_from_catalog: (
+        bool  # indicates if the node was missing from the catalog.json
+    )
 
     owner: Optional[str]
 
     columns: List[DBTColumn] = field(default_factory=list)
     upstream_nodes: List[str] = field(default_factory=list)  # list of upstream dbt_name
     upstream_cll: List[DBTColumnLineageInfo] = field(default_factory=list)
-    raw_sql_parsing_result: Optional[
-        SqlParsingResult
-    ] = None  # only set for nodes that don't depend on ephemeral models
+    raw_sql_parsing_result: Optional[SqlParsingResult] = (
+        None  # only set for nodes that don't depend on ephemeral models
+    )
     cll_debug_info: Optional[SqlParsingDebugInfo] = None
 
     meta: Dict[str, Any] = field(default_factory=dict)
@@ -869,10 +871,10 @@ def create_test_entity_mcps(
                                 "platform": DBT_PLATFORM,
                                 "name": node.dbt_name,
                                 "instance": self.config.platform_instance,
+                                # Ideally we'd include the env unconditionally. However, we started out
+                                # not including env in the guid, so we need to maintain backwards compatibility
+                                # with existing PROD assertions.
                                 **(
-                                    # Ideally we'd include the env unconditionally. However, we started out
-                                    # not including env in the guid, so we need to maintain backwards compatibility
-                                    # with existing PROD assertions.
                                     {"env": self.config.env}
                                     if self.config.env != mce_builder.DEFAULT_ENV
                                     and self.config.include_env_in_assertion_guid
@@ -1927,7 +1929,7 @@ def _translate_dbt_name_to_upstream_urn(dbt_name: str) -> str:
                             else None
                         ),
                     )
-                    for downstream, upstreams in itertools.groupby(
+                    for downstream, upstreams in groupby_unsorted(
                         node.upstream_cll, lambda x: x.downstream_col
                     )
                 ]
diff --git a/metadata-ingestion/src/datahub/ingestion/source/dremio/dremio_api.py b/metadata-ingestion/src/datahub/ingestion/source/dremio/dremio_api.py
index d913b7e42065d2..cf2d9670400ca5 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/dremio/dremio_api.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/dremio/dremio_api.py
@@ -181,7 +181,7 @@ def authenticate(self, connection_args: "DremioSourceConfig") -> None:
             return
 
         # On-prem Dremio authentication (PAT or Basic Auth)
-        for retry in range(1, self._retry_count + 1):
+        for _ in range(1, self._retry_count + 1):
             try:
                 if connection_args.authentication_method == "PAT":
                     self.session.headers.update(
@@ -191,9 +191,9 @@ def authenticate(self, connection_args: "DremioSourceConfig") -> None:
                     )
                     return
                 else:
-                    assert (
-                        connection_args.username and connection_args.password
-                    ), "Username and password are required for authentication"
+                    assert connection_args.username and connection_args.password, (
+                        "Username and password are required for authentication"
+                    )
                     host = connection_args.hostname
                     port = connection_args.port
                     protocol = "https" if connection_args.tls else "http"
diff --git a/metadata-ingestion/src/datahub/ingestion/source/dremio/dremio_datahub_source_mapping.py b/metadata-ingestion/src/datahub/ingestion/source/dremio/dremio_datahub_source_mapping.py
index e5d6b8e40fb3d8..482647f8d77da1 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/dremio/dremio_datahub_source_mapping.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/dremio/dremio_datahub_source_mapping.py
@@ -101,9 +101,9 @@ def add_mapping(
         Add a new source type if not in the map (e.g., Dremio ARP).
         """
         dremio_source_type = dremio_source_type.upper()
-        DremioToDataHubSourceTypeMapping.SOURCE_TYPE_MAPPING[
-            dremio_source_type
-        ] = datahub_source_type
+        DremioToDataHubSourceTypeMapping.SOURCE_TYPE_MAPPING[dremio_source_type] = (
+            datahub_source_type
+        )
 
         if category:
             if category.lower() == "file_object_storage":
diff --git a/metadata-ingestion/src/datahub/ingestion/source/elastic_search.py b/metadata-ingestion/src/datahub/ingestion/source/elastic_search.py
index 99aa5f54f6a576..ce1c60dcafdd46 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/elastic_search.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/elastic_search.py
@@ -111,10 +111,10 @@ class ElasticToSchemaFieldConverter:
 
     @staticmethod
     def get_column_type(elastic_column_type: str) -> SchemaFieldDataType:
-        type_class: Optional[
-            Type
-        ] = ElasticToSchemaFieldConverter._field_type_to_schema_field_type.get(
-            elastic_column_type
+        type_class: Optional[Type] = (
+            ElasticToSchemaFieldConverter._field_type_to_schema_field_type.get(
+                elastic_column_type
+            )
         )
         if type_class is None:
             logger.warning(
diff --git a/metadata-ingestion/src/datahub/ingestion/source/fivetran/config.py b/metadata-ingestion/src/datahub/ingestion/source/fivetran/config.py
index 86826ae7bedc09..b048b6157fbb9f 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/fivetran/config.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/fivetran/config.py
@@ -167,6 +167,10 @@ class PlatformDetail(ConfigModel):
         description="The database that all assets produced by this connector belong to. "
         "For destinations, this defaults to the fivetran log config's database.",
     )
+    include_schema_in_urn: bool = pydantic.Field(
+        default=True,
+        description="Include schema in the dataset URN. In some cases, the schema is not relevant to the dataset URN and Fivetran sets it to the source and destination table names in the connector.",
+    )
 
 
 class FivetranSourceConfig(StatefulIngestionConfigBase, DatasetSourceConfigMixin):
diff --git a/metadata-ingestion/src/datahub/ingestion/source/fivetran/fivetran.py b/metadata-ingestion/src/datahub/ingestion/source/fivetran/fivetran.py
index d8ebbe5b63d1ae..a3469696699f17 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/fivetran/fivetran.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/fivetran/fivetran.py
@@ -119,21 +119,31 @@ def _extend_lineage(self, connector: Connector, datajob: DataJob) -> Dict[str, s
             )
 
         for lineage in connector.lineage:
+            source_table = (
+                lineage.source_table
+                if source_details.include_schema_in_urn
+                else lineage.source_table.split(".", 1)[1]
+            )
             input_dataset_urn = DatasetUrn.create_from_ids(
                 platform_id=source_details.platform,
                 table_name=(
-                    f"{source_details.database.lower()}.{lineage.source_table}"
+                    f"{source_details.database.lower()}.{source_table}"
                     if source_details.database
-                    else lineage.source_table
+                    else source_table
                 ),
                 env=source_details.env,
                 platform_instance=source_details.platform_instance,
             )
             input_dataset_urn_list.append(input_dataset_urn)
 
+            destination_table = (
+                lineage.destination_table
+                if destination_details.include_schema_in_urn
+                else lineage.destination_table.split(".", 1)[1]
+            )
             output_dataset_urn = DatasetUrn.create_from_ids(
                 platform_id=destination_details.platform,
-                table_name=f"{destination_details.database.lower()}.{lineage.destination_table}",
+                table_name=f"{destination_details.database.lower()}.{destination_table}",
                 env=destination_details.env,
                 platform_instance=destination_details.platform_instance,
             )
@@ -176,12 +186,12 @@ def _extend_lineage(self, connector: Connector, datajob: DataJob) -> Dict[str, s
             **{
                 f"source.{k}": str(v)
                 for k, v in source_details.dict().items()
-                if v is not None
+                if v is not None and not isinstance(v, bool)
             },
             **{
                 f"destination.{k}": str(v)
                 for k, v in destination_details.dict().items()
-                if v is not None
+                if v is not None and not isinstance(v, bool)
             },
         )
 
diff --git a/metadata-ingestion/src/datahub/ingestion/source/gc/datahub_gc.py b/metadata-ingestion/src/datahub/ingestion/source/gc/datahub_gc.py
index b4cc5423277c5a..5c03b873c5505e 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/gc/datahub_gc.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/gc/datahub_gc.py
@@ -292,6 +292,7 @@ def revoke_expired_tokens(self) -> None:
             tokens = list_access_tokens.get("tokens", [])
             total = list_access_tokens.get("total", 0)
             if tokens == []:
+                # Due to a server bug we cannot rely on just total
                 break
             for token in tokens:
                 self.report.expired_tokens_revoked += 1
diff --git a/metadata-ingestion/src/datahub/ingestion/source/gc/soft_deleted_entity_cleanup.py b/metadata-ingestion/src/datahub/ingestion/source/gc/soft_deleted_entity_cleanup.py
index b3d6e8514c02f2..ffcd9218a2103c 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/gc/soft_deleted_entity_cleanup.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/gc/soft_deleted_entity_cleanup.py
@@ -99,6 +99,7 @@ class SoftDeletedEntitiesCleanupConfig(ConfigModel):
 
 @dataclass
 class SoftDeletedEntitiesReport(SourceReport):
+    num_calls_made: Dict[str, int] = field(default_factory=dict)
     num_entities_found: Dict[str, int] = field(default_factory=dict)
     num_soft_deleted_entity_processed: int = 0
     num_soft_deleted_retained_due_to_age: int = 0
@@ -154,9 +155,9 @@ def _update_report(self, urn: str, entity_type: str) -> None:
             current_count = self.report.num_hard_deleted_by_type.get(entity_type, 0)
             self.report.num_hard_deleted_by_type[entity_type] = current_count + 1
             if entity_type not in self.report.sample_hard_deleted_aspects_by_type:
-                self.report.sample_hard_deleted_aspects_by_type[
-                    entity_type
-                ] = LossyList()
+                self.report.sample_hard_deleted_aspects_by_type[entity_type] = (
+                    LossyList()
+                )
             self.report.sample_hard_deleted_aspects_by_type[entity_type].append(urn)
 
     def delete_entity(self, urn: str) -> None:
@@ -242,6 +243,11 @@ def _get_soft_deleted(self, graphql_query: str, entity_type: str) -> Iterable[st
 
         while True:
             try:
+                if entity_type not in self.report.num_calls_made:
+                    self.report.num_calls_made[entity_type] = 1
+                else:
+                    self.report.num_calls_made[entity_type] += 1
+                self._print_report()
                 result = self.ctx.graph.execute_graphql(
                     graphql_query,
                     {
@@ -270,7 +276,13 @@ def _get_soft_deleted(self, graphql_query: str, entity_type: str) -> Iterable[st
                 )
                 break
             scroll_across_entities = result.get("scrollAcrossEntities")
-            if not scroll_across_entities or not scroll_across_entities.get("count"):
+            if not scroll_across_entities:
+                break
+            search_results = scroll_across_entities.get("searchResults")
+            count = scroll_across_entities.get("count")
+            if not count or not search_results:
+                # Due to a server bug we cannot rely on just count as it was returning response like this
+                # {'count': 1, 'nextScrollId': None, 'searchResults': []}
                 break
             if entity_type == "DATA_PROCESS_INSTANCE":
                 # Temp workaround. See note in beginning of the function
@@ -282,7 +294,7 @@ def _get_soft_deleted(self, graphql_query: str, entity_type: str) -> Iterable[st
             self.report.num_entities_found[entity_type] += scroll_across_entities.get(
                 "count"
             )
-            for query in scroll_across_entities.get("searchResults"):
+            for query in search_results:
                 yield query["entity"]["urn"]
 
     def _get_urns(self) -> Iterable[str]:
diff --git a/metadata-ingestion/src/datahub/ingestion/source/gcs/gcs_source.py b/metadata-ingestion/src/datahub/ingestion/source/gcs/gcs_source.py
index 18838af9bdf85f..95490147d48204 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/gcs/gcs_source.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/gcs/gcs_source.py
@@ -88,6 +88,7 @@ def __init__(self, config: GCSSourceConfig, ctx: PipelineContext):
         super().__init__(config, ctx)
         self.config = config
         self.report = GCSSourceReport()
+        self.platform: str = PLATFORM_GCS
         self.s3_source = self.create_equivalent_s3_source(ctx)
 
     @classmethod
@@ -135,14 +136,15 @@ def create_equivalent_s3_path_specs(self):
 
     def create_equivalent_s3_source(self, ctx: PipelineContext) -> S3Source:
         config = self.create_equivalent_s3_config()
-        return self.s3_source_overrides(S3Source(config, ctx))
+        return self.s3_source_overrides(S3Source(config, PipelineContext(ctx.run_id)))
 
     def s3_source_overrides(self, source: S3Source) -> S3Source:
         source.source_config.platform = PLATFORM_GCS
 
         source.is_s3_platform = lambda: True  # type: ignore
-        source.create_s3_path = lambda bucket_name, key: unquote(f"s3://{bucket_name}/{key}")  # type: ignore
-
+        source.create_s3_path = lambda bucket_name, key: unquote(  # type: ignore
+            f"s3://{bucket_name}/{key}"
+        )
         return source
 
     def get_workunit_processors(self) -> List[Optional[MetadataWorkUnitProcessor]]:
diff --git a/metadata-ingestion/src/datahub/ingestion/source/ge_data_profiler.py b/metadata-ingestion/src/datahub/ingestion/source/ge_data_profiler.py
index c2c43a0e805b24..bde26f97bf271f 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/ge_data_profiler.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/ge_data_profiler.py
@@ -267,7 +267,6 @@ def _is_single_row_query_method(query: Any) -> bool:
         "get_column_max",
         "get_column_mean",
         "get_column_stdev",
-        "get_column_stdev",
         "get_column_nonnull_count",
         "get_column_unique_count",
     }
@@ -328,7 +327,7 @@ def _is_single_row_query_method(query: Any) -> bool:
 
 
 def _run_with_query_combiner(
-    method: Callable[Concatenate["_SingleDatasetProfiler", P], None]
+    method: Callable[Concatenate["_SingleDatasetProfiler", P], None],
 ) -> Callable[Concatenate["_SingleDatasetProfiler", P], None]:
     @functools.wraps(method)
     def inner(
@@ -1538,9 +1537,7 @@ def create_bigquery_temp_table(
         query_job: Optional["google.cloud.bigquery.job.query.QueryJob"] = (
             # In google-cloud-bigquery 3.15.0, the _query_job attribute was
             # made public and renamed to query_job.
-            cursor.query_job
-            if hasattr(cursor, "query_job")
-            else cursor._query_job  # type: ignore[attr-defined]
+            cursor.query_job if hasattr(cursor, "query_job") else cursor._query_job  # type: ignore[attr-defined]
         )
         assert query_job
         temp_destination_table = query_job.destination
diff --git a/metadata-ingestion/src/datahub/ingestion/source/ge_profiling_config.py b/metadata-ingestion/src/datahub/ingestion/source/ge_profiling_config.py
index 42d0def0a46e7d..93142a347ca0e6 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/ge_profiling_config.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/ge_profiling_config.py
@@ -220,9 +220,9 @@ def ensure_field_level_settings_are_normalized(
                         )
                     values[field_level_metric] = False
 
-            assert (
-                max_num_fields_to_profile is None
-            ), f"{max_num_fields_to_profile_key} should be set to None"
+            assert max_num_fields_to_profile is None, (
+                f"{max_num_fields_to_profile_key} should be set to None"
+            )
 
         # Disable expensive queries.
         if values.get("turn_off_expensive_profiling_metrics"):
diff --git a/metadata-ingestion/src/datahub/ingestion/source/iceberg/iceberg.py b/metadata-ingestion/src/datahub/ingestion/source/iceberg/iceberg.py
index 8101f0110509e3..9a62ee2dab52f4 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/iceberg/iceberg.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/iceberg/iceberg.py
@@ -296,9 +296,9 @@ def _create_iceberg_workunit(
                 custom_properties["snapshot-id"] = str(
                     table.current_snapshot().snapshot_id
                 )
-                custom_properties[
-                    "manifest-list"
-                ] = table.current_snapshot().manifest_list
+                custom_properties["manifest-list"] = (
+                    table.current_snapshot().manifest_list
+                )
             dataset_properties = DatasetPropertiesClass(
                 name=table.name()[-1],
                 description=table.metadata.properties.get("comment", None),
diff --git a/metadata-ingestion/src/datahub/ingestion/source/identity/azure_ad.py b/metadata-ingestion/src/datahub/ingestion/source/identity/azure_ad.py
index 885b6514779cc4..edb9b7b8bd5264 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/identity/azure_ad.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/identity/azure_ad.py
@@ -354,9 +354,9 @@ def get_workunits_internal(self) -> Iterable[MetadataWorkUnit]:
                     yield MetadataWorkUnit(id=group_status_wu_id, mcp=group_status_mcp)
 
         # Populate GroupMembership Aspects for CorpUsers
-        datahub_corp_user_urn_to_group_membership: Dict[
-            str, GroupMembershipClass
-        ] = defaultdict(lambda: GroupMembershipClass(groups=[]))
+        datahub_corp_user_urn_to_group_membership: Dict[str, GroupMembershipClass] = (
+            defaultdict(lambda: GroupMembershipClass(groups=[]))
+        )
         if (
             self.config.ingest_group_membership
             and len(self.selected_azure_ad_groups) > 0
diff --git a/metadata-ingestion/src/datahub/ingestion/source/identity/okta.py b/metadata-ingestion/src/datahub/ingestion/source/identity/okta.py
index dda81b0e34a8d2..5452fbcd3f053b 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/identity/okta.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/identity/okta.py
@@ -344,9 +344,9 @@ def get_workunits_internal(self) -> Iterable[MetadataWorkUnit]:
                 ).as_workunit()
 
         # Step 2: Populate GroupMembership Aspects for CorpUsers
-        datahub_corp_user_urn_to_group_membership: Dict[
-            str, GroupMembershipClass
-        ] = defaultdict(lambda: GroupMembershipClass(groups=[]))
+        datahub_corp_user_urn_to_group_membership: Dict[str, GroupMembershipClass] = (
+            defaultdict(lambda: GroupMembershipClass(groups=[]))
+        )
         if self.config.ingest_group_membership and okta_groups is not None:
             # Fetch membership for each group.
             for okta_group in okta_groups:
diff --git a/metadata-ingestion/src/datahub/ingestion/source/kafka/kafka.py b/metadata-ingestion/src/datahub/ingestion/source/kafka/kafka.py
index fa842a15ba7328..9f15eda1501f11 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/kafka/kafka.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/kafka/kafka.py
@@ -419,10 +419,10 @@ def _extract_record(
             custom_props = self.build_custom_properties(
                 topic, topic_detail, extra_topic_config
             )
-            schema_name: Optional[
-                str
-            ] = self.schema_registry_client._get_subject_for_topic(
-                topic, is_key_schema=False
+            schema_name: Optional[str] = (
+                self.schema_registry_client._get_subject_for_topic(
+                    topic, is_key_schema=False
+                )
             )
             if schema_name is not None:
                 custom_props["Schema Name"] = schema_name
@@ -610,11 +610,13 @@ def fetch_extra_topic_details(self, topics: List[str]) -> Dict[str, dict]:
 
     def fetch_topic_configurations(self, topics: List[str]) -> Dict[str, dict]:
         logger.info("Fetching config details for all topics")
-        configs: Dict[
-            ConfigResource, concurrent.futures.Future
-        ] = self.admin_client.describe_configs(
-            resources=[ConfigResource(ConfigResource.Type.TOPIC, t) for t in topics],
-            request_timeout=self.source_config.connection.client_timeout_seconds,
+        configs: Dict[ConfigResource, concurrent.futures.Future] = (
+            self.admin_client.describe_configs(
+                resources=[
+                    ConfigResource(ConfigResource.Type.TOPIC, t) for t in topics
+                ],
+                request_timeout=self.source_config.connection.client_timeout_seconds,
+            )
         )
         logger.debug("Waiting for config details futures to complete")
         concurrent.futures.wait(configs.values())
diff --git a/metadata-ingestion/src/datahub/ingestion/source/kafka_connect/kafka_connect.py b/metadata-ingestion/src/datahub/ingestion/source/kafka_connect/kafka_connect.py
index 72be864fc30a1c..9edfce5855f430 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/kafka_connect/kafka_connect.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/kafka_connect/kafka_connect.py
@@ -110,9 +110,8 @@ def get_connectors_manifest(self) -> Iterable[ConnectorManifest]:
             connector_manifest = self._get_connector_manifest(
                 connector_name, connector_url
             )
-            if (
-                connector_manifest is None
-                or not self.config.connector_patterns.allowed(connector_manifest.name)
+            if connector_manifest is None or not self.config.connector_patterns.allowed(
+                connector_manifest.name
             ):
                 self.report.report_dropped(connector_name)
                 continue
diff --git a/metadata-ingestion/src/datahub/ingestion/source/kafka_connect/sink_connectors.py b/metadata-ingestion/src/datahub/ingestion/source/kafka_connect/sink_connectors.py
index 2790460c8e6019..10255ed544b812 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/kafka_connect/sink_connectors.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/kafka_connect/sink_connectors.py
@@ -199,9 +199,9 @@ def get_parser(
             transforms.append(transform)
             for key in self.connector_manifest.config.keys():
                 if key.startswith(f"transforms.{name}."):
-                    transform[
-                        key.replace(f"transforms.{name}.", "")
-                    ] = self.connector_manifest.config[key]
+                    transform[key.replace(f"transforms.{name}.", "")] = (
+                        self.connector_manifest.config[key]
+                    )
 
         if "defaultDataset" in connector_manifest.config:
             defaultDataset = connector_manifest.config["defaultDataset"]
diff --git a/metadata-ingestion/src/datahub/ingestion/source/kafka_connect/source_connectors.py b/metadata-ingestion/src/datahub/ingestion/source/kafka_connect/source_connectors.py
index 7b3b6e551a0a1f..5e64d4e161e3ea 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/kafka_connect/source_connectors.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/kafka_connect/source_connectors.py
@@ -123,9 +123,9 @@ def get_parser(
             transforms.append(transform)
             for key in self.connector_manifest.config.keys():
                 if key.startswith(f"transforms.{name}."):
-                    transform[
-                        key.replace(f"transforms.{name}.", "")
-                    ] = self.connector_manifest.config[key]
+                    transform[key.replace(f"transforms.{name}.", "")] = (
+                        self.connector_manifest.config[key]
+                    )
 
         return self.JdbcParser(
             db_connection_url,
diff --git a/metadata-ingestion/src/datahub/ingestion/source/looker/looker_common.py b/metadata-ingestion/src/datahub/ingestion/source/looker/looker_common.py
index 1183916e9b3fef..abe9b5684f8f1f 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/looker/looker_common.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/looker/looker_common.py
@@ -596,9 +596,9 @@ class LookerUtil:
 
     @staticmethod
     def _extract_view_from_field(field: str) -> str:
-        assert (
-            field.count(".") == 1
-        ), f"Error: A field must be prefixed by a view name, field is: {field}"
+        assert field.count(".") == 1, (
+            f"Error: A field must be prefixed by a view name, field is: {field}"
+        )
         return field.split(".")[0]
 
     @staticmethod
@@ -815,9 +815,9 @@ class LookerExplore:
     project_name: Optional[str] = None
     label: Optional[str] = None
     description: Optional[str] = None
-    upstream_views: Optional[
-        List[ProjectInclude]
-    ] = None  # captures the view name(s) this explore is derived from
+    upstream_views: Optional[List[ProjectInclude]] = (
+        None  # captures the view name(s) this explore is derived from
+    )
     upstream_views_file_path: Dict[str, Optional[str]] = dataclasses_field(
         default_factory=dict
     )  # view_name is key and file_path is value. A single file may contains multiple views
@@ -889,7 +889,7 @@ def from_dict(
                     upstream_views.extend(parsed_explore.upstream_views or [])
                 else:
                     logger.warning(
-                        f'Could not find extended explore {extended_explore} for explore {dict["name"]} in model {model_name}'
+                        f"Could not find extended explore {extended_explore} for explore {dict['name']} in model {model_name}"
                     )
         else:
             # we only fallback to the view_names list if this is not an extended explore
@@ -903,7 +903,7 @@ def from_dict(
                 )
                 if not info:
                     logger.warning(
-                        f'Could not resolve view {view_name} for explore {dict["name"]} in model {model_name}'
+                        f"Could not resolve view {view_name} for explore {dict['name']} in model {model_name}"
                     )
                 else:
                     upstream_views.append(
@@ -935,9 +935,9 @@ def from_api(  # noqa: C901
         try:
             explore = client.lookml_model_explore(model, explore_name)
             views: Set[str] = set()
-            lkml_fields: List[
-                LookmlModelExploreField
-            ] = explore_field_set_to_lkml_fields(explore)
+            lkml_fields: List[LookmlModelExploreField] = (
+                explore_field_set_to_lkml_fields(explore)
+            )
 
             if explore.view_name is not None and explore.view_name != explore.name:
                 # explore is not named after a view and is instead using a from field, which is modeled as view_name.
@@ -1034,9 +1034,9 @@ def from_api(  # noqa: C901
                         if measure_field.name is None:
                             continue
                         else:
-                            field_name_vs_raw_explore_field[
-                                measure_field.name
-                            ] = measure_field
+                            field_name_vs_raw_explore_field[measure_field.name] = (
+                                measure_field
+                            )
 
                             view_fields.append(
                                 ViewField(
@@ -1072,11 +1072,11 @@ def from_api(  # noqa: C901
             if view_project_map:
                 logger.debug(f"views and their projects: {view_project_map}")
 
-            upstream_views_file_path: Dict[
-                str, Optional[str]
-            ] = create_upstream_views_file_path_map(
-                lkml_fields=lkml_fields,
-                view_names=views,
+            upstream_views_file_path: Dict[str, Optional[str]] = (
+                create_upstream_views_file_path_map(
+                    lkml_fields=lkml_fields,
+                    view_names=views,
+                )
             )
             if upstream_views_file_path:
                 logger.debug(f"views and their file-paths: {upstream_views_file_path}")
diff --git a/metadata-ingestion/src/datahub/ingestion/source/looker/looker_config.py b/metadata-ingestion/src/datahub/ingestion/source/looker/looker_config.py
index 4e9d0f68928a45..3ed3186399588e 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/looker/looker_config.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/looker/looker_config.py
@@ -166,9 +166,9 @@ def _get_generic_definition(
         # e.g. spark1 or hive2 or druid_18
         platform = re.sub(r"[0-9]+", "", dialect_name.split("_")[0])
 
-    assert (
-        platform is not None
-    ), f"Failed to extract a valid platform from connection {looker_connection}"
+    assert platform is not None, (
+        f"Failed to extract a valid platform from connection {looker_connection}"
+    )
     db = looker_connection.database
     schema = looker_connection.schema  # ok for this to be None
     return platform, db, schema
diff --git a/metadata-ingestion/src/datahub/ingestion/source/looker/looker_source.py b/metadata-ingestion/src/datahub/ingestion/source/looker/looker_source.py
index 8487d5113bc1d3..2f1fcd378d40fb 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/looker/looker_source.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/looker/looker_source.py
@@ -250,9 +250,9 @@ def _set_test_connection_capability(
 
     @staticmethod
     def _extract_view_from_field(field: str) -> str:
-        assert (
-            field.count(".") == 1
-        ), f"Error: A field must be prefixed by a view name, field is: {field}"
+        assert field.count(".") == 1, (
+            f"Error: A field must be prefixed by a view name, field is: {field}"
+        )
         return field.split(".")[0]
 
     def _get_views_from_fields(self, fields: List[str]) -> List[str]:
@@ -610,12 +610,12 @@ def _get_folder_browse_path_v2_entries(
     def _create_platform_instance_aspect(
         self,
     ) -> DataPlatformInstance:
-        assert (
-            self.source_config.platform_name
-        ), "Platform name is not set in the configuration."
-        assert (
-            self.source_config.platform_instance
-        ), "Platform instance is not set in the configuration."
+        assert self.source_config.platform_name, (
+            "Platform name is not set in the configuration."
+        )
+        assert self.source_config.platform_instance, (
+            "Platform instance is not set in the configuration."
+        )
 
         return DataPlatformInstance(
             platform=builder.make_data_platform_urn(self.source_config.platform_name),
@@ -1016,9 +1016,9 @@ def _make_dashboard_and_chart_mces(
         yield from chart_events
 
         # Step 2: Emit metadata events for the Dashboard itself.
-        chart_urns: Set[
-            str
-        ] = set()  # Collect the unique child chart urns for dashboard input lineage.
+        chart_urns: Set[str] = (
+            set()
+        )  # Collect the unique child chart urns for dashboard input lineage.
         for chart_event in chart_events:
             chart_event_urn = self._extract_event_urn(chart_event)
             if chart_event_urn:
@@ -1538,20 +1538,20 @@ def extract_independent_looks(self) -> Iterable[MetadataWorkUnit]:
                     }
                 )
 
-            dashboard_element: Optional[
-                LookerDashboardElement
-            ] = self._get_looker_dashboard_element(
-                DashboardElement(
-                    id=f"looks_{look.id}",  # to avoid conflict with non-standalone looks (element.id prefixes),
-                    # we add the "looks_" prefix to look.id.
-                    title=look.title,
-                    subtitle_text=look.description,
-                    look_id=look.id,
-                    dashboard_id=None,  # As this is an independent look
-                    look=LookWithQuery(
-                        query=query, folder=look.folder, user_id=look.user_id
+            dashboard_element: Optional[LookerDashboardElement] = (
+                self._get_looker_dashboard_element(
+                    DashboardElement(
+                        id=f"looks_{look.id}",  # to avoid conflict with non-standalone looks (element.id prefixes),
+                        # we add the "looks_" prefix to look.id.
+                        title=look.title,
+                        subtitle_text=look.description,
+                        look_id=look.id,
+                        dashboard_id=None,  # As this is an independent look
+                        look=LookWithQuery(
+                            query=query, folder=look.folder, user_id=look.user_id
+                        ),
                     ),
-                ),
+                )
             )
 
             if dashboard_element is not None:
diff --git a/metadata-ingestion/src/datahub/ingestion/source/looker/looker_template_language.py b/metadata-ingestion/src/datahub/ingestion/source/looker/looker_template_language.py
index 6d49d57e077435..2bcae4d46b8d52 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/looker/looker_template_language.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/looker/looker_template_language.py
@@ -33,9 +33,9 @@
 
 
 class SpecialVariable:
-    SPECIAL_VARIABLE_PATTERN: ClassVar[
-        str
-    ] = r"\b\w+(\.\w+)*\._(is_selected|in_query|is_filtered)\b"
+    SPECIAL_VARIABLE_PATTERN: ClassVar[str] = (
+        r"\b\w+(\.\w+)*\._(is_selected|in_query|is_filtered)\b"
+    )
     liquid_variable: dict
 
     def __init__(self, liquid_variable):
diff --git a/metadata-ingestion/src/datahub/ingestion/source/looker/looker_usage.py b/metadata-ingestion/src/datahub/ingestion/source/looker/looker_usage.py
index 098d7d73a3da84..05806840b5c954 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/looker/looker_usage.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/looker/looker_usage.py
@@ -257,9 +257,9 @@ def _process_entity_timeseries_rows(
 
         for row in rows:
             logger.debug(row)
-            entity_stat_aspect[
-                self.get_entity_stat_key(row)
-            ] = self.to_entity_timeseries_stat_aspect(row)
+            entity_stat_aspect[self.get_entity_stat_key(row)] = (
+                self.to_entity_timeseries_stat_aspect(row)
+            )
 
         return entity_stat_aspect
 
@@ -385,10 +385,8 @@ def generate_usage_stat_mcps(self) -> Iterable[MetadataChangeProposalWrapper]:
         entity_rows: List[Dict] = self._execute_query(
             entity_query_with_filters, "entity_query"
         )
-        entity_usage_stat: Dict[
-            Tuple[str, str], Any
-        ] = self._process_entity_timeseries_rows(
-            entity_rows
+        entity_usage_stat: Dict[Tuple[str, str], Any] = (
+            self._process_entity_timeseries_rows(entity_rows)
         )  # Any type to pass mypy unbound Aspect type error
 
         user_wise_query_with_filters: LookerQuery = self._append_filters(
diff --git a/metadata-ingestion/src/datahub/ingestion/source/looker/lookml_concept_context.py b/metadata-ingestion/src/datahub/ingestion/source/looker/lookml_concept_context.py
index 103f4175a9ccff..4e38165bb56286 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/looker/lookml_concept_context.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/looker/lookml_concept_context.py
@@ -38,16 +38,16 @@ def merge_parent_and_child_fields(
     # Create a map field-name vs field
     child_field_map: dict = {}
     for field in child_fields:
-        assert (
-            NAME in field
-        ), "A lookml view must have a name field"  # name is required field of lookml field array
+        assert NAME in field, (
+            "A lookml view must have a name field"
+        )  # name is required field of lookml field array
 
         child_field_map[field[NAME]] = field
 
     for field in parent_fields:
-        assert (
-            NAME in field
-        ), "A lookml view must have a name field"  # name is required field of lookml field array
+        assert NAME in field, (
+            "A lookml view must have a name field"
+        )  # name is required field of lookml field array
 
         if field[NAME] in child_field_map:
             # Fields defined in the child view take higher precedence.
diff --git a/metadata-ingestion/src/datahub/ingestion/source/looker/lookml_source.py b/metadata-ingestion/src/datahub/ingestion/source/looker/lookml_source.py
index c7d3724472d3c8..a8575c84b510d5 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/looker/lookml_source.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/looker/lookml_source.py
@@ -482,14 +482,14 @@ def get_project_name(self, model_name: str) -> str:
         if self.source_config.project_name is not None:
             return self.source_config.project_name
 
-        assert (
-            self.looker_client is not None
-        ), "Failed to find a configured Looker API client"
+        assert self.looker_client is not None, (
+            "Failed to find a configured Looker API client"
+        )
         try:
             model = self.looker_client.lookml_model(model_name, fields="project_name")
-            assert (
-                model.project_name is not None
-            ), f"Failed to find a project name for model {model_name}"
+            assert model.project_name is not None, (
+                f"Failed to find a project name for model {model_name}"
+            )
             return model.project_name
         except SDKError:
             raise ValueError(
@@ -541,9 +541,9 @@ def get_workunits_internal(self) -> Iterable[MetadataWorkUnit]:
                 self.reporter.git_clone_latency = datetime.now() - start_time
                 self.source_config.base_folder = checkout_dir.resolve()
 
-            self.base_projects_folder[
-                BASE_PROJECT_NAME
-            ] = self.source_config.base_folder
+            self.base_projects_folder[BASE_PROJECT_NAME] = (
+                self.source_config.base_folder
+            )
 
             visited_projects: Set[str] = set()
 
@@ -641,9 +641,9 @@ def _recursively_check_manifests(
                     repo_url=remote_project.url,
                 )
 
-                self.base_projects_folder[
-                    remote_project.name
-                ] = p_checkout_dir.resolve()
+                self.base_projects_folder[remote_project.name] = (
+                    p_checkout_dir.resolve()
+                )
                 repo = p_cloner.get_last_repo_cloned()
                 assert repo
                 remote_git_info = GitInfo(
@@ -930,9 +930,7 @@ def get_internal_workunits(self) -> Iterable[MetadataWorkUnit]:  # noqa: C901
                                         logger.warning(
                                             f"view {maybe_looker_view.id.view_name} from model {model_name}, connection {model.connection} was previously processed via model {prev_model_name}, connection {prev_model_connection} and will likely lead to incorrect lineage to the underlying tables"
                                         )
-                                        if (
-                                            not self.source_config.emit_reachable_views_only
-                                        ):
+                                        if not self.source_config.emit_reachable_views_only:
                                             logger.warning(
                                                 "Consider enabling the `emit_reachable_views_only` flag to handle this case."
                                             )
diff --git a/metadata-ingestion/src/datahub/ingestion/source/looker/view_upstream.py b/metadata-ingestion/src/datahub/ingestion/source/looker/view_upstream.py
index 971181e4300d69..f77eebb3cdd8cb 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/looker/view_upstream.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/looker/view_upstream.py
@@ -484,11 +484,11 @@ def __init__(
         )
 
     def __get_upstream_dataset_urn(self) -> List[str]:
-        current_view_id: Optional[
-            LookerViewId
-        ] = self.looker_view_id_cache.get_looker_view_id(
-            view_name=self.view_context.name(),
-            base_folder_path=self.view_context.base_folder_path,
+        current_view_id: Optional[LookerViewId] = (
+            self.looker_view_id_cache.get_looker_view_id(
+                view_name=self.view_context.name(),
+                base_folder_path=self.view_context.base_folder_path,
+            )
         )
 
         # Current view will always be present in cache. assert  will silence the lint
diff --git a/metadata-ingestion/src/datahub/ingestion/source/mlflow.py b/metadata-ingestion/src/datahub/ingestion/source/mlflow.py
index b0b04dff20bffc..02125db83d2582 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/mlflow.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/mlflow.py
@@ -172,10 +172,10 @@ def _get_mlflow_registered_models(self) -> Iterable[RegisteredModel]:
         """
         Get all Registered Models in MLflow Model Registry.
         """
-        registered_models: Iterable[
-            RegisteredModel
-        ] = self._traverse_mlflow_search_func(
-            search_func=self.client.search_registered_models,
+        registered_models: Iterable[RegisteredModel] = (
+            self._traverse_mlflow_search_func(
+                search_func=self.client.search_registered_models,
+            )
         )
         return registered_models
 
diff --git a/metadata-ingestion/src/datahub/ingestion/source/mode.py b/metadata-ingestion/src/datahub/ingestion/source/mode.py
index 68ecc5d8694ac5..bf0a33e423446a 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/mode.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/mode.py
@@ -893,11 +893,11 @@ def normalize_mode_query(self, query: str) -> str:
                         jinja_params[key] = parameters[key].get("default", "")
 
                 normalized_query = re.sub(
-                    r"{% form %}(.*){% endform %}",
-                    "",
-                    query,
-                    0,
-                    re.MULTILINE | re.DOTALL,
+                    pattern=r"{% form %}(.*){% endform %}",
+                    repl="",
+                    string=query,
+                    count=0,
+                    flags=re.MULTILINE | re.DOTALL,
                 )
 
             # Wherever we don't resolve the jinja params, we replace it with NULL
diff --git a/metadata-ingestion/src/datahub/ingestion/source/mongodb.py b/metadata-ingestion/src/datahub/ingestion/source/mongodb.py
index bbc4897d227bac..ad8487c1a759ec 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/mongodb.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/mongodb.py
@@ -288,7 +288,9 @@ def __init__(self, ctx: PipelineContext, config: MongoDBConfig):
 
         # See https://pymongo.readthedocs.io/en/stable/examples/datetimes.html#handling-out-of-range-datetimes
         self.mongo_client = MongoClient(
-            self.config.connect_uri, datetime_conversion="DATETIME_AUTO", **options  # type: ignore
+            self.config.connect_uri,
+            datetime_conversion="DATETIME_AUTO",
+            **options,  # type: ignore
         )
 
         # This cheaply tests the connection. For details, see
@@ -470,9 +472,9 @@ def _infer_schema_metadata(
             )
             # Add this information to the custom properties so user can know they are looking at downsampled schema
             dataset_properties.customProperties["schema.downsampled"] = "True"
-            dataset_properties.customProperties[
-                "schema.totalFields"
-            ] = f"{collection_schema_size}"
+            dataset_properties.customProperties["schema.totalFields"] = (
+                f"{collection_schema_size}"
+            )
 
         logger.debug(f"Size of collection fields = {len(collection_fields)}")
         # append each schema field (sort so output is consistent)
diff --git a/metadata-ingestion/src/datahub/ingestion/source/neo4j/neo4j_source.py b/metadata-ingestion/src/datahub/ingestion/source/neo4j/neo4j_source.py
index 2c9107b967e4f8..8cdd4b17733e01 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/neo4j/neo4j_source.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/neo4j/neo4j_source.py
@@ -286,7 +286,7 @@ def get_workunits_internal(self) -> Iterable[MetadataWorkUnit]:
         df = self.get_neo4j_metadata(
             "CALL apoc.meta.schema() YIELD value UNWIND keys(value) AS key RETURN key, value[key] AS value;"
         )
-        for index, row in df.iterrows():
+        for _, row in df.iterrows():
             try:
                 yield MetadataWorkUnit(
                     id=row["key"],
diff --git a/metadata-ingestion/src/datahub/ingestion/source/nifi.py b/metadata-ingestion/src/datahub/ingestion/source/nifi.py
index 7f446f6d1c2718..52b1386e21d85a 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/nifi.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/nifi.py
@@ -184,9 +184,9 @@ def validator_site_url_to_site_name(cls, values):
 
     @validator("site_url")
     def validator_site_url(cls, site_url: str) -> str:
-        assert site_url.startswith(
-            ("http://", "https://")
-        ), "site_url must start with http:// or https://"
+        assert site_url.startswith(("http://", "https://")), (
+            "site_url must start with http:// or https://"
+        )
 
         if not site_url.endswith("/"):
             site_url = site_url + "/"
@@ -487,9 +487,7 @@ def rest_api_base_url(self):
     def get_report(self) -> SourceReport:
         return self.report
 
-    def update_flow(
-        self, pg_flow_dto: Dict, recursion_level: int = 0
-    ) -> None:  # noqa: C901
+    def update_flow(self, pg_flow_dto: Dict, recursion_level: int = 0) -> None:  # noqa: C901
         """
         Update self.nifi_flow with contents of the input process group `pg_flow_dto`
         """
@@ -548,16 +546,16 @@ def update_flow(
         for inputPort in flow_dto.get("inputPorts", []):
             component = inputPort.get("component")
             if inputPort.get("allowRemoteAccess"):
-                self.nifi_flow.remotely_accessible_ports[
-                    component.get("id")
-                ] = NifiComponent(
-                    component.get("id"),
-                    component.get("name"),
-                    component.get("type"),
-                    component.get("parentGroupId"),
-                    NifiType.INPUT_PORT,
-                    comments=component.get("comments"),
-                    status=component.get("status", {}).get("runStatus"),
+                self.nifi_flow.remotely_accessible_ports[component.get("id")] = (
+                    NifiComponent(
+                        component.get("id"),
+                        component.get("name"),
+                        component.get("type"),
+                        component.get("parentGroupId"),
+                        NifiType.INPUT_PORT,
+                        comments=component.get("comments"),
+                        status=component.get("status", {}).get("runStatus"),
+                    )
                 )
                 logger.debug(f"Adding remotely accessible port {component.get('id')}")
             else:
@@ -576,16 +574,16 @@ def update_flow(
         for outputPort in flow_dto.get("outputPorts", []):
             component = outputPort.get("component")
             if outputPort.get("allowRemoteAccess"):
-                self.nifi_flow.remotely_accessible_ports[
-                    component.get("id")
-                ] = NifiComponent(
-                    component.get("id"),
-                    component.get("name"),
-                    component.get("type"),
-                    component.get("parentGroupId"),
-                    NifiType.OUTPUT_PORT,
-                    comments=component.get("comments"),
-                    status=component.get("status", {}).get("runStatus"),
+                self.nifi_flow.remotely_accessible_ports[component.get("id")] = (
+                    NifiComponent(
+                        component.get("id"),
+                        component.get("name"),
+                        component.get("type"),
+                        component.get("parentGroupId"),
+                        NifiType.OUTPUT_PORT,
+                        comments=component.get("comments"),
+                        status=component.get("status", {}).get("runStatus"),
+                    )
                 )
                 logger.debug(f"Adding remotely accessible port {component.get('id')}")
             else:
diff --git a/metadata-ingestion/src/datahub/ingestion/source/openapi.py b/metadata-ingestion/src/datahub/ingestion/source/openapi.py
index 8289265483d598..2075e999ea1d0e 100755
--- a/metadata-ingestion/src/datahub/ingestion/source/openapi.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/openapi.py
@@ -101,16 +101,16 @@ def get_swagger(self) -> Dict:
                 # details there once, and then use that session for all requests.
                 self.token = f"Bearer {self.bearer_token}"
             else:
-                assert (
-                    "url_complement" in self.get_token.keys()
-                ), "When 'request_type' is set to 'get', an url_complement is needed for the request."
+                assert "url_complement" in self.get_token.keys(), (
+                    "When 'request_type' is set to 'get', an url_complement is needed for the request."
+                )
                 if self.get_token["request_type"] == "get":
-                    assert (
-                        "{username}" in self.get_token["url_complement"]
-                    ), "we expect the keyword {username} to be present in the url"
-                    assert (
-                        "{password}" in self.get_token["url_complement"]
-                    ), "we expect the keyword {password} to be present in the url"
+                    assert "{username}" in self.get_token["url_complement"], (
+                        "we expect the keyword {username} to be present in the url"
+                    )
+                    assert "{password}" in self.get_token["url_complement"], (
+                        "we expect the keyword {password} to be present in the url"
+                    )
                     url4req = self.get_token["url_complement"].replace(
                         "{username}", self.username
                     )
diff --git a/metadata-ingestion/src/datahub/ingestion/source/powerbi/config.py b/metadata-ingestion/src/datahub/ingestion/source/powerbi/config.py
index b49d40a0c7eb6a..14beab6bc9391e 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/powerbi/config.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/powerbi/config.py
@@ -225,9 +225,9 @@ def report_charts_dropped(self, view: str) -> None:
 def default_for_dataset_type_mapping() -> Dict[str, str]:
     dict_: dict = {}
     for item in SupportedDataPlatform:
-        dict_[
-            item.value.powerbi_data_platform_name
-        ] = item.value.datahub_data_platform_name
+        dict_[item.value.powerbi_data_platform_name] = (
+            item.value.datahub_data_platform_name
+        )
 
     return dict_
 
@@ -303,15 +303,15 @@ class PowerBiDashboardSourceConfig(
     # Dataset type mapping PowerBI support many type of data-sources. Here user needs to define what type of PowerBI
     # DataSource needs to be mapped to corresponding DataHub Platform DataSource. For example, PowerBI `Snowflake` is
     # mapped to DataHub `snowflake` PowerBI `PostgreSQL` is mapped to DataHub `postgres` and so on.
-    dataset_type_mapping: Union[
-        Dict[str, str], Dict[str, PlatformDetail]
-    ] = pydantic.Field(
-        default_factory=default_for_dataset_type_mapping,
-        description="[deprecated] Use server_to_platform_instance instead. Mapping of PowerBI datasource type to "
-        "DataHub supported datasources."
-        "You can configured platform instance for dataset lineage. "
-        "See Quickstart Recipe for mapping",
-        hidden_from_docs=True,
+    dataset_type_mapping: Union[Dict[str, str], Dict[str, PlatformDetail]] = (
+        pydantic.Field(
+            default_factory=default_for_dataset_type_mapping,
+            description="[deprecated] Use server_to_platform_instance instead. Mapping of PowerBI datasource type to "
+            "DataHub supported datasources."
+            "You can configured platform instance for dataset lineage. "
+            "See Quickstart Recipe for mapping",
+            hidden_from_docs=True,
+        )
     )
     # PowerBI datasource's server to platform instance mapping
     server_to_platform_instance: Dict[
diff --git a/metadata-ingestion/src/datahub/ingestion/source/powerbi/m_query/parser.py b/metadata-ingestion/src/datahub/ingestion/source/powerbi/m_query/parser.py
index 2a5de7494920b2..759fc6d7dadfba 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/powerbi/m_query/parser.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/powerbi/m_query/parser.py
@@ -128,17 +128,17 @@ def get_upstream_tables(
     reporter.m_query_parse_successes += 1
 
     try:
-        lineage: List[
-            datahub.ingestion.source.powerbi.m_query.data_classes.Lineage
-        ] = resolver.MQueryResolver(
-            table=table,
-            parse_tree=parse_tree,
-            reporter=reporter,
-            parameters=parameters,
-        ).resolve_to_lineage(
-            ctx=ctx,
-            config=config,
-            platform_instance_resolver=platform_instance_resolver,
+        lineage: List[datahub.ingestion.source.powerbi.m_query.data_classes.Lineage] = (
+            resolver.MQueryResolver(
+                table=table,
+                parse_tree=parse_tree,
+                reporter=reporter,
+                parameters=parameters,
+            ).resolve_to_lineage(
+                ctx=ctx,
+                config=config,
+                platform_instance_resolver=platform_instance_resolver,
+            )
         )
 
         if lineage:
diff --git a/metadata-ingestion/src/datahub/ingestion/source/powerbi/m_query/pattern_handler.py b/metadata-ingestion/src/datahub/ingestion/source/powerbi/m_query/pattern_handler.py
index 63520bd731de86..54b810650f5854 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/powerbi/m_query/pattern_handler.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/powerbi/m_query/pattern_handler.py
@@ -170,8 +170,7 @@ def create_reference_table(
         logger.debug(f"Processing arguments {arguments}")
 
         if (
-            len(arguments)
-            >= 4  # [0] is warehouse FQDN.
+            len(arguments) >= 4  # [0] is warehouse FQDN.
             # [1] is endpoint, we are not using it.
             # [2] is "Catalog" key
             # [3] is catalog's value
@@ -215,16 +214,16 @@ def parse_custom_sql(
             native_sql_parser.remove_special_characters(query)
         )
 
-        parsed_result: Optional[
-            "SqlParsingResult"
-        ] = native_sql_parser.parse_custom_sql(
-            ctx=self.ctx,
-            query=query,
-            platform=self.get_platform_pair().datahub_data_platform_name,
-            platform_instance=platform_detail.platform_instance,
-            env=platform_detail.env,
-            database=database,
-            schema=schema,
+        parsed_result: Optional["SqlParsingResult"] = (
+            native_sql_parser.parse_custom_sql(
+                ctx=self.ctx,
+                query=query,
+                platform=self.get_platform_pair().datahub_data_platform_name,
+                platform_instance=platform_detail.platform_instance,
+                env=platform_detail.env,
+                database=database,
+                schema=schema,
+            )
         )
 
         if parsed_result is None:
@@ -410,9 +409,9 @@ def create_lineage(
             f"Processing Databrick data-access function detail {data_access_func_detail}"
         )
         table_detail: Dict[str, str] = {}
-        temp_accessor: Optional[
-            IdentifierAccessor
-        ] = data_access_func_detail.identifier_accessor
+        temp_accessor: Optional[IdentifierAccessor] = (
+            data_access_func_detail.identifier_accessor
+        )
 
         while temp_accessor:
             # Condition to handle databricks M-query pattern where table, schema and database all are present in
@@ -647,11 +646,13 @@ def create_lineage(
         db_name: str = data_access_func_detail.identifier_accessor.items["Name"]  # type: ignore
         # Second is schema name
         schema_name: str = cast(
-            IdentifierAccessor, data_access_func_detail.identifier_accessor.next  # type: ignore
+            IdentifierAccessor,
+            data_access_func_detail.identifier_accessor.next,  # type: ignore
         ).items["Name"]
         # Third is table name
         table_name: str = cast(
-            IdentifierAccessor, data_access_func_detail.identifier_accessor.next.next  # type: ignore
+            IdentifierAccessor,
+            data_access_func_detail.identifier_accessor.next.next,  # type: ignore
         ).items["Name"]
 
         qualified_table_name: str = f"{db_name}.{schema_name}.{table_name}"
@@ -768,10 +769,13 @@ def get_db_name(self, data_access_tokens: List[str]) -> Optional[str]:
         ):  # database name is explicitly set
             return database
 
-        return get_next_item(  # database name is set in Name argument
-            data_access_tokens, "Name"
-        ) or get_next_item(  # If both above arguments are not available, then try Catalog
-            data_access_tokens, "Catalog"
+        return (
+            get_next_item(  # database name is set in Name argument
+                data_access_tokens, "Name"
+            )
+            or get_next_item(  # If both above arguments are not available, then try Catalog
+                data_access_tokens, "Catalog"
+            )
         )
 
     def create_lineage(
@@ -819,9 +823,7 @@ def create_lineage(
             values=tree_function.remove_whitespaces_from_list(
                 tree_function.token_values(flat_argument_list[1])
             ),
-        )[
-            0
-        ]  # Remove any whitespaces and double quotes character
+        )[0]  # Remove any whitespaces and double quotes character
 
         server = tree_function.strip_char_from_list([data_access_tokens[2]])[0]
 
diff --git a/metadata-ingestion/src/datahub/ingestion/source/powerbi/m_query/resolver.py b/metadata-ingestion/src/datahub/ingestion/source/powerbi/m_query/resolver.py
index 2756a113d1ef0c..42963c08d992d1 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/powerbi/m_query/resolver.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/powerbi/m_query/resolver.py
@@ -188,9 +188,9 @@ def _process_invoke_expression(
         # - The inner function Table.TransformColumnTypes takes #"Removed Columns1"
         #   (a table reference) as its first argument
         # - Its result is then passed as the first argument to Table.SplitColumn
-        second_invoke_expression: Optional[
-            Tree
-        ] = tree_function.first_invoke_expression_func(first_argument)
+        second_invoke_expression: Optional[Tree] = (
+            tree_function.first_invoke_expression_func(first_argument)
+        )
         if second_invoke_expression:
             # 1. The First argument is function call
             # 2. That function's first argument references next table variable
@@ -304,14 +304,14 @@ def internal(
                 logger.debug(v_statement.pretty())
                 return None
 
-            invoke_expression: Optional[
-                Tree
-            ] = tree_function.first_invoke_expression_func(rh_tree)
+            invoke_expression: Optional[Tree] = (
+                tree_function.first_invoke_expression_func(rh_tree)
+            )
 
             if invoke_expression is not None:
-                result: Union[
-                    DataAccessFunctionDetail, List[str], None
-                ] = self._process_invoke_expression(invoke_expression)
+                result: Union[DataAccessFunctionDetail, List[str], None] = (
+                    self._process_invoke_expression(invoke_expression)
+                )
                 if result is None:
                     return None  # No need to process some un-expected grammar found while processing invoke_expression
                 if isinstance(result, DataAccessFunctionDetail):
@@ -368,9 +368,9 @@ def resolve_to_lineage(
             return lineage
 
         # Parse M-Query and use output_variable as root of tree and create instance of DataAccessFunctionDetail
-        table_links: List[
-            DataAccessFunctionDetail
-        ] = self.create_data_access_functional_detail(output_variable)
+        table_links: List[DataAccessFunctionDetail] = (
+            self.create_data_access_functional_detail(output_variable)
+        )
 
         # Each item is data-access function
         for f_detail in table_links:
@@ -390,7 +390,7 @@ def resolve_to_lineage(
 
             # From supported_resolver enum get respective handler like AmazonRedshift or Snowflake or Oracle or NativeQuery and create instance of it
             # & also pass additional information that will be need to generate lineage
-            pattern_handler: (AbstractLineage) = supported_resolver.handler()(
+            pattern_handler: AbstractLineage = supported_resolver.handler()(
                 ctx=ctx,
                 table=self.table,
                 config=config,
diff --git a/metadata-ingestion/src/datahub/ingestion/source/powerbi/powerbi.py b/metadata-ingestion/src/datahub/ingestion/source/powerbi/powerbi.py
index 044946a5d308d1..89184d37569ffc 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/powerbi/powerbi.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/powerbi/powerbi.py
@@ -945,9 +945,9 @@ def to_datahub_work_units(
         # Convert tiles to charts
         ds_mcps, chart_mcps = self.to_datahub_chart(dashboard.tiles, workspace)
         # Lets convert dashboard to datahub dashboard
-        dashboard_mcps: List[
-            MetadataChangeProposalWrapper
-        ] = self.to_datahub_dashboard_mcp(dashboard, workspace, chart_mcps, user_mcps)
+        dashboard_mcps: List[MetadataChangeProposalWrapper] = (
+            self.to_datahub_dashboard_mcp(dashboard, workspace, chart_mcps, user_mcps)
+        )
 
         # Now add MCPs in sequence
         mcps.extend(ds_mcps)
@@ -1322,14 +1322,14 @@ def extract_independent_datasets(
                     context=",".join(
                         [
                             dataset.name
-                            for dataset in workspace.independent_datasets
+                            for dataset in workspace.independent_datasets.values()
                             if dataset.name
                         ]
                     ),
                 )
             return
 
-        for dataset in workspace.independent_datasets:
+        for dataset in workspace.independent_datasets.values():
             yield from auto_workunit(
                 stream=self.mapper.to_datahub_dataset(
                     dataset=dataset,
@@ -1440,7 +1440,7 @@ def get_workspace_workunit(
 
         yield from auto_workunit(self.emit_app(workspace=workspace))
 
-        for dashboard in workspace.dashboards:
+        for dashboard in workspace.dashboards.values():
             try:
                 # Fetch PowerBi users for dashboards
                 dashboard.users = self.powerbi_client.get_dashboard_users(dashboard)
@@ -1459,7 +1459,7 @@ def get_workspace_workunit(
                 if wu is not None:
                     yield wu
 
-        for report in workspace.reports:
+        for report in workspace.reports.values():
             for work_unit in self.mapper.report_to_datahub_work_units(
                 report, workspace
             ):
@@ -1472,9 +1472,9 @@ def get_workspace_workunit(
     def _get_dashboard_patch_work_unit(
         self, work_unit: MetadataWorkUnit
     ) -> Optional[MetadataWorkUnit]:
-        dashboard_info_aspect: Optional[
-            DashboardInfoClass
-        ] = work_unit.get_aspect_of_type(DashboardInfoClass)
+        dashboard_info_aspect: Optional[DashboardInfoClass] = (
+            work_unit.get_aspect_of_type(DashboardInfoClass)
+        )
 
         if dashboard_info_aspect and self.source_config.patch_metadata:
             return convert_dashboard_info_to_patch(
diff --git a/metadata-ingestion/src/datahub/ingestion/source/powerbi/rest_api_wrapper/data_classes.py b/metadata-ingestion/src/datahub/ingestion/source/powerbi/rest_api_wrapper/data_classes.py
index 70e4ee68f53515..f3e968896b3c96 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/powerbi/rest_api_wrapper/data_classes.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/powerbi/rest_api_wrapper/data_classes.py
@@ -71,13 +71,13 @@ class Workspace:
     id: str
     name: str
     type: str  # This is used as a subtype of the Container entity.
-    dashboards: List["Dashboard"]
-    reports: List["Report"]
-    datasets: Dict[str, "PowerBIDataset"]
-    report_endorsements: Dict[str, List[str]]
-    dashboard_endorsements: Dict[str, List[str]]
+    dashboards: Dict[str, "Dashboard"]  # key = dashboard id
+    reports: Dict[str, "Report"]  # key = report id
+    datasets: Dict[str, "PowerBIDataset"]  # key = dataset id
+    report_endorsements: Dict[str, List[str]]  # key = report id
+    dashboard_endorsements: Dict[str, List[str]]  # key = dashboard id
     scan_result: dict
-    independent_datasets: List["PowerBIDataset"]
+    independent_datasets: Dict[str, "PowerBIDataset"]  # key = dataset id
     app: Optional["App"]
 
     def get_urn_part(self, workspace_id_as_urn_part: Optional[bool] = False) -> str:
diff --git a/metadata-ingestion/src/datahub/ingestion/source/powerbi/rest_api_wrapper/data_resolver.py b/metadata-ingestion/src/datahub/ingestion/source/powerbi/rest_api_wrapper/data_resolver.py
index 161975fa635fdb..927840c44bf0b0 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/powerbi/rest_api_wrapper/data_resolver.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/powerbi/rest_api_wrapper/data_resolver.py
@@ -425,9 +425,9 @@ def itr_pages(
 
             response.raise_for_status()
 
-            assert (
-                Constant.VALUE in response.json()
-            ), "'value' key is not present in paginated response"
+            assert Constant.VALUE in response.json(), (
+                "'value' key is not present in paginated response"
+            )
 
             if not response.json()[Constant.VALUE]:  # if it is an empty list then break
                 break
@@ -447,13 +447,13 @@ def get_app(
         if raw_app is None:
             return None
 
-        assert (
-            Constant.ID in raw_app
-        ), f"{Constant.ID} is required field not present in server response"
+        assert Constant.ID in raw_app, (
+            f"{Constant.ID} is required field not present in server response"
+        )
 
-        assert (
-            Constant.NAME in raw_app
-        ), f"{Constant.NAME} is required field not present in server response"
+        assert Constant.NAME in raw_app, (
+            f"{Constant.NAME} is required field not present in server response"
+        )
 
         return App(
             id=raw_app[Constant.ID],
diff --git a/metadata-ingestion/src/datahub/ingestion/source/powerbi/rest_api_wrapper/powerbi_api.py b/metadata-ingestion/src/datahub/ingestion/source/powerbi/rest_api_wrapper/powerbi_api.py
index 5ae333430a78bc..93d734652f0ff1 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/powerbi/rest_api_wrapper/powerbi_api.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/powerbi/rest_api_wrapper/powerbi_api.py
@@ -96,7 +96,7 @@ def log_http_error(self, message: str) -> Any:
             url: str = e.request.url if e.request else "URL not available"
             self.reporter.warning(
                 title="Metadata API Timeout",
-                message=f"Metadata endpoints are not reachable. Check network connectivity to PowerBI Service.",
+                message="Metadata endpoints are not reachable. Check network connectivity to PowerBI Service.",
                 context=f"url={url}",
             )
 
@@ -173,7 +173,7 @@ def _get_entity_users(
                 entity=entity_name,
                 entity_id=entity_id,
             )
-        except:  # It will catch all type of exception
+        except Exception:
             e = self.log_http_error(
                 message=f"Unable to fetch users for {entity_name}({entity_id})."
             )
@@ -193,15 +193,18 @@ def get_dashboard_users(self, dashboard: Dashboard) -> List[User]:
     def get_report_users(self, workspace_id: str, report_id: str) -> List[User]:
         return self._get_entity_users(workspace_id, Constant.REPORTS, report_id)
 
-    def get_reports(self, workspace: Workspace) -> List[Report]:
+    def get_reports(self, workspace: Workspace) -> Dict[str, Report]:
         """
         Fetch the report from PowerBi for the given Workspace
         """
-        reports: List[Report] = []
+        reports: Dict[str, Report] = {}
         try:
-            reports = self._get_resolver().get_reports(workspace)
+            reports = {
+                report.id: report
+                for report in self._get_resolver().get_reports(workspace)
+            }
             # Fill Report dataset
-            for report in reports:
+            for report in reports.values():
                 if report.dataset_id:
                     report.dataset = self.dataset_registry.get(report.dataset_id)
                     if report.dataset is None:
@@ -210,7 +213,7 @@ def get_reports(self, workspace: Workspace) -> List[Report]:
                             message="A cross-workspace reference that failed to be resolved. Please ensure that no global workspace is being filtered out due to the workspace_id_pattern.",
                             context=f"report-name: {report.name} and dataset-id: {report.dataset_id}",
                         )
-        except:
+        except Exception:
             self.log_http_error(
                 message=f"Unable to fetch reports for workspace {workspace.name}"
             )
@@ -222,7 +225,7 @@ def fill_ownership() -> None:
                 )
                 return
 
-            for report in reports:
+            for report in reports.values():
                 report.users = self.get_report_users(
                     workspace_id=workspace.id, report_id=report.id
                 )
@@ -234,7 +237,7 @@ def fill_tags() -> None:
                 )
                 return
 
-            for report in reports:
+            for report in reports.values():
                 report.tags = workspace.report_endorsements.get(report.id, [])
 
         fill_ownership()
@@ -260,7 +263,7 @@ def get_workspaces(self) -> List[Workspace]:
 
             groups = self._get_resolver().get_groups(filter_=filter_)
 
-        except:
+        except Exception:
             self.log_http_error(message="Unable to fetch list of workspaces")
             # raise  # we want this exception to bubble up
 
@@ -270,12 +273,12 @@ def get_workspaces(self) -> List[Workspace]:
                 name=workspace[Constant.NAME],
                 type=workspace[Constant.TYPE],
                 datasets={},
-                dashboards=[],
-                reports=[],
+                dashboards={},
+                reports={},
                 report_endorsements={},
                 dashboard_endorsements={},
                 scan_result={},
-                independent_datasets=[],
+                independent_datasets={},
                 app=None,  # It will be populated in _fill_metadata_from_scan_result method
             )
             for workspace in groups
@@ -292,7 +295,7 @@ def get_modified_workspaces(self) -> List[str]:
             modified_workspace_ids = self.__admin_api_resolver.get_modified_workspaces(
                 self.__config.modified_since
             )
-        except:
+        except Exception:
             self.log_http_error(message="Unable to fetch list of modified workspaces.")
 
         return modified_workspace_ids
@@ -303,8 +306,8 @@ def _get_scan_result(self, workspace_ids: List[str]) -> Any:
             scan_id = self.__admin_api_resolver.create_scan_job(
                 workspace_ids=workspace_ids
             )
-        except:
-            e = self.log_http_error(message=f"Unable to fetch get scan result.")
+        except Exception:
+            e = self.log_http_error(message="Unable to fetch get scan result.")
             if data_resolver.is_permission_error(cast(Exception, e)):
                 logger.warning(
                     "Dataset lineage can not be ingestion because this user does not have access to the PowerBI Admin "
@@ -561,12 +564,12 @@ def _fill_metadata_from_scan_result(
                 name=workspace_metadata[Constant.NAME],
                 type=workspace_metadata[Constant.TYPE],
                 datasets={},
-                dashboards=[],
-                reports=[],
+                dashboards={},
+                reports={},
                 report_endorsements={},
                 dashboard_endorsements={},
                 scan_result={},
-                independent_datasets=[],
+                independent_datasets={},
                 app=None,  # It is getting set from scan-result
             )
             cur_workspace.scan_result = workspace_metadata
@@ -597,25 +600,28 @@ def _fill_metadata_from_scan_result(
     def _fill_independent_datasets(self, workspace: Workspace) -> None:
         reachable_datasets: List[str] = []
         # Find out reachable datasets
-        for dashboard in workspace.dashboards:
+        for dashboard in workspace.dashboards.values():
             for tile in dashboard.tiles:
                 if tile.dataset is not None:
                     reachable_datasets.append(tile.dataset.id)
 
-        for report in workspace.reports:
+        for report in workspace.reports.values():
             if report.dataset is not None:
                 reachable_datasets.append(report.dataset.id)
 
         # Set datasets not present in reachable_datasets
         for dataset in workspace.datasets.values():
             if dataset.id not in reachable_datasets:
-                workspace.independent_datasets.append(dataset)
+                workspace.independent_datasets[dataset.id] = dataset
 
     def _fill_regular_metadata_detail(self, workspace: Workspace) -> None:
         def fill_dashboards() -> None:
-            workspace.dashboards = self._get_resolver().get_dashboards(workspace)
+            workspace.dashboards = {
+                dashboard.id: dashboard
+                for dashboard in self._get_resolver().get_dashboards(workspace)
+            }
             # set tiles of Dashboard
-            for dashboard in workspace.dashboards:
+            for dashboard in workspace.dashboards.values():
                 dashboard.tiles = self._get_resolver().get_tiles(
                     workspace, dashboard=dashboard
                 )
@@ -644,7 +650,7 @@ def fill_dashboard_tags() -> None:
                     "Skipping tag retrieval for dashboard as extract_endorsements_to_tags is set to false"
                 )
                 return
-            for dashboard in workspace.dashboards:
+            for dashboard in workspace.dashboards.values():
                 dashboard.tags = workspace.dashboard_endorsements.get(dashboard.id, [])
 
         if self.__config.extract_dashboards:
diff --git a/metadata-ingestion/src/datahub/ingestion/source/qlik_sense/qlik_api.py b/metadata-ingestion/src/datahub/ingestion/source/qlik_sense/qlik_api.py
index 31b0731aaa751c..10b062c98c147f 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/qlik_sense/qlik_api.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/qlik_sense/qlik_api.py
@@ -156,7 +156,7 @@ def _get_sheet(
                 )
                 if chart:
                     if not chart.title:
-                        chart.title = f"Object {i+1} of Sheet '{sheet.title}'"
+                        chart.title = f"Object {i + 1} of Sheet '{sheet.title}'"
                     sheet.charts.append(chart)
                 websocket_connection.handle.pop()
             return sheet
diff --git a/metadata-ingestion/src/datahub/ingestion/source/redshift/config.py b/metadata-ingestion/src/datahub/ingestion/source/redshift/config.py
index cad48eaf1c2375..932ada0a908b28 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/redshift/config.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/redshift/config.py
@@ -178,9 +178,9 @@ class RedshiftConfig(
     @root_validator(pre=True)
     def check_email_is_set_on_usage(cls, values):
         if values.get("include_usage_statistics"):
-            assert (
-                "email_domain" in values and values["email_domain"]
-            ), "email_domain needs to be set if usage is enabled"
+            assert "email_domain" in values and values["email_domain"], (
+                "email_domain needs to be set if usage is enabled"
+            )
         return values
 
     @root_validator(skip_on_failure=True)
diff --git a/metadata-ingestion/src/datahub/ingestion/source/redshift/query.py b/metadata-ingestion/src/datahub/ingestion/source/redshift/query.py
index 71a20890d35e88..62f7d0a3901c7a 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/redshift/query.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/redshift/query.py
@@ -797,61 +797,91 @@ def stl_scan_based_lineage_query(
         db_name: str, start_time: datetime, end_time: datetime
     ) -> str:
         return """
-            SELECT
-                distinct cluster,
-                target_schema,
-                target_table,
-                username,
-                source_schema,
-                source_table,
-                query_text AS ddl,
-                start_time AS timestamp
-            FROM
-            (
-                SELECT
-                    sti.schema AS target_schema,
-                    sti.table AS target_table,
-                    sti.database AS cluster,
-                    qi.table_id AS target_table_id,
-                    qi.query_id AS query_id,
-                    qi.start_time AS start_time
-                FROM
-                    SYS_QUERY_DETAIL qi
-                    JOIN
-                    SVV_TABLE_INFO sti on sti.table_id = qi.table_id
-                WHERE
-                    start_time >= '{start_time}' and
-                    start_time < '{end_time}' and
-                    cluster = '{db_name}' and
-                    step_name = 'insert'
-            ) AS target_tables
-            JOIN
-            (
+            WITH queries AS (
                 SELECT
-                    sti.schema AS source_schema,
-                    sti.table AS source_table,
-                    qs.table_id AS source_table_id,
-                    qs.query_id AS query_id,
-                    sui.user_name AS username,
-                    LISTAGG(qt."text") WITHIN GROUP (ORDER BY sequence) AS query_text
+                    sti.database as cluster,
+                    sti.schema AS "schema",
+                    sti.table AS "table",
+                    qs.table_id AS table_id,
+                    qs.query_id as query_id,
+                    qs.step_name as step_name,
+                    sui.user_name as username,
+                    source,
+                    MIN(qs.start_time) as "timestamp" -- multiple duplicate records with start_time increasing slightly by miliseconds
                 FROM
                     SYS_QUERY_DETAIL qs
                     JOIN
                     SVV_TABLE_INFO sti ON sti.table_id = qs.table_id
                     LEFT JOIN
-                    SYS_QUERY_TEXT qt ON qt.query_id = qs.query_id
-                    LEFT JOIN
                     SVV_USER_INFO sui ON qs.user_id = sui.user_id
                 WHERE
-                    qs.step_name = 'scan' AND
-                    qs.source = 'Redshift(local)' AND
-                    qt.sequence < 16 AND -- See https://stackoverflow.com/questions/72770890/redshift-result-size-exceeds-listagg-limit-on-svl-statementtext
-                    sti.database = '{db_name}' AND -- this was required to not retrieve some internal redshift tables, try removing to see what happens
-                    sui.user_name <> 'rdsdb' -- not entirely sure about this filter
-                GROUP BY sti.schema, sti.table, qs.table_id, qs.query_id, sui.user_name
-            ) AS source_tables ON target_tables.query_id = source_tables.query_id
-            WHERE source_tables.source_table_id <> target_tables.target_table_id
-            ORDER BY cluster, target_schema, target_table, start_time ASC
+                    cluster = '{db_name}' AND
+                    qs.user_id <> 1 AND -- this is user 'rdsdb'
+                    qs.start_time >= '{start_time}' AND
+                    qs.start_time < '{end_time}'
+                GROUP BY cluster, "schema", "table", qs.table_id, query_id, step_name, username, source -- to be sure we are not making duplicates ourselves the list of group by must match whatever we use in "group by" and "where" of subsequent queries ("cluster" is already set to single value in this query)
+            ),
+            unique_query_text AS (
+                SELECT
+                    query_id,
+                    sequence,
+                    text
+                FROM (
+                    SELECT
+                        query_id,
+                        "sequence",
+                        text,
+                        ROW_NUMBER() OVER (
+                        PARTITION BY query_id, sequence
+                        ) as rn
+                    FROM SYS_QUERY_TEXT
+                    )
+                WHERE rn = 1
+            ),
+            scan_queries AS (
+                SELECT
+                    "schema" as source_schema,
+                    "table" as source_table,
+                    table_id as source_table_id,
+                    queries.query_id as query_id,
+                    username,
+                    LISTAGG(qt."text") WITHIN GROUP (ORDER BY sequence) AS query_text
+                FROM
+                    "queries" LEFT JOIN
+                    unique_query_text qt ON qt.query_id = queries.query_id
+                WHERE
+                    source = 'Redshift(local)' AND
+                    step_name = 'scan' AND
+                    qt.sequence < 16 -- truncating query to not exceed Redshift limit on LISTAGG function (each sequence has at most 4k characters, limit is 64k, divided by 4k gives 16, starts count from 0)
+                GROUP BY source_schema, source_table, source_table_id, queries.query_id, username
+            ),
+            insert_queries AS (
+                SELECT
+                    "schema" as target_schema,
+                    "table" as target_table,
+                    table_id as target_table_id,
+                    query_id,
+                    cluster,
+                    min("timestamp") as "timestamp"
+                FROM
+                    queries
+                WHERE
+                    step_name = 'insert'
+                GROUP BY cluster, target_schema, target_table, target_table_id, query_id
+            )
+            SELECT
+                cluster,
+                target_schema,
+                target_table,
+                username,
+                source_schema,
+                source_table,
+                query_text AS ddl,
+                "timestamp"
+            FROM scan_queries
+                JOIN insert_queries on insert_queries.query_id = scan_queries.query_id
+            WHERE source_table_id <> target_table_id
+            ORDER BY cluster, target_schema, target_table, "timestamp" ASC;
                     """.format(
             # We need the original database name for filtering
             db_name=db_name,
diff --git a/metadata-ingestion/src/datahub/ingestion/source/redshift/redshift.py b/metadata-ingestion/src/datahub/ingestion/source/redshift/redshift.py
index 9bfca941ce48fb..cce282c71056a2 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/redshift/redshift.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/redshift/redshift.py
@@ -305,13 +305,13 @@ def test_connection(config_dict: dict) -> TestConnectionReport:
             test_report.capability_report = {}
             try:
                 RedshiftDataDictionary.get_schemas(connection, database=config.database)
-                test_report.capability_report[
-                    SourceCapability.SCHEMA_METADATA
-                ] = CapabilityReport(capable=True)
+                test_report.capability_report[SourceCapability.SCHEMA_METADATA] = (
+                    CapabilityReport(capable=True)
+                )
             except Exception as e:
-                test_report.capability_report[
-                    SourceCapability.SCHEMA_METADATA
-                ] = CapabilityReport(capable=False, failure_reason=str(e))
+                test_report.capability_report[SourceCapability.SCHEMA_METADATA] = (
+                    CapabilityReport(capable=False, failure_reason=str(e))
+                )
 
         except Exception as e:
             test_report.basic_connectivity = CapabilityReport(
@@ -947,9 +947,9 @@ def cache_tables_and_views(self, connection, database):
     def get_all_tables(
         self,
     ) -> Dict[str, Dict[str, List[Union[RedshiftView, RedshiftTable]]]]:
-        all_tables: Dict[
-            str, Dict[str, List[Union[RedshiftView, RedshiftTable]]]
-        ] = defaultdict(dict)
+        all_tables: Dict[str, Dict[str, List[Union[RedshiftView, RedshiftTable]]]] = (
+            defaultdict(dict)
+        )
         for db in set().union(self.db_tables, self.db_views):
             tables = self.db_tables.get(db, {})
             views = self.db_views.get(db, {})
@@ -967,9 +967,9 @@ def extract_usage(
         all_tables: Dict[str, Dict[str, List[Union[RedshiftView, RedshiftTable]]]],
     ) -> Iterable[MetadataWorkUnit]:
         with PerfTimer() as timer:
-            redundant_usage_run_skip_handler: Optional[
-                RedundantUsageRunSkipHandler
-            ] = None
+            redundant_usage_run_skip_handler: Optional[RedundantUsageRunSkipHandler] = (
+                None
+            )
             if self.config.enable_stateful_usage_ingestion:
                 redundant_usage_run_skip_handler = RedundantUsageRunSkipHandler(
                     source=self,
diff --git a/metadata-ingestion/src/datahub/ingestion/source/redshift/usage.py b/metadata-ingestion/src/datahub/ingestion/source/redshift/usage.py
index d66a1ee18be40f..a5758bdd825702 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/redshift/usage.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/redshift/usage.py
@@ -199,10 +199,10 @@ def _get_workunits_internal(
                 end_time=self.end_time.strftime(REDSHIFT_DATETIME_FORMAT),
                 database=self.config.database,
             )
-            access_events_iterable: Iterable[
-                RedshiftAccessEvent
-            ] = self._gen_access_events_from_history_query(
-                query, connection=self.connection, all_tables=all_tables
+            access_events_iterable: Iterable[RedshiftAccessEvent] = (
+                self._gen_access_events_from_history_query(
+                    query, connection=self.connection, all_tables=all_tables
+                )
             )
 
             aggregated_events: AggregatedAccessEvents = self._aggregate_access_events(
@@ -225,10 +225,10 @@ def _gen_operation_aspect_workunits(
             start_time=self.start_time.strftime(REDSHIFT_DATETIME_FORMAT),
             end_time=self.end_time.strftime(REDSHIFT_DATETIME_FORMAT),
         )
-        access_events_iterable: Iterable[
-            RedshiftAccessEvent
-        ] = self._gen_access_events_from_history_query(
-            query, connection, all_tables=all_tables
+        access_events_iterable: Iterable[RedshiftAccessEvent] = (
+            self._gen_access_events_from_history_query(
+                query, connection, all_tables=all_tables
+            )
         )
 
         # Generate operation aspect work units from the access events
diff --git a/metadata-ingestion/src/datahub/ingestion/source/s3/datalake_profiler_config.py b/metadata-ingestion/src/datahub/ingestion/source/s3/datalake_profiler_config.py
index 89c092875e4490..58e930eb6e809c 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/s3/datalake_profiler_config.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/s3/datalake_profiler_config.py
@@ -85,8 +85,8 @@ def ensure_field_level_settings_are_normalized(
                 if field_level_metric.startswith("include_field_"):
                     values.setdefault(field_level_metric, False)
 
-            assert (
-                max_num_fields_to_profile is None
-            ), f"{max_num_fields_to_profile_key} should be set to None"
+            assert max_num_fields_to_profile is None, (
+                f"{max_num_fields_to_profile_key} should be set to None"
+            )
 
         return values
diff --git a/metadata-ingestion/src/datahub/ingestion/source/s3/source.py b/metadata-ingestion/src/datahub/ingestion/source/s3/source.py
index 989d0d734352a2..b5bc08912ee946 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/s3/source.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/s3/source.py
@@ -40,7 +40,6 @@
     get_bucket_name,
     get_bucket_relative_path,
     get_key_prefix,
-    group_s3_objects_by_dirname,
     strip_s3_prefix,
 )
 from datahub.ingestion.source.data_lake_common.data_lake_utils import ContainerWUCreator
@@ -73,6 +72,7 @@
     _Aspect,
 )
 from datahub.telemetry import stats, telemetry
+from datahub.utilities.groupby import groupby_unsorted
 from datahub.utilities.perf_timer import PerfTimer
 
 if TYPE_CHECKING:
@@ -868,7 +868,11 @@ def get_folder_info(
         """
         partitions: List[Folder] = []
         s3_objects = bucket.objects.filter(Prefix=prefix).page_size(PAGE_SIZE)
-        for key, group in group_s3_objects_by_dirname(s3_objects).items():
+        grouped_s3_objects_by_dirname = groupby_unsorted(
+            s3_objects,
+            key=lambda obj: obj.key.rsplit("/", 1)[0],
+        )
+        for key, group in grouped_s3_objects_by_dirname:
             file_size = 0
             creation_time = None
             modification_time = None
@@ -1124,7 +1128,7 @@ def get_workunits_internal(self) -> Iterable[MetadataWorkUnit]:
                                 table_data.table_path
                             ].timestamp = table_data.timestamp
 
-                for guid, table_data in table_dict.items():
+                for _, table_data in table_dict.items():
                     yield from self.ingest_table(table_data, path_spec)
 
             if not self.source_config.is_profiling_enabled():
diff --git a/metadata-ingestion/src/datahub/ingestion/source/salesforce.py b/metadata-ingestion/src/datahub/ingestion/source/salesforce.py
index 88679efdf5fc31..66e0e6b741d1ff 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/salesforce.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/salesforce.py
@@ -236,12 +236,12 @@ def __init__(self, config: SalesforceConfig, ctx: PipelineContext) -> None:
         try:
             if self.config.auth is SalesforceAuthType.DIRECT_ACCESS_TOKEN:
                 logger.debug("Access Token Provided in Config")
-                assert (
-                    self.config.access_token is not None
-                ), "Config access_token is required for DIRECT_ACCESS_TOKEN auth"
-                assert (
-                    self.config.instance_url is not None
-                ), "Config instance_url is required for DIRECT_ACCESS_TOKEN auth"
+                assert self.config.access_token is not None, (
+                    "Config access_token is required for DIRECT_ACCESS_TOKEN auth"
+                )
+                assert self.config.instance_url is not None, (
+                    "Config instance_url is required for DIRECT_ACCESS_TOKEN auth"
+                )
 
                 self.sf = Salesforce(
                     instance_url=self.config.instance_url,
@@ -250,15 +250,15 @@ def __init__(self, config: SalesforceConfig, ctx: PipelineContext) -> None:
                 )
             elif self.config.auth is SalesforceAuthType.USERNAME_PASSWORD:
                 logger.debug("Username/Password Provided in Config")
-                assert (
-                    self.config.username is not None
-                ), "Config username is required for USERNAME_PASSWORD auth"
-                assert (
-                    self.config.password is not None
-                ), "Config password is required for USERNAME_PASSWORD auth"
-                assert (
-                    self.config.security_token is not None
-                ), "Config security_token is required for USERNAME_PASSWORD auth"
+                assert self.config.username is not None, (
+                    "Config username is required for USERNAME_PASSWORD auth"
+                )
+                assert self.config.password is not None, (
+                    "Config password is required for USERNAME_PASSWORD auth"
+                )
+                assert self.config.security_token is not None, (
+                    "Config security_token is required for USERNAME_PASSWORD auth"
+                )
 
                 self.sf = Salesforce(
                     username=self.config.username,
@@ -269,15 +269,15 @@ def __init__(self, config: SalesforceConfig, ctx: PipelineContext) -> None:
 
             elif self.config.auth is SalesforceAuthType.JSON_WEB_TOKEN:
                 logger.debug("Json Web Token provided in the config")
-                assert (
-                    self.config.username is not None
-                ), "Config username is required for JSON_WEB_TOKEN auth"
-                assert (
-                    self.config.consumer_key is not None
-                ), "Config consumer_key is required for JSON_WEB_TOKEN auth"
-                assert (
-                    self.config.private_key is not None
-                ), "Config private_key is required for JSON_WEB_TOKEN auth"
+                assert self.config.username is not None, (
+                    "Config username is required for JSON_WEB_TOKEN auth"
+                )
+                assert self.config.consumer_key is not None, (
+                    "Config consumer_key is required for JSON_WEB_TOKEN auth"
+                )
+                assert self.config.private_key is not None, (
+                    "Config private_key is required for JSON_WEB_TOKEN auth"
+                )
 
                 self.sf = Salesforce(
                     username=self.config.username,
@@ -439,7 +439,8 @@ def get_platform_instance_workunit(self, datasetUrn: str) -> MetadataWorkUnit:
         dataPlatformInstance = DataPlatformInstanceClass(
             builder.make_data_platform_urn(self.platform),
             instance=builder.make_dataplatform_instance_urn(
-                self.platform, self.config.platform_instance  # type:ignore
+                self.platform,
+                self.config.platform_instance,  # type:ignore
             ),
         )
 
diff --git a/metadata-ingestion/src/datahub/ingestion/source/schema/json_schema.py b/metadata-ingestion/src/datahub/ingestion/source/schema/json_schema.py
index 635e894d18c7e5..a50e99393fdc27 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/schema/json_schema.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/schema/json_schema.py
@@ -354,7 +354,7 @@ def get_workunits_internal(self) -> Iterable[MetadataWorkUnit]:
             browse_prefix = f"/{self.config.env.lower()}/{self.config.platform}/{self.config.platform_instance}"
 
         if os.path.isdir(self.config.path):
-            for root, dirs, files in os.walk(self.config.path, topdown=False):
+            for root, _, files in os.walk(self.config.path, topdown=False):
                 for file_name in [f for f in files if f.endswith(".json")]:
                     try:
                         yield from self._load_one_file(
diff --git a/metadata-ingestion/src/datahub/ingestion/source/sigma/sigma.py b/metadata-ingestion/src/datahub/ingestion/source/sigma/sigma.py
index e96eeb58d96efe..0468792f44aabb 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/sigma/sigma.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/sigma/sigma.py
@@ -477,9 +477,9 @@ def _gen_elements_workunit(
                     upstream_dataset_urns
                     and dataset_urn not in self.dataset_upstream_urn_mapping
                 ):
-                    self.dataset_upstream_urn_mapping[
-                        dataset_urn
-                    ] = upstream_dataset_urns
+                    self.dataset_upstream_urn_mapping[dataset_urn] = (
+                        upstream_dataset_urns
+                    )
 
             element_input_fields = [
                 InputFieldClass(
diff --git a/metadata-ingestion/src/datahub/ingestion/source/sigma/sigma_api.py b/metadata-ingestion/src/datahub/ingestion/source/sigma/sigma_api.py
index 3e88f43142ede6..6762302ebe57c7 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/sigma/sigma_api.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/sigma/sigma_api.py
@@ -126,9 +126,9 @@ def fill_workspaces(self) -> None:
                 response.raise_for_status()
                 response_dict = response.json()
                 for workspace_dict in response_dict[Constant.ENTRIES]:
-                    self.workspaces[
-                        workspace_dict[Constant.WORKSPACEID]
-                    ] = Workspace.parse_obj(workspace_dict)
+                    self.workspaces[workspace_dict[Constant.WORKSPACEID]] = (
+                        Workspace.parse_obj(workspace_dict)
+                    )
                 if response_dict[Constant.NEXTPAGE]:
                     url = f"{workspace_url}&page={response_dict[Constant.NEXTPAGE]}"
                 else:
@@ -147,9 +147,9 @@ def _get_users(self) -> Dict[str, str]:
                 response.raise_for_status()
                 response_dict = response.json()
                 for user_dict in response_dict[Constant.ENTRIES]:
-                    users[
-                        user_dict[Constant.MEMBERID]
-                    ] = f"{user_dict[Constant.FIRSTNAME]}_{user_dict[Constant.LASTNAME]}"
+                    users[user_dict[Constant.MEMBERID]] = (
+                        f"{user_dict[Constant.FIRSTNAME]}_{user_dict[Constant.LASTNAME]}"
+                    )
                 if response_dict[Constant.NEXTPAGE]:
                     url = f"{members_url}&page={response_dict[Constant.NEXTPAGE]}"
                 else:
@@ -327,10 +327,12 @@ def get_page_elements(self, workbook: Workbook, page: Page) -> List[Element]:
             response.raise_for_status()
             for i, element_dict in enumerate(response.json()[Constant.ENTRIES]):
                 if not element_dict.get(Constant.NAME):
-                    element_dict[Constant.NAME] = f"Element {i+1} of Page '{page.name}'"
-                element_dict[
-                    Constant.URL
-                ] = f"{workbook.url}?:nodeId={element_dict[Constant.ELEMENTID]}&:fullScreen=true"
+                    element_dict[Constant.NAME] = (
+                        f"Element {i + 1} of Page '{page.name}'"
+                    )
+                element_dict[Constant.URL] = (
+                    f"{workbook.url}?:nodeId={element_dict[Constant.ELEMENTID]}&:fullScreen=true"
+                )
                 element = Element.parse_obj(element_dict)
                 if (
                     self.config.extract_lineage
diff --git a/metadata-ingestion/src/datahub/ingestion/source/slack/slack.py b/metadata-ingestion/src/datahub/ingestion/source/slack/slack.py
index 00b783a89774cc..6d147e27405818 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/slack/slack.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/slack/slack.py
@@ -5,6 +5,8 @@
 
 from pydantic import Field, SecretStr
 from slack_sdk import WebClient
+from tenacity import retry, wait_exponential
+from tenacity.before_sleep import before_sleep_log
 
 import datahub.emitter.mce_builder as builder
 from datahub.configuration.common import ConfigModel
@@ -294,6 +296,10 @@ def populate_slack_id_from_email(self, user_obj: CorpUser) -> None:
                 return
             raise e
 
+    @retry(
+        wait=wait_exponential(multiplier=2, min=4, max=60),
+        before_sleep=before_sleep_log(logger, logging.ERROR, True),
+    )
     def get_user_to_be_updated(self) -> Iterable[CorpUser]:
         graphql_query = textwrap.dedent(
             """
diff --git a/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_config.py b/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_config.py
index b14e51a982082c..5f732e2621656f 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_config.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_config.py
@@ -384,18 +384,20 @@ def validate_shares(
                     assert all(
                         consumer.platform_instance != share_details.platform_instance
                         for consumer in share_details.consumers
-                    ), "Share's platform_instance can not be same as consumer's platform instance. Self-sharing not supported in Snowflake."
+                    ), (
+                        "Share's platform_instance can not be same as consumer's platform instance. Self-sharing not supported in Snowflake."
+                    )
 
                 databases_included_in_share.append(shared_db)
                 databases_created_from_share.extend(share_details.consumers)
 
             for db_from_share in databases_created_from_share:
-                assert (
-                    db_from_share not in databases_included_in_share
-                ), "Database included in a share can not be present as consumer in any share."
-                assert (
-                    databases_created_from_share.count(db_from_share) == 1
-                ), "Same database can not be present as consumer in more than one share."
+                assert db_from_share not in databases_included_in_share, (
+                    "Database included in a share can not be present as consumer in any share."
+                )
+                assert databases_created_from_share.count(db_from_share) == 1, (
+                    "Same database can not be present as consumer in more than one share."
+                )
 
         return shares
 
diff --git a/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_connection.py b/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_connection.py
index 2239338972d9be..2854a99198d62b 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_connection.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_connection.py
@@ -250,9 +250,9 @@ def get_connect_args(self) -> dict:
             if self.private_key is not None:
                 pkey_bytes = self.private_key.replace("\\n", "\n").encode()
             else:
-                assert (
-                    self.private_key_path
-                ), "missing required private key path to read key from"
+                assert self.private_key_path, (
+                    "missing required private key path to read key from"
+                )
                 with open(self.private_key_path, "rb") as key:
                     pkey_bytes = key.read()
 
@@ -284,9 +284,9 @@ def get_options(self) -> dict:
         return self.options
 
     def get_oauth_connection(self) -> NativeSnowflakeConnection:
-        assert (
-            self.oauth_config
-        ), "oauth_config should be provided if using oauth based authentication"
+        assert self.oauth_config, (
+            "oauth_config should be provided if using oauth based authentication"
+        )
         generator = OAuthTokenGenerator(
             client_id=self.oauth_config.client_id,
             authority_url=self.oauth_config.authority_url,
diff --git a/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_queries.py b/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_queries.py
index b82734cbbe84ea..69d0b62a8edfdf 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_queries.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_queries.py
@@ -623,7 +623,7 @@ def _build_enriched_query_log_query(
         query_history.start_time >= to_timestamp_ltz({start_time_millis}, 3)
         AND query_history.start_time < to_timestamp_ltz({end_time_millis}, 3)
         AND execution_status = 'SUCCESS'
-        AND {users_filter or 'TRUE'}
+        AND {users_filter or "TRUE"}
 )
 , deduplicated_queries as (
     SELECT
@@ -651,7 +651,7 @@ def _build_enriched_query_log_query(
     WHERE
         query_start_time >= to_timestamp_ltz({start_time_millis}, 3)
         AND query_start_time < to_timestamp_ltz({end_time_millis}, 3)
-        AND {users_filter or 'TRUE'}
+        AND {users_filter or "TRUE"}
         AND query_id IN (
             SELECT query_id FROM deduplicated_queries
         )
diff --git a/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_schema.py b/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_schema.py
index d165be3f3cc656..173024aec0cf38 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_schema.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_schema.py
@@ -142,9 +142,9 @@ def __init__(self) -> None:
         )
 
         # self._table_tags[<database_name>][<schema_name>][<table_name>] = list of tags applied to table
-        self._table_tags: Dict[
-            str, Dict[str, Dict[str, List[SnowflakeTag]]]
-        ] = defaultdict(lambda: defaultdict(lambda: defaultdict(list)))
+        self._table_tags: Dict[str, Dict[str, Dict[str, List[SnowflakeTag]]]] = (
+            defaultdict(lambda: defaultdict(lambda: defaultdict(list)))
+        )
 
         # self._column_tags[<database_name>][<schema_name>][<table_name>][<column_name>] = list of tags applied to column
         self._column_tags: Dict[
diff --git a/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_schema_gen.py b/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_schema_gen.py
index 393e4d3c96d51f..04bc51f1ebd3f5 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_schema_gen.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_schema_gen.py
@@ -194,9 +194,9 @@ def __init__(
             config, self.data_dictionary, self.report
         )
         self.profiler: Optional[SnowflakeProfiler] = profiler
-        self.snowsight_url_builder: Optional[
-            SnowsightUrlBuilder
-        ] = snowsight_url_builder
+        self.snowsight_url_builder: Optional[SnowsightUrlBuilder] = (
+            snowsight_url_builder
+        )
 
         # These are populated as side-effects of get_workunits_internal.
         self.databases: List[SnowflakeDatabase] = []
@@ -267,9 +267,9 @@ def get_databases(self) -> Optional[List[SnowflakeDatabase]]:
             )
             return None
         else:
-            ischema_databases: List[
-                SnowflakeDatabase
-            ] = self.get_databases_from_ischema(databases)
+            ischema_databases: List[SnowflakeDatabase] = (
+                self.get_databases_from_ischema(databases)
+            )
 
             if len(ischema_databases) == 0:
                 self.structured_reporter.failure(
@@ -491,15 +491,25 @@ def fetch_secure_view_definition(
         try:
             view_definitions = self.data_dictionary.get_secure_view_definitions()
             return view_definitions[db_name][schema_name][table_name]
+        except KeyError:
+            # Received secure view definitions but the view is not present in results
+            self.structured_reporter.info(
+                title="Secure view definition not found",
+                message="Lineage will be missing for the view.",
+                context=f"{db_name}.{schema_name}.{table_name}",
+            )
+            return None
         except Exception as e:
-            if isinstance(e, SnowflakePermissionError):
-                error_msg = (
-                    "Failed to get secure views definitions. Please check permissions."
-                )
-            else:
-                error_msg = "Failed to get secure views definitions"
+            action_msg = (
+                "Please check permissions."
+                if isinstance(e, SnowflakePermissionError)
+                else ""
+            )
+
             self.structured_reporter.warning(
-                error_msg,
+                title="Failed to get secure views definitions",
+                message=f"Lineage will be missing for the view. {action_msg}",
+                context=f"{db_name}.{schema_name}.{table_name}",
                 exc=e,
             )
             return None
diff --git a/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_tag.py b/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_tag.py
index 75567cc3da8830..597e7bee4d4cc0 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_tag.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_tag.py
@@ -38,9 +38,9 @@ def _get_tags_on_object_without_propagation(
         table_name: Optional[str],
     ) -> List[SnowflakeTag]:
         if db_name not in self.tag_cache:
-            self.tag_cache[
-                db_name
-            ] = self.data_dictionary.get_tags_for_database_without_propagation(db_name)
+            self.tag_cache[db_name] = (
+                self.data_dictionary.get_tags_for_database_without_propagation(db_name)
+            )
 
         if domain == SnowflakeObjectDomain.DATABASE:
             return self.tag_cache[db_name].get_database_tags(db_name)
@@ -130,10 +130,10 @@ def get_column_tags_for_table(
         temp_column_tags: Dict[str, List[SnowflakeTag]] = {}
         if self.config.extract_tags == TagOption.without_lineage:
             if db_name not in self.tag_cache:
-                self.tag_cache[
-                    db_name
-                ] = self.data_dictionary.get_tags_for_database_without_propagation(
-                    db_name
+                self.tag_cache[db_name] = (
+                    self.data_dictionary.get_tags_for_database_without_propagation(
+                        db_name
+                    )
                 )
             temp_column_tags = self.tag_cache[db_name].get_column_tags_for_table(
                 table_name, schema_name, db_name
diff --git a/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_usage_v2.py b/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_usage_v2.py
index 85e4071aec07df..edd13ee48326bb 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_usage_v2.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_usage_v2.py
@@ -549,9 +549,9 @@ def parse_event_objects(self, event_dict: Dict) -> None:
         ):
             # NOTE: Generated emails may be incorrect, as email may be different than
             # username@email_domain
-            event_dict[
-                "EMAIL"
-            ] = f'{event_dict["USER_NAME"]}@{self.config.email_domain}'.lower()
+            event_dict["EMAIL"] = (
+                f"{event_dict['USER_NAME']}@{self.config.email_domain}".lower()
+            )
 
         if not event_dict["EMAIL"]:
             self.report.rows_missing_email += 1
diff --git a/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_utils.py b/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_utils.py
index 885bee1ccdb908..030edfde4ca1da 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_utils.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_utils.py
@@ -21,8 +21,7 @@
 class SnowflakeStructuredReportMixin(abc.ABC):
     @property
     @abc.abstractmethod
-    def structured_reporter(self) -> SourceReport:
-        ...
+    def structured_reporter(self) -> SourceReport: ...
 
 
 class SnowsightUrlBuilder:
diff --git a/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_v2.py b/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_v2.py
index b8afd145727400..b4ef2180d71d45 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_v2.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_v2.py
@@ -211,9 +211,9 @@ def __init__(self, ctx: PipelineContext, config: SnowflakeV2Config):
 
         self.usage_extractor: Optional[SnowflakeUsageExtractor] = None
         if self.config.include_usage_stats or self.config.include_operational_stats:
-            redundant_usage_run_skip_handler: Optional[
-                RedundantUsageRunSkipHandler
-            ] = None
+            redundant_usage_run_skip_handler: Optional[RedundantUsageRunSkipHandler] = (
+                None
+            )
             if self.config.enable_stateful_usage_ingestion:
                 redundant_usage_run_skip_handler = RedundantUsageRunSkipHandler(
                     source=self,
@@ -296,7 +296,16 @@ class SnowflakePrivilege:
 
         _report: Dict[Union[SourceCapability, str], CapabilityReport] = dict()
         privileges: List[SnowflakePrivilege] = []
-        capabilities: List[SourceCapability] = [c.capability for c in SnowflakeV2Source.get_capabilities() if c.capability not in (SourceCapability.PLATFORM_INSTANCE, SourceCapability.DOMAINS, SourceCapability.DELETION_DETECTION)]  # type: ignore
+        capabilities: List[SourceCapability] = [
+            c.capability
+            for c in SnowflakeV2Source.get_capabilities()  # type: ignore
+            if c.capability
+            not in (
+                SourceCapability.PLATFORM_INSTANCE,
+                SourceCapability.DOMAINS,
+                SourceCapability.DELETION_DETECTION,
+            )
+        ]
 
         cur = conn.query("select current_role()")
         current_role = [row["CURRENT_ROLE()"] for row in cur][0]
diff --git a/metadata-ingestion/src/datahub/ingestion/source/sql/athena.py b/metadata-ingestion/src/datahub/ingestion/source/sql/athena.py
index 6f7decc79b1df2..cfc43454b51fad 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/sql/athena.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/sql/athena.py
@@ -104,9 +104,7 @@ def get_view_definition(self, connection, view_name, schema=None, **kw):
             return "\n".join([r for r in res])
 
     @typing.no_type_check
-    def _get_column_type(
-        self, type_: Union[str, Dict[str, Any]]
-    ) -> TypeEngine:  # noqa: C901
+    def _get_column_type(self, type_: Union[str, Dict[str, Any]]) -> TypeEngine:  # noqa: C901
         """Derives the data type of the Athena column.
 
         This method is overwritten to extend the behavior of PyAthena.
diff --git a/metadata-ingestion/src/datahub/ingestion/source/sql/clickhouse.py b/metadata-ingestion/src/datahub/ingestion/source/sql/clickhouse.py
index aeb21e88d04437..a2db116cf20912 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/sql/clickhouse.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/sql/clickhouse.py
@@ -53,7 +53,6 @@
 )
 from datahub.metadata.schema_classes import (
     DatasetLineageTypeClass,
-    DatasetPropertiesClass,
     DatasetSnapshotClass,
     UpstreamClass,
 )
@@ -218,9 +217,7 @@ def _get_all_table_comments_and_properties(self, connection, **kw):
              , comment
              , {properties_clause} AS properties
           FROM system.tables
-         WHERE name NOT LIKE '.inner%'""".format(
-            properties_clause=properties_clause
-        )
+         WHERE name NOT LIKE '.inner%'""".format(properties_clause=properties_clause)
     )
 
     all_table_comments: Dict[Tuple[str, str], Dict[str, Any]] = {}
@@ -268,7 +265,7 @@ def _get_table_or_view_names(self, relkind, connection, schema=None, **kw):
     info_cache = kw.get("info_cache")
     all_relations = self._get_all_relation_info(connection, info_cache=info_cache)
     relation_names = []
-    for key, relation in all_relations.items():
+    for _, relation in all_relations.items():
         if relation.database == schema and relation.relkind == relkind:
             relation_names.append(relation.relname)
     return relation_names
@@ -301,9 +298,7 @@ def _get_schema_column_info(self, connection, schema=None, **kw):
              , comment
           FROM system.columns
          WHERE {schema_clause}
-         ORDER BY database, table, position""".format(
-                    schema_clause=schema_clause
-                )
+         ORDER BY database, table, position""".format(schema_clause=schema_clause)
             )
         )
     )
@@ -422,41 +417,11 @@ def get_workunits_internal(self) -> Iterable[Union[MetadataWorkUnit, SqlWorkUnit
                 dataset_snapshot: DatasetSnapshotClass = wu.metadata.proposedSnapshot
                 assert dataset_snapshot
 
-                lineage_mcp, lineage_properties_aspect = self.get_lineage_mcp(
-                    wu.metadata.proposedSnapshot.urn
-                )
+                lineage_mcp = self.get_lineage_mcp(wu.metadata.proposedSnapshot.urn)
 
                 if lineage_mcp is not None:
                     yield lineage_mcp.as_workunit()
 
-                if lineage_properties_aspect:
-                    aspects = dataset_snapshot.aspects
-                    if aspects is None:
-                        aspects = []
-
-                    dataset_properties_aspect: Optional[DatasetPropertiesClass] = None
-
-                    for aspect in aspects:
-                        if isinstance(aspect, DatasetPropertiesClass):
-                            dataset_properties_aspect = aspect
-
-                    if dataset_properties_aspect is None:
-                        dataset_properties_aspect = DatasetPropertiesClass()
-                        aspects.append(dataset_properties_aspect)
-
-                    custom_properties = (
-                        {
-                            **dataset_properties_aspect.customProperties,
-                            **lineage_properties_aspect.customProperties,
-                        }
-                        if dataset_properties_aspect.customProperties
-                        else lineage_properties_aspect.customProperties
-                    )
-                    dataset_properties_aspect.customProperties = custom_properties
-                    dataset_snapshot.aspects = aspects
-
-                    dataset_snapshot.aspects.append(dataset_properties_aspect)
-
             # Emit the work unit from super.
             yield wu
 
@@ -474,7 +439,7 @@ def _get_all_tables(self) -> Set[str]:
         logger.debug(f"sql_alchemy_url={url}")
         engine = create_engine(url, **self.config.options)
         for db_row in engine.execute(text(all_tables_query)):
-            all_tables_set.add(f'{db_row["database"]}.{db_row["table_name"]}')
+            all_tables_set.add(f"{db_row['database']}.{db_row['table_name']}")
 
         return all_tables_set
 
@@ -503,7 +468,7 @@ def _populate_lineage_map(
 
         try:
             for db_row in engine.execute(text(query)):
-                dataset_name = f'{db_row["target_schema"]}.{db_row["target_table"]}'
+                dataset_name = f"{db_row['target_schema']}.{db_row['target_table']}"
                 if not self.config.database_pattern.allowed(
                     db_row["target_schema"]
                 ) or not self.config.table_pattern.allowed(dataset_name):
@@ -512,7 +477,7 @@ def _populate_lineage_map(
 
                 # Target
                 target_path = (
-                    f'{self.config.platform_instance+"." if self.config.platform_instance else ""}'
+                    f"{self.config.platform_instance + '.' if self.config.platform_instance else ''}"
                     f"{dataset_name}"
                 )
                 target = LineageItem(
@@ -525,7 +490,7 @@ def _populate_lineage_map(
 
                 # Source
                 platform = LineageDatasetPlatform.CLICKHOUSE
-                path = f'{db_row["source_schema"]}.{db_row["source_table"]}'
+                path = f"{db_row['source_schema']}.{db_row['source_table']}"
 
                 sources = [
                     LineageDataset(
@@ -552,9 +517,7 @@ def _populate_lineage_map(
                         target.dataset.path
                     ].upstreams = self._lineage_map[
                         target.dataset.path
-                    ].upstreams.union(
-                        target.upstreams
-                    )
+                    ].upstreams.union(target.upstreams)
 
                 else:
                     self._lineage_map[target.dataset.path] = target
@@ -662,19 +625,16 @@ def _populate_lineage(self) -> None:
 
     def get_lineage_mcp(
         self, dataset_urn: str
-    ) -> Tuple[
-        Optional[MetadataChangeProposalWrapper], Optional[DatasetPropertiesClass]
-    ]:
+    ) -> Optional[MetadataChangeProposalWrapper]:
         dataset_key = mce_builder.dataset_urn_to_key(dataset_urn)
         if dataset_key is None:
-            return None, None
+            return None
 
         if not self._lineage_map:
             self._populate_lineage()
         assert self._lineage_map is not None
 
         upstream_lineage: List[UpstreamClass] = []
-        custom_properties: Dict[str, str] = {}
 
         if dataset_key.name in self._lineage_map:
             item = self._lineage_map[dataset_key.name]
@@ -690,16 +650,12 @@ def get_lineage_mcp(
                 )
                 upstream_lineage.append(upstream_table)
 
-        properties = None
-        if custom_properties:
-            properties = DatasetPropertiesClass(customProperties=custom_properties)
-
         if not upstream_lineage:
-            return None, properties
+            return None
 
         mcp = MetadataChangeProposalWrapper(
             entityUrn=dataset_urn,
             aspect=UpstreamLineage(upstreams=upstream_lineage),
         )
 
-        return mcp, properties
+        return mcp
diff --git a/metadata-ingestion/src/datahub/ingestion/source/sql/hive_metastore.py b/metadata-ingestion/src/datahub/ingestion/source/sql/hive_metastore.py
index 60ecbaf38838a6..82b22d896698a8 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/sql/hive_metastore.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/sql/hive_metastore.py
@@ -2,7 +2,6 @@
 import json
 import logging
 from collections import namedtuple
-from itertools import groupby
 from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
 
 from pydantic.dataclasses import dataclass
@@ -58,6 +57,7 @@
     SubTypesClass,
     ViewPropertiesClass,
 )
+from datahub.utilities.groupby import groupby_unsorted
 from datahub.utilities.hive_schema_to_avro import get_schema_fields_for_hive_column
 from datahub.utilities.str_enum import StrEnum
 
@@ -490,7 +490,7 @@ def loop_tables(
 
         iter_res = self._alchemy_client.execute_query(statement)
 
-        for key, group in groupby(iter_res, self._get_table_key):
+        for key, group in groupby_unsorted(iter_res, self._get_table_key):
             schema_name = (
                 f"{db_name}.{key.schema}"
                 if self.config.include_catalog_name_in_ids
@@ -647,7 +647,7 @@ def get_hive_view_columns(self, inspector: Inspector) -> Iterable[ViewDataset]:
         )
 
         iter_res = self._alchemy_client.execute_query(statement)
-        for key, group in groupby(iter_res, self._get_table_key):
+        for key, group in groupby_unsorted(iter_res, self._get_table_key):
             db_name = self.get_db_name(inspector)
 
             schema_name = (
diff --git a/metadata-ingestion/src/datahub/ingestion/source/sql/mssql/job_models.py b/metadata-ingestion/src/datahub/ingestion/source/sql/mssql/job_models.py
index d3941e7add0fd0..9cd485f6d9a4b1 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/sql/mssql/job_models.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/sql/mssql/job_models.py
@@ -156,7 +156,7 @@ class MSSQLDataJob:
     entity: Union[StoredProcedure, JobStep]
     type: str = "dataJob"
     source: str = "mssql"
-    external_url: str = ""
+    external_url: Optional[str] = None
     description: Optional[str] = None
     status: Optional[str] = None
     incoming: List[str] = field(default_factory=list)
@@ -228,7 +228,7 @@ class MSSQLDataFlow:
     entity: Union[MSSQLJob, MSSQLProceduresContainer]
     type: str = "dataFlow"
     source: str = "mssql"
-    external_url: str = ""
+    external_url: Optional[str] = None
     flow_properties: Dict[str, str] = field(default_factory=dict)
 
     def add_property(
diff --git a/metadata-ingestion/src/datahub/ingestion/source/sql/mssql/source.py b/metadata-ingestion/src/datahub/ingestion/source/sql/mssql/source.py
index a2338f14196d77..2d0c4dc9f6f6f9 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/sql/mssql/source.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/sql/mssql/source.py
@@ -11,6 +11,7 @@
 from sqlalchemy.engine.reflection import Inspector
 from sqlalchemy.exc import ProgrammingError, ResourceClosedError
 
+import datahub.metadata.schema_classes as models
 from datahub.configuration.common import AllowDenyPattern
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.ingestion.api.common import PipelineContext
@@ -49,21 +50,15 @@
     make_sqlalchemy_uri,
 )
 from datahub.ingestion.source.sql.sql_report import SQLSourceReport
-from datahub.metadata.schema_classes import (
-    BooleanTypeClass,
-    NumberTypeClass,
-    StringTypeClass,
-    UnionTypeClass,
-)
 from datahub.utilities.file_backed_collections import FileBackedList
 
 logger: logging.Logger = logging.getLogger(__name__)
 
-register_custom_type(sqlalchemy.dialects.mssql.BIT, BooleanTypeClass)
-register_custom_type(sqlalchemy.dialects.mssql.MONEY, NumberTypeClass)
-register_custom_type(sqlalchemy.dialects.mssql.SMALLMONEY, NumberTypeClass)
-register_custom_type(sqlalchemy.dialects.mssql.SQL_VARIANT, UnionTypeClass)
-register_custom_type(sqlalchemy.dialects.mssql.UNIQUEIDENTIFIER, StringTypeClass)
+register_custom_type(sqlalchemy.dialects.mssql.BIT, models.BooleanTypeClass)
+register_custom_type(sqlalchemy.dialects.mssql.MONEY, models.NumberTypeClass)
+register_custom_type(sqlalchemy.dialects.mssql.SMALLMONEY, models.NumberTypeClass)
+register_custom_type(sqlalchemy.dialects.mssql.SQL_VARIANT, models.UnionTypeClass)
+register_custom_type(sqlalchemy.dialects.mssql.UNIQUEIDENTIFIER, models.StringTypeClass)
 
 
 class SQLServerConfig(BasicSQLAlchemyConfig):
@@ -651,6 +646,26 @@ def construct_job_workunits(
                 entityUrn=data_job.urn,
                 aspect=data_job.as_datajob_input_output_aspect,
             ).as_workunit()
+
+        if (
+            self.config.include_stored_procedures_code
+            and isinstance(data_job.entity, StoredProcedure)
+            and data_job.entity.code is not None
+        ):
+            yield MetadataChangeProposalWrapper(
+                entityUrn=data_job.urn,
+                aspect=models.DataTransformLogicClass(
+                    transforms=[
+                        models.DataTransformClass(
+                            queryStatement=models.QueryStatementClass(
+                                value=data_job.entity.code,
+                                language=models.QueryLanguageClass.SQL,
+                            ),
+                        )
+                    ]
+                ),
+            ).as_workunit()
+
         # TODO: Add SubType when it appear
 
     def construct_flow_workunits(
diff --git a/metadata-ingestion/src/datahub/ingestion/source/sql/oracle.py b/metadata-ingestion/src/datahub/ingestion/source/sql/oracle.py
index 52db3cd11a759d..ac568c58af6c68 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/sql/oracle.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/sql/oracle.py
@@ -234,9 +234,7 @@ def get_columns(
                     WHERE col.table_name = id.table_name
                     AND col.column_name = id.column_name
                     AND col.owner = id.owner
-                ) AS identity_options""".format(
-                dblink=dblink
-            )
+                ) AS identity_options""".format(dblink=dblink)
         else:
             identity_cols = "NULL as default_on_null, NULL as identity_options"
 
diff --git a/metadata-ingestion/src/datahub/ingestion/source/sql/sql_generic_profiler.py b/metadata-ingestion/src/datahub/ingestion/source/sql/sql_generic_profiler.py
index c91be9b494c006..664735053f1852 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/sql/sql_generic_profiler.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/sql/sql_generic_profiler.py
@@ -278,8 +278,7 @@ def is_dataset_eligible_for_profiling(
 
         if self.config.profiling.profile_table_size_limit is not None and (
             size_in_bytes is not None
-            and size_in_bytes / (2**30)
-            > self.config.profiling.profile_table_size_limit
+            and size_in_bytes / (2**30) > self.config.profiling.profile_table_size_limit
         ):
             self.report.profiling_skipped_size_limit[schema_name] += 1
             logger.debug(
diff --git a/metadata-ingestion/src/datahub/ingestion/source/sql/sql_types.py b/metadata-ingestion/src/datahub/ingestion/source/sql/sql_types.py
index 1acf962d7c4750..6b9b8bdeb39e97 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/sql/sql_types.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/sql/sql_types.py
@@ -384,7 +384,6 @@ def resolve_snowflake_modified_type(type_string: str) -> Any:
     "varchar": StringType,
     "char": StringType,
     "varbinary": BytesType,
-    "json": RecordType,
     "date": DateType,
     "time": TimeType,
     "timestamp": TimeType,
diff --git a/metadata-ingestion/src/datahub/ingestion/source/sql/teradata.py b/metadata-ingestion/src/datahub/ingestion/source/sql/teradata.py
index 5b76fe41d92e97..e6319f668ecb8c 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/sql/teradata.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/sql/teradata.py
@@ -3,7 +3,6 @@
 from dataclasses import dataclass
 from datetime import datetime
 from functools import lru_cache
-from itertools import groupby
 from typing import (
     Any,
     Dict,
@@ -59,6 +58,7 @@
 from datahub.metadata.schema_classes import SchemaMetadataClass
 from datahub.sql_parsing.schema_resolver import SchemaResolver
 from datahub.sql_parsing.sqlglot_lineage import sqlglot_lineage
+from datahub.utilities.groupby import groupby_unsorted
 
 logger: logging.Logger = logging.getLogger(__name__)
 
@@ -286,7 +286,7 @@ def grouper(fk_row):
 
     # TODO: Check if there's a better way
     fk_dicts = list()
-    for constraint_info, constraint_cols in groupby(res, grouper):
+    for constraint_info, constraint_cols in groupby_unsorted(res, grouper):
         fk_dict = {
             "name": str(constraint_info["name"]),
             "constrained_columns": list(),
@@ -599,7 +599,12 @@ def __init__(self, config: TeradataConfig, ctx: PipelineContext):
             setattr(  # noqa: B010
                 TeradataDialect,
                 "get_columns",
-                lambda self, connection, table_name, schema=None, use_qvci=self.config.use_qvci, **kw: optimized_get_columns(
+                lambda self,
+                connection,
+                table_name,
+                schema=None,
+                use_qvci=self.config.use_qvci,
+                **kw: optimized_get_columns(
                     self,
                     connection,
                     table_name,
@@ -613,7 +618,11 @@ def __init__(self, config: TeradataConfig, ctx: PipelineContext):
             setattr(  # noqa: B010
                 TeradataDialect,
                 "get_pk_constraint",
-                lambda self, connection, table_name, schema=None, **kw: optimized_get_pk_constraint(
+                lambda self,
+                connection,
+                table_name,
+                schema=None,
+                **kw: optimized_get_pk_constraint(
                     self, connection, table_name, schema, **kw
                 ),
             )
@@ -621,7 +630,11 @@ def __init__(self, config: TeradataConfig, ctx: PipelineContext):
             setattr(  # noqa: B010
                 TeradataDialect,
                 "get_foreign_keys",
-                lambda self, connection, table_name, schema=None, **kw: optimized_get_foreign_keys(
+                lambda self,
+                connection,
+                table_name,
+                schema=None,
+                **kw: optimized_get_foreign_keys(
                     self, connection, table_name, schema, **kw
                 ),
             )
diff --git a/metadata-ingestion/src/datahub/ingestion/source/state/profiling_state_handler.py b/metadata-ingestion/src/datahub/ingestion/source/state/profiling_state_handler.py
index 9883bc2b8e9b0b..6080ddadb65e40 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/state/profiling_state_handler.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/state/profiling_state_handler.py
@@ -41,9 +41,9 @@ def __init__(
         run_id: str,
     ):
         self.state_provider = source.state_provider
-        self.stateful_ingestion_config: Optional[
-            ProfilingStatefulIngestionConfig
-        ] = config.stateful_ingestion
+        self.stateful_ingestion_config: Optional[ProfilingStatefulIngestionConfig] = (
+            config.stateful_ingestion
+        )
         self.pipeline_name = pipeline_name
         self.run_id = run_id
         self.checkpointing_enabled: bool = (
diff --git a/metadata-ingestion/src/datahub/ingestion/source/state/redundant_run_skip_handler.py b/metadata-ingestion/src/datahub/ingestion/source/state/redundant_run_skip_handler.py
index 8630a959d3f6a3..e4a2646f6ccd3c 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/state/redundant_run_skip_handler.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/state/redundant_run_skip_handler.py
@@ -48,9 +48,9 @@ def __init__(
     ):
         self.source = source
         self.state_provider = source.state_provider
-        self.stateful_ingestion_config: Optional[
-            StatefulIngestionConfig
-        ] = config.stateful_ingestion
+        self.stateful_ingestion_config: Optional[StatefulIngestionConfig] = (
+            config.stateful_ingestion
+        )
         self.pipeline_name = pipeline_name
         self.run_id = run_id
         self._job_id = self._init_job_id()
@@ -145,8 +145,7 @@ def should_skip_this_run(
             )
 
             logger.debug(
-                f"{self.job_id} : Last run start, end times:"
-                f"({last_run_time_window})"
+                f"{self.job_id} : Last run start, end times:({last_run_time_window})"
             )
 
             # If current run's time window is subset of last run's time window, then skip.
@@ -212,8 +211,7 @@ def suggest_run_time_window(
             )
 
         self.log(
-            "Adjusted start, end times: "
-            f"({suggested_start_time}, {suggested_end_time})"
+            f"Adjusted start, end times: ({suggested_start_time}, {suggested_end_time})"
         )
         return (suggested_start_time, suggested_end_time)
 
diff --git a/metadata-ingestion/src/datahub/ingestion/source/state/stale_entity_removal_handler.py b/metadata-ingestion/src/datahub/ingestion/source/state/stale_entity_removal_handler.py
index d4fcbf09924e9e..017d78bc1abf8d 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/state/stale_entity_removal_handler.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/state/stale_entity_removal_handler.py
@@ -111,9 +111,9 @@ def __init__(
         self.state_type_class = state_type_class
         self.pipeline_name = pipeline_name
         self.run_id = run_id
-        self.stateful_ingestion_config: Optional[
-            StatefulStaleMetadataRemovalConfig
-        ] = config.stateful_ingestion
+        self.stateful_ingestion_config: Optional[StatefulStaleMetadataRemovalConfig] = (
+            config.stateful_ingestion
+        )
         self.checkpointing_enabled: bool = (
             True
             if (
diff --git a/metadata-ingestion/src/datahub/ingestion/source/state_provider/datahub_ingestion_checkpointing_provider.py b/metadata-ingestion/src/datahub/ingestion/source/state_provider/datahub_ingestion_checkpointing_provider.py
index 8f4a53ffc3ed58..1f5a651fc64a79 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/state_provider/datahub_ingestion_checkpointing_provider.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/state_provider/datahub_ingestion_checkpointing_provider.py
@@ -70,20 +70,20 @@ def get_latest_checkpoint(
             self.orchestrator_name, pipeline_name, job_name
         )
 
-        latest_checkpoint: Optional[
-            DatahubIngestionCheckpointClass
-        ] = self.graph.get_latest_timeseries_value(
-            entity_urn=data_job_urn,
-            aspect_type=DatahubIngestionCheckpointClass,
-            filter_criteria_map={
-                "pipelineName": pipeline_name,
-            },
+        latest_checkpoint: Optional[DatahubIngestionCheckpointClass] = (
+            self.graph.get_latest_timeseries_value(
+                entity_urn=data_job_urn,
+                aspect_type=DatahubIngestionCheckpointClass,
+                filter_criteria_map={
+                    "pipelineName": pipeline_name,
+                },
+            )
         )
         if latest_checkpoint:
             logger.debug(
                 f"The last committed ingestion checkpoint for pipelineName:'{pipeline_name}',"
                 f" job_name:'{job_name}' found with start_time:"
-                f" {datetime.utcfromtimestamp(latest_checkpoint.timestampMillis/1000)}"
+                f" {datetime.utcfromtimestamp(latest_checkpoint.timestampMillis / 1000)}"
             )
             return latest_checkpoint
         else:
diff --git a/metadata-ingestion/src/datahub/ingestion/source/state_provider/file_ingestion_checkpointing_provider.py b/metadata-ingestion/src/datahub/ingestion/source/state_provider/file_ingestion_checkpointing_provider.py
index a37774773b84d7..55f0903b9c91c7 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/state_provider/file_ingestion_checkpointing_provider.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/state_provider/file_ingestion_checkpointing_provider.py
@@ -67,7 +67,7 @@ def get_latest_checkpoint(
             logger.debug(
                 f"The last committed ingestion checkpoint for pipelineName:'{pipeline_name}',"
                 f" job_name:'{job_name}' found with start_time:"
-                f" {datetime.utcfromtimestamp(latest_checkpoint.timestampMillis/1000)}"
+                f" {datetime.utcfromtimestamp(latest_checkpoint.timestampMillis / 1000)}"
             )
             return latest_checkpoint
         else:
diff --git a/metadata-ingestion/src/datahub/ingestion/source/tableau/tableau.py b/metadata-ingestion/src/datahub/ingestion/source/tableau/tableau.py
index ea3fb6c979a19c..36e62c928744ef 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/tableau/tableau.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/tableau/tableau.py
@@ -174,6 +174,8 @@
 from datahub.utilities.stats_collections import TopKDict
 from datahub.utilities.urns.dataset_urn import DatasetUrn
 
+DEFAULT_PAGE_SIZE = 10
+
 try:
     # On earlier versions of the tableauserverclient, the NonXMLResponseError
     # was thrown when reauthentication was necessary. We'll keep both exceptions
@@ -279,9 +281,9 @@ def get_tableau_auth(
         return authentication
 
     def make_tableau_client(self, site: str) -> Server:
-        authentication: Union[
-            TableauAuth, PersonalAccessTokenAuth
-        ] = self.get_tableau_auth(site)
+        authentication: Union[TableauAuth, PersonalAccessTokenAuth] = (
+            self.get_tableau_auth(site)
+        )
         try:
             server = Server(
                 self.connect_uri,
@@ -342,11 +344,140 @@ class PermissionIngestionConfig(ConfigModel):
     )
 
 
+class TableauPageSizeConfig(ConfigModel):
+    """
+    Configuration for setting page sizes for different Tableau metadata objects.
+
+    Some considerations:
+    - All have default values, so no setting is mandatory.
+    - In general, with the `effective_` methods, if not specifically set fine-grained metrics fallback to `page_size`
+    or correlate with `page_size`.
+
+    Measuring the impact of changing these values can be done by looking at the
+    `num_(filter_|paginated_)?queries_by_connection_type` metrics in the report.
+    """
+
+    page_size: int = Field(
+        default=DEFAULT_PAGE_SIZE,
+        description="[advanced] Number of metadata objects (e.g. CustomSQLTable, PublishedDatasource, etc) to query at a time using the Tableau API.",
+    )
+
+    database_server_page_size: Optional[int] = Field(
+        default=None,
+        description="[advanced] Number of database servers to query at a time using the Tableau API; fallbacks to `page_size` if not set.",
+    )
+
+    @property
+    def effective_database_server_page_size(self) -> int:
+        return self.database_server_page_size or self.page_size
+
+    # We've found that even with a small workbook page size (e.g. 10), the Tableau API often
+    # returns warnings like this:
+    # {
+    # 	'message': 'Showing partial results. The request exceeded the 20000 node limit. Use pagination, additional filtering, or both in the query to adjust results.',
+    # 	'extensions': {
+    # 		'severity': 'WARNING',
+    # 		'code': 'NODE_LIMIT_EXCEEDED',
+    # 		'properties': {
+    # 			'nodeLimit': 20000
+    # 		}
+    # 	}
+    # }
+    # Reducing the page size for the workbook queries helps to avoid this.
+    workbook_page_size: Optional[int] = Field(
+        default=1,
+        description="[advanced] Number of workbooks to query at a time using the Tableau API; defaults to `1` and fallbacks to `page_size` if not set.",
+    )
+
+    @property
+    def effective_workbook_page_size(self) -> int:
+        return self.workbook_page_size or self.page_size
+
+    sheet_page_size: Optional[int] = Field(
+        default=None,
+        description="[advanced] Number of sheets to query at a time using the Tableau API; fallbacks to `page_size` if not set.",
+    )
+
+    @property
+    def effective_sheet_page_size(self) -> int:
+        return self.sheet_page_size or self.page_size
+
+    dashboard_page_size: Optional[int] = Field(
+        default=None,
+        description="[advanced] Number of dashboards to query at a time using the Tableau API; fallbacks to `page_size` if not set.",
+    )
+
+    @property
+    def effective_dashboard_page_size(self) -> int:
+        return self.dashboard_page_size or self.page_size
+
+    embedded_datasource_page_size: Optional[int] = Field(
+        default=None,
+        description="[advanced] Number of embedded datasources to query at a time using the Tableau API; fallbacks to `page_size` if not set.",
+    )
+
+    @property
+    def effective_embedded_datasource_page_size(self) -> int:
+        return self.embedded_datasource_page_size or self.page_size
+
+    # Since the field upstream query was separated from the embedded datasource queries into an independent query,
+    # the number of queries increased significantly and so the execution time.
+    # To increase the batching and so reduce the number of queries, we can increase the page size for that
+    # particular case.
+    #
+    # That's why unless specifically set, we will effectively use 10 times the page size as the default page size.
+    embedded_datasource_field_upstream_page_size: Optional[int] = Field(
+        default=None,
+        description="[advanced] Number of upstream fields to query at a time for embedded datasources using the Tableau API; fallbacks to `page_size` * 10 if not set.",
+    )
+
+    @property
+    def effective_embedded_datasource_field_upstream_page_size(self) -> int:
+        return self.embedded_datasource_field_upstream_page_size or self.page_size * 10
+
+    published_datasource_page_size: Optional[int] = Field(
+        default=None,
+        description="[advanced] Number of published datasources to query at a time using the Tableau API; fallbacks to `page_size` if not set.",
+    )
+
+    @property
+    def effective_published_datasource_page_size(self) -> int:
+        return self.published_datasource_page_size or self.page_size
+
+    published_datasource_field_upstream_page_size: Optional[int] = Field(
+        default=None,
+        description="[advanced] Number of upstream fields to query at a time for published datasources using the Tableau API; fallbacks to `page_size` * 10 if not set.",
+    )
+
+    @property
+    def effective_published_datasource_field_upstream_page_size(self) -> int:
+        return self.published_datasource_field_upstream_page_size or self.page_size * 10
+
+    custom_sql_table_page_size: Optional[int] = Field(
+        default=None,
+        description="[advanced] Number of custom sql datasources to query at a time using the Tableau API; fallbacks to `page_size` if not set.",
+    )
+
+    @property
+    def effective_custom_sql_table_page_size(self) -> int:
+        return self.custom_sql_table_page_size or self.page_size
+
+    database_table_page_size: Optional[int] = Field(
+        default=None,
+        description="[advanced] Number of database tables to query at a time using the Tableau API; fallbacks to `page_size` if not set.",
+    )
+
+    @property
+    def effective_database_table_page_size(self) -> int:
+        return self.database_table_page_size or self.page_size
+
+
 class TableauConfig(
     DatasetLineageProviderConfigBase,
     StatefulIngestionConfigBase,
     DatasetSourceConfigMixin,
     TableauConnectionConfig,
+    TableauPageSizeConfig,
 ):
     projects: Optional[List[str]] = Field(
         default=["default"],
@@ -396,29 +527,6 @@ class TableauConfig(
         description="Ingest details for tables external to (not embedded in) tableau as entities.",
     )
 
-    page_size: int = Field(
-        default=10,
-        description="[advanced] Number of metadata objects (e.g. CustomSQLTable, PublishedDatasource, etc) to query at a time using the Tableau API.",
-    )
-
-    # We've found that even with a small workbook page size (e.g. 10), the Tableau API often
-    # returns warnings like this:
-    # {
-    # 	'message': 'Showing partial results. The request exceeded the 20000 node limit. Use pagination, additional filtering, or both in the query to adjust results.',
-    # 	'extensions': {
-    # 		'severity': 'WARNING',
-    # 		'code': 'NODE_LIMIT_EXCEEDED',
-    # 		'properties': {
-    # 			'nodeLimit': 20000
-    # 		}
-    # 	}
-    # }
-    # Reducing the page size for the workbook queries helps to avoid this.
-    workbook_page_size: int = Field(
-        default=1,
-        description="[advanced] Number of workbooks to query at a time using the Tableau API.",
-    )
-
     env: str = Field(
         default=builder.DEFAULT_ENV,
         description="Environment to use in namespace when constructing URNs.",
@@ -527,7 +635,7 @@ def projects_backward_compatibility(cls, values: Dict) -> Dict:
         project_path_pattern = values.get("project_path_pattern")
         if project_pattern is None and project_path_pattern is None and projects:
             logger.warning(
-                "projects is deprecated, please use " "project_path_pattern instead."
+                "projects is deprecated, please use project_path_pattern instead."
             )
             logger.info("Initializing project_pattern from projects")
             values["project_pattern"] = AllowDenyPattern(
@@ -600,18 +708,18 @@ class DatabaseTable:
     """
 
     urn: str
-    id: Optional[
-        str
-    ] = None  # is not None only for tables that came from Tableau metadata
+    id: Optional[str] = (
+        None  # is not None only for tables that came from Tableau metadata
+    )
     num_cols: Optional[int] = None
 
-    paths: Optional[
-        Set[str]
-    ] = None  # maintains all browse paths encountered for this table
+    paths: Optional[Set[str]] = (
+        None  # maintains all browse paths encountered for this table
+    )
 
-    parsed_columns: Optional[
-        Set[str]
-    ] = None  # maintains all columns encountered for this table during parsing SQL queries
+    parsed_columns: Optional[Set[str]] = (
+        None  # maintains all columns encountered for this table during parsing SQL queries
+    )
 
     def update_table(
         self,
@@ -700,6 +808,23 @@ class TableauSourceReport(
         default_factory=(lambda: defaultdict(int))
     )
 
+    # Counters for tracking the number of queries made to get_connection_objects method
+    # by connection type (static and short set of keys):
+    # - num_queries_by_connection_type: total number of queries
+    # - num_filter_queries_by_connection_type: number of paginated queries due to splitting query filters
+    # - num_paginated_queries_by_connection_type: total number of queries due to Tableau pagination
+    # These counters are useful to understand the impact of changing the page size.
+
+    num_queries_by_connection_type: Dict[str, int] = dataclass_field(
+        default_factory=(lambda: defaultdict(int))
+    )
+    num_filter_queries_by_connection_type: Dict[str, int] = dataclass_field(
+        default_factory=(lambda: defaultdict(int))
+    )
+    num_paginated_queries_by_connection_type: Dict[str, int] = dataclass_field(
+        default_factory=(lambda: defaultdict(int))
+    )
+
 
 def report_user_role(report: TableauSourceReport, server: Server) -> None:
     title: str = "Insufficient Permissions"
@@ -994,7 +1119,9 @@ def maybe_parse_hostname():
             return server_connection
 
         for database_server in self.get_connection_objects(
-            database_servers_graphql_query, c.DATABASE_SERVERS_CONNECTION
+            query=database_servers_graphql_query,
+            connection_type=c.DATABASE_SERVERS_CONNECTION,
+            page_size=self.config.effective_database_server_page_size,
         ):
             database_server_id = database_server.get(c.ID)
             server_connection = database_server.get(c.HOST_NAME)
@@ -1020,23 +1147,36 @@ def fetch_projects():
                 )
             # Set parent project name
             for _project_id, project in all_project_map.items():
-                if (
-                    project.parent_id is not None
-                    and project.parent_id in all_project_map
-                ):
+                if project.parent_id is None:
+                    continue
+
+                if project.parent_id in all_project_map:
                     project.parent_name = all_project_map[project.parent_id].name
+                else:
+                    self.report.warning(
+                        title="Incomplete project hierarchy",
+                        message="Project details missing. Child projects will be ingested without reference to their parent project. We generally need Site Administrator Explorer permissions to extract the complete project hierarchy.",
+                        context=f"Missing {project.parent_id}, referenced by {project.id} {project.project_name}",
+                    )
+                    project.parent_id = None
+
+            # Post-condition
+            assert all(
+                [
+                    ((project.parent_id is None) == (project.parent_name is None))
+                    and (
+                        project.parent_id is None
+                        or project.parent_id in all_project_map
+                    )
+                    for project in all_project_map.values()
+                ]
+            ), "Parent project id and name should be consistent"
 
         def set_project_path():
             def form_path(project_id: str) -> List[str]:
                 cur_proj = all_project_map[project_id]
                 ancestors = [cur_proj.name]
                 while cur_proj.parent_id is not None:
-                    if cur_proj.parent_id not in all_project_map:
-                        self.report.warning(
-                            "project-issue",
-                            f"Parent project {cur_proj.parent_id} not found. We need Site Administrator Explorer permissions.",
-                        )
-                        break
                     cur_proj = all_project_map[cur_proj.parent_id]
                     ancestors = [cur_proj.name, *ancestors]
                 return ancestors
@@ -1420,22 +1560,30 @@ def get_connection_objects(
         self,
         query: str,
         connection_type: str,
+        page_size: int,
         query_filter: dict = {},
-        page_size_override: Optional[int] = None,
     ) -> Iterable[dict]:
         query_filter = optimize_query_filter(query_filter)
 
         # Calls the get_connection_object_page function to get the objects,
         # and automatically handles pagination.
-        page_size = page_size_override or self.config.page_size
 
         filter_pages = get_filter_pages(query_filter, page_size)
+        self.report.num_queries_by_connection_type[connection_type] += 1
+        self.report.num_filter_queries_by_connection_type[connection_type] += len(
+            filter_pages
+        )
+
         for filter_page in filter_pages:
             has_next_page = 1
             current_cursor: Optional[str] = None
             while has_next_page:
                 filter_: str = make_filter(filter_page)
 
+                self.report.num_paginated_queries_by_connection_type[
+                    connection_type
+                ] += 1
+
                 self.report.num_expected_tableau_metadata_queries += 1
                 (
                     connection_objects,
@@ -1463,10 +1611,10 @@ def emit_workbooks(self) -> Iterable[MetadataWorkUnit]:
             projects = {c.PROJECT_NAME_WITH_IN: project_names}
 
             for workbook in self.get_connection_objects(
-                workbook_graphql_query,
-                c.WORKBOOKS_CONNECTION,
-                projects,
-                page_size_override=self.config.workbook_page_size,
+                query=workbook_graphql_query,
+                connection_type=c.WORKBOOKS_CONNECTION,
+                query_filter=projects,
+                page_size=self.config.effective_workbook_page_size,
             ):
                 # This check is needed as we are using projectNameWithin which return project as per project name so if
                 # user want to ingest only nested project C from A->B->C then tableau might return more than one Project
@@ -1921,9 +2069,10 @@ def emit_custom_sql_datasources(self) -> Iterable[MetadataWorkUnit]:
 
         custom_sql_connection = list(
             self.get_connection_objects(
-                custom_sql_graphql_query,
-                c.CUSTOM_SQL_TABLE_CONNECTION,
-                custom_sql_filter,
+                query=custom_sql_graphql_query,
+                connection_type=c.CUSTOM_SQL_TABLE_CONNECTION,
+                query_filter=custom_sql_filter,
+                page_size=self.config.effective_custom_sql_table_page_size,
             )
         )
 
@@ -2174,8 +2323,7 @@ def _get_datasource_project_luid(self, ds: dict) -> Optional[str]:
             c.EMBEDDED_DATA_SOURCE,
         ):
             logger.debug(
-                f"datasource {ds.get(c.NAME)} type {ds.get(c.TYPE_NAME)} is "
-                f"unsupported"
+                f"datasource {ds.get(c.NAME)} type {ds.get(c.TYPE_NAME)} is unsupported"
             )
             return None
 
@@ -2357,9 +2505,9 @@ def parse_custom_sql(
     def _enrich_database_tables_with_parsed_schemas(
         self, parsing_result: SqlParsingResult
     ) -> None:
-        in_tables_schemas: Dict[
-            str, Set[str]
-        ] = transform_parsing_result_to_in_tables_schemas(parsing_result)
+        in_tables_schemas: Dict[str, Set[str]] = (
+            transform_parsing_result_to_in_tables_schemas(parsing_result)
+        )
 
         if not in_tables_schemas:
             logger.info("Unable to extract table schema from parsing result")
@@ -2632,6 +2780,7 @@ def update_datasource_for_field_upstream(
         self,
         datasource: dict,
         field_upstream_query: str,
+        page_size: int,
     ) -> dict:
         # Collect field ids to fetch field upstreams
         field_ids: List[str] = []
@@ -2642,9 +2791,10 @@ def update_datasource_for_field_upstream(
         # Fetch field upstreams and arrange them in map
         field_vs_upstream: Dict[str, dict] = {}
         for field_upstream in self.get_connection_objects(
-            field_upstream_query,
-            c.FIELDS_CONNECTION,
-            {c.ID_WITH_IN: field_ids},
+            query=field_upstream_query,
+            connection_type=c.FIELDS_CONNECTION,
+            query_filter={c.ID_WITH_IN: field_ids},
+            page_size=page_size,
         ):
             if field_upstream.get(c.ID):
                 field_id = field_upstream[c.ID]
@@ -2667,13 +2817,15 @@ def emit_published_datasources(self) -> Iterable[MetadataWorkUnit]:
         datasource_filter = {c.ID_WITH_IN: self.datasource_ids_being_used}
 
         for datasource in self.get_connection_objects(
-            published_datasource_graphql_query,
-            c.PUBLISHED_DATA_SOURCES_CONNECTION,
-            datasource_filter,
+            query=published_datasource_graphql_query,
+            connection_type=c.PUBLISHED_DATA_SOURCES_CONNECTION,
+            query_filter=datasource_filter,
+            page_size=self.config.effective_published_datasource_page_size,
         ):
             datasource = self.update_datasource_for_field_upstream(
                 datasource=datasource,
                 field_upstream_query=datasource_upstream_fields_graphql_query,
+                page_size=self.config.effective_published_datasource_field_upstream_page_size,
             )
 
             yield from self.emit_datasource(datasource)
@@ -2689,11 +2841,12 @@ def emit_upstream_tables(self) -> Iterable[MetadataWorkUnit]:
             c.ID_WITH_IN: list(tableau_database_table_id_to_urn_map.keys())
         }
 
-        # Emmitting tables that came from Tableau metadata
+        # Emitting tables that came from Tableau metadata
         for tableau_table in self.get_connection_objects(
-            database_tables_graphql_query,
-            c.DATABASE_TABLES_CONNECTION,
-            tables_filter,
+            query=database_tables_graphql_query,
+            connection_type=c.DATABASE_TABLES_CONNECTION,
+            query_filter=tables_filter,
+            page_size=self.config.effective_database_table_page_size,
         ):
             database_table = self.database_tables[
                 tableau_database_table_id_to_urn_map[tableau_table[c.ID]]
@@ -2882,9 +3035,10 @@ def emit_sheets(self) -> Iterable[MetadataWorkUnit]:
         sheets_filter = {c.ID_WITH_IN: self.sheet_ids}
 
         for sheet in self.get_connection_objects(
-            sheet_graphql_query,
-            c.SHEETS_CONNECTION,
-            sheets_filter,
+            query=sheet_graphql_query,
+            connection_type=c.SHEETS_CONNECTION,
+            query_filter=sheets_filter,
+            page_size=self.config.effective_sheet_page_size,
         ):
             if self.config.ingest_hidden_assets or not self._is_hidden_view(sheet):
                 yield from self.emit_sheets_as_charts(sheet, sheet.get(c.WORKBOOK))
@@ -3202,9 +3356,10 @@ def emit_dashboards(self) -> Iterable[MetadataWorkUnit]:
         dashboards_filter = {c.ID_WITH_IN: self.dashboard_ids}
 
         for dashboard in self.get_connection_objects(
-            dashboard_graphql_query,
-            c.DASHBOARDS_CONNECTION,
-            dashboards_filter,
+            query=dashboard_graphql_query,
+            connection_type=c.DASHBOARDS_CONNECTION,
+            query_filter=dashboards_filter,
+            page_size=self.config.effective_dashboard_page_size,
         ):
             if self.config.ingest_hidden_assets or not self._is_hidden_view(dashboard):
                 yield from self.emit_dashboard(dashboard, dashboard.get(c.WORKBOOK))
@@ -3349,13 +3504,15 @@ def emit_embedded_datasources(self) -> Iterable[MetadataWorkUnit]:
         datasource_filter = {c.ID_WITH_IN: self.embedded_datasource_ids_being_used}
 
         for datasource in self.get_connection_objects(
-            embedded_datasource_graphql_query,
-            c.EMBEDDED_DATA_SOURCES_CONNECTION,
-            datasource_filter,
+            query=embedded_datasource_graphql_query,
+            connection_type=c.EMBEDDED_DATA_SOURCES_CONNECTION,
+            query_filter=datasource_filter,
+            page_size=self.config.effective_embedded_datasource_page_size,
         ):
             datasource = self.update_datasource_for_field_upstream(
                 datasource=datasource,
                 field_upstream_query=datasource_upstream_fields_graphql_query,
+                page_size=self.config.effective_embedded_datasource_field_upstream_page_size,
             )
             yield from self.emit_datasource(
                 datasource,
@@ -3414,25 +3571,25 @@ def emit_project_in_topological_order(
 
             generated_project_keys.add(project_key.guid())
 
-            parent_project_key: Optional[
-                Union[ProjectKey, SiteKey]
-            ] = None  # It is going
+            parent_project_key: Optional[Union[ProjectKey, SiteKey]] = (
+                None  # It is going
+            )
             # to be used as a parent container key for the current tableau project
 
             if project_.parent_id is not None:
                 # Go to the parent project as we need to generate container first for parent
                 parent_project_key = self.gen_project_key(project_.parent_id)
 
-                parent_tableau_project: Optional[
-                    TableauProject
-                ] = self.tableau_project_registry.get(project_.parent_id)
+                parent_tableau_project: Optional[TableauProject] = (
+                    self.tableau_project_registry.get(project_.parent_id)
+                )
 
                 if (
                     parent_tableau_project is None
                 ):  # It is not in project registry because of project_pattern
-                    assert (
-                        project_.parent_name
-                    ), f"project {project_.name} should not be null"
+                    assert project_.parent_name, (
+                        f"project {project_.name} should not be null"
+                    )
                     parent_tableau_project = TableauProject(
                         id=project_.parent_id,
                         name=project_.parent_name,
@@ -3460,7 +3617,7 @@ def emit_project_in_topological_order(
                 parent_container_key=parent_project_key,
             )
 
-        for id_, project in self.tableau_project_registry.items():
+        for project in self.tableau_project_registry.values():
             logger.debug(
                 f"project {project.name} and it's parent {project.parent_name} and parent id {project.parent_id}"
             )
@@ -3524,16 +3681,16 @@ def ingest_tableau_site(self):
             if self.config.extract_usage_stats:
                 with PerfTimer() as timer:
                     self._populate_usage_stat_registry()
-                    self.report.extract_usage_stats_timer[
-                        self.site_content_url
-                    ] = timer.elapsed_seconds(digits=2)
+                    self.report.extract_usage_stats_timer[self.site_content_url] = (
+                        timer.elapsed_seconds(digits=2)
+                    )
 
             if self.config.permission_ingestion:
                 with PerfTimer() as timer:
                     self._fetch_groups()
-                    self.report.fetch_groups_timer[
-                        self.site_content_url
-                    ] = timer.elapsed_seconds(digits=2)
+                    self.report.fetch_groups_timer[self.site_content_url] = (
+                        timer.elapsed_seconds(digits=2)
+                    )
 
             # Populate the map of database names and database hostnames to be used later to map
             # databases to platform instances.
@@ -3546,9 +3703,9 @@ def ingest_tableau_site(self):
 
             with PerfTimer() as timer:
                 self._populate_projects_registry()
-                self.report.populate_projects_registry_timer[
-                    self.site_content_url
-                ] = timer.elapsed_seconds(digits=2)
+                self.report.populate_projects_registry_timer[self.site_content_url] = (
+                    timer.elapsed_seconds(digits=2)
+                )
 
             if self.config.add_site_container:
                 yield from self.emit_site_container()
@@ -3556,23 +3713,23 @@ def ingest_tableau_site(self):
 
             with PerfTimer() as timer:
                 yield from self.emit_workbooks()
-                self.report.emit_workbooks_timer[
-                    self.site_content_url
-                ] = timer.elapsed_seconds(digits=2)
+                self.report.emit_workbooks_timer[self.site_content_url] = (
+                    timer.elapsed_seconds(digits=2)
+                )
 
             if self.sheet_ids:
                 with PerfTimer() as timer:
                     yield from self.emit_sheets()
-                    self.report.emit_sheets_timer[
-                        self.site_content_url
-                    ] = timer.elapsed_seconds(digits=2)
+                    self.report.emit_sheets_timer[self.site_content_url] = (
+                        timer.elapsed_seconds(digits=2)
+                    )
 
             if self.dashboard_ids:
                 with PerfTimer() as timer:
                     yield from self.emit_dashboards()
-                    self.report.emit_dashboards_timer[
-                        self.site_content_url
-                    ] = timer.elapsed_seconds(digits=2)
+                    self.report.emit_dashboards_timer[self.site_content_url] = (
+                        timer.elapsed_seconds(digits=2)
+                    )
 
             if self.embedded_datasource_ids_being_used:
                 with PerfTimer() as timer:
@@ -3598,6 +3755,6 @@ def ingest_tableau_site(self):
             if self.database_tables:
                 with PerfTimer() as timer:
                     yield from self.emit_upstream_tables()
-                    self.report.emit_upstream_tables_timer[
-                        self.site_content_url
-                    ] = timer.elapsed_seconds(digits=2)
+                    self.report.emit_upstream_tables_timer[self.site_content_url] = (
+                        timer.elapsed_seconds(digits=2)
+                    )
diff --git a/metadata-ingestion/src/datahub/ingestion/source/tableau/tableau_common.py b/metadata-ingestion/src/datahub/ingestion/source/tableau/tableau_common.py
index 8f9d81eb9a18c1..5d5103330fe302 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/tableau/tableau_common.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/tableau/tableau_common.py
@@ -642,8 +642,11 @@ class TableauUpstreamReference:
 
     @classmethod
     def create(
-        cls, d: dict, default_schema_map: Optional[Dict[str, str]] = None
+        cls, d: Dict, default_schema_map: Optional[Dict[str, str]] = None
     ) -> "TableauUpstreamReference":
+        if d is None:
+            raise ValueError("TableauUpstreamReference.create: d is None")
+
         # Values directly from `table` object from Tableau
         database_dict = (
             d.get(c.DATABASE) or {}
@@ -717,7 +720,7 @@ def parse_full_name(full_name: Optional[str]) -> Optional[List[str]]:
         #  schema
 
         # TODO: Validate the startswith check. Currently required for our integration tests
-        if full_name is None or not full_name.startswith("["):
+        if full_name is None:
             return None
 
         return full_name.replace("[", "").replace("]", "").split(".")
diff --git a/metadata-ingestion/src/datahub/ingestion/source/unity/config.py b/metadata-ingestion/src/datahub/ingestion/source/unity/config.py
index 1fdce3aa1e2d34..6c3f7a51294797 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/unity/config.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/unity/config.py
@@ -254,7 +254,9 @@ class UnityCatalogSourceConfig(
     )
 
     # TODO: Remove `type:ignore` by refactoring config
-    profiling: Union[UnityCatalogGEProfilerConfig, UnityCatalogAnalyzeProfilerConfig] = Field(  # type: ignore
+    profiling: Union[
+        UnityCatalogGEProfilerConfig, UnityCatalogAnalyzeProfilerConfig
+    ] = Field(  # type: ignore
         default=UnityCatalogGEProfilerConfig(),
         description="Data profiling configuration",
         discriminator="method",
diff --git a/metadata-ingestion/src/datahub/ingestion/source/unity/proxy.py b/metadata-ingestion/src/datahub/ingestion/source/unity/proxy.py
index 9b96953794dcd5..86e577febf454c 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/unity/proxy.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/unity/proxy.py
@@ -26,7 +26,7 @@
 )
 from databricks.sdk.service.workspace import ObjectType
 
-import datahub
+from datahub._version import nice_version_name
 from datahub.emitter.mce_builder import parse_ts_millis
 from datahub.ingestion.source.unity.hive_metastore_proxy import HiveMetastoreProxy
 from datahub.ingestion.source.unity.proxy_profiling import (
@@ -103,7 +103,7 @@ def __init__(
             host=workspace_url,
             token=personal_access_token,
             product="datahub",
-            product_version=datahub.nice_version_name(),
+            product_version=nice_version_name(),
         )
         self.warehouse_id = warehouse_id or ""
         self.report = report
@@ -363,7 +363,7 @@ def _escape_sequence(value: str) -> str:
 
     @staticmethod
     def _create_metastore(
-        obj: Union[GetMetastoreSummaryResponse, MetastoreInfo]
+        obj: Union[GetMetastoreSummaryResponse, MetastoreInfo],
     ) -> Optional[Metastore]:
         if not obj.name:
             return None
diff --git a/metadata-ingestion/src/datahub/ingestion/source/unity/source.py b/metadata-ingestion/src/datahub/ingestion/source/unity/source.py
index 43bd788f809c3e..29562eaf3ce5b1 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/unity/source.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/unity/source.py
@@ -205,9 +205,9 @@ def __init__(self, ctx: PipelineContext, config: UnityCatalogSourceConfig):
         self.table_refs: Set[TableReference] = set()
         self.view_refs: Set[TableReference] = set()
         self.notebooks: FileBackedDict[Notebook] = FileBackedDict()
-        self.view_definitions: FileBackedDict[
-            Tuple[TableReference, str]
-        ] = FileBackedDict()
+        self.view_definitions: FileBackedDict[Tuple[TableReference, str]] = (
+            FileBackedDict()
+        )
 
         # Global map of tables, for profiling
         self.tables: FileBackedDict[Table] = FileBackedDict()
diff --git a/metadata-ingestion/src/datahub/ingestion/source/unity/usage.py b/metadata-ingestion/src/datahub/ingestion/source/unity/usage.py
index 718818d9b347bf..2e9f7fc00c8784 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/unity/usage.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/unity/usage.py
@@ -103,7 +103,9 @@ def _get_workunits_internal(
                                 query, table_info
                             )
                         for source_table in table_info.source_tables:
-                            with self.report.usage_perf_report.aggregator_add_event_timer:
+                            with (
+                                self.report.usage_perf_report.aggregator_add_event_timer
+                            ):
                                 self.usage_aggregator.aggregate_event(
                                     resource=source_table,
                                     start_time=query.start_time,
diff --git a/metadata-ingestion/src/datahub/ingestion/source/usage/clickhouse_usage.py b/metadata-ingestion/src/datahub/ingestion/source/usage/clickhouse_usage.py
index 4c2e4d42c440e8..2e1e315c4df956 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/usage/clickhouse_usage.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/usage/clickhouse_usage.py
@@ -213,15 +213,15 @@ def _get_joined_access_event(self, events):
     def _aggregate_access_events(
         self, events: List[ClickHouseJoinedAccessEvent]
     ) -> Dict[datetime, Dict[ClickHouseTableRef, AggregatedDataset]]:
-        datasets: Dict[
-            datetime, Dict[ClickHouseTableRef, AggregatedDataset]
-        ] = collections.defaultdict(dict)
+        datasets: Dict[datetime, Dict[ClickHouseTableRef, AggregatedDataset]] = (
+            collections.defaultdict(dict)
+        )
 
         for event in events:
             floored_ts = get_time_bucket(event.starttime, self.config.bucket_duration)
 
             resource = (
-                f'{self.config.platform_instance+"." if self.config.platform_instance else ""}'
+                f"{self.config.platform_instance + '.' if self.config.platform_instance else ''}"
                 f"{event.database}.{event.table}"
             )
 
diff --git a/metadata-ingestion/src/datahub/ingestion/source/usage/starburst_trino_usage.py b/metadata-ingestion/src/datahub/ingestion/source/usage/starburst_trino_usage.py
index 6ded11027c83a8..e4138696186416 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/usage/starburst_trino_usage.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/usage/starburst_trino_usage.py
@@ -235,9 +235,9 @@ def _get_joined_access_event(self, events):
     def _aggregate_access_events(
         self, events: List[TrinoJoinedAccessEvent]
     ) -> Dict[datetime, Dict[TrinoTableRef, AggregatedDataset]]:
-        datasets: Dict[
-            datetime, Dict[TrinoTableRef, AggregatedDataset]
-        ] = collections.defaultdict(dict)
+        datasets: Dict[datetime, Dict[TrinoTableRef, AggregatedDataset]] = (
+            collections.defaultdict(dict)
+        )
 
         for event in events:
             floored_ts = get_time_bucket(event.starttime, self.config.bucket_duration)
diff --git a/metadata-ingestion/src/datahub/ingestion/source/usage/usage_common.py b/metadata-ingestion/src/datahub/ingestion/source/usage/usage_common.py
index 95c2345232a1ee..73e7e415e2b9eb 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/usage/usage_common.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/usage/usage_common.py
@@ -89,7 +89,7 @@ def make_usage_workunit(
     top_sql_queries: Optional[List[str]] = None
     if query_freq is not None:
         if top_n_queries < len(query_freq):
-            logger.warn(
+            logger.warning(
                 f"Top N query limit exceeded on {str(resource)}.  Max number of queries {top_n_queries} <  {len(query_freq)}. Truncating top queries to {top_n_queries}."
             )
             query_freq = query_freq[0:top_n_queries]
diff --git a/metadata-ingestion/src/datahub/ingestion/transformer/add_dataset_dataproduct.py b/metadata-ingestion/src/datahub/ingestion/transformer/add_dataset_dataproduct.py
index bb1c297513de10..b4dc8835f9fba9 100644
--- a/metadata-ingestion/src/datahub/ingestion/transformer/add_dataset_dataproduct.py
+++ b/metadata-ingestion/src/datahub/ingestion/transformer/add_dataset_dataproduct.py
@@ -80,10 +80,10 @@ def handle_end_of_stream(
                         ).add_asset(container_urn)
                         data_products_container[data_product_urn] = container_product
                     else:
-                        data_products_container[
-                            data_product_urn
-                        ] = data_products_container[data_product_urn].add_asset(
-                            container_urn
+                        data_products_container[data_product_urn] = (
+                            data_products_container[data_product_urn].add_asset(
+                                container_urn
+                            )
                         )
 
         mcps: List[
diff --git a/metadata-ingestion/src/datahub/ingestion/transformer/add_dataset_properties.py b/metadata-ingestion/src/datahub/ingestion/transformer/add_dataset_properties.py
index 668f6ed7abe074..4b9b4c9e6f5da6 100644
--- a/metadata-ingestion/src/datahub/ingestion/transformer/add_dataset_properties.py
+++ b/metadata-ingestion/src/datahub/ingestion/transformer/add_dataset_properties.py
@@ -61,9 +61,9 @@ def _merge_with_server_properties(
     ) -> Optional[DatasetPropertiesClass]:
         assert dataset_properties_aspect
 
-        server_dataset_properties_aspect: Optional[
-            DatasetPropertiesClass
-        ] = graph.get_dataset_properties(entity_urn)
+        server_dataset_properties_aspect: Optional[DatasetPropertiesClass] = (
+            graph.get_dataset_properties(entity_urn)
+        )
         # No need to take any action if server properties is None or there is not customProperties in server properties
         if (
             server_dataset_properties_aspect is None
diff --git a/metadata-ingestion/src/datahub/ingestion/transformer/add_dataset_schema_tags.py b/metadata-ingestion/src/datahub/ingestion/transformer/add_dataset_schema_tags.py
index ba3b6508daaecd..d2687ebc5e76f6 100644
--- a/metadata-ingestion/src/datahub/ingestion/transformer/add_dataset_schema_tags.py
+++ b/metadata-ingestion/src/datahub/ingestion/transformer/add_dataset_schema_tags.py
@@ -89,9 +89,9 @@ def transform_aspect(
         server_field_map: dict = {}
         if self.config.semantics == TransformerSemantics.PATCH:
             assert self.ctx.graph
-            server_schema_metadata_aspect: Optional[
-                SchemaMetadataClass
-            ] = self.ctx.graph.get_schema_metadata(entity_urn=entity_urn)
+            server_schema_metadata_aspect: Optional[SchemaMetadataClass] = (
+                self.ctx.graph.get_schema_metadata(entity_urn=entity_urn)
+            )
             if server_schema_metadata_aspect is not None:
                 if not schema_metadata_aspect:
                     schema_metadata_aspect = server_schema_metadata_aspect
diff --git a/metadata-ingestion/src/datahub/ingestion/transformer/add_dataset_schema_terms.py b/metadata-ingestion/src/datahub/ingestion/transformer/add_dataset_schema_terms.py
index a7e92d4bd7edbd..d17a39bee6cfbf 100644
--- a/metadata-ingestion/src/datahub/ingestion/transformer/add_dataset_schema_terms.py
+++ b/metadata-ingestion/src/datahub/ingestion/transformer/add_dataset_schema_terms.py
@@ -108,9 +108,9 @@ def transform_aspect(
         ] = {}  # Map to cache server field objects, where fieldPath is key
         if self.config.semantics == TransformerSemantics.PATCH:
             assert self.ctx.graph
-            server_schema_metadata_aspect: Optional[
-                SchemaMetadataClass
-            ] = self.ctx.graph.get_schema_metadata(entity_urn=entity_urn)
+            server_schema_metadata_aspect: Optional[SchemaMetadataClass] = (
+                self.ctx.graph.get_schema_metadata(entity_urn=entity_urn)
+            )
             if server_schema_metadata_aspect is not None:
                 if not schema_metadata_aspect:
                     schema_metadata_aspect = server_schema_metadata_aspect
diff --git a/metadata-ingestion/src/datahub/ingestion/transformer/dataset_domain_based_on_tags.py b/metadata-ingestion/src/datahub/ingestion/transformer/dataset_domain_based_on_tags.py
index 7be8069e1b0852..bb2f318dcac8b8 100644
--- a/metadata-ingestion/src/datahub/ingestion/transformer/dataset_domain_based_on_tags.py
+++ b/metadata-ingestion/src/datahub/ingestion/transformer/dataset_domain_based_on_tags.py
@@ -60,10 +60,10 @@ def transform_aspect(
                 domain_aspect.domains.extend(mapped_domains.domains)
                 if self.config.semantics == TransformerSemantics.PATCH:
                     # Try merging with server-side domains
-                    patch_domain_aspect: Optional[
-                        DomainsClass
-                    ] = AddDatasetDomain._merge_with_server_domains(
-                        self.ctx.graph, entity_urn, domain_aspect
+                    patch_domain_aspect: Optional[DomainsClass] = (
+                        AddDatasetDomain._merge_with_server_domains(
+                            self.ctx.graph, entity_urn, domain_aspect
+                        )
                     )
                     return cast(Optional[Aspect], patch_domain_aspect)
                 return cast(Optional[Aspect], domain_aspect)
diff --git a/metadata-ingestion/src/datahub/ingestion/transformer/extract_ownership_from_tags.py b/metadata-ingestion/src/datahub/ingestion/transformer/extract_ownership_from_tags.py
index 212e018dd64fb7..32707dcd3a372f 100644
--- a/metadata-ingestion/src/datahub/ingestion/transformer/extract_ownership_from_tags.py
+++ b/metadata-ingestion/src/datahub/ingestion/transformer/extract_ownership_from_tags.py
@@ -141,9 +141,9 @@ def transform_aspect(
                 else:
                     owner_type = get_owner_type(self.config.owner_type)
                     if owner_type == OwnershipTypeClass.CUSTOM:
-                        assert (
-                            self.config.owner_type_urn is not None
-                        ), "owner_type_urn must be set if owner_type is CUSTOM"
+                        assert self.config.owner_type_urn is not None, (
+                            "owner_type_urn must be set if owner_type is CUSTOM"
+                        )
 
                     owners.append(
                         OwnerClass(
diff --git a/metadata-ingestion/src/datahub/ingestion/transformer/tags_to_terms.py b/metadata-ingestion/src/datahub/ingestion/transformer/tags_to_terms.py
index 7e6125079f16e3..65cf2ac3614ae0 100644
--- a/metadata-ingestion/src/datahub/ingestion/transformer/tags_to_terms.py
+++ b/metadata-ingestion/src/datahub/ingestion/transformer/tags_to_terms.py
@@ -92,9 +92,9 @@ def transform_aspect(
         in_global_tags_aspect: Optional[GlobalTagsClass] = self.ctx.graph.get_tags(
             entity_urn
         )
-        in_schema_metadata_aspect: Optional[
-            SchemaMetadataClass
-        ] = self.ctx.graph.get_schema_metadata(entity_urn)
+        in_schema_metadata_aspect: Optional[SchemaMetadataClass] = (
+            self.ctx.graph.get_schema_metadata(entity_urn)
+        )
 
         if in_global_tags_aspect is None and in_schema_metadata_aspect is None:
             return cast(Aspect, in_glossary_terms)
@@ -134,10 +134,10 @@ def transform_aspect(
         )
 
         if self.config.semantics == TransformerSemantics.PATCH:
-            patch_glossary_terms: Optional[
-                GlossaryTermsClass
-            ] = TagsToTermMapper._merge_with_server_glossary_terms(
-                self.ctx.graph, entity_urn, out_glossary_terms
+            patch_glossary_terms: Optional[GlossaryTermsClass] = (
+                TagsToTermMapper._merge_with_server_glossary_terms(
+                    self.ctx.graph, entity_urn, out_glossary_terms
+                )
             )
             return cast(Optional[Aspect], patch_glossary_terms)
         else:
diff --git a/metadata-ingestion/src/datahub/integrations/assertion/snowflake/compiler.py b/metadata-ingestion/src/datahub/integrations/assertion/snowflake/compiler.py
index 8d2ae2960ebd05..e32f1ddc3943ae 100644
--- a/metadata-ingestion/src/datahub/integrations/assertion/snowflake/compiler.py
+++ b/metadata-ingestion/src/datahub/integrations/assertion/snowflake/compiler.py
@@ -61,17 +61,17 @@ def __init__(self, output_dir: str, extras: Dict[str, str]) -> None:
     def create(
         cls, output_dir: str, extras: Dict[str, str]
     ) -> "SnowflakeAssertionCompiler":
-        assert os.path.exists(
-            output_dir
-        ), f"Specified location {output_dir} does not exist."
+        assert os.path.exists(output_dir), (
+            f"Specified location {output_dir} does not exist."
+        )
 
-        assert os.path.isdir(
-            output_dir
-        ), f"Specified location {output_dir} is not a folder."
+        assert os.path.isdir(output_dir), (
+            f"Specified location {output_dir} is not a folder."
+        )
 
-        assert any(
-            x.upper() == DMF_SCHEMA_PROPERTY_KEY for x in extras
-        ), "Must specify value for DMF schema using -x DMF_SCHEMA=<db.schema>"
+        assert any(x.upper() == DMF_SCHEMA_PROPERTY_KEY for x in extras), (
+            "Must specify value for DMF schema using -x DMF_SCHEMA=<db.schema>"
+        )
 
         return SnowflakeAssertionCompiler(output_dir, extras)
 
@@ -232,6 +232,6 @@ def get_dmf_schedule(trigger: AssertionTrigger) -> str:
     elif isinstance(trigger.trigger, CronTrigger):
         return f"USING CRON {trigger.trigger.cron} {trigger.trigger.timezone}"
     elif isinstance(trigger.trigger, IntervalTrigger):
-        return f"{trigger.trigger.interval.seconds/60} MIN"
+        return f"{trigger.trigger.interval.seconds / 60} MIN"
     else:
         raise ValueError(f"Unsupported trigger type {type(trigger.trigger)}")
diff --git a/metadata-ingestion/src/datahub/lite/duckdb_lite.py b/metadata-ingestion/src/datahub/lite/duckdb_lite.py
index 89317383520923..fe025842822b13 100644
--- a/metadata-ingestion/src/datahub/lite/duckdb_lite.py
+++ b/metadata-ingestion/src/datahub/lite/duckdb_lite.py
@@ -163,9 +163,9 @@ def write(
 
                 if "properties" not in writeable_dict["systemMetadata"]:
                     writeable_dict["systemMetadata"]["properties"] = {}
-                writeable_dict["systemMetadata"]["properties"][
-                    "sysVersion"
-                ] = new_version
+                writeable_dict["systemMetadata"]["properties"]["sysVersion"] = (
+                    new_version
+                )
                 if needs_write:
                     self.duckdb_client.execute(
                         query="INSERT INTO metadata_aspect_v2 VALUES (?, ?, ?, ?, ?, ?)",
@@ -208,9 +208,9 @@ def write(
                             "lastObserved": writeable.systemMetadata.lastObserved
                         }
                     else:
-                        system_metadata[
-                            "lastObserved"
-                        ] = writeable.systemMetadata.lastObserved
+                        system_metadata["lastObserved"] = (
+                            writeable.systemMetadata.lastObserved
+                        )
                     self.duckdb_client.execute(
                         query="UPDATE metadata_aspect_v2 SET system_metadata = ? WHERE urn = ? AND aspect_name = ? AND version = 0",
                         parameters=[
@@ -497,9 +497,9 @@ def get_all_entities(
             aspect_name = r[1]
             aspect_payload = json.loads(r[2])
             if typed:
-                assert (
-                    aspect_name in ASPECT_MAP
-                ), f"Missing aspect name {aspect_name} in the registry"
+                assert aspect_name in ASPECT_MAP, (
+                    f"Missing aspect name {aspect_name} in the registry"
+                )
                 try:
                     aspect_payload = ASPECT_MAP[aspect_name].from_obj(
                         post_json_transform(aspect_payload)
@@ -531,7 +531,9 @@ def get_all_aspects(self) -> Iterable[MetadataChangeProposalWrapper]:
         for r in results.fetchall():
             urn = r[0]
             aspect_name = r[1]
-            aspect_metadata = ASPECT_MAP[aspect_name].from_obj(post_json_transform(json.loads(r[2])))  # type: ignore
+            aspect_metadata = ASPECT_MAP[aspect_name].from_obj(
+                post_json_transform(json.loads(r[2]))
+            )  # type: ignore
             system_metadata = SystemMetadataClass.from_obj(json.loads(r[3]))
             mcp = MetadataChangeProposalWrapper(
                 entityUrn=urn,
diff --git a/metadata-ingestion/src/datahub/secret/datahub_secrets_client.py b/metadata-ingestion/src/datahub/secret/datahub_secrets_client.py
index c60aeff5db2f3f..6cd542c03a4260 100644
--- a/metadata-ingestion/src/datahub/secret/datahub_secrets_client.py
+++ b/metadata-ingestion/src/datahub/secret/datahub_secrets_client.py
@@ -11,34 +11,25 @@ class DataHubSecretsClient:
     def __init__(self, graph: DataHubGraph):
         self.graph = graph
 
+    def _cleanup_secret_name(self, secret_names: List[str]) -> List[str]:
+        """Remove empty strings from the list of secret names."""
+        return [secret_name for secret_name in secret_names if secret_name]
+
     def get_secret_values(self, secret_names: List[str]) -> Dict[str, Optional[str]]:
         if len(secret_names) == 0:
             return {}
 
-        request_json = {
-            "query": """query getSecretValues($input: GetSecretValuesInput!) {\n
-                getSecretValues(input: $input) {\n
-                    name\n
-                    value\n
-                }\n
+        res_data = self.graph.execute_graphql(
+            query="""query getSecretValues($input: GetSecretValuesInput!) {
+                getSecretValues(input: $input) {
+                    name
+                    value
+                }
             }""",
-            "variables": {"input": {"secrets": secret_names}},
-        }
-        # TODO: Use graph.execute_graphql() instead.
-
-        # Fetch secrets using GraphQL API f
-        response = self.graph._session.post(
-            f"{self.graph.config.server}/api/graphql", json=request_json
+            variables={"input": {"secrets": self._cleanup_secret_name(secret_names)}},
         )
-        response.raise_for_status()
-
-        # Verify response
-        res_data = response.json()
-        if "errors" in res_data:
-            raise Exception("Failed to retrieve secrets from DataHub.")
-
         # Convert list of name, value secret pairs into a dict and return
-        secret_value_list = res_data["data"]["getSecretValues"]
+        secret_value_list = res_data["getSecretValues"]
         secret_value_dict = dict()
         for secret_value in secret_value_list:
             secret_value_dict[secret_value["name"]] = secret_value["value"]
diff --git a/metadata-ingestion/src/datahub/secret/secret_common.py b/metadata-ingestion/src/datahub/secret/secret_common.py
index 2f7a584d875380..a116c70407af23 100644
--- a/metadata-ingestion/src/datahub/secret/secret_common.py
+++ b/metadata-ingestion/src/datahub/secret/secret_common.py
@@ -2,10 +2,7 @@
 import logging
 from typing import List
 
-from datahub.configuration.config_loader import (
-    list_referenced_env_variables,
-    resolve_env_variables,
-)
+from datahub.configuration.config_loader import EnvResolver
 from datahub.secret.secret_store import SecretStore
 
 logger = logging.getLogger(__name__)
@@ -42,18 +39,27 @@ def resolve_secrets(secret_names: List[str], secret_stores: List[SecretStore]) -
     return final_secret_values
 
 
-def resolve_recipe(recipe: str, secret_stores: List[SecretStore]) -> dict:
+def resolve_recipe(
+    recipe: str, secret_stores: List[SecretStore], strict_env_syntax: bool = True
+) -> dict:
+    # Note: the default for `strict_env_syntax` is normally False, but here we override
+    # it to be true. Particularly when fetching secrets from external secret stores, we
+    # want to be more careful about not over-fetching secrets.
+
     json_recipe_raw = json.loads(recipe)
 
     # 1. Extract all secrets needing resolved.
-    secrets_to_resolve = list_referenced_env_variables(json_recipe_raw)
+    secrets_to_resolve = EnvResolver.list_referenced_variables(
+        json_recipe_raw, strict_env_syntax=strict_env_syntax
+    )
 
     # 2. Resolve secret values
     secret_values_dict = resolve_secrets(list(secrets_to_resolve), secret_stores)
 
     # 3. Substitute secrets into recipe file
-    json_recipe_resolved = resolve_env_variables(
-        json_recipe_raw, environ=secret_values_dict
+    resolver = EnvResolver(
+        environ=secret_values_dict, strict_env_syntax=strict_env_syntax
     )
+    json_recipe_resolved = resolver.resolve(json_recipe_raw)
 
     return json_recipe_resolved
diff --git a/metadata-ingestion/src/datahub/specific/aspect_helpers/custom_properties.py b/metadata-ingestion/src/datahub/specific/aspect_helpers/custom_properties.py
index 1fd1585a913581..4b8b4d0bc99bc0 100644
--- a/metadata-ingestion/src/datahub/specific/aspect_helpers/custom_properties.py
+++ b/metadata-ingestion/src/datahub/specific/aspect_helpers/custom_properties.py
@@ -9,8 +9,7 @@
 class HasCustomPropertiesPatch(MetadataPatchProposal):
     @classmethod
     @abstractmethod
-    def _custom_properties_location(self) -> Tuple[str, PatchPath]:
-        ...
+    def _custom_properties_location(self) -> Tuple[str, PatchPath]: ...
 
     def add_custom_property(self, key: str, value: str) -> Self:
         """Add a custom property to the entity.
diff --git a/metadata-ingestion/src/datahub/sql_parsing/schema_resolver.py b/metadata-ingestion/src/datahub/sql_parsing/schema_resolver.py
index 6aa10381a883ef..55b026a144c6d5 100644
--- a/metadata-ingestion/src/datahub/sql_parsing/schema_resolver.py
+++ b/metadata-ingestion/src/datahub/sql_parsing/schema_resolver.py
@@ -33,14 +33,11 @@ class GraphQLSchemaMetadata(TypedDict):
 
 class SchemaResolverInterface(Protocol):
     @property
-    def platform(self) -> str:
-        ...
+    def platform(self) -> str: ...
 
-    def includes_temp_tables(self) -> bool:
-        ...
+    def includes_temp_tables(self) -> bool: ...
 
-    def resolve_table(self, table: _TableName) -> Tuple[str, Optional[SchemaInfo]]:
-        ...
+    def resolve_table(self, table: _TableName) -> Tuple[str, Optional[SchemaInfo]]: ...
 
     def __hash__(self) -> int:
         # Mainly to make lru_cache happy in methods that accept a schema resolver.
@@ -232,8 +229,7 @@ def convert_graphql_schema_metadata_to_info(
         return {
             get_simple_field_path_from_v2_field_path(field["fieldPath"]): (
                 # The actual types are more of a "nice to have".
-                field["nativeDataType"]
-                or "str"
+                field["nativeDataType"] or "str"
             )
             for field in schema["fields"]
             # TODO: We can't generate lineage to columns nested within structs yet.
@@ -289,8 +285,7 @@ def _convert_schema_field_list_to_info(
     return {
         get_simple_field_path_from_v2_field_path(col.fieldPath): (
             # The actual types are more of a "nice to have".
-            col.nativeDataType
-            or "str"
+            col.nativeDataType or "str"
         )
         for col in schema_fields
         # TODO: We can't generate lineage to columns nested within structs yet.
diff --git a/metadata-ingestion/src/datahub/sql_parsing/sql_parsing_aggregator.py b/metadata-ingestion/src/datahub/sql_parsing/sql_parsing_aggregator.py
index 25b63ffac45f96..893e89f177094f 100644
--- a/metadata-ingestion/src/datahub/sql_parsing/sql_parsing_aggregator.py
+++ b/metadata-ingestion/src/datahub/sql_parsing/sql_parsing_aggregator.py
@@ -2,7 +2,6 @@
 import dataclasses
 import enum
 import functools
-import itertools
 import json
 import logging
 import os
@@ -63,6 +62,7 @@
     FileBackedDict,
     FileBackedList,
 )
+from datahub.utilities.groupby import groupby_unsorted
 from datahub.utilities.lossy_collections import LossyDict, LossyList
 from datahub.utilities.ordered_set import OrderedSet
 from datahub.utilities.perf_timer import PerfTimer
@@ -284,6 +284,7 @@ class SqlAggregatorReport(Report):
 
     # Queries.
     num_queries_entities_generated: int = 0
+    num_queries_used_in_lineage: Optional[int] = None
     num_queries_skipped_due_to_filters: int = 0
 
     # Usage-related.
@@ -681,10 +682,10 @@ def add_known_lineage_mapping(
         query_id = self._known_lineage_query_id()
 
         # Generate CLL if schema of downstream is known
-        column_lineage: List[
-            ColumnLineageInfo
-        ] = self._generate_identity_column_lineage(
-            upstream_urn=upstream_urn, downstream_urn=downstream_urn
+        column_lineage: List[ColumnLineageInfo] = (
+            self._generate_identity_column_lineage(
+                upstream_urn=upstream_urn, downstream_urn=downstream_urn
+            )
         )
 
         # Register the query.
@@ -1043,9 +1044,9 @@ def _make_schema_resolver_for_session(
             temp_table_schemas: Dict[str, Optional[List[models.SchemaFieldClass]]] = {}
             for temp_table_urn, query_ids in self._temp_lineage_map[session_id].items():
                 for query_id in query_ids:
-                    temp_table_schemas[
-                        temp_table_urn
-                    ] = self._inferred_temp_schemas.get(query_id)
+                    temp_table_schemas[temp_table_urn] = (
+                        self._inferred_temp_schemas.get(query_id)
+                    )
                     if temp_table_schemas:
                         break
 
@@ -1072,9 +1073,9 @@ def _process_view_definition(
             schema_resolver=self._schema_resolver,
         )
         if parsed.debug_info.error:
-            self.report.views_parse_failures[
-                view_urn
-            ] = f"{parsed.debug_info.error} on query: {view_definition.view_definition[:100]}"
+            self.report.views_parse_failures[view_urn] = (
+                f"{parsed.debug_info.error} on query: {view_definition.view_definition[:100]}"
+            )
         if parsed.debug_info.table_error:
             self.report.num_views_failed += 1
             return  # we can't do anything with this query
@@ -1200,6 +1201,7 @@ def gen_metadata(self) -> Iterable[MetadataChangeProposalWrapper]:
         queries_generated: Set[QueryId] = set()
 
         yield from self._gen_lineage_mcps(queries_generated)
+        self.report.num_queries_used_in_lineage = len(queries_generated)
         yield from self._gen_usage_statistics_mcps()
         yield from self._gen_operation_mcps(queries_generated)
         yield from self._gen_remaining_queries(queries_generated)
@@ -1312,8 +1314,8 @@ def _gen_lineage_for_downstream(
         upstream_aspect.fineGrainedLineages = []
         for downstream_column, all_upstream_columns in cll.items():
             # Group by query ID.
-            for query_id, upstream_columns_for_query in itertools.groupby(
-                sorted(all_upstream_columns.items(), key=lambda x: x[1]),
+            for query_id, upstream_columns_for_query in groupby_unsorted(
+                all_upstream_columns.items(),
                 key=lambda x: x[1],
             ):
                 upstream_columns = [x[0] for x in upstream_columns_for_query]
@@ -1581,9 +1583,9 @@ def _recurse_into_query(
                                     temp_query_lineage_info
                                 )
                             else:
-                                temp_upstream_queries[
-                                    upstream
-                                ] = temp_query_lineage_info
+                                temp_upstream_queries[upstream] = (
+                                    temp_query_lineage_info
+                                )
 
             # Compute merged upstreams.
             new_upstreams = OrderedSet[UrnStr]()
@@ -1663,9 +1665,9 @@ def _recurse_into_query(
         composed_of_queries_truncated: LossyList[str] = LossyList()
         for query_id in composed_of_queries:
             composed_of_queries_truncated.append(query_id)
-        self.report.queries_with_temp_upstreams[
-            composite_query_id
-        ] = composed_of_queries_truncated
+        self.report.queries_with_temp_upstreams[composite_query_id] = (
+            composed_of_queries_truncated
+        )
 
         merged_query_text = ";\n\n".join(
             [q.formatted_query_string for q in ordered_queries]
diff --git a/metadata-ingestion/src/datahub/sql_parsing/sqlglot_lineage.py b/metadata-ingestion/src/datahub/sql_parsing/sqlglot_lineage.py
index bf28ab0e7b229b..c825deeccd9592 100644
--- a/metadata-ingestion/src/datahub/sql_parsing/sqlglot_lineage.py
+++ b/metadata-ingestion/src/datahub/sql_parsing/sqlglot_lineage.py
@@ -442,9 +442,9 @@ def _create_table_ddl_cll(
 ) -> List[_ColumnLineageInfo]:
     column_lineage: List[_ColumnLineageInfo] = []
 
-    assert (
-        output_table is not None
-    ), "output_table must be set for create DDL statements"
+    assert output_table is not None, (
+        "output_table must be set for create DDL statements"
+    )
 
     create_schema: sqlglot.exp.Schema = statement.this
     sqlglot_columns = create_schema.expressions
diff --git a/metadata-ingestion/src/datahub/sql_parsing/sqlglot_utils.py b/metadata-ingestion/src/datahub/sql_parsing/sqlglot_utils.py
index 57a5cc3c9a6574..5b12c64a831666 100644
--- a/metadata-ingestion/src/datahub/sql_parsing/sqlglot_utils.py
+++ b/metadata-ingestion/src/datahub/sql_parsing/sqlglot_utils.py
@@ -404,7 +404,7 @@ def replace_cte_refs(node: sqlglot.exp.Expression) -> sqlglot.exp.Expression:
         if new_statement == statement:
             if iteration > 1:
                 logger.debug(
-                    f"Required {iteration+1} iterations to detach and eliminate all CTEs"
+                    f"Required {iteration + 1} iterations to detach and eliminate all CTEs"
                 )
             break
         statement = new_statement
diff --git a/metadata-ingestion/src/datahub/telemetry/stats.py b/metadata-ingestion/src/datahub/telemetry/stats.py
index bf98bd72b574ce..d6835e49de56aa 100644
--- a/metadata-ingestion/src/datahub/telemetry/stats.py
+++ b/metadata-ingestion/src/datahub/telemetry/stats.py
@@ -5,8 +5,7 @@
 
 
 class SupportsLT(Protocol):
-    def __lt__(self, __other: Any) -> Any:
-        ...
+    def __lt__(self, __other: Any) -> Any: ...
 
 
 _SupportsComparisonT = TypeVar("_SupportsComparisonT", bound=SupportsLT)
diff --git a/metadata-ingestion/src/datahub/telemetry/telemetry.py b/metadata-ingestion/src/datahub/telemetry/telemetry.py
index 22b2cb6a101af9..c64c133fbf456f 100644
--- a/metadata-ingestion/src/datahub/telemetry/telemetry.py
+++ b/metadata-ingestion/src/datahub/telemetry/telemetry.py
@@ -12,7 +12,7 @@
 from mixpanel import Consumer, Mixpanel
 from typing_extensions import ParamSpec
 
-import datahub as datahub_package
+from datahub._version import __version__, nice_version_name
 from datahub.cli.config_utils import DATAHUB_ROOT_FOLDER
 from datahub.cli.env_utils import get_boolean_env_variable
 from datahub.configuration.common import ExceptionWithProps
@@ -106,7 +106,7 @@
 
 def _default_telemetry_properties() -> Dict[str, Any]:
     return {
-        "datahub_version": datahub_package.nice_version_name(),
+        "datahub_version": nice_version_name(),
         "python_version": platform.python_version(),
         "os": platform.system(),
         "arch": platform.machine(),
@@ -132,7 +132,7 @@ def __init__(self):
                 sentry_sdk.init(
                     dsn=SENTRY_DSN,
                     environment=SENTRY_ENVIRONMENT,
-                    release=datahub_package.__version__,
+                    release=__version__,
                 )
             except Exception as e:
                 # We need to print initialization errors to stderr, since logger is not initialized yet
@@ -277,7 +277,7 @@ def init_capture_exception(self) -> None:
                 "environment",
                 {
                     "environment": SENTRY_ENVIRONMENT,
-                    "datahub_version": datahub_package.nice_version_name(),
+                    "datahub_version": nice_version_name(),
                     "os": platform.system(),
                     "python_version": platform.python_version(),
                 },
diff --git a/metadata-ingestion/src/datahub/testing/check_imports.py b/metadata-ingestion/src/datahub/testing/check_imports.py
index e4bf07882b36ae..b65b3aa90dca35 100644
--- a/metadata-ingestion/src/datahub/testing/check_imports.py
+++ b/metadata-ingestion/src/datahub/testing/check_imports.py
@@ -1,4 +1,5 @@
 import pathlib
+import re
 from typing import List
 
 
@@ -32,3 +33,30 @@ def ensure_no_indirect_model_imports(dirs: List[pathlib.Path]) -> None:
                                 f"Disallowed import found in {file}: `{line.rstrip()}`. "
                                 f"Import from {replacement} instead."
                             )
+
+
+def ban_direct_datahub_imports(dirs: List[pathlib.Path]) -> None:
+    # We also want to ban all direct imports of datahub.
+    # The base `datahub` package is used to export public-facing classes.
+    # If we import it directly, we'll likely end up with circular imports.
+
+    banned_strings = [
+        r"^import datahub[\s$]",
+        r"^from datahub import",
+    ]
+    ignored_files = {
+        __file__,
+    }
+    for dir in dirs:
+        for file in dir.rglob("*.py"):
+            if str(file) in ignored_files:
+                continue
+
+            file_contents = file.read_text()
+
+            for banned_string in banned_strings:
+                if re.search(banned_string, file_contents, re.MULTILINE):
+                    raise ValueError(
+                        f"Disallowed bare datahub import found in {file}. "
+                        f"Do not import datahub directly; instead import from the underlying file."
+                    )
diff --git a/metadata-ingestion/src/datahub/testing/mcp_diff.py b/metadata-ingestion/src/datahub/testing/mcp_diff.py
index 5e669a718e9ad3..b58afc10148edc 100644
--- a/metadata-ingestion/src/datahub/testing/mcp_diff.py
+++ b/metadata-ingestion/src/datahub/testing/mcp_diff.py
@@ -246,7 +246,7 @@ def pretty(self, verbose: bool = False) -> str:
         for urn in self.aspect_changes.keys() - self.urns_added - self.urns_removed:
             aspect_map = self.aspect_changes[urn]
             s.append(f"Urn changed, {urn}:")
-            for aspect_name, aspect_diffs in aspect_map.items():
+            for aspect_diffs in aspect_map.values():
                 for i, ga in aspect_diffs.aspects_added.items():
                     s.append(self.report_aspect(ga, i, "added"))
                     if verbose:
diff --git a/metadata-ingestion/src/datahub/upgrade/upgrade.py b/metadata-ingestion/src/datahub/upgrade/upgrade.py
index fb14514588e5fc..276f4ccd54a4a1 100644
--- a/metadata-ingestion/src/datahub/upgrade/upgrade.py
+++ b/metadata-ingestion/src/datahub/upgrade/upgrade.py
@@ -10,7 +10,7 @@
 from packaging.version import Version
 from pydantic import BaseModel
 
-from datahub import __version__
+from datahub._version import __version__
 from datahub.cli.config_utils import load_client_config
 from datahub.ingestion.graph.client import DataHubGraph
 from datahub.utilities.perf_timer import PerfTimer
@@ -93,11 +93,11 @@ async def get_github_stats():
     async with aiohttp.ClientSession(
         headers={"Accept": "application/vnd.github.v3+json"}
     ) as session:
-        gh_url = "https://api.github.com/repos/datahub-project/datahub/releases"
+        gh_url = "https://api.github.com/repos/datahub-project/datahub/releases/latest"
         async with session.get(gh_url) as gh_response:
             gh_response_json = await gh_response.json()
-            latest_server_version = Version(gh_response_json[0].get("tag_name"))
-            latest_server_date = gh_response_json[0].get("published_at")
+            latest_server_version = Version(gh_response_json.get("tag_name"))
+            latest_server_date = gh_response_json.get("published_at")
             return (latest_server_version, latest_server_date)
 
 
diff --git a/metadata-ingestion/src/datahub/utilities/file_backed_collections.py b/metadata-ingestion/src/datahub/utilities/file_backed_collections.py
index fb028605c35b77..79da90ba20ea9f 100644
--- a/metadata-ingestion/src/datahub/utilities/file_backed_collections.py
+++ b/metadata-ingestion/src/datahub/utilities/file_backed_collections.py
@@ -224,9 +224,9 @@ class FileBackedDict(MutableMapping[str, _VT], Closeable, Generic[_VT]):
     _use_sqlite_on_conflict: bool = field(repr=False, default=True)
 
     def __post_init__(self) -> None:
-        assert (
-            self.cache_eviction_batch_size > 0
-        ), "cache_eviction_batch_size must be positive"
+        assert self.cache_eviction_batch_size > 0, (
+            "cache_eviction_batch_size must be positive"
+        )
 
         for reserved_column in ("key", "value", "rowid"):
             if reserved_column in self.extra_columns:
@@ -261,7 +261,7 @@ def __post_init__(self) -> None:
                 rowid INTEGER PRIMARY KEY AUTOINCREMENT,
                 key TEXT UNIQUE,
                 value BLOB
-                {''.join(f', {column_name} BLOB' for column_name in self.extra_columns.keys())}
+                {"".join(f", {column_name} BLOB" for column_name in self.extra_columns.keys())}
             )"""
         )
 
@@ -316,12 +316,12 @@ def _prune_cache(self, num_items_to_prune: int) -> None:
                 f"""INSERT INTO {self.tablename} (
                     key,
                     value
-                    {''.join(f', {column_name}' for column_name in self.extra_columns.keys())}
+                    {"".join(f", {column_name}" for column_name in self.extra_columns.keys())}
                 )
-                VALUES ({', '.join(['?'] *(2 + len(self.extra_columns)))})
+                VALUES ({", ".join(["?"] * (2 + len(self.extra_columns)))})
                 ON CONFLICT (key) DO UPDATE SET
                     value = excluded.value
-                    {''.join(f', {column_name} = excluded.{column_name}' for column_name in self.extra_columns.keys())}
+                    {"".join(f", {column_name} = excluded.{column_name}" for column_name in self.extra_columns.keys())}
                 """,
                 items_to_write,
             )
@@ -332,16 +332,16 @@ def _prune_cache(self, num_items_to_prune: int) -> None:
                         f"""INSERT INTO {self.tablename} (
                             key,
                             value
-                            {''.join(f', {column_name}' for column_name in self.extra_columns.keys())}
+                            {"".join(f", {column_name}" for column_name in self.extra_columns.keys())}
                         )
-                        VALUES ({', '.join(['?'] *(2 + len(self.extra_columns)))})""",
+                        VALUES ({", ".join(["?"] * (2 + len(self.extra_columns)))})""",
                         item,
                     )
                 except sqlite3.IntegrityError:
                     self._conn.execute(
                         f"""UPDATE {self.tablename} SET
                             value = ?
-                            {''.join(f', {column_name} = ?' for column_name in self.extra_columns.keys())}
+                            {"".join(f", {column_name} = ?" for column_name in self.extra_columns.keys())}
                         WHERE key = ?""",
                         (*item[1:], item[0]),
                     )
diff --git a/metadata-ingestion/src/datahub/utilities/groupby.py b/metadata-ingestion/src/datahub/utilities/groupby.py
new file mode 100644
index 00000000000000..463bca44c1e49c
--- /dev/null
+++ b/metadata-ingestion/src/datahub/utilities/groupby.py
@@ -0,0 +1,17 @@
+import collections
+from typing import Callable, Iterable, Tuple, TypeVar
+
+T = TypeVar("T")
+K = TypeVar("K")
+
+
+def groupby_unsorted(
+    iterable: Iterable[T], key: Callable[[T], K]
+) -> Iterable[Tuple[K, Iterable[T]]]:
+    """The default itertools.groupby() requires that the iterable is already sorted by the key.
+    This method is similar to groupby() but without the pre-sorted requirement."""
+
+    values = collections.defaultdict(list)
+    for v in iterable:
+        values[key(v)].append(v)
+    return values.items()
diff --git a/metadata-ingestion/src/datahub/utilities/hive_schema_to_avro.py b/metadata-ingestion/src/datahub/utilities/hive_schema_to_avro.py
index e98fe42c1d56ce..fccd8dd8a60c35 100644
--- a/metadata-ingestion/src/datahub/utilities/hive_schema_to_avro.py
+++ b/metadata-ingestion/src/datahub/utilities/hive_schema_to_avro.py
@@ -142,10 +142,10 @@ def _parse_struct_fields_string(s: str, **kwargs: Any) -> Dict[str, object]:
                 fields.append({"name": field_name, "type": field_type})
 
         if kwargs.get("ustruct_seqn") is not None:
-            struct_name = f'__structn_{kwargs["ustruct_seqn"]}_{str(uuid.uuid4()).replace("-", "")}'
+            struct_name = f"__structn_{kwargs['ustruct_seqn']}_{str(uuid.uuid4()).replace('-', '')}"
 
         else:
-            struct_name = f'__struct_{str(uuid.uuid4()).replace("-", "")}'
+            struct_name = f"__struct_{str(uuid.uuid4()).replace('-', '')}"
         return {
             "type": "record",
             "name": struct_name,
diff --git a/metadata-ingestion/src/datahub/utilities/logging_manager.py b/metadata-ingestion/src/datahub/utilities/logging_manager.py
index 926b8782fbf119..a5fd20fef307d0 100644
--- a/metadata-ingestion/src/datahub/utilities/logging_manager.py
+++ b/metadata-ingestion/src/datahub/utilities/logging_manager.py
@@ -130,9 +130,9 @@ def _formatMessageColor(self, record: logging.LogRecord) -> str:
         # Mimic our default format, but with color.
         message_fg = self.MESSAGE_COLORS.get(record.levelname)
         return (
-            f'{click.style(f"[{self.formatTime(record, self.datefmt)}]", fg="green", dim=True)} '
+            f"{click.style(f'[{self.formatTime(record, self.datefmt)}]', fg='green', dim=True)} "
             f"{click.style(f'{record.levelname:8}', fg=message_fg)} "
-            f'{click.style(f"{{{record.name}:{record.lineno}}}", fg="blue", dim=True)} - '
+            f"{click.style(f'{{{record.name}:{record.lineno}}}', fg='blue', dim=True)} - "
             f"{click.style(record.getMessage(), fg=message_fg)}"
         )
 
diff --git a/metadata-ingestion/src/datahub/utilities/lossy_collections.py b/metadata-ingestion/src/datahub/utilities/lossy_collections.py
index f71aad51ab0b6b..31d6d0eb842d04 100644
--- a/metadata-ingestion/src/datahub/utilities/lossy_collections.py
+++ b/metadata-ingestion/src/datahub/utilities/lossy_collections.py
@@ -151,9 +151,9 @@ def __str__(self) -> str:
     def as_obj(self) -> Dict[Union[_KT, str], Union[_VT, str]]:
         base_dict: Dict[Union[_KT, str], Union[_VT, str]] = super().copy()  # type: ignore
         if self.sampled:
-            base_dict[
-                "sampled"
-            ] = f"{len(self.keys())} sampled of at most {self.total_key_count()} entries."
+            base_dict["sampled"] = (
+                f"{len(self.keys())} sampled of at most {self.total_key_count()} entries."
+            )
         return base_dict
 
     def total_key_count(self) -> int:
diff --git a/metadata-ingestion/src/datahub/utilities/mapping.py b/metadata-ingestion/src/datahub/utilities/mapping.py
index 17023c7b388e76..96870fc6fcd378 100644
--- a/metadata-ingestion/src/datahub/utilities/mapping.py
+++ b/metadata-ingestion/src/datahub/utilities/mapping.py
@@ -349,9 +349,9 @@ def convert_to_aspects(self, operation_map: Dict[str, list]) -> Dict[str, Any]:
                         elements=[institutional_memory_element]
                     )
 
-                    aspect_map[
-                        Constants.ADD_DOC_LINK_OPERATION
-                    ] = institutional_memory_aspect
+                    aspect_map[Constants.ADD_DOC_LINK_OPERATION] = (
+                        institutional_memory_aspect
+                    )
                 else:
                     raise Exception(
                         f"Expected 1 item of type list for the documentation_link meta_mapping config,"
diff --git a/metadata-ingestion/src/datahub/utilities/memory_footprint.py b/metadata-ingestion/src/datahub/utilities/memory_footprint.py
index fa9e64cbf47380..f5f5fdf54ff9b7 100644
--- a/metadata-ingestion/src/datahub/utilities/memory_footprint.py
+++ b/metadata-ingestion/src/datahub/utilities/memory_footprint.py
@@ -1,7 +1,7 @@
 from collections import deque
 from itertools import chain
 from sys import getsizeof
-from typing import Any, Callable
+from typing import Any, Iterator
 
 
 def total_size(o: Any, handlers: Any = {}) -> int:
@@ -15,7 +15,8 @@ def total_size(o: Any, handlers: Any = {}) -> int:
     Based on https://github.com/ActiveState/recipe-577504-compute-mem-footprint/blob/master/recipe.py
     """
 
-    dict_handler: Callable[[Any], chain[Any]] = lambda d: chain.from_iterable(d.items())
+    def dict_handler(d: dict) -> Iterator[Any]:
+        return chain.from_iterable(d.items())
 
     all_handlers = {
         tuple: iter,
diff --git a/metadata-ingestion/src/datahub/utilities/serialized_lru_cache.py b/metadata-ingestion/src/datahub/utilities/serialized_lru_cache.py
index b5f490720340ce..bdfe4285065522 100644
--- a/metadata-ingestion/src/datahub/utilities/serialized_lru_cache.py
+++ b/metadata-ingestion/src/datahub/utilities/serialized_lru_cache.py
@@ -41,7 +41,9 @@ def decorator(func: Callable[_F, _T]) -> Callable[_F, _T]:
         def wrapper(*args: _F.args, **kwargs: _F.kwargs) -> _T:
             # We need a type ignore here because there's no way for us to require that
             # the args and kwargs are hashable while using ParamSpec.
-            key: _Key = cachetools.keys.hashkey(*args, **{k: v for k, v in kwargs.items() if "cache_exclude" not in k})  # type: ignore
+            key: _Key = cachetools.keys.hashkey(
+                *args, **{k: v for k, v in kwargs.items() if "cache_exclude" not in k}
+            )  # type: ignore
 
             with cache_lock:
                 if key in cache:
diff --git a/metadata-ingestion/src/datahub/utilities/sqlalchemy_query_combiner.py b/metadata-ingestion/src/datahub/utilities/sqlalchemy_query_combiner.py
index 11c04082ee7ad5..cf92336c68cdf6 100644
--- a/metadata-ingestion/src/datahub/utilities/sqlalchemy_query_combiner.py
+++ b/metadata-ingestion/src/datahub/utilities/sqlalchemy_query_combiner.py
@@ -160,12 +160,12 @@ class SQLAlchemyQueryCombiner:
     _greenlets_by_thread_lock: threading.Lock = dataclasses.field(
         default_factory=lambda: threading.Lock()
     )
-    _queries_by_thread: Dict[
-        greenlet.greenlet, Dict[str, _QueryFuture]
-    ] = dataclasses.field(default_factory=lambda: collections.defaultdict(dict))
-    _greenlets_by_thread: Dict[
-        greenlet.greenlet, Set[greenlet.greenlet]
-    ] = dataclasses.field(default_factory=lambda: collections.defaultdict(set))
+    _queries_by_thread: Dict[greenlet.greenlet, Dict[str, _QueryFuture]] = (
+        dataclasses.field(default_factory=lambda: collections.defaultdict(dict))
+    )
+    _greenlets_by_thread: Dict[greenlet.greenlet, Set[greenlet.greenlet]] = (
+        dataclasses.field(default_factory=lambda: collections.defaultdict(set))
+    )
 
     @staticmethod
     def _generate_sql_safe_identifier() -> str:
diff --git a/metadata-ingestion/src/datahub/utilities/sqllineage_patch.py b/metadata-ingestion/src/datahub/utilities/sqllineage_patch.py
index afcd02478ae687..4c237d02727f72 100644
--- a/metadata-ingestion/src/datahub/utilities/sqllineage_patch.py
+++ b/metadata-ingestion/src/datahub/utilities/sqllineage_patch.py
@@ -8,7 +8,7 @@
 
 # Patch based on sqllineage v1.3.3
 def end_of_query_cleanup_patch(self, holder: SubQueryLineageHolder) -> None:  # type: ignore
-    for i, tbl in enumerate(self.tables):
+    for tbl in self.tables:
         holder.add_read(tbl)
     self.union_barriers.append((len(self.columns), len(self.tables)))
 
diff --git a/metadata-ingestion/src/datahub/utilities/stats_collections.py b/metadata-ingestion/src/datahub/utilities/stats_collections.py
index 09a9490abc0fbe..c0bd9d058e5d37 100644
--- a/metadata-ingestion/src/datahub/utilities/stats_collections.py
+++ b/metadata-ingestion/src/datahub/utilities/stats_collections.py
@@ -48,7 +48,9 @@ def as_obj(self) -> Dict[_KT, _VT]:
                 total_value: Union[_VT, str] = sum(trimmed_dict.values())  # type: ignore
             except Exception:
                 total_value = ""
-            trimmed_dict[f"... top {self.top_k} of total {len(self)} entries"] = total_value  # type: ignore
+            trimmed_dict[f"... top {self.top_k} of total {len(self)} entries"] = (  # type: ignore
+                total_value  # type: ignore
+            )
             return trimmed_dict
 
 
diff --git a/metadata-ingestion/src/datahub/utilities/urns/_urn_base.py b/metadata-ingestion/src/datahub/utilities/urns/_urn_base.py
index e8e22cd85ac9ff..55144ebf61b197 100644
--- a/metadata-ingestion/src/datahub/utilities/urns/_urn_base.py
+++ b/metadata-ingestion/src/datahub/utilities/urns/_urn_base.py
@@ -6,6 +6,7 @@
 from deprecated import deprecated
 from typing_extensions import Self
 
+from datahub._codegen.aspect import _Aspect
 from datahub.utilities.urns.error import InvalidUrnError
 
 URN_TYPES: Dict[str, Type["_SpecificUrn"]] = {}
@@ -270,7 +271,7 @@ def make_form_urn(form: str) -> str:
 
 
 class _SpecificUrn(Urn):
-    ENTITY_TYPE: str = ""
+    ENTITY_TYPE: ClassVar[str] = ""
 
     def __init_subclass__(cls) -> None:
         # Validate the subclass.
@@ -286,7 +287,10 @@ def __init_subclass__(cls) -> None:
         return super().__init_subclass__()
 
     @classmethod
-    def underlying_key_aspect_type(cls) -> Type:
+    def underlying_key_aspect_type(cls) -> Type[_Aspect]:
+        raise NotImplementedError()
+
+    def to_key_aspect(self) -> _Aspect:
         raise NotImplementedError()
 
     @classmethod
diff --git a/metadata-ingestion/src/datahub/utilities/urns/urn_iter.py b/metadata-ingestion/src/datahub/utilities/urns/urn_iter.py
index f0e4c6f5ee14a1..d792e0bba649dd 100644
--- a/metadata-ingestion/src/datahub/utilities/urns/urn_iter.py
+++ b/metadata-ingestion/src/datahub/utilities/urns/urn_iter.py
@@ -21,7 +21,7 @@ def _add_prefix_to_paths(
 
 
 def list_urns_with_path(
-    model: Union[DictWrapper, MetadataChangeProposalWrapper]
+    model: Union[DictWrapper, MetadataChangeProposalWrapper],
 ) -> List[Tuple[str, _Path]]:
     """List urns in the given model with their paths.
 
@@ -145,7 +145,7 @@ def lowercase_dataset_urns(
         MetadataChangeEventClass,
         MetadataChangeProposalClass,
         MetadataChangeProposalWrapper,
-    ]
+    ],
 ) -> None:
     def modify_urn(urn: str) -> str:
         if guess_entity_type(urn) == "dataset":
diff --git a/metadata-ingestion/tests/integration/azure_ad/test_azure_ad.py b/metadata-ingestion/tests/integration/azure_ad/test_azure_ad.py
index 024bb62bbe9ce9..5bb078a368dd50 100644
--- a/metadata-ingestion/tests/integration/azure_ad/test_azure_ad.py
+++ b/metadata-ingestion/tests/integration/azure_ad/test_azure_ad.py
@@ -98,7 +98,9 @@ def load_test_resources(test_resources_dir):
     with azure_ad_nested_group_json_file.open() as azure_ad_nested_group_json:
         reference_nested_group = json.loads(azure_ad_nested_group_json.read())
 
-    with azure_ad_nested_groups_members_json_file.open() as azure_ad_nested_groups_users_json:
+    with (
+        azure_ad_nested_groups_members_json_file.open()
+    ) as azure_ad_nested_groups_users_json:
         reference_nested_groups_users = json.loads(
             azure_ad_nested_groups_users_json.read()
         )
diff --git a/metadata-ingestion/tests/integration/dremio/test_dremio.py b/metadata-ingestion/tests/integration/dremio/test_dremio.py
index 401f487d8a14b8..c286746c68b79d 100644
--- a/metadata-ingestion/tests/integration/dremio/test_dremio.py
+++ b/metadata-ingestion/tests/integration/dremio/test_dremio.py
@@ -190,9 +190,9 @@ def create_mysql_source(headers):
         "type": "MYSQL",
     }
     response = requests.post(url, headers=headers, data=json.dumps(payload))
-    assert (
-        response.status_code == 200
-    ), f"Failed to add mysql datasource: {response.text}"
+    assert response.status_code == 200, (
+        f"Failed to add mysql datasource: {response.text}"
+    )
 
 
 def upload_dataset(headers):
@@ -537,9 +537,9 @@ def test_dremio_platform_instance_urns(
 
         # Check dataset URN structure
         if mce["entityType"] == "dataset" and "entityUrn" in mce:
-            assert (
-                "test-platform.dremio" in mce["entityUrn"]
-            ), f"Platform instance missing in dataset URN: {mce['entityUrn']}"
+            assert "test-platform.dremio" in mce["entityUrn"], (
+                f"Platform instance missing in dataset URN: {mce['entityUrn']}"
+            )
 
         # Check aspects for both datasets and containers
         if "aspectName" in mce:
@@ -558,9 +558,9 @@ def test_dremio_platform_instance_urns(
 
                 instance = aspect_json["instance"]
                 expected_instance = "urn:li:dataPlatformInstance:(urn:li:dataPlatform:dremio,test-platform)"
-                assert (
-                    instance == expected_instance
-                ), f"Invalid platform instance format: {instance}"
+                assert instance == expected_instance, (
+                    f"Invalid platform instance format: {instance}"
+                )
 
     # Verify against golden file
     mce_helpers.check_golden_file(
diff --git a/metadata-ingestion/tests/integration/fivetran/fivetran_snowflake_empty_connection_user_golden.json b/metadata-ingestion/tests/integration/fivetran/fivetran_snowflake_empty_connection_user_golden.json
index 7934153051c60b..cf77ab39cbcbe7 100644
--- a/metadata-ingestion/tests/integration/fivetran/fivetran_snowflake_empty_connection_user_golden.json
+++ b/metadata-ingestion/tests/integration/fivetran/fivetran_snowflake_empty_connection_user_golden.json
@@ -595,6 +595,323 @@
         "lastRunId": "no-run-id-provided"
     }
 },
+{
+    "entityType": "dataFlow",
+    "entityUrn": "urn:li:dataFlow:(fivetran,my_confluent_cloud_connector_id,PROD)",
+    "changeType": "UPSERT",
+    "aspectName": "dataFlowInfo",
+    "aspect": {
+        "json": {
+            "customProperties": {},
+            "name": "confluent_cloud",
+            "env": "PROD"
+        }
+    },
+    "systemMetadata": {
+        "lastObserved": 1654621200000,
+        "runId": "powerbi-test",
+        "lastRunId": "no-run-id-provided"
+    }
+},
+{
+    "entityType": "dataFlow",
+    "entityUrn": "urn:li:dataFlow:(fivetran,my_confluent_cloud_connector_id,PROD)",
+    "changeType": "UPSERT",
+    "aspectName": "status",
+    "aspect": {
+        "json": {
+            "removed": false
+        }
+    },
+    "systemMetadata": {
+        "lastObserved": 1654621200000,
+        "runId": "powerbi-test",
+        "lastRunId": "no-run-id-provided"
+    }
+},
+{
+    "entityType": "dataFlow",
+    "entityUrn": "urn:li:dataFlow:(fivetran,my_confluent_cloud_connector_id,PROD)",
+    "changeType": "UPSERT",
+    "aspectName": "ownership",
+    "aspect": {
+        "json": {
+            "owners": [],
+            "ownerTypes": {},
+            "lastModified": {
+                "time": 0,
+                "actor": "urn:li:corpuser:fivetran"
+            }
+        }
+    },
+    "systemMetadata": {
+        "lastObserved": 1654621200000,
+        "runId": "powerbi-test",
+        "lastRunId": "no-run-id-provided"
+    }
+},
+{
+    "entityType": "dataFlow",
+    "entityUrn": "urn:li:dataFlow:(fivetran,my_confluent_cloud_connector_id,PROD)",
+    "changeType": "UPSERT",
+    "aspectName": "globalTags",
+    "aspect": {
+        "json": {
+            "tags": []
+        }
+    },
+    "systemMetadata": {
+        "lastObserved": 1654621200000,
+        "runId": "powerbi-test",
+        "lastRunId": "no-run-id-provided"
+    }
+},
+{
+    "entityType": "dataJob",
+    "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(fivetran,my_confluent_cloud_connector_id,PROD),my_confluent_cloud_connector_id)",
+    "changeType": "UPSERT",
+    "aspectName": "dataJobInfo",
+    "aspect": {
+        "json": {
+            "customProperties": {
+                "connector_id": "my_confluent_cloud_connector_id",
+                "connector_type": "confluent_cloud",
+                "paused": "False",
+                "sync_frequency": "1440",
+                "destination_id": "interval_unconstitutional",
+                "source.platform": "kafka",
+                "source.env": "PROD",
+                "source.database": "kafka_prod",
+                "destination.platform": "snowflake",
+                "destination.env": "PROD",
+                "destination.database": "test_database"
+            },
+            "name": "confluent_cloud",
+            "type": {
+                "string": "COMMAND"
+            },
+            "env": "PROD"
+        }
+    },
+    "systemMetadata": {
+        "lastObserved": 1654621200000,
+        "runId": "powerbi-test",
+        "lastRunId": "no-run-id-provided"
+    }
+},
+{
+    "entityType": "dataJob",
+    "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(fivetran,my_confluent_cloud_connector_id,PROD),my_confluent_cloud_connector_id)",
+    "changeType": "UPSERT",
+    "aspectName": "status",
+    "aspect": {
+        "json": {
+            "removed": false
+        }
+    },
+    "systemMetadata": {
+        "lastObserved": 1654621200000,
+        "runId": "powerbi-test",
+        "lastRunId": "no-run-id-provided"
+    }
+},
+{
+    "entityType": "dataJob",
+    "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(fivetran,my_confluent_cloud_connector_id,PROD),my_confluent_cloud_connector_id)",
+    "changeType": "UPSERT",
+    "aspectName": "dataJobInputOutput",
+    "aspect": {
+        "json": {
+            "inputDatasets": [
+                "urn:li:dataset:(urn:li:dataPlatform:kafka,kafka_prod.my-source-topic,PROD)"
+            ],
+            "outputDatasets": [
+                "urn:li:dataset:(urn:li:dataPlatform:snowflake,test_database.confluent_cloud.my-destination-topic,PROD)"
+            ],
+            "inputDatajobs": [],
+            "fineGrainedLineages": []
+        }
+    },
+    "systemMetadata": {
+        "lastObserved": 1654621200000,
+        "runId": "powerbi-test",
+        "lastRunId": "no-run-id-provided"
+    }
+},
+{
+    "entityType": "dataJob",
+    "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(fivetran,my_confluent_cloud_connector_id,PROD),my_confluent_cloud_connector_id)",
+    "changeType": "UPSERT",
+    "aspectName": "ownership",
+    "aspect": {
+        "json": {
+            "owners": [],
+            "ownerTypes": {},
+            "lastModified": {
+                "time": 0,
+                "actor": "urn:li:corpuser:fivetran"
+            }
+        }
+    },
+    "systemMetadata": {
+        "lastObserved": 1654621200000,
+        "runId": "powerbi-test",
+        "lastRunId": "no-run-id-provided"
+    }
+},
+{
+    "entityType": "dataJob",
+    "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(fivetran,my_confluent_cloud_connector_id,PROD),my_confluent_cloud_connector_id)",
+    "changeType": "UPSERT",
+    "aspectName": "globalTags",
+    "aspect": {
+        "json": {
+            "tags": []
+        }
+    },
+    "systemMetadata": {
+        "lastObserved": 1654621200000,
+        "runId": "powerbi-test",
+        "lastRunId": "no-run-id-provided"
+    }
+},
+{
+    "entityType": "dataProcessInstance",
+    "entityUrn": "urn:li:dataProcessInstance:08bddbc8007d76bfbbb8e231d1c65290",
+    "changeType": "UPSERT",
+    "aspectName": "dataProcessInstanceProperties",
+    "aspect": {
+        "json": {
+            "customProperties": {},
+            "name": "d9a03d6-eded-4422-a46a-163266e58244",
+            "type": "BATCH_SCHEDULED",
+            "created": {
+                "time": 1695191853000,
+                "actor": "urn:li:corpuser:datahub"
+            }
+        }
+    },
+    "systemMetadata": {
+        "lastObserved": 1654621200000,
+        "runId": "powerbi-test",
+        "lastRunId": "no-run-id-provided"
+    }
+},
+{
+    "entityType": "dataProcessInstance",
+    "entityUrn": "urn:li:dataProcessInstance:08bddbc8007d76bfbbb8e231d1c65290",
+    "changeType": "UPSERT",
+    "aspectName": "dataProcessInstanceRelationships",
+    "aspect": {
+        "json": {
+            "parentTemplate": "urn:li:dataJob:(urn:li:dataFlow:(fivetran,my_confluent_cloud_connector_id,PROD),my_confluent_cloud_connector_id)",
+            "upstreamInstances": []
+        }
+    },
+    "systemMetadata": {
+        "lastObserved": 1654621200000,
+        "runId": "powerbi-test",
+        "lastRunId": "no-run-id-provided"
+    }
+},
+{
+    "entityType": "dataProcessInstance",
+    "entityUrn": "urn:li:dataProcessInstance:08bddbc8007d76bfbbb8e231d1c65290",
+    "changeType": "UPSERT",
+    "aspectName": "dataProcessInstanceInput",
+    "aspect": {
+        "json": {
+            "inputs": [
+                "urn:li:dataset:(urn:li:dataPlatform:kafka,kafka_prod.my-source-topic,PROD)"
+            ]
+        }
+    },
+    "systemMetadata": {
+        "lastObserved": 1654621200000,
+        "runId": "powerbi-test",
+        "lastRunId": "no-run-id-provided"
+    }
+},
+{
+    "entityType": "dataProcessInstance",
+    "entityUrn": "urn:li:dataProcessInstance:08bddbc8007d76bfbbb8e231d1c65290",
+    "changeType": "UPSERT",
+    "aspectName": "dataProcessInstanceOutput",
+    "aspect": {
+        "json": {
+            "outputs": [
+                "urn:li:dataset:(urn:li:dataPlatform:snowflake,test_database.confluent_cloud.my-destination-topic,PROD)"
+            ]
+        }
+    },
+    "systemMetadata": {
+        "lastObserved": 1654621200000,
+        "runId": "powerbi-test",
+        "lastRunId": "no-run-id-provided"
+    }
+},
+{
+    "entityType": "dataProcessInstance",
+    "entityUrn": "urn:li:dataProcessInstance:08bddbc8007d76bfbbb8e231d1c65290",
+    "changeType": "UPSERT",
+    "aspectName": "dataProcessInstanceRunEvent",
+    "aspect": {
+        "json": {
+            "timestampMillis": 1695191853000,
+            "partitionSpec": {
+                "partition": "FULL_TABLE_SNAPSHOT",
+                "type": "FULL_TABLE"
+            },
+            "status": "STARTED"
+        }
+    },
+    "systemMetadata": {
+        "lastObserved": 1654621200000,
+        "runId": "powerbi-test",
+        "lastRunId": "no-run-id-provided"
+    }
+},
+{
+    "entityType": "dataProcessInstance",
+    "entityUrn": "urn:li:dataProcessInstance:08bddbc8007d76bfbbb8e231d1c65290",
+    "changeType": "UPSERT",
+    "aspectName": "dataProcessInstanceRunEvent",
+    "aspect": {
+        "json": {
+            "timestampMillis": 1695191885000,
+            "partitionSpec": {
+                "partition": "FULL_TABLE_SNAPSHOT",
+                "type": "FULL_TABLE"
+            },
+            "status": "COMPLETE",
+            "result": {
+                "type": "SUCCESS",
+                "nativeResultType": "fivetran"
+            }
+        }
+    },
+    "systemMetadata": {
+        "lastObserved": 1654621200000,
+        "runId": "powerbi-test",
+        "lastRunId": "no-run-id-provided"
+    }
+},
+{
+    "entityType": "dataProcessInstance",
+    "entityUrn": "urn:li:dataProcessInstance:08bddbc8007d76bfbbb8e231d1c65290",
+    "changeType": "UPSERT",
+    "aspectName": "status",
+    "aspect": {
+        "json": {
+            "removed": false
+        }
+    },
+    "systemMetadata": {
+        "lastObserved": 1654621200000,
+        "runId": "powerbi-test",
+        "lastRunId": "no-run-id-provided"
+    }
+},
 {
     "entityType": "dataProcessInstance",
     "entityUrn": "urn:li:dataProcessInstance:be36f55c13ec4e313c7510770e50784a",
diff --git a/metadata-ingestion/tests/integration/fivetran/fivetran_snowflake_golden.json b/metadata-ingestion/tests/integration/fivetran/fivetran_snowflake_golden.json
index c46bd8fb65d876..cec1ef1b6221d3 100644
--- a/metadata-ingestion/tests/integration/fivetran/fivetran_snowflake_golden.json
+++ b/metadata-ingestion/tests/integration/fivetran/fivetran_snowflake_golden.json
@@ -603,6 +603,331 @@
         "lastRunId": "no-run-id-provided"
     }
 },
+{
+    "entityType": "dataFlow",
+    "entityUrn": "urn:li:dataFlow:(fivetran,my_confluent_cloud_connector_id,PROD)",
+    "changeType": "UPSERT",
+    "aspectName": "dataFlowInfo",
+    "aspect": {
+        "json": {
+            "customProperties": {},
+            "name": "confluent_cloud",
+            "env": "PROD"
+        }
+    },
+    "systemMetadata": {
+        "lastObserved": 1654621200000,
+        "runId": "powerbi-test",
+        "lastRunId": "no-run-id-provided"
+    }
+},
+{
+    "entityType": "dataFlow",
+    "entityUrn": "urn:li:dataFlow:(fivetran,my_confluent_cloud_connector_id,PROD)",
+    "changeType": "UPSERT",
+    "aspectName": "status",
+    "aspect": {
+        "json": {
+            "removed": false
+        }
+    },
+    "systemMetadata": {
+        "lastObserved": 1654621200000,
+        "runId": "powerbi-test",
+        "lastRunId": "no-run-id-provided"
+    }
+},
+{
+    "entityType": "dataFlow",
+    "entityUrn": "urn:li:dataFlow:(fivetran,my_confluent_cloud_connector_id,PROD)",
+    "changeType": "UPSERT",
+    "aspectName": "ownership",
+    "aspect": {
+        "json": {
+            "owners": [],
+            "ownerTypes": {},
+            "lastModified": {
+                "time": 0,
+                "actor": "urn:li:corpuser:fivetran"
+            }
+        }
+    },
+    "systemMetadata": {
+        "lastObserved": 1654621200000,
+        "runId": "powerbi-test",
+        "lastRunId": "no-run-id-provided"
+    }
+},
+{
+    "entityType": "dataFlow",
+    "entityUrn": "urn:li:dataFlow:(fivetran,my_confluent_cloud_connector_id,PROD)",
+    "changeType": "UPSERT",
+    "aspectName": "globalTags",
+    "aspect": {
+        "json": {
+            "tags": []
+        }
+    },
+    "systemMetadata": {
+        "lastObserved": 1654621200000,
+        "runId": "powerbi-test",
+        "lastRunId": "no-run-id-provided"
+    }
+},
+{
+    "entityType": "dataJob",
+    "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(fivetran,my_confluent_cloud_connector_id,PROD),my_confluent_cloud_connector_id)",
+    "changeType": "UPSERT",
+    "aspectName": "dataJobInfo",
+    "aspect": {
+        "json": {
+            "customProperties": {
+                "connector_id": "my_confluent_cloud_connector_id",
+                "connector_type": "confluent_cloud",
+                "paused": "False",
+                "sync_frequency": "1440",
+                "destination_id": "my_confluent_cloud_connector_id",
+                "source.platform": "kafka",
+                "source.env": "PROD",
+                "source.database": "kafka_prod",
+                "destination.platform": "kafka",
+                "destination.env": "PROD",
+                "destination.database": "kafka_prod"
+            },
+            "name": "confluent_cloud",
+            "type": {
+                "string": "COMMAND"
+            },
+            "env": "PROD"
+        }
+    },
+    "systemMetadata": {
+        "lastObserved": 1654621200000,
+        "runId": "powerbi-test",
+        "lastRunId": "no-run-id-provided"
+    }
+},
+{
+    "entityType": "dataJob",
+    "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(fivetran,my_confluent_cloud_connector_id,PROD),my_confluent_cloud_connector_id)",
+    "changeType": "UPSERT",
+    "aspectName": "status",
+    "aspect": {
+        "json": {
+            "removed": false
+        }
+    },
+    "systemMetadata": {
+        "lastObserved": 1654621200000,
+        "runId": "powerbi-test",
+        "lastRunId": "no-run-id-provided"
+    }
+},
+{
+    "entityType": "dataJob",
+    "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(fivetran,my_confluent_cloud_connector_id,PROD),my_confluent_cloud_connector_id)",
+    "changeType": "UPSERT",
+    "aspectName": "dataJobInputOutput",
+    "aspect": {
+        "json": {
+            "inputDatasets": [
+                "urn:li:dataset:(urn:li:dataPlatform:kafka,kafka_prod.my-source-topic,PROD)"
+            ],
+            "outputDatasets": [
+                "urn:li:dataset:(urn:li:dataPlatform:kafka,kafka_prod.my-destination-topic,PROD)"
+            ],
+            "inputDatajobs": [],
+            "fineGrainedLineages": []
+        }
+    },
+    "systemMetadata": {
+        "lastObserved": 1654621200000,
+        "runId": "powerbi-test",
+        "lastRunId": "no-run-id-provided"
+    }
+},
+{
+    "entityType": "dataJob",
+    "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(fivetran,my_confluent_cloud_connector_id,PROD),my_confluent_cloud_connector_id)",
+    "changeType": "UPSERT",
+    "aspectName": "ownership",
+    "aspect": {
+        "json": {
+            "owners": [
+                {
+                    "owner": "urn:li:corpuser:abc.xyz@email.com",
+                    "type": "DEVELOPER",
+                    "source": {
+                        "type": "SERVICE"
+                    }
+                }
+            ],
+            "ownerTypes": {},
+            "lastModified": {
+                "time": 0,
+                "actor": "urn:li:corpuser:fivetran"
+            }
+        }
+    },
+    "systemMetadata": {
+        "lastObserved": 1654621200000,
+        "runId": "powerbi-test",
+        "lastRunId": "no-run-id-provided"
+    }
+},
+{
+    "entityType": "dataJob",
+    "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(fivetran,my_confluent_cloud_connector_id,PROD),my_confluent_cloud_connector_id)",
+    "changeType": "UPSERT",
+    "aspectName": "globalTags",
+    "aspect": {
+        "json": {
+            "tags": []
+        }
+    },
+    "systemMetadata": {
+        "lastObserved": 1654621200000,
+        "runId": "powerbi-test",
+        "lastRunId": "no-run-id-provided"
+    }
+},
+{
+    "entityType": "dataProcessInstance",
+    "entityUrn": "urn:li:dataProcessInstance:08bddbc8007d76bfbbb8e231d1c65290",
+    "changeType": "UPSERT",
+    "aspectName": "dataProcessInstanceProperties",
+    "aspect": {
+        "json": {
+            "customProperties": {},
+            "name": "d9a03d6-eded-4422-a46a-163266e58244",
+            "type": "BATCH_SCHEDULED",
+            "created": {
+                "time": 1695191853000,
+                "actor": "urn:li:corpuser:datahub"
+            }
+        }
+    },
+    "systemMetadata": {
+        "lastObserved": 1654621200000,
+        "runId": "powerbi-test",
+        "lastRunId": "no-run-id-provided"
+    }
+},
+{
+    "entityType": "dataProcessInstance",
+    "entityUrn": "urn:li:dataProcessInstance:08bddbc8007d76bfbbb8e231d1c65290",
+    "changeType": "UPSERT",
+    "aspectName": "dataProcessInstanceRelationships",
+    "aspect": {
+        "json": {
+            "parentTemplate": "urn:li:dataJob:(urn:li:dataFlow:(fivetran,my_confluent_cloud_connector_id,PROD),my_confluent_cloud_connector_id)",
+            "upstreamInstances": []
+        }
+    },
+    "systemMetadata": {
+        "lastObserved": 1654621200000,
+        "runId": "powerbi-test",
+        "lastRunId": "no-run-id-provided"
+    }
+},
+{
+    "entityType": "dataProcessInstance",
+    "entityUrn": "urn:li:dataProcessInstance:08bddbc8007d76bfbbb8e231d1c65290",
+    "changeType": "UPSERT",
+    "aspectName": "dataProcessInstanceInput",
+    "aspect": {
+        "json": {
+            "inputs": [
+                "urn:li:dataset:(urn:li:dataPlatform:kafka,kafka_prod.my-source-topic,PROD)"
+            ]
+        }
+    },
+    "systemMetadata": {
+        "lastObserved": 1654621200000,
+        "runId": "powerbi-test",
+        "lastRunId": "no-run-id-provided"
+    }
+},
+{
+    "entityType": "dataProcessInstance",
+    "entityUrn": "urn:li:dataProcessInstance:08bddbc8007d76bfbbb8e231d1c65290",
+    "changeType": "UPSERT",
+    "aspectName": "dataProcessInstanceOutput",
+    "aspect": {
+        "json": {
+            "outputs": [
+                "urn:li:dataset:(urn:li:dataPlatform:kafka,kafka_prod.my-destination-topic,PROD)"
+            ]
+        }
+    },
+    "systemMetadata": {
+        "lastObserved": 1654621200000,
+        "runId": "powerbi-test",
+        "lastRunId": "no-run-id-provided"
+    }
+},
+{
+    "entityType": "dataProcessInstance",
+    "entityUrn": "urn:li:dataProcessInstance:08bddbc8007d76bfbbb8e231d1c65290",
+    "changeType": "UPSERT",
+    "aspectName": "dataProcessInstanceRunEvent",
+    "aspect": {
+        "json": {
+            "timestampMillis": 1695191853000,
+            "partitionSpec": {
+                "partition": "FULL_TABLE_SNAPSHOT",
+                "type": "FULL_TABLE"
+            },
+            "status": "STARTED"
+        }
+    },
+    "systemMetadata": {
+        "lastObserved": 1654621200000,
+        "runId": "powerbi-test",
+        "lastRunId": "no-run-id-provided"
+    }
+},
+{
+    "entityType": "dataProcessInstance",
+    "entityUrn": "urn:li:dataProcessInstance:08bddbc8007d76bfbbb8e231d1c65290",
+    "changeType": "UPSERT",
+    "aspectName": "dataProcessInstanceRunEvent",
+    "aspect": {
+        "json": {
+            "timestampMillis": 1695191885000,
+            "partitionSpec": {
+                "partition": "FULL_TABLE_SNAPSHOT",
+                "type": "FULL_TABLE"
+            },
+            "status": "COMPLETE",
+            "result": {
+                "type": "SUCCESS",
+                "nativeResultType": "fivetran"
+            }
+        }
+    },
+    "systemMetadata": {
+        "lastObserved": 1654621200000,
+        "runId": "powerbi-test",
+        "lastRunId": "no-run-id-provided"
+    }
+},
+{
+    "entityType": "dataProcessInstance",
+    "entityUrn": "urn:li:dataProcessInstance:08bddbc8007d76bfbbb8e231d1c65290",
+    "changeType": "UPSERT",
+    "aspectName": "status",
+    "aspect": {
+        "json": {
+            "removed": false
+        }
+    },
+    "systemMetadata": {
+        "lastObserved": 1654621200000,
+        "runId": "powerbi-test",
+        "lastRunId": "no-run-id-provided"
+    }
+},
 {
     "entityType": "dataProcessInstance",
     "entityUrn": "urn:li:dataProcessInstance:be36f55c13ec4e313c7510770e50784a",
diff --git a/metadata-ingestion/tests/integration/fivetran/test_fivetran.py b/metadata-ingestion/tests/integration/fivetran/test_fivetran.py
index 2e6c2b1370d166..9f53ae2382d406 100644
--- a/metadata-ingestion/tests/integration/fivetran/test_fivetran.py
+++ b/metadata-ingestion/tests/integration/fivetran/test_fivetran.py
@@ -32,6 +32,15 @@
         "sync_frequency": 1440,
         "destination_id": "interval_unconstitutional",
     },
+    {
+        "connector_id": "my_confluent_cloud_connector_id",
+        "connecting_user_id": "reapply_phone",
+        "connector_type_id": "confluent_cloud",
+        "connector_name": "confluent_cloud",
+        "paused": False,
+        "sync_frequency": 1440,
+        "destination_id": "my_confluent_cloud_connector_id",
+    },
 ]
 
 
@@ -45,7 +54,7 @@ def default_query_results(
     elif query == fivetran_log_query.get_connectors_query():
         return connector_query_results
     elif query == fivetran_log_query.get_table_lineage_query(
-        connector_ids=["calendar_elected"]
+        connector_ids=["calendar_elected", "my_confluent_cloud_connector_id"]
     ):
         return [
             {
@@ -66,9 +75,18 @@ def default_query_results(
                 "destination_table_name": "company",
                 "destination_schema_name": "postgres_public",
             },
+            {
+                "connector_id": "my_confluent_cloud_connector_id",
+                "source_table_id": "10042",
+                "source_table_name": "my-source-topic",
+                "source_schema_name": "confluent_cloud",
+                "destination_table_id": "7781",
+                "destination_table_name": "my-destination-topic",
+                "destination_schema_name": "confluent_cloud",
+            },
         ]
     elif query == fivetran_log_query.get_column_lineage_query(
-        connector_ids=["calendar_elected"]
+        connector_ids=["calendar_elected", "my_confluent_cloud_connector_id"]
     ):
         return [
             {
@@ -107,7 +125,7 @@ def default_query_results(
         ]
     elif query == fivetran_log_query.get_sync_logs_query(
         syncs_interval=7,
-        connector_ids=["calendar_elected"],
+        connector_ids=["calendar_elected", "my_confluent_cloud_connector_id"],
     ):
         return [
             {
@@ -131,6 +149,13 @@ def default_query_results(
                 "end_time": datetime.datetime(2023, 10, 3, 14, 36, 29, 678000),
                 "end_message_data": '"{\\"reason\\":\\"java.lang.RuntimeException: FATAL: too many connections for role \\\\\\"hxwraqld\\\\\\"\\",\\"taskType\\":\\"reconnect\\",\\"status\\":\\"FAILURE_WITH_TASK\\"}"',
             },
+            {
+                "connector_id": "my_confluent_cloud_connector_id",
+                "sync_id": "d9a03d6-eded-4422-a46a-163266e58244",
+                "start_time": datetime.datetime(2023, 9, 20, 6, 37, 32, 606000),
+                "end_time": datetime.datetime(2023, 9, 20, 6, 38, 5, 56000),
+                "end_message_data": '"{\\"status\\":\\"SUCCESSFUL\\"}"',
+            },
         ]
     # Unreachable code
     raise Exception(f"Unknown query {query}")
@@ -172,19 +197,30 @@ def test_fivetran_with_snowflake_dest(pytestconfig, tmp_path):
                             },
                         },
                         "connector_patterns": {
-                            "allow": [
-                                "postgres",
-                            ]
+                            "allow": ["postgres", "confluent_cloud"]
                         },
                         "destination_patterns": {
                             "allow": [
                                 "interval_unconstitutional",
+                                "my_confluent_cloud_connector_id",
                             ]
                         },
                         "sources_to_platform_instance": {
                             "calendar_elected": {
                                 "database": "postgres_db",
                                 "env": "DEV",
+                            },
+                            "my_confluent_cloud_connector_id": {
+                                "platform": "kafka",
+                                "include_schema_in_urn": False,
+                                "database": "kafka_prod",
+                            },
+                        },
+                        "destination_to_platform_instance": {
+                            "my_confluent_cloud_connector_id": {
+                                "platform": "kafka",
+                                "include_schema_in_urn": False,
+                                "database": "kafka_prod",
                             }
                         },
                     },
@@ -234,6 +270,15 @@ def test_fivetran_with_snowflake_dest_and_null_connector_user(pytestconfig, tmp_
                 "sync_frequency": 1440,
                 "destination_id": "interval_unconstitutional",
             },
+            {
+                "connector_id": "my_confluent_cloud_connector_id",
+                "connecting_user_id": None,
+                "connector_type_id": "confluent_cloud",
+                "connector_name": "confluent_cloud",
+                "paused": False,
+                "sync_frequency": 1440,
+                "destination_id": "interval_unconstitutional",
+            },
         ]
 
         connection_magic_mock.execute.side_effect = partial(
@@ -261,9 +306,7 @@ def test_fivetran_with_snowflake_dest_and_null_connector_user(pytestconfig, tmp_
                             },
                         },
                         "connector_patterns": {
-                            "allow": [
-                                "postgres",
-                            ]
+                            "allow": ["postgres", "confluent_cloud"]
                         },
                         "destination_patterns": {
                             "allow": [
@@ -275,6 +318,18 @@ def test_fivetran_with_snowflake_dest_and_null_connector_user(pytestconfig, tmp_
                                 "platform": "postgres",
                                 "env": "DEV",
                                 "database": "postgres_db",
+                            },
+                            "my_confluent_cloud_connector_id": {
+                                "platform": "kafka",
+                                "database": "kafka_prod",
+                                "include_schema_in_urn": False,
+                            },
+                        },
+                        "destination_to_platform_instance": {
+                            "my_confluent_cloud_connector_id": {
+                                "platform": "kafka",
+                                "database": "kafka_prod",
+                                "include_schema_in_urn": False,
                             }
                         },
                     },
diff --git a/metadata-ingestion/tests/integration/grafana/test_grafana.py b/metadata-ingestion/tests/integration/grafana/test_grafana.py
index 6eb6b0b8509263..cbac965884365d 100644
--- a/metadata-ingestion/tests/integration/grafana/test_grafana.py
+++ b/metadata-ingestion/tests/integration/grafana/test_grafana.py
@@ -120,7 +120,7 @@ def test_grafana_dashboard(loaded_grafana, pytestconfig, tmp_path, test_resource
         time.sleep(5)
         resp = requests.get(url)
         if resp.status_code == 200:
-            logging.info(f"Grafana started after waiting {i*5} seconds")
+            logging.info(f"Grafana started after waiting {i * 5} seconds")
             break
     else:
         pytest.fail("Grafana did not start in time")
@@ -131,12 +131,12 @@ def test_grafana_dashboard(loaded_grafana, pytestconfig, tmp_path, test_resource
     assert resp.status_code == 200, "Failed to load default dashboard"
     dashboard = resp.json()
 
-    assert (
-        dashboard["dashboard"]["title"] == "Default Dashboard"
-    ), "Default dashboard title mismatch"
-    assert any(
-        panel["type"] == "text" for panel in dashboard["dashboard"]["panels"]
-    ), "Default dashboard missing text panel"
+    assert dashboard["dashboard"]["title"] == "Default Dashboard", (
+        "Default dashboard title mismatch"
+    )
+    assert any(panel["type"] == "text" for panel in dashboard["dashboard"]["panels"]), (
+        "Default dashboard missing text panel"
+    )
 
     # Verify the output. (You can add further checks here if needed)
     logging.info("Default dashboard verified successfully")
@@ -153,7 +153,7 @@ def test_grafana_ingest(
         time.sleep(5)
         resp = requests.get(url)
         if resp.status_code == 200:
-            logging.info(f"Grafana started after waiting {i*5} seconds")
+            logging.info(f"Grafana started after waiting {i * 5} seconds")
             break
     else:
         pytest.fail("Grafana did not start in time")
diff --git a/metadata-ingestion/tests/integration/kafka-connect/test_kafka_connect.py b/metadata-ingestion/tests/integration/kafka-connect/test_kafka_connect.py
index d0f4fc35fc03eb..d8c98b12951f5d 100644
--- a/metadata-ingestion/tests/integration/kafka-connect/test_kafka_connect.py
+++ b/metadata-ingestion/tests/integration/kafka-connect/test_kafka_connect.py
@@ -482,9 +482,9 @@ def test_kafka_connect_ingest_stateful(
             "mysql_source1",
             "mysql_source2",
         ]
-        pipeline_run1_config["sink"]["config"][
-            "filename"
-        ] = f"{tmp_path}/{output_file_name}"
+        pipeline_run1_config["sink"]["config"]["filename"] = (
+            f"{tmp_path}/{output_file_name}"
+        )
         pipeline_run1 = Pipeline.create(pipeline_run1_config)
         pipeline_run1.run()
         pipeline_run1.raise_from_status()
@@ -506,14 +506,16 @@ def test_kafka_connect_ingest_stateful(
         mock_datahub_graph,
     ) as mock_checkpoint:
         mock_checkpoint.return_value = mock_datahub_graph
-        pipeline_run2_config: Dict[str, Dict[str, Dict[str, Any]]] = dict(base_pipeline_config)  # type: ignore
+        pipeline_run2_config: Dict[str, Dict[str, Dict[str, Any]]] = dict(
+            base_pipeline_config  # type: ignore
+        )
         # Set the special properties for this run
         pipeline_run1_config["source"]["config"]["connector_patterns"]["allow"] = [
             "mysql_source1",
         ]
-        pipeline_run2_config["sink"]["config"][
-            "filename"
-        ] = f"{tmp_path}/{output_file_deleted_name}"
+        pipeline_run2_config["sink"]["config"]["filename"] = (
+            f"{tmp_path}/{output_file_deleted_name}"
+        )
         pipeline_run2 = Pipeline.create(pipeline_run2_config)
         pipeline_run2.run()
         pipeline_run2.raise_from_status()
diff --git a/metadata-ingestion/tests/integration/looker/test_looker.py b/metadata-ingestion/tests/integration/looker/test_looker.py
index c96bcc729a95da..bbcc6332539c02 100644
--- a/metadata-ingestion/tests/integration/looker/test_looker.py
+++ b/metadata-ingestion/tests/integration/looker/test_looker.py
@@ -1096,9 +1096,9 @@ def test_file_path_in_view_naming_pattern(
 ):
     mocked_client = mock.MagicMock()
     new_recipe = get_default_recipe(output_file_path=f"{tmp_path}/looker_mces.json")
-    new_recipe["source"]["config"][
-        "view_naming_pattern"
-    ] = "{project}.{file_path}.view.{name}"
+    new_recipe["source"]["config"]["view_naming_pattern"] = (
+        "{project}.{file_path}.view.{name}"
+    )
 
     with mock.patch(
         "datahub.ingestion.source.state_provider.datahub_ingestion_checkpointing_provider.DataHubGraph",
diff --git a/metadata-ingestion/tests/integration/lookml/test_lookml.py b/metadata-ingestion/tests/integration/lookml/test_lookml.py
index 940e7f36675f79..d803b8498104fd 100644
--- a/metadata-ingestion/tests/integration/lookml/test_lookml.py
+++ b/metadata-ingestion/tests/integration/lookml/test_lookml.py
@@ -101,13 +101,13 @@ def test_lookml_refinement_ingest(pytestconfig, tmp_path, mock_time):
     )
     new_recipe["source"]["config"]["process_refinements"] = True
 
-    new_recipe["source"]["config"][
-        "view_naming_pattern"
-    ] = "{project}.{file_path}.view.{name}"
+    new_recipe["source"]["config"]["view_naming_pattern"] = (
+        "{project}.{file_path}.view.{name}"
+    )
 
-    new_recipe["source"]["config"][
-        "view_browse_pattern"
-    ] = "/{env}/{platform}/{project}/{file_path}/views"
+    new_recipe["source"]["config"]["view_browse_pattern"] = (
+        "/{env}/{platform}/{project}/{file_path}/views"
+    )
 
     pipeline = Pipeline.create(new_recipe)
     pipeline.run()
diff --git a/metadata-ingestion/tests/integration/nifi/test_nifi.py b/metadata-ingestion/tests/integration/nifi/test_nifi.py
index b992de058879ef..924e854a47e4eb 100644
--- a/metadata-ingestion/tests/integration/nifi/test_nifi.py
+++ b/metadata-ingestion/tests/integration/nifi/test_nifi.py
@@ -72,7 +72,7 @@ def test_nifi_ingest_standalone(
             status = next(s for s in statuses if s["name"] == "FetchS3Object")
 
             if status["aggregateSnapshot"]["flowFilesOut"] >= 1:
-                logging.info(f"Waited for time {i*5} seconds")
+                logging.info(f"Waited for time {i * 5} seconds")
                 break
 
     # Run the metadata ingestion pipeline.
@@ -124,7 +124,7 @@ def test_nifi_ingest_cluster(loaded_nifi, pytestconfig, tmp_path, test_resources
             statuses = [pg["status"] for pg in pgs]
             status = next(s for s in statuses if s["name"] == "Cluster_Site_S3_to_S3")
             if status["aggregateSnapshot"]["flowFilesSent"] >= 1:
-                logging.info(f"Waited for time {i*5} seconds")
+                logging.info(f"Waited for time {i * 5} seconds")
                 break
     test_resources_dir = pytestconfig.rootpath / "tests/integration/nifi"
     # Run the metadata ingestion pipeline.
diff --git a/metadata-ingestion/tests/integration/powerbi/test_m_parser.py b/metadata-ingestion/tests/integration/powerbi/test_m_parser.py
index 832d00d9c54702..0d85d370265cae 100644
--- a/metadata-ingestion/tests/integration/powerbi/test_m_parser.py
+++ b/metadata-ingestion/tests/integration/powerbi/test_m_parser.py
@@ -765,14 +765,14 @@ def test_sqlglot_parser():
         }
     )
 
-    lineage: List[
-        datahub.ingestion.source.powerbi.m_query.data_classes.Lineage
-    ] = parser.get_upstream_tables(
-        table,
-        reporter,
-        ctx=ctx,
-        config=config,
-        platform_instance_resolver=platform_instance_resolver,
+    lineage: List[datahub.ingestion.source.powerbi.m_query.data_classes.Lineage] = (
+        parser.get_upstream_tables(
+            table,
+            reporter,
+            ctx=ctx,
+            config=config,
+            platform_instance_resolver=platform_instance_resolver,
+        )
     )
 
     data_platform_tables: List[DataPlatformTable] = lineage[0].upstreams
@@ -814,9 +814,9 @@ def test_sqlglot_parser():
 def test_databricks_multi_cloud():
     q = M_QUERIES[25]
 
-    lineage: List[
-        datahub.ingestion.source.powerbi.m_query.data_classes.Lineage
-    ] = get_data_platform_tables_with_dummy_table(q=q)
+    lineage: List[datahub.ingestion.source.powerbi.m_query.data_classes.Lineage] = (
+        get_data_platform_tables_with_dummy_table(q=q)
+    )
 
     assert len(lineage) == 1
 
@@ -833,9 +833,9 @@ def test_databricks_multi_cloud():
 def test_databricks_catalog_pattern_1():
     q = M_QUERIES[26]
 
-    lineage: List[
-        datahub.ingestion.source.powerbi.m_query.data_classes.Lineage
-    ] = get_data_platform_tables_with_dummy_table(q=q)
+    lineage: List[datahub.ingestion.source.powerbi.m_query.data_classes.Lineage] = (
+        get_data_platform_tables_with_dummy_table(q=q)
+    )
 
     assert len(lineage) == 1
 
@@ -904,14 +904,14 @@ def test_sqlglot_parser_2():
         }
     )
 
-    lineage: List[
-        datahub.ingestion.source.powerbi.m_query.data_classes.Lineage
-    ] = parser.get_upstream_tables(
-        table,
-        reporter,
-        ctx=ctx,
-        config=config,
-        platform_instance_resolver=platform_instance_resolver,
+    lineage: List[datahub.ingestion.source.powerbi.m_query.data_classes.Lineage] = (
+        parser.get_upstream_tables(
+            table,
+            reporter,
+            ctx=ctx,
+            config=config,
+            platform_instance_resolver=platform_instance_resolver,
+        )
     )
 
     data_platform_tables: List[DataPlatformTable] = lineage[0].upstreams
@@ -965,9 +965,9 @@ def test_databricks_regular_case_with_view():
 def test_snowflake_double_double_quotes():
     q = M_QUERIES[30]
 
-    lineage: List[
-        datahub.ingestion.source.powerbi.m_query.data_classes.Lineage
-    ] = get_data_platform_tables_with_dummy_table(q=q)
+    lineage: List[datahub.ingestion.source.powerbi.m_query.data_classes.Lineage] = (
+        get_data_platform_tables_with_dummy_table(q=q)
+    )
 
     assert len(lineage) == 1
 
@@ -984,9 +984,9 @@ def test_snowflake_double_double_quotes():
 def test_databricks_multicloud():
     q = M_QUERIES[31]
 
-    lineage: List[
-        datahub.ingestion.source.powerbi.m_query.data_classes.Lineage
-    ] = get_data_platform_tables_with_dummy_table(q=q)
+    lineage: List[datahub.ingestion.source.powerbi.m_query.data_classes.Lineage] = (
+        get_data_platform_tables_with_dummy_table(q=q)
+    )
 
     assert len(lineage) == 1
 
@@ -1003,9 +1003,9 @@ def test_databricks_multicloud():
 def test_snowflake_multi_function_call():
     q = M_QUERIES[32]
 
-    lineage: List[
-        datahub.ingestion.source.powerbi.m_query.data_classes.Lineage
-    ] = get_data_platform_tables_with_dummy_table(q=q)
+    lineage: List[datahub.ingestion.source.powerbi.m_query.data_classes.Lineage] = (
+        get_data_platform_tables_with_dummy_table(q=q)
+    )
 
     assert len(lineage) == 1
 
@@ -1022,9 +1022,9 @@ def test_snowflake_multi_function_call():
 def test_mssql_drop_with_select():
     q = M_QUERIES[33]
 
-    lineage: List[
-        datahub.ingestion.source.powerbi.m_query.data_classes.Lineage
-    ] = get_data_platform_tables_with_dummy_table(q=q)
+    lineage: List[datahub.ingestion.source.powerbi.m_query.data_classes.Lineage] = (
+        get_data_platform_tables_with_dummy_table(q=q)
+    )
 
     assert len(lineage) == 1
 
@@ -1070,23 +1070,23 @@ def test_unsupported_data_platform():
     )  # type :ignore
 
     is_entry_present: bool = False
-    for key, entry in info_entries.items():
+    for entry in info_entries.values():
         if entry.title == "Non-Data Platform Expression":
             is_entry_present = True
             break
 
-    assert (
-        is_entry_present
-    ), 'Info message "Non-Data Platform Expression" should be present in reporter'
+    assert is_entry_present, (
+        'Info message "Non-Data Platform Expression" should be present in reporter'
+    )
 
 
 def test_empty_string_in_m_query():
     # TRIM(TRIM(TRIM(AGENT_NAME, '\"\"'), '+'), '\\'') is in Query
     q = "let\n  Source = Value.NativeQuery(Snowflake.Databases(\"bu10758.ap-unknown-2.fakecomputing.com\",\"operations_analytics_warehouse_prod\",[Role=\"OPERATIONS_ANALYTICS_MEMBER\"]){[Name=\"OPERATIONS_ANALYTICS\"]}[Data], \"select #(lf)UPPER(REPLACE(AGENT_NAME,'-','')) AS CLIENT_DIRECTOR,#(lf)TRIM(TRIM(TRIM(AGENT_NAME, '\"\"'), '+'), '\\'') AS TRIM_AGENT_NAME,#(lf)TIER,#(lf)UPPER(MANAGER),#(lf)TEAM_TYPE,#(lf)DATE_TARGET,#(lf)MONTHID,#(lf)TARGET_TEAM,#(lf)SELLER_EMAIL,#(lf)concat((UPPER(REPLACE(AGENT_NAME,'-',''))), MONTHID) as AGENT_KEY,#(lf)UNIT_TARGET AS SME_Quota,#(lf)AMV_TARGET AS Revenue_Quota,#(lf)SERVICE_QUOTA,#(lf)BL_TARGET,#(lf)SOFTWARE_QUOTA as Software_Quota#(lf)#(lf)from OPERATIONS_ANALYTICS.TRANSFORMED_PROD.V_SME_UNIT_TARGETS inner join OPERATIONS_ANALYTICS.TRANSFORMED_PROD.V_SME_UNIT #(lf)#(lf)where YEAR_TARGET >= 2022#(lf)and TEAM_TYPE = 'Accounting'#(lf)and TARGET_TEAM = 'Enterprise'#(lf)AND TIER = 'Client Director'\", null, [EnableFolding=true])\nin\n    Source"
 
-    lineage: List[
-        datahub.ingestion.source.powerbi.m_query.data_classes.Lineage
-    ] = get_data_platform_tables_with_dummy_table(q=q)
+    lineage: List[datahub.ingestion.source.powerbi.m_query.data_classes.Lineage] = (
+        get_data_platform_tables_with_dummy_table(q=q)
+    )
 
     assert len(lineage) == 1
 
@@ -1108,9 +1108,9 @@ def test_double_quotes_in_alias():
     # SELECT CAST(sales_date AS DATE) AS \"\"Date\"\" in query
     q = 'let \n Source = Sql.Database("abc.com", "DB", [Query="SELECT CAST(sales_date AS DATE) AS ""Date"",#(lf) SUM(cshintrpret) / 60.0      AS ""Total Order All Items"",#(lf)#(tab)#(tab)#(tab)  SUM(cshintrpret) / 60.0 - LAG(SUM(cshintrpret) / 60.0, 1) OVER (ORDER BY CAST(sales_date AS DATE)) AS ""Total minute difference"",#(lf)#(tab)#(tab)#(tab)  SUM(sale_price)  / 60.0 - LAG(SUM(sale_price)  / 60.0, 1) OVER (ORDER BY CAST(sales_date AS DATE)) AS ""Normal minute difference""#(lf)        FROM   [DB].[dbo].[sales_t]#(lf)        WHERE  sales_date >= GETDATE() - 365#(lf)        GROUP  BY CAST(sales_date AS DATE),#(lf)#(tab)#(tab)CAST(sales_date AS TIME);"]) \n in \n Source'
 
-    lineage: List[
-        datahub.ingestion.source.powerbi.m_query.data_classes.Lineage
-    ] = get_data_platform_tables_with_dummy_table(q=q)
+    lineage: List[datahub.ingestion.source.powerbi.m_query.data_classes.Lineage] = (
+        get_data_platform_tables_with_dummy_table(q=q)
+    )
 
     assert len(lineage) == 1
 
@@ -1163,14 +1163,14 @@ def test_m_query_timeout(mock_get_lark_parser):
     )  # type :ignore
 
     is_entry_present: bool = False
-    for key, entry in warn_entries.items():
+    for entry in warn_entries.values():
         if entry.title == "M-Query Parsing Timeout":
             is_entry_present = True
             break
 
-    assert (
-        is_entry_present
-    ), 'Warning message "M-Query Parsing Timeout" should be present in reporter'
+    assert is_entry_present, (
+        'Warning message "M-Query Parsing Timeout" should be present in reporter'
+    )
 
 
 def test_comments_in_m_query():
diff --git a/metadata-ingestion/tests/integration/powerbi/test_powerbi.py b/metadata-ingestion/tests/integration/powerbi/test_powerbi.py
index 739be7cc8408dd..571c0092dcc8c1 100644
--- a/metadata-ingestion/tests/integration/powerbi/test_powerbi.py
+++ b/metadata-ingestion/tests/integration/powerbi/test_powerbi.py
@@ -828,9 +828,9 @@ def dataset_type_mapping_set_to_all_platform(pipeline: Pipeline) -> None:
     # Generate default dataset_type_mapping and compare it with source_config.dataset_type_mapping
     default_dataset_type_mapping: dict = {}
     for item in SupportedDataPlatform:
-        default_dataset_type_mapping[
-            item.value.powerbi_data_platform_name
-        ] = item.value.datahub_data_platform_name
+        default_dataset_type_mapping[item.value.powerbi_data_platform_name] = (
+            item.value.datahub_data_platform_name
+        )
 
     assert default_dataset_type_mapping == source_config.dataset_type_mapping
 
@@ -978,16 +978,16 @@ def validate_pipeline(pipeline: Pipeline) -> None:
         name="demo-workspace",
         type="Workspace",
         datasets={},
-        dashboards=[],
-        reports=[],
+        dashboards={},
+        reports={},
         report_endorsements={},
         dashboard_endorsements={},
         scan_result={},
-        independent_datasets=[],
+        independent_datasets={},
         app=None,
     )
     # Fetch actual reports
-    reports: List[Report] = cast(
+    reports: Dict[str, Report] = cast(
         PowerBiDashboardSource, pipeline.source
     ).powerbi_client.get_reports(workspace=mock_workspace)
 
@@ -1020,8 +1020,8 @@ def validate_pipeline(pipeline: Pipeline) -> None:
             "pages": [],
         },
     ]
-    expected_reports: List[Report] = [
-        Report(
+    expected_reports: Dict[str, Report] = {
+        report[Constant.ID]: Report(
             id=report[Constant.ID],
             name=report[Constant.NAME],
             type=ReportType.PowerBIReport,
@@ -1045,14 +1045,15 @@ def validate_pipeline(pipeline: Pipeline) -> None:
             dataset=mock_workspace.datasets.get(report[Constant.DATASET_ID]),
         )
         for report in mock_reports
-    ]
+    }
     # Compare actual and expected reports
     for i in range(2):
-        assert reports[i].id == expected_reports[i].id
-        assert reports[i].name == expected_reports[i].name
-        assert reports[i].description == expected_reports[i].description
-        assert reports[i].dataset == expected_reports[i].dataset
-        assert reports[i].pages == expected_reports[i].pages
+        report_id = mock_reports[i][Constant.ID]
+        assert reports[report_id].id == expected_reports[report_id].id
+        assert reports[report_id].name == expected_reports[report_id].name
+        assert reports[report_id].description == expected_reports[report_id].description
+        assert reports[report_id].dataset == expected_reports[report_id].dataset
+        assert reports[report_id].pages == expected_reports[report_id].pages
 
 
 @freeze_time(FROZEN_TIME)
@@ -1438,14 +1439,14 @@ def test_powerbi_cross_workspace_reference_info_message(
 
     is_entry_present: bool = False
     # Printing INFO entries
-    for key, entry in info_entries.items():
+    for entry in info_entries.values():
         if entry.title == "Missing Lineage For Tile":
             is_entry_present = True
             break
 
-    assert (
-        is_entry_present
-    ), 'Info message "Missing Lineage For Tile" should be present in reporter'
+    assert is_entry_present, (
+        'Info message "Missing Lineage For Tile" should be present in reporter'
+    )
 
     test_resources_dir = pytestconfig.rootpath / "tests/integration/powerbi"
 
@@ -1563,11 +1564,11 @@ def test_powerbi_app_ingest_info_message(
 
     is_entry_present: bool = False
     # Printing INFO entries
-    for key, entry in info_entries.items():
+    for entry in info_entries.values():
         if entry.title == "App Ingestion Is Disabled":
             is_entry_present = True
             break
 
-    assert (
-        is_entry_present
-    ), "The extract_app flag should be set to false by default. We need to keep this flag as false until all GMS instances are updated to the latest release."
+    assert is_entry_present, (
+        "The extract_app flag should be set to false by default. We need to keep this flag as false until all GMS instances are updated to the latest release."
+    )
diff --git a/metadata-ingestion/tests/integration/salesforce/test_salesforce.py b/metadata-ingestion/tests/integration/salesforce/test_salesforce.py
index 89a37a372df843..9e68ff22a767e2 100644
--- a/metadata-ingestion/tests/integration/salesforce/test_salesforce.py
+++ b/metadata-ingestion/tests/integration/salesforce/test_salesforce.py
@@ -89,15 +89,15 @@ def test_latest_version(mock_sdk):
     )
     SalesforceSource(config=config, ctx=Mock())
     calls = mock_sf._call_salesforce.mock_calls
-    assert (
-        len(calls) == 1
-    ), "We didn't specify version but source didn't call SF API to get the latest one"
-    assert calls[0].ends_with(
-        "/services/data"
-    ), "Source didn't call proper SF API endpoint to get all versions"
-    assert (
-        mock_sf.sf_version == "54.0"
-    ), "API version was not correctly set (see versions_responses.json)"
+    assert len(calls) == 1, (
+        "We didn't specify version but source didn't call SF API to get the latest one"
+    )
+    assert calls[0].ends_with("/services/data"), (
+        "Source didn't call proper SF API endpoint to get all versions"
+    )
+    assert mock_sf.sf_version == "54.0", (
+        "API version was not correctly set (see versions_responses.json)"
+    )
 
 
 @mock.patch("datahub.ingestion.source.salesforce.Salesforce")
@@ -133,12 +133,12 @@ def test_custom_version(mock_sdk):
     SalesforceSource(config=config, ctx=Mock())
 
     calls = mock_sf._call_salesforce.mock_calls
-    assert (
-        len(calls) == 0
-    ), "Source called API to get all versions even though we specified proper version"
-    assert (
-        mock_sdk.call_args.kwargs["version"] == "46.0"
-    ), "API client object was not correctly initialized with the custom version"
+    assert len(calls) == 0, (
+        "Source called API to get all versions even though we specified proper version"
+    )
+    assert mock_sdk.call_args.kwargs["version"] == "46.0", (
+        "API client object was not correctly initialized with the custom version"
+    )
 
 
 @freeze_time(FROZEN_TIME)
diff --git a/metadata-ingestion/tests/integration/snowflake/test_snowflake_failures.py b/metadata-ingestion/tests/integration/snowflake/test_snowflake_failures.py
index de6e996a52642b..4cb6cec4906efa 100644
--- a/metadata-ingestion/tests/integration/snowflake/test_snowflake_failures.py
+++ b/metadata-ingestion/tests/integration/snowflake/test_snowflake_failures.py
@@ -260,3 +260,63 @@ def test_snowflake_missing_snowflake_operations_permission_causes_pipeline_failu
         assert "usage-permission-error" in [
             failure.message for failure in pipeline.source.get_report().failures
         ]
+
+
+@freeze_time(FROZEN_TIME)
+def test_snowflake_missing_snowflake_secure_view_definitions_raises_pipeline_info(
+    pytestconfig,
+    snowflake_pipeline_config,
+):
+    with mock.patch("snowflake.connector.connect") as mock_connect:
+        sf_connection = mock.MagicMock()
+        sf_cursor = mock.MagicMock()
+        mock_connect.return_value = sf_connection
+        sf_connection.cursor.return_value = sf_cursor
+
+        # Empty secure view definitions
+        sf_cursor.execute.side_effect = query_permission_response_override(
+            default_query_results,
+            [snowflake_query.SnowflakeQuery.get_secure_view_definitions()],
+            [],
+        )
+        pipeline = Pipeline(snowflake_pipeline_config)
+        pipeline.run()
+
+        pipeline.raise_from_status(raise_warnings=True)
+        assert pipeline.source.get_report().infos.as_obj() == [
+            {
+                "title": "Secure view definition not found",
+                "message": "Lineage will be missing for the view.",
+                "context": ["TEST_DB.TEST_SCHEMA.VIEW_1"],
+            }
+        ]
+
+
+@freeze_time(FROZEN_TIME)
+def test_snowflake_failed_secure_view_definitions_query_raises_pipeline_warning(
+    pytestconfig,
+    snowflake_pipeline_config,
+):
+    with mock.patch("snowflake.connector.connect") as mock_connect:
+        sf_connection = mock.MagicMock()
+        sf_cursor = mock.MagicMock()
+        mock_connect.return_value = sf_connection
+        sf_connection.cursor.return_value = sf_cursor
+
+        # Error in getting secure view definitions
+        sf_cursor.execute.side_effect = query_permission_error_override(
+            default_query_results,
+            [snowflake_query.SnowflakeQuery.get_secure_view_definitions()],
+            "Database 'SNOWFLAKE' does not exist or not authorized.",
+        )
+        pipeline = Pipeline(snowflake_pipeline_config)
+        pipeline.run()
+        assert pipeline.source.get_report().warnings.as_obj() == [
+            {
+                "title": "Failed to get secure views definitions",
+                "message": "Lineage will be missing for the view. Please check permissions.",
+                "context": [
+                    "TEST_DB.TEST_SCHEMA.VIEW_1 <class 'datahub.ingestion.source.snowflake.snowflake_connection.SnowflakePermissionError'>: Database 'SNOWFLAKE' does not exist or not authorized."
+                ],
+            }
+        ]
diff --git a/metadata-ingestion/tests/integration/sql_server/golden_files/golden_mces_mssql_no_db_to_file.json b/metadata-ingestion/tests/integration/sql_server/golden_files/golden_mces_mssql_no_db_to_file.json
index 0d9386dcda0cdb..0b478e13b86d2a 100644
--- a/metadata-ingestion/tests/integration/sql_server/golden_files/golden_mces_mssql_no_db_to_file.json
+++ b/metadata-ingestion/tests/integration/sql_server/golden_files/golden_mces_mssql_no_db_to_file.json
@@ -95,7 +95,6 @@
     "aspect": {
         "json": {
             "customProperties": {},
-            "externalUrl": "",
             "name": "Weekly Demo Data Backup"
         }
     },
@@ -113,17 +112,16 @@
     "aspect": {
         "json": {
             "customProperties": {
-                "job_id": "a06cfdca-b65e-42de-8db2-8c21c183c5dd",
+                "job_id": "f5a6c120-500a-4300-9b21-0c3225af1f80",
                 "job_name": "Weekly Demo Data Backup",
                 "description": "No description available.",
-                "date_created": "2024-12-26 12:03:35.420000",
-                "date_modified": "2024-12-26 12:03:35.590000",
+                "date_created": "2024-12-30 19:59:24.690000",
+                "date_modified": "2024-12-30 19:59:24.690000",
                 "step_id": "1",
                 "step_name": "Set database to read only",
                 "subsystem": "TSQL",
                 "command": "ALTER DATABASE DemoData SET READ_ONLY"
             },
-            "externalUrl": "",
             "name": "Weekly Demo Data Backup",
             "type": {
                 "string": "MSSQL_JOB_STEP"
@@ -2259,7 +2257,6 @@
     "aspect": {
         "json": {
             "customProperties": {},
-            "externalUrl": "",
             "name": "DemoData.Foo.stored_procedures"
         }
     },
@@ -2282,10 +2279,9 @@
                 "code": "CREATE PROCEDURE [Foo].[Proc.With.SpecialChar] @ID INT\nAS\n    SELECT @ID AS ThatDB;\n",
                 "input parameters": "['@ID']",
                 "parameter @ID": "{'type': 'int'}",
-                "date_created": "2024-12-26 12:03:35.230000",
-                "date_modified": "2024-12-26 12:03:35.230000"
+                "date_created": "2024-12-30 19:59:24.690000",
+                "date_modified": "2024-12-30 19:59:24.690000"
             },
-            "externalUrl": "",
             "name": "DemoData.Foo.Proc.With.SpecialChar",
             "type": {
                 "string": "MSSQL_STORED_PROCEDURE"
@@ -2298,6 +2294,29 @@
         "lastRunId": "no-run-id-provided"
     }
 },
+{
+    "entityType": "dataJob",
+    "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(mssql,DemoData.Foo.stored_procedures,PROD),Proc.With.SpecialChar)",
+    "changeType": "UPSERT",
+    "aspectName": "dataTransformLogic",
+    "aspect": {
+        "json": {
+            "transforms": [
+                {
+                    "queryStatement": {
+                        "value": "CREATE PROCEDURE [Foo].[Proc.With.SpecialChar] @ID INT\nAS\n    SELECT @ID AS ThatDB;\n",
+                        "language": "SQL"
+                    }
+                }
+            ]
+        }
+    },
+    "systemMetadata": {
+        "lastObserved": 1615443388097,
+        "runId": "mssql-test",
+        "lastRunId": "no-run-id-provided"
+    }
+},
 {
     "entityType": "dataJob",
     "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(mssql,DemoData.Foo.stored_procedures,PROD),NewProc)",
@@ -2310,10 +2329,9 @@
                 "depending_on_procedure": "{}",
                 "code": "CREATE PROCEDURE [Foo].[NewProc]\n    AS\n    BEGIN\n        --insert into items table from salesreason table\n        insert into Foo.Items (ID, ItemName)\n        SELECT TempID, Name\n        FROM Foo.SalesReason;\n\n\n       IF OBJECT_ID('Foo.age_dist', 'U') IS NULL\n       BEGIN\n            -- Create and populate if table doesn't exist\n            SELECT Age, COUNT(*) as Count\n            INTO Foo.age_dist\n            FROM Foo.Persons\n            GROUP BY Age\n        END\n        ELSE\n        BEGIN\n            -- Update existing table\n            TRUNCATE TABLE Foo.age_dist;\n\n            INSERT INTO Foo.age_dist (Age, Count)\n            SELECT Age, COUNT(*) as Count\n            FROM Foo.Persons\n            GROUP BY Age\n        END\n\n        SELECT ID, Age INTO #TEMPTABLE FROM NewData.FooNew.PersonsNew\n        \n        UPDATE DemoData.Foo.Persons\n        SET Age = t.Age\n        FROM DemoData.Foo.Persons p\n        JOIN #TEMPTABLE t ON p.ID = t.ID\n\n    END\n",
                 "input parameters": "[]",
-                "date_created": "2024-12-26 12:03:35.237000",
-                "date_modified": "2024-12-26 12:03:35.237000"
+                "date_created": "2024-12-30 19:59:24.690000",
+                "date_modified": "2024-12-30 19:59:24.690000"
             },
-            "externalUrl": "",
             "name": "DemoData.Foo.NewProc",
             "type": {
                 "string": "MSSQL_STORED_PROCEDURE"
@@ -2326,6 +2344,29 @@
         "lastRunId": "no-run-id-provided"
     }
 },
+{
+    "entityType": "dataJob",
+    "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(mssql,DemoData.Foo.stored_procedures,PROD),NewProc)",
+    "changeType": "UPSERT",
+    "aspectName": "dataTransformLogic",
+    "aspect": {
+        "json": {
+            "transforms": [
+                {
+                    "queryStatement": {
+                        "value": "CREATE PROCEDURE [Foo].[NewProc]\n    AS\n    BEGIN\n        --insert into items table from salesreason table\n        insert into Foo.Items (ID, ItemName)\n        SELECT TempID, Name\n        FROM Foo.SalesReason;\n\n\n       IF OBJECT_ID('Foo.age_dist', 'U') IS NULL\n       BEGIN\n            -- Create and populate if table doesn't exist\n            SELECT Age, COUNT(*) as Count\n            INTO Foo.age_dist\n            FROM Foo.Persons\n            GROUP BY Age\n        END\n        ELSE\n        BEGIN\n            -- Update existing table\n            TRUNCATE TABLE Foo.age_dist;\n\n            INSERT INTO Foo.age_dist (Age, Count)\n            SELECT Age, COUNT(*) as Count\n            FROM Foo.Persons\n            GROUP BY Age\n        END\n\n        SELECT ID, Age INTO #TEMPTABLE FROM NewData.FooNew.PersonsNew\n        \n        UPDATE DemoData.Foo.Persons\n        SET Age = t.Age\n        FROM DemoData.Foo.Persons p\n        JOIN #TEMPTABLE t ON p.ID = t.ID\n\n    END\n",
+                        "language": "SQL"
+                    }
+                }
+            ]
+        }
+    },
+    "systemMetadata": {
+        "lastObserved": 1615443388097,
+        "runId": "mssql-test",
+        "lastRunId": "no-run-id-provided"
+    }
+},
 {
     "entityType": "container",
     "entityUrn": "urn:li:container:250ce23f940485303fa5e5d4f5194975",
@@ -4928,7 +4969,7 @@
                 "actor": "urn:li:corpuser:_ingestion"
             },
             "lastModified": {
-                "time": 1735214618898,
+                "time": 1735588784503,
                 "actor": "urn:li:corpuser:_ingestion"
             }
         }
@@ -5051,7 +5092,7 @@
                 "actor": "urn:li:corpuser:_ingestion"
             },
             "lastModified": {
-                "time": 1735214618906,
+                "time": 1735588784511,
                 "actor": "urn:li:corpuser:_ingestion"
             }
         }
diff --git a/metadata-ingestion/tests/integration/sql_server/golden_files/golden_mces_mssql_no_db_with_filter.json b/metadata-ingestion/tests/integration/sql_server/golden_files/golden_mces_mssql_no_db_with_filter.json
index 07098f0161fc3d..c1748ff13ac93c 100644
--- a/metadata-ingestion/tests/integration/sql_server/golden_files/golden_mces_mssql_no_db_with_filter.json
+++ b/metadata-ingestion/tests/integration/sql_server/golden_files/golden_mces_mssql_no_db_with_filter.json
@@ -95,7 +95,6 @@
     "aspect": {
         "json": {
             "customProperties": {},
-            "externalUrl": "",
             "name": "Weekly Demo Data Backup"
         }
     },
@@ -113,17 +112,16 @@
     "aspect": {
         "json": {
             "customProperties": {
-                "job_id": "a06cfdca-b65e-42de-8db2-8c21c183c5dd",
+                "job_id": "f5a6c120-500a-4300-9b21-0c3225af1f80",
                 "job_name": "Weekly Demo Data Backup",
                 "description": "No description available.",
-                "date_created": "2024-12-26 12:03:35.420000",
-                "date_modified": "2024-12-26 12:03:35.590000",
+                "date_created": "2024-12-30 19:59:24.690000",
+                "date_modified": "2024-12-30 19:59:24.690000",
                 "step_id": "1",
                 "step_name": "Set database to read only",
                 "subsystem": "TSQL",
                 "command": "ALTER DATABASE DemoData SET READ_ONLY"
             },
-            "externalUrl": "",
             "name": "Weekly Demo Data Backup",
             "type": {
                 "string": "MSSQL_JOB_STEP"
@@ -2259,7 +2257,6 @@
     "aspect": {
         "json": {
             "customProperties": {},
-            "externalUrl": "",
             "name": "DemoData.Foo.stored_procedures"
         }
     },
@@ -2282,10 +2279,9 @@
                 "code": "CREATE PROCEDURE [Foo].[Proc.With.SpecialChar] @ID INT\nAS\n    SELECT @ID AS ThatDB;\n",
                 "input parameters": "['@ID']",
                 "parameter @ID": "{'type': 'int'}",
-                "date_created": "2024-12-26 12:03:35.230000",
-                "date_modified": "2024-12-26 12:03:35.230000"
+                "date_created": "2024-12-30 19:59:24.690000",
+                "date_modified": "2024-12-30 19:59:24.690000"
             },
-            "externalUrl": "",
             "name": "DemoData.Foo.Proc.With.SpecialChar",
             "type": {
                 "string": "MSSQL_STORED_PROCEDURE"
@@ -2298,6 +2294,29 @@
         "lastRunId": "no-run-id-provided"
     }
 },
+{
+    "entityType": "dataJob",
+    "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(mssql,DemoData.Foo.stored_procedures,PROD),Proc.With.SpecialChar)",
+    "changeType": "UPSERT",
+    "aspectName": "dataTransformLogic",
+    "aspect": {
+        "json": {
+            "transforms": [
+                {
+                    "queryStatement": {
+                        "value": "CREATE PROCEDURE [Foo].[Proc.With.SpecialChar] @ID INT\nAS\n    SELECT @ID AS ThatDB;\n",
+                        "language": "SQL"
+                    }
+                }
+            ]
+        }
+    },
+    "systemMetadata": {
+        "lastObserved": 1615443388097,
+        "runId": "mssql-test",
+        "lastRunId": "no-run-id-provided"
+    }
+},
 {
     "entityType": "container",
     "entityUrn": "urn:li:container:250ce23f940485303fa5e5d4f5194975",
@@ -2675,7 +2694,7 @@
                 "actor": "urn:li:corpuser:_ingestion"
             },
             "lastModified": {
-                "time": 1735214621644,
+                "time": 1735588789629,
                 "actor": "urn:li:corpuser:_ingestion"
             }
         }
diff --git a/metadata-ingestion/tests/integration/sql_server/golden_files/golden_mces_mssql_to_file.json b/metadata-ingestion/tests/integration/sql_server/golden_files/golden_mces_mssql_to_file.json
index bf30448469c309..cbfc374decab2f 100644
--- a/metadata-ingestion/tests/integration/sql_server/golden_files/golden_mces_mssql_to_file.json
+++ b/metadata-ingestion/tests/integration/sql_server/golden_files/golden_mces_mssql_to_file.json
@@ -102,7 +102,6 @@
     "aspect": {
         "json": {
             "customProperties": {},
-            "externalUrl": "",
             "name": "Weekly Demo Data Backup"
         }
     },
@@ -137,17 +136,16 @@
     "aspect": {
         "json": {
             "customProperties": {
-                "job_id": "a06cfdca-b65e-42de-8db2-8c21c183c5dd",
+                "job_id": "f5a6c120-500a-4300-9b21-0c3225af1f80",
                 "job_name": "Weekly Demo Data Backup",
                 "description": "No description available.",
-                "date_created": "2024-12-26 12:03:35.420000",
-                "date_modified": "2024-12-26 12:03:35.590000",
+                "date_created": "2024-12-30 19:59:24.690000",
+                "date_modified": "2024-12-30 19:59:24.690000",
                 "step_id": "1",
                 "step_name": "Set database to read only",
                 "subsystem": "TSQL",
                 "command": "ALTER DATABASE DemoData SET READ_ONLY"
             },
-            "externalUrl": "",
             "name": "Weekly Demo Data Backup",
             "type": {
                 "string": "MSSQL_JOB_STEP"
@@ -2492,7 +2490,6 @@
     "aspect": {
         "json": {
             "customProperties": {},
-            "externalUrl": "",
             "name": "DemoData.Foo.stored_procedures"
         }
     },
@@ -2532,10 +2529,9 @@
                 "code": "CREATE PROCEDURE [Foo].[Proc.With.SpecialChar] @ID INT\nAS\n    SELECT @ID AS ThatDB;\n",
                 "input parameters": "['@ID']",
                 "parameter @ID": "{'type': 'int'}",
-                "date_created": "2024-12-26 12:03:35.230000",
-                "date_modified": "2024-12-26 12:03:35.230000"
+                "date_created": "2024-12-30 19:59:24.690000",
+                "date_modified": "2024-12-30 19:59:24.690000"
             },
-            "externalUrl": "",
             "name": "DemoData.Foo.Proc.With.SpecialChar",
             "type": {
                 "string": "MSSQL_STORED_PROCEDURE"
@@ -2565,6 +2561,29 @@
         "lastRunId": "no-run-id-provided"
     }
 },
+{
+    "entityType": "dataJob",
+    "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(mssql,my-instance.DemoData.Foo.stored_procedures,PROD),Proc.With.SpecialChar)",
+    "changeType": "UPSERT",
+    "aspectName": "dataTransformLogic",
+    "aspect": {
+        "json": {
+            "transforms": [
+                {
+                    "queryStatement": {
+                        "value": "CREATE PROCEDURE [Foo].[Proc.With.SpecialChar] @ID INT\nAS\n    SELECT @ID AS ThatDB;\n",
+                        "language": "SQL"
+                    }
+                }
+            ]
+        }
+    },
+    "systemMetadata": {
+        "lastObserved": 1615443388097,
+        "runId": "mssql-test",
+        "lastRunId": "no-run-id-provided"
+    }
+},
 {
     "entityType": "dataJob",
     "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(mssql,my-instance.DemoData.Foo.stored_procedures,PROD),NewProc)",
@@ -2577,10 +2596,9 @@
                 "depending_on_procedure": "{}",
                 "code": "CREATE PROCEDURE [Foo].[NewProc]\n    AS\n    BEGIN\n        --insert into items table from salesreason table\n        insert into Foo.Items (ID, ItemName)\n        SELECT TempID, Name\n        FROM Foo.SalesReason;\n\n\n       IF OBJECT_ID('Foo.age_dist', 'U') IS NULL\n       BEGIN\n            -- Create and populate if table doesn't exist\n            SELECT Age, COUNT(*) as Count\n            INTO Foo.age_dist\n            FROM Foo.Persons\n            GROUP BY Age\n        END\n        ELSE\n        BEGIN\n            -- Update existing table\n            TRUNCATE TABLE Foo.age_dist;\n\n            INSERT INTO Foo.age_dist (Age, Count)\n            SELECT Age, COUNT(*) as Count\n            FROM Foo.Persons\n            GROUP BY Age\n        END\n\n        SELECT ID, Age INTO #TEMPTABLE FROM NewData.FooNew.PersonsNew\n        \n        UPDATE DemoData.Foo.Persons\n        SET Age = t.Age\n        FROM DemoData.Foo.Persons p\n        JOIN #TEMPTABLE t ON p.ID = t.ID\n\n    END\n",
                 "input parameters": "[]",
-                "date_created": "2024-12-26 12:03:35.237000",
-                "date_modified": "2024-12-26 12:03:35.237000"
+                "date_created": "2024-12-30 19:59:24.690000",
+                "date_modified": "2024-12-30 19:59:24.690000"
             },
-            "externalUrl": "",
             "name": "DemoData.Foo.NewProc",
             "type": {
                 "string": "MSSQL_STORED_PROCEDURE"
@@ -2610,6 +2628,29 @@
         "lastRunId": "no-run-id-provided"
     }
 },
+{
+    "entityType": "dataJob",
+    "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(mssql,my-instance.DemoData.Foo.stored_procedures,PROD),NewProc)",
+    "changeType": "UPSERT",
+    "aspectName": "dataTransformLogic",
+    "aspect": {
+        "json": {
+            "transforms": [
+                {
+                    "queryStatement": {
+                        "value": "CREATE PROCEDURE [Foo].[NewProc]\n    AS\n    BEGIN\n        --insert into items table from salesreason table\n        insert into Foo.Items (ID, ItemName)\n        SELECT TempID, Name\n        FROM Foo.SalesReason;\n\n\n       IF OBJECT_ID('Foo.age_dist', 'U') IS NULL\n       BEGIN\n            -- Create and populate if table doesn't exist\n            SELECT Age, COUNT(*) as Count\n            INTO Foo.age_dist\n            FROM Foo.Persons\n            GROUP BY Age\n        END\n        ELSE\n        BEGIN\n            -- Update existing table\n            TRUNCATE TABLE Foo.age_dist;\n\n            INSERT INTO Foo.age_dist (Age, Count)\n            SELECT Age, COUNT(*) as Count\n            FROM Foo.Persons\n            GROUP BY Age\n        END\n\n        SELECT ID, Age INTO #TEMPTABLE FROM NewData.FooNew.PersonsNew\n        \n        UPDATE DemoData.Foo.Persons\n        SET Age = t.Age\n        FROM DemoData.Foo.Persons p\n        JOIN #TEMPTABLE t ON p.ID = t.ID\n\n    END\n",
+                        "language": "SQL"
+                    }
+                }
+            ]
+        }
+    },
+    "systemMetadata": {
+        "lastObserved": 1615443388097,
+        "runId": "mssql-test",
+        "lastRunId": "no-run-id-provided"
+    }
+},
 {
     "entityType": "container",
     "entityUrn": "urn:li:container:5631370915311469374ef3cb5f0ebbf0",
@@ -3005,7 +3046,7 @@
                 "actor": "urn:li:corpuser:_ingestion"
             },
             "lastModified": {
-                "time": 1735214620908,
+                "time": 1735588787786,
                 "actor": "urn:li:corpuser:_ingestion"
             }
         }
diff --git a/metadata-ingestion/tests/integration/sql_server/golden_files/golden_mces_mssql_with_lower_case_urn.json b/metadata-ingestion/tests/integration/sql_server/golden_files/golden_mces_mssql_with_lower_case_urn.json
index ff27989d71de1b..70f5784d6d0b18 100644
--- a/metadata-ingestion/tests/integration/sql_server/golden_files/golden_mces_mssql_with_lower_case_urn.json
+++ b/metadata-ingestion/tests/integration/sql_server/golden_files/golden_mces_mssql_with_lower_case_urn.json
@@ -95,7 +95,6 @@
     "aspect": {
         "json": {
             "customProperties": {},
-            "externalUrl": "",
             "name": "Weekly Demo Data Backup"
         }
     },
@@ -113,17 +112,16 @@
     "aspect": {
         "json": {
             "customProperties": {
-                "job_id": "a06cfdca-b65e-42de-8db2-8c21c183c5dd",
+                "job_id": "f5a6c120-500a-4300-9b21-0c3225af1f80",
                 "job_name": "Weekly Demo Data Backup",
                 "description": "No description available.",
-                "date_created": "2024-12-26 12:03:35.420000",
-                "date_modified": "2024-12-26 12:03:35.590000",
+                "date_created": "2024-12-30 19:59:24.690000",
+                "date_modified": "2024-12-30 19:59:24.690000",
                 "step_id": "1",
                 "step_name": "Set database to read only",
                 "subsystem": "TSQL",
                 "command": "ALTER DATABASE DemoData SET READ_ONLY"
             },
-            "externalUrl": "",
             "name": "Weekly Demo Data Backup",
             "type": {
                 "string": "MSSQL_JOB_STEP"
@@ -2259,7 +2257,6 @@
     "aspect": {
         "json": {
             "customProperties": {},
-            "externalUrl": "",
             "name": "DemoData.Foo.stored_procedures"
         }
     },
@@ -2282,10 +2279,9 @@
                 "code": "CREATE PROCEDURE [Foo].[Proc.With.SpecialChar] @ID INT\nAS\n    SELECT @ID AS ThatDB;\n",
                 "input parameters": "['@ID']",
                 "parameter @ID": "{'type': 'int'}",
-                "date_created": "2024-12-26 12:03:35.230000",
-                "date_modified": "2024-12-26 12:03:35.230000"
+                "date_created": "2024-12-30 19:59:24.690000",
+                "date_modified": "2024-12-30 19:59:24.690000"
             },
-            "externalUrl": "",
             "name": "DemoData.Foo.Proc.With.SpecialChar",
             "type": {
                 "string": "MSSQL_STORED_PROCEDURE"
@@ -2298,6 +2294,29 @@
         "lastRunId": "no-run-id-provided"
     }
 },
+{
+    "entityType": "dataJob",
+    "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(mssql,DemoData.Foo.stored_procedures,PROD),Proc.With.SpecialChar)",
+    "changeType": "UPSERT",
+    "aspectName": "dataTransformLogic",
+    "aspect": {
+        "json": {
+            "transforms": [
+                {
+                    "queryStatement": {
+                        "value": "CREATE PROCEDURE [Foo].[Proc.With.SpecialChar] @ID INT\nAS\n    SELECT @ID AS ThatDB;\n",
+                        "language": "SQL"
+                    }
+                }
+            ]
+        }
+    },
+    "systemMetadata": {
+        "lastObserved": 1615443388097,
+        "runId": "mssql-test",
+        "lastRunId": "no-run-id-provided"
+    }
+},
 {
     "entityType": "dataJob",
     "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(mssql,DemoData.Foo.stored_procedures,PROD),NewProc)",
@@ -2310,10 +2329,9 @@
                 "depending_on_procedure": "{}",
                 "code": "CREATE PROCEDURE [Foo].[NewProc]\n    AS\n    BEGIN\n        --insert into items table from salesreason table\n        insert into Foo.Items (ID, ItemName)\n        SELECT TempID, Name\n        FROM Foo.SalesReason;\n\n\n       IF OBJECT_ID('Foo.age_dist', 'U') IS NULL\n       BEGIN\n            -- Create and populate if table doesn't exist\n            SELECT Age, COUNT(*) as Count\n            INTO Foo.age_dist\n            FROM Foo.Persons\n            GROUP BY Age\n        END\n        ELSE\n        BEGIN\n            -- Update existing table\n            TRUNCATE TABLE Foo.age_dist;\n\n            INSERT INTO Foo.age_dist (Age, Count)\n            SELECT Age, COUNT(*) as Count\n            FROM Foo.Persons\n            GROUP BY Age\n        END\n\n        SELECT ID, Age INTO #TEMPTABLE FROM NewData.FooNew.PersonsNew\n        \n        UPDATE DemoData.Foo.Persons\n        SET Age = t.Age\n        FROM DemoData.Foo.Persons p\n        JOIN #TEMPTABLE t ON p.ID = t.ID\n\n    END\n",
                 "input parameters": "[]",
-                "date_created": "2024-12-26 12:03:35.237000",
-                "date_modified": "2024-12-26 12:03:35.237000"
+                "date_created": "2024-12-30 19:59:24.690000",
+                "date_modified": "2024-12-30 19:59:24.690000"
             },
-            "externalUrl": "",
             "name": "DemoData.Foo.NewProc",
             "type": {
                 "string": "MSSQL_STORED_PROCEDURE"
@@ -2326,6 +2344,29 @@
         "lastRunId": "no-run-id-provided"
     }
 },
+{
+    "entityType": "dataJob",
+    "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(mssql,DemoData.Foo.stored_procedures,PROD),NewProc)",
+    "changeType": "UPSERT",
+    "aspectName": "dataTransformLogic",
+    "aspect": {
+        "json": {
+            "transforms": [
+                {
+                    "queryStatement": {
+                        "value": "CREATE PROCEDURE [Foo].[NewProc]\n    AS\n    BEGIN\n        --insert into items table from salesreason table\n        insert into Foo.Items (ID, ItemName)\n        SELECT TempID, Name\n        FROM Foo.SalesReason;\n\n\n       IF OBJECT_ID('Foo.age_dist', 'U') IS NULL\n       BEGIN\n            -- Create and populate if table doesn't exist\n            SELECT Age, COUNT(*) as Count\n            INTO Foo.age_dist\n            FROM Foo.Persons\n            GROUP BY Age\n        END\n        ELSE\n        BEGIN\n            -- Update existing table\n            TRUNCATE TABLE Foo.age_dist;\n\n            INSERT INTO Foo.age_dist (Age, Count)\n            SELECT Age, COUNT(*) as Count\n            FROM Foo.Persons\n            GROUP BY Age\n        END\n\n        SELECT ID, Age INTO #TEMPTABLE FROM NewData.FooNew.PersonsNew\n        \n        UPDATE DemoData.Foo.Persons\n        SET Age = t.Age\n        FROM DemoData.Foo.Persons p\n        JOIN #TEMPTABLE t ON p.ID = t.ID\n\n    END\n",
+                        "language": "SQL"
+                    }
+                }
+            ]
+        }
+    },
+    "systemMetadata": {
+        "lastObserved": 1615443388097,
+        "runId": "mssql-test",
+        "lastRunId": "no-run-id-provided"
+    }
+},
 {
     "entityType": "container",
     "entityUrn": "urn:li:container:250ce23f940485303fa5e5d4f5194975",
@@ -4978,7 +5019,7 @@
                 "actor": "urn:li:corpuser:_ingestion"
             },
             "lastModified": {
-                "time": 1735214622805,
+                "time": 1735588791954,
                 "actor": "urn:li:corpuser:_ingestion"
             }
         }
@@ -5125,7 +5166,7 @@
                 "actor": "urn:li:corpuser:_ingestion"
             },
             "lastModified": {
-                "time": 1735214622810,
+                "time": 1735588791966,
                 "actor": "urn:li:corpuser:_ingestion"
             }
         }
diff --git a/metadata-ingestion/tests/integration/sql_server/test_sql_server.py b/metadata-ingestion/tests/integration/sql_server/test_sql_server.py
index b969f77b4c3c18..7fab5fc7dae1ba 100644
--- a/metadata-ingestion/tests/integration/sql_server/test_sql_server.py
+++ b/metadata-ingestion/tests/integration/sql_server/test_sql_server.py
@@ -57,7 +57,7 @@ def test_mssql_ingest(mssql_runner, pytestconfig, tmp_path, mock_time, config_fi
         pytestconfig,
         output_path=tmp_path / "mssql_mces.json",
         golden_path=test_resources_dir
-        / f"golden_files/golden_mces_{config_file.replace('yml','json')}",
+        / f"golden_files/golden_mces_{config_file.replace('yml', 'json')}",
         ignore_paths=[
             r"root\[\d+\]\['aspect'\]\['json'\]\['customProperties'\]\['job_id'\]",
             r"root\[\d+\]\['aspect'\]\['json'\]\['customProperties'\]\['date_created'\]",
diff --git a/metadata-ingestion/tests/integration/tableau/test_tableau_ingest.py b/metadata-ingestion/tests/integration/tableau/test_tableau_ingest.py
index d7868038a40aa1..32b1ef2ed1f835 100644
--- a/metadata-ingestion/tests/integration/tableau/test_tableau_ingest.py
+++ b/metadata-ingestion/tests/integration/tableau/test_tableau_ingest.py
@@ -69,7 +69,7 @@
     "projects": ["default", "Project 2", "Samples"],
     "extract_project_hierarchy": False,
     "page_size": 1000,
-    "workbook_page_size": 1000,
+    "workbook_page_size": None,
     "ingest_tags": True,
     "ingest_owner": True,
     "ingest_tables_external": True,
@@ -645,7 +645,7 @@ def test_tableau_ingest_with_platform_instance(
         "platform_instance": "acryl_site1",
         "projects": ["default", "Project 2"],
         "page_size": 1000,
-        "workbook_page_size": 1000,
+        "workbook_page_size": None,
         "ingest_tags": True,
         "ingest_owner": True,
         "ingest_tables_external": True,
diff --git a/metadata-ingestion/tests/integration/unity/test_unity_catalog_ingest.py b/metadata-ingestion/tests/integration/unity/test_unity_catalog_ingest.py
index b8b0563a1d24e5..9c7b86a275f6d0 100644
--- a/metadata-ingestion/tests/integration/unity/test_unity_catalog_ingest.py
+++ b/metadata-ingestion/tests/integration/unity/test_unity_catalog_ingest.py
@@ -205,55 +205,57 @@ def register_mock_data(workspace_client):
         ),
     ]
 
-    workspace_client.tables.get = lambda *args, **kwargs: databricks.sdk.service.catalog.TableInfo.from_dict(
-        {
-            "name": "quickstart_table",
-            "catalog_name": "quickstart_catalog",
-            "schema_name": "quickstart_schema",
-            "table_type": "MANAGED",
-            "data_source_format": "DELTA",
-            "columns": [
-                {
-                    "name": "columnA",
-                    "type_text": "int",
-                    "type_json": '{"name":"columnA","type":"integer","nullable":true,"metadata":{}}',
-                    "type_name": "INT",
-                    "type_precision": 0,
-                    "type_scale": 0,
-                    "position": 0,
-                    "nullable": True,
-                },
-                {
-                    "name": "columnB",
-                    "type_text": "string",
-                    "type_json": '{"name":"columnB","type":"string","nullable":true,"metadata":{}}',
-                    "type_name": "STRING",
-                    "type_precision": 0,
-                    "type_scale": 0,
-                    "position": 1,
-                    "nullable": True,
+    workspace_client.tables.get = (
+        lambda *args, **kwargs: databricks.sdk.service.catalog.TableInfo.from_dict(
+            {
+                "name": "quickstart_table",
+                "catalog_name": "quickstart_catalog",
+                "schema_name": "quickstart_schema",
+                "table_type": "MANAGED",
+                "data_source_format": "DELTA",
+                "columns": [
+                    {
+                        "name": "columnA",
+                        "type_text": "int",
+                        "type_json": '{"name":"columnA","type":"integer","nullable":true,"metadata":{}}',
+                        "type_name": "INT",
+                        "type_precision": 0,
+                        "type_scale": 0,
+                        "position": 0,
+                        "nullable": True,
+                    },
+                    {
+                        "name": "columnB",
+                        "type_text": "string",
+                        "type_json": '{"name":"columnB","type":"string","nullable":true,"metadata":{}}',
+                        "type_name": "STRING",
+                        "type_precision": 0,
+                        "type_scale": 0,
+                        "position": 1,
+                        "nullable": True,
+                    },
+                ],
+                "storage_location": "s3://db-02eec1f70bfe4115445be9fdb1aac6ac-s3-root-bucket/metastore/2c983545-d403-4f87-9063-5b7e3b6d3736/tables/cff27aa1-1c6a-4d78-b713-562c660c2896",
+                "owner": "account users",
+                "properties": {
+                    "delta.lastCommitTimestamp": "1666185711000",
+                    "delta.lastUpdateVersion": "1",
+                    "delta.minReaderVersion": "1",
+                    "delta.minWriterVersion": "2",
+                    "spark.sql.statistics.numRows": "10",
+                    "spark.sql.statistics.totalSize": "512",
                 },
-            ],
-            "storage_location": "s3://db-02eec1f70bfe4115445be9fdb1aac6ac-s3-root-bucket/metastore/2c983545-d403-4f87-9063-5b7e3b6d3736/tables/cff27aa1-1c6a-4d78-b713-562c660c2896",
-            "owner": "account users",
-            "properties": {
-                "delta.lastCommitTimestamp": "1666185711000",
-                "delta.lastUpdateVersion": "1",
-                "delta.minReaderVersion": "1",
-                "delta.minWriterVersion": "2",
-                "spark.sql.statistics.numRows": "10",
-                "spark.sql.statistics.totalSize": "512",
-            },
-            "generation": 2,
-            "metastore_id": "2c983545-d403-4f87-9063-5b7e3b6d3736",
-            "full_name": "quickstart_catalog.quickstart_schema.quickstart_table",
-            "data_access_configuration_id": "00000000-0000-0000-0000-000000000000",
-            "created_at": 1666185698688,
-            "created_by": "abc@acryl.io",
-            "updated_at": 1666186049633,
-            "updated_by": "abc@acryl.io",
-            "table_id": "cff27aa1-1c6a-4d78-b713-562c660c2896",
-        }
+                "generation": 2,
+                "metastore_id": "2c983545-d403-4f87-9063-5b7e3b6d3736",
+                "full_name": "quickstart_catalog.quickstart_schema.quickstart_table",
+                "data_access_configuration_id": "00000000-0000-0000-0000-000000000000",
+                "created_at": 1666185698688,
+                "created_by": "abc@acryl.io",
+                "updated_at": 1666186049633,
+                "updated_by": "abc@acryl.io",
+                "table_id": "cff27aa1-1c6a-4d78-b713-562c660c2896",
+            }
+        )
     )
 
     workspace_client.service_principals.list.return_value = [
@@ -437,9 +439,7 @@ def test_ingestion(pytestconfig, tmp_path, requests_mock):
         "datahub.ingestion.source.unity.proxy.WorkspaceClient"
     ) as mock_client, patch.object(
         HiveMetastoreProxy, "get_inspector"
-    ) as get_inspector, patch.object(
-        HiveMetastoreProxy, "_execute_sql"
-    ) as execute_sql:
+    ) as get_inspector, patch.object(HiveMetastoreProxy, "_execute_sql") as execute_sql:
         workspace_client: mock.MagicMock = mock.MagicMock()
         mock_client.return_value = workspace_client
         register_mock_data(workspace_client)
diff --git a/metadata-ingestion/tests/performance/data_generation.py b/metadata-ingestion/tests/performance/data_generation.py
index fcff13edf59363..266c0d9af03224 100644
--- a/metadata-ingestion/tests/performance/data_generation.py
+++ b/metadata-ingestion/tests/performance/data_generation.py
@@ -198,7 +198,7 @@ def generate_queries(
 
     all_tables = seed_metadata.tables + seed_metadata.views
     users = [f"user_{i}@xyz.com" for i in range(num_users)]
-    for i in range(num_selects):  # Pure SELECT statements
+    for _ in range(num_selects):  # Pure SELECT statements
         tables = _sample_list(all_tables, tables_per_select)
         all_columns = [
             FieldAccess(column, table) for table in tables for column in table.columns
@@ -213,7 +213,7 @@ def generate_queries(
             fields_accessed=_sample_list(all_columns, columns_per_select),
         )
 
-    for i in range(num_operations):
+    for _ in range(num_operations):
         modified_table = random.choice(seed_metadata.tables)
         n_col = len(modified_table.columns)
         num_columns_modified = NormalDistribution(n_col / 2, n_col / 2)
diff --git a/metadata-ingestion/tests/performance/databricks/generator.py b/metadata-ingestion/tests/performance/databricks/generator.py
index 29df325d856a1a..b11771e55b2c9e 100644
--- a/metadata-ingestion/tests/performance/databricks/generator.py
+++ b/metadata-ingestion/tests/performance/databricks/generator.py
@@ -167,7 +167,7 @@ def _generate_insert_lineage(table: Table, upstream: Table) -> str:
 def _generate_view_definition(view: View) -> str:
     from_statement = f"FROM {_quote_table(view.upstreams[0])} t0"
     join_statement = " ".join(
-        f"JOIN {_quote_table(upstream)} t{i+1} ON t0.id = t{i+1}.id"
+        f"JOIN {_quote_table(upstream)} t{i + 1} ON t0.id = t{i + 1}.id"
         for i, upstream in enumerate(view.upstreams[1:])
     )
     return f"CREATE VIEW {_quote_table(view)} AS SELECT * {from_statement} {join_statement} {view.definition}"
diff --git a/metadata-ingestion/tests/test_helpers/mce_helpers.py b/metadata-ingestion/tests/test_helpers/mce_helpers.py
index 0105e6d596970b..d70a440dab0657 100644
--- a/metadata-ingestion/tests/test_helpers/mce_helpers.py
+++ b/metadata-ingestion/tests/test_helpers/mce_helpers.py
@@ -300,9 +300,9 @@ def assert_for_each_entity(
     for urn, aspect_val in aspect_map.items():
         if aspect_val is not None:
             for f in aspect_field_matcher:
-                assert aspect_field_matcher[f] == _get_element(
-                    aspect_val, [f]
-                ), f"urn: {urn} -> Field {f} must match value {aspect_field_matcher[f]}, found {_get_element(aspect_val, [f])}"
+                assert aspect_field_matcher[f] == _get_element(aspect_val, [f]), (
+                    f"urn: {urn} -> Field {f} must match value {aspect_field_matcher[f]}, found {_get_element(aspect_val, [f])}"
+                )
             success.append(urn)
         elif urn not in exception_urns:
             print(f"Adding {urn} to failures")
@@ -361,9 +361,9 @@ def assert_entity_mcp_aspect(
             assert mcp.aspect
             aspect_val = mcp.aspect.to_obj()
             for f in aspect_field_matcher:
-                assert aspect_field_matcher[f] == _get_element(
-                    aspect_val, [f]
-                ), f"urn: {mcp.entityUrn} -> Field {f} must match value {aspect_field_matcher[f]}, found {_get_element(aspect_val, [f])}"
+                assert aspect_field_matcher[f] == _get_element(aspect_val, [f]), (
+                    f"urn: {mcp.entityUrn} -> Field {f} must match value {aspect_field_matcher[f]}, found {_get_element(aspect_val, [f])}"
+                )
                 matches = matches + 1
     return matches
 
diff --git a/metadata-ingestion/tests/test_helpers/state_helpers.py b/metadata-ingestion/tests/test_helpers/state_helpers.py
index f68aef742fc730..c469db6ce8cf80 100644
--- a/metadata-ingestion/tests/test_helpers/state_helpers.py
+++ b/metadata-ingestion/tests/test_helpers/state_helpers.py
@@ -104,7 +104,7 @@ def monkey_patch_get_latest_timeseries_value(
 
 @pytest.fixture
 def mock_datahub_graph_instance(
-    mock_datahub_graph: Callable[[DatahubClientConfig], DataHubGraph]
+    mock_datahub_graph: Callable[[DatahubClientConfig], DataHubGraph],
 ) -> DataHubGraph:
     return mock_datahub_graph(DatahubClientConfig(server="http://fake.domain.local"))
 
diff --git a/metadata-ingestion/tests/unit/api/entities/dataproducts/test_dataproduct.py b/metadata-ingestion/tests/unit/api/entities/dataproducts/test_dataproduct.py
index e796f0b3f37219..dad7662d9ad00b 100644
--- a/metadata-ingestion/tests/unit/api/entities/dataproducts/test_dataproduct.py
+++ b/metadata-ingestion/tests/unit/api/entities/dataproducts/test_dataproduct.py
@@ -26,7 +26,7 @@ def base_entity_metadata():
 
 @pytest.fixture
 def base_mock_graph(
-    base_entity_metadata: Dict[str, Dict[str, Any]]
+    base_entity_metadata: Dict[str, Dict[str, Any]],
 ) -> MockDataHubGraph:
     return MockDataHubGraph(entity_graph=base_entity_metadata)
 
diff --git a/metadata-ingestion/tests/unit/api/source_helpers/test_ensure_aspect_size.py b/metadata-ingestion/tests/unit/api/source_helpers/test_ensure_aspect_size.py
index bdf1e0a2e0e860..8a45efb46893ae 100644
--- a/metadata-ingestion/tests/unit/api/source_helpers/test_ensure_aspect_size.py
+++ b/metadata-ingestion/tests/unit/api/source_helpers/test_ensure_aspect_size.py
@@ -159,9 +159,9 @@ def test_ensure_size_of_proper_dataset_profile(processor):
     processor.ensure_dataset_profile_size(
         "urn:li:dataset:(s3, dummy_dataset, DEV)", profile
     )
-    assert orig_repr == json.dumps(
-        profile.to_obj()
-    ), "Aspect was modified in case where workunit processor should have been no-op"
+    assert orig_repr == json.dumps(profile.to_obj()), (
+        "Aspect was modified in case where workunit processor should have been no-op"
+    )
 
 
 @freeze_time("2023-01-02 00:00:00")
@@ -177,9 +177,9 @@ def test_ensure_size_of_too_big_schema_metadata(processor):
     # +100kb is completely arbitrary, but we are truncating the aspect based on schema fields size only, not total taken
     # by other parameters of the aspect - it is reasonable approach though - schema fields is the only field in schema
     # metadata which can be expected to grow out of control
-    assert (
-        len(json.dumps(schema.to_obj())) < INGEST_MAX_PAYLOAD_BYTES + 100000
-    ), "Aspect exceeded acceptable size"
+    assert len(json.dumps(schema.to_obj())) < INGEST_MAX_PAYLOAD_BYTES + 100000, (
+        "Aspect exceeded acceptable size"
+    )
 
 
 @freeze_time("2023-01-02 00:00:00")
@@ -189,9 +189,9 @@ def test_ensure_size_of_proper_schema_metadata(processor):
     processor.ensure_schema_metadata_size(
         "urn:li:dataset:(s3, dummy_dataset, DEV)", schema
     )
-    assert orig_repr == json.dumps(
-        schema.to_obj()
-    ), "Aspect was modified in case where workunit processor should have been no-op"
+    assert orig_repr == json.dumps(schema.to_obj()), (
+        "Aspect was modified in case where workunit processor should have been no-op"
+    )
 
 
 @freeze_time("2023-01-02 00:00:00")
@@ -214,9 +214,9 @@ def test_ensure_size_of_too_big_dataset_profile(processor):
     )
     assert expected_profile.fieldProfiles
     expected_profile.fieldProfiles.insert(4, reduced_field)
-    assert json.dumps(profile.to_obj()) == json.dumps(
-        expected_profile.to_obj()
-    ), "Field 'big' was not properly removed from aspect due to its size"
+    assert json.dumps(profile.to_obj()) == json.dumps(expected_profile.to_obj()), (
+        "Field 'big' was not properly removed from aspect due to its size"
+    )
 
 
 @freeze_time("2023-01-02 00:00:00")
diff --git a/metadata-ingestion/tests/unit/bigquery/test_bigquery_lineage.py b/metadata-ingestion/tests/unit/bigquery/test_bigquery_lineage.py
index a1981ccf767916..f494ed78211dcf 100644
--- a/metadata-ingestion/tests/unit/bigquery/test_bigquery_lineage.py
+++ b/metadata-ingestion/tests/unit/bigquery/test_bigquery_lineage.py
@@ -232,9 +232,9 @@ def fake_schema_metadata(entity_urn: str) -> models.SchemaMetadataClass:
         else []
         for fine_grained_lineage in upstream_lineage.fineGrainedLineages
     ]
-    assert all(
-        urn in expected_schema_field_urns for urn in actual_schema_field_urns
-    ), "Some expected column URNs are missing from fine grained lineage."
+    assert all(urn in expected_schema_field_urns for urn in actual_schema_field_urns), (
+        "Some expected column URNs are missing from fine grained lineage."
+    )
 
 
 def test_lineage_for_external_bq_table_no_column_lineage(mock_datahub_graph_instance):
@@ -286,9 +286,9 @@ def fake_schema_metadata(entity_urn: str) -> Optional[models.SchemaMetadataClass
     assert len(upstream_lineage.upstreams) == 3
     # Extracting dataset URNs from upstream_lineage.upstreams
     actual_dataset_urns = [upstream.dataset for upstream in upstream_lineage.upstreams]
-    assert all(
-        urn in actual_dataset_urns for urn in expected_dataset_urns
-    ), "Some expected dataset URNs are missing from upstream lineage."
+    assert all(urn in actual_dataset_urns for urn in expected_dataset_urns), (
+        "Some expected dataset URNs are missing from upstream lineage."
+    )
     assert upstream_lineage.fineGrainedLineages is None
 
 
diff --git a/metadata-ingestion/tests/unit/cli/assertion/test_compile.py b/metadata-ingestion/tests/unit/cli/assertion/test_compile.py
index 47253b5b0d71ea..0a1870d83212e8 100644
--- a/metadata-ingestion/tests/unit/cli/assertion/test_compile.py
+++ b/metadata-ingestion/tests/unit/cli/assertion/test_compile.py
@@ -37,6 +37,6 @@ def test_compile_assertion_config_spec_for_snowflake(pytestconfig, tmp_path):
 
     for file_name in output_file_names:
         assert os.path.exists(tmp_path / file_name)
-        assert filecmp.cmp(
-            golden_file_path / file_name, tmp_path / file_name
-        ), f"{file_name} is not as expected"
+        assert filecmp.cmp(golden_file_path / file_name, tmp_path / file_name), (
+            f"{file_name} is not as expected"
+        )
diff --git a/metadata-ingestion/tests/unit/config/test_config_loader.py b/metadata-ingestion/tests/unit/config/test_config_loader.py
index 25ee289ec4e4e7..43781acd7f80c0 100644
--- a/metadata-ingestion/tests/unit/config/test_config_loader.py
+++ b/metadata-ingestion/tests/unit/config/test_config_loader.py
@@ -10,6 +10,7 @@
 
 from datahub.configuration.common import ConfigurationError
 from datahub.configuration.config_loader import (
+    EnvResolver,
     list_referenced_env_variables,
     load_config_file,
 )
@@ -138,6 +139,28 @@ def test_load_success(pytestconfig, filename, golden_config, env, referenced_env
         # TODO check referenced env vars
 
 
+def test_load_strict_env_syntax() -> None:
+    config = {
+        "foo": "${BAR}",
+        "baz": "$BAZ",
+        "qux": "qux$QUX",
+    }
+    assert EnvResolver.list_referenced_variables(
+        config,
+        strict_env_syntax=True,
+    ) == {"BAR"}
+
+    assert EnvResolver(
+        environ={
+            "BAR": "bar",
+        }
+    ).resolve(config) == {
+        "foo": "bar",
+        "baz": "$BAZ",
+        "qux": "qux$QUX",
+    }
+
+
 @pytest.mark.parametrize(
     "filename,env,error_type",
     [
diff --git a/metadata-ingestion/tests/unit/redshift/test_redshift_lineage.py b/metadata-ingestion/tests/unit/redshift/test_redshift_lineage.py
index 941d13be0a6139..27045dfc656cbe 100644
--- a/metadata-ingestion/tests/unit/redshift/test_redshift_lineage.py
+++ b/metadata-ingestion/tests/unit/redshift/test_redshift_lineage.py
@@ -262,8 +262,7 @@ def test_collapse_temp_lineage():
     lineage_item: LineageItem = lineage_extractor._lineage_map[target_urn]
 
     assert list(lineage_item.upstreams)[0].urn == (
-        "urn:li:dataset:(urn:li:dataPlatform:redshift,"
-        "test.public.player_activity,PROD)"
+        "urn:li:dataset:(urn:li:dataPlatform:redshift,test.public.player_activity,PROD)"
     )
 
     assert lineage_item.cll is not None
@@ -276,8 +275,7 @@ def test_collapse_temp_lineage():
     assert lineage_item.cll[0].downstream.column == "price"
 
     assert lineage_item.cll[0].upstreams[0].table == (
-        "urn:li:dataset:(urn:li:dataPlatform:redshift,"
-        "test.public.player_activity,PROD)"
+        "urn:li:dataset:(urn:li:dataPlatform:redshift,test.public.player_activity,PROD)"
     )
 
     assert lineage_item.cll[0].upstreams[0].column == "price"
@@ -441,8 +439,7 @@ def test_collapse_temp_recursive_cll_lineage():
     )
 
     assert target_dataset_cll[0].upstreams[0].table == (
-        "urn:li:dataset:(urn:li:dataPlatform:redshift,"
-        "dev.public.player_activity,PROD)"
+        "urn:li:dataset:(urn:li:dataPlatform:redshift,dev.public.player_activity,PROD)"
     )
     assert target_dataset_cll[0].upstreams[0].column == "price"
 
@@ -638,8 +635,7 @@ def test_collapse_temp_recursive_with_compex_column_cll_lineage():
     )
 
     assert target_dataset_cll[0].upstreams[0].table == (
-        "urn:li:dataset:(urn:li:dataPlatform:redshift,"
-        "dev.public.player_activity,PROD)"
+        "urn:li:dataset:(urn:li:dataPlatform:redshift,dev.public.player_activity,PROD)"
     )
     assert target_dataset_cll[0].upstreams[0].column == "price"
     assert target_dataset_cll[0].upstreams[1].column == "tax"
diff --git a/metadata-ingestion/tests/unit/s3/test_s3_util.py b/metadata-ingestion/tests/unit/s3/test_s3_util.py
deleted file mode 100644
index 7850d65ca8b01f..00000000000000
--- a/metadata-ingestion/tests/unit/s3/test_s3_util.py
+++ /dev/null
@@ -1,29 +0,0 @@
-from unittest.mock import Mock
-
-from datahub.ingestion.source.aws.s3_util import group_s3_objects_by_dirname
-
-
-def test_group_s3_objects_by_dirname():
-    s3_objects = [
-        Mock(key="/dir1/file1.txt"),
-        Mock(key="/dir2/file2.txt"),
-        Mock(key="/dir1/file3.txt"),
-    ]
-
-    grouped_objects = group_s3_objects_by_dirname(s3_objects)
-
-    assert len(grouped_objects) == 2
-    assert grouped_objects["/dir1"] == [s3_objects[0], s3_objects[2]]
-    assert grouped_objects["/dir2"] == [s3_objects[1]]
-
-
-def test_group_s3_objects_by_dirname_files_in_root_directory():
-    s3_objects = [
-        Mock(key="file1.txt"),
-        Mock(key="file2.txt"),
-    ]
-
-    grouped_objects = group_s3_objects_by_dirname(s3_objects)
-
-    assert len(grouped_objects) == 1
-    assert grouped_objects["/"] == s3_objects
diff --git a/metadata-ingestion/tests/unit/sdk/test_mcp_builder.py b/metadata-ingestion/tests/unit/sdk/test_mcp_builder.py
index e304edb24789cd..8032df632c549e 100644
--- a/metadata-ingestion/tests/unit/sdk/test_mcp_builder.py
+++ b/metadata-ingestion/tests/unit/sdk/test_mcp_builder.py
@@ -86,6 +86,70 @@ def test_guid_generators():
     assert guid == guid_datahub
 
 
+def _assert_eq(a: builder.ContainerKey, b: builder.ContainerKey) -> None:
+    assert a == b
+    assert a.__class__ is b.__class__
+    assert a.guid() == b.guid()
+    assert a.property_dict() == b.property_dict()
+
+
+def test_parent_key() -> None:
+    schema_key = builder.SchemaKey(
+        database="test",
+        schema="Test",
+        platform="mysql",
+        instance="TestInstance",
+        env="DEV",
+    )
+    db_key = schema_key.parent_key()
+    assert db_key is not None
+    _assert_eq(
+        db_key,
+        builder.DatabaseKey(
+            database="test",
+            platform="mysql",
+            instance="TestInstance",
+            env="DEV",
+        ),
+    )
+
+    assert db_key.parent_key() is None
+
+
+def test_parent_key_with_backcompat_env_as_instance() -> None:
+    schema_key = builder.SchemaKey(
+        database="test",
+        schema="Test",
+        platform="mysql",
+        instance=None,
+        env="PROD",
+        backcompat_env_as_instance=True,
+    )
+    db_key = schema_key.parent_key()
+    assert db_key is not None
+    _assert_eq(
+        db_key,
+        builder.DatabaseKey(
+            database="test",
+            platform="mysql",
+            instance=None,
+            env="PROD",
+            backcompat_env_as_instance=True,
+        ),
+    )
+
+
+def test_parent_key_on_container_key() -> None:
+    # In general, people shouldn't be calling parent_key() on ContainerKey directly.
+    # But just in case, we should make sure it works.
+    container_key = builder.ContainerKey(
+        platform="bigquery",
+        name="test",
+        env="DEV",
+    )
+    assert container_key.parent_key() is None
+
+
 def test_entity_supports_aspect():
     assert builder.entity_supports_aspect("dataset", StatusClass)
     assert not builder.entity_supports_aspect("telemetry", StatusClass)
diff --git a/metadata-ingestion/tests/unit/serde/test_codegen.py b/metadata-ingestion/tests/unit/serde/test_codegen.py
index b49f7153129136..13fcf3d919cc03 100644
--- a/metadata-ingestion/tests/unit/serde/test_codegen.py
+++ b/metadata-ingestion/tests/unit/serde/test_codegen.py
@@ -156,9 +156,9 @@ def _err(msg: str) -> None:
                         f"entity {entity_type}: aspect {aspect_name} is missing from the entity registry"
                     )
 
-    assert (
-        not errors
-    ), f'To fix these errors, run "UPDATE_ENTITY_REGISTRY=true pytest {__file__}"'
+    assert not errors, (
+        f'To fix these errors, run "UPDATE_ENTITY_REGISTRY=true pytest {__file__}"'
+    )
 
 
 def test_enum_options():
diff --git a/metadata-ingestion/tests/unit/sql_parsing/aggregator_goldens/test_overlapping_inserts.json b/metadata-ingestion/tests/unit/sql_parsing/aggregator_goldens/test_overlapping_inserts.json
index fcbe0ec5aeb839..fa693381854972 100644
--- a/metadata-ingestion/tests/unit/sql_parsing/aggregator_goldens/test_overlapping_inserts.json
+++ b/metadata-ingestion/tests/unit/sql_parsing/aggregator_goldens/test_overlapping_inserts.json
@@ -38,26 +38,26 @@
                 {
                     "upstreamType": "FIELD_SET",
                     "upstreams": [
-                        "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:redshift,dev.public.upstream1,PROD),a)"
+                        "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:redshift,dev.public.upstream2,PROD),a)"
                     ],
                     "downstreamType": "FIELD",
                     "downstreams": [
                         "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:redshift,dev.public.downstream,PROD),a)"
                     ],
                     "confidenceScore": 0.2,
-                    "query": "urn:li:query:377a73bbf094c8b176b15157c24242cdfc7a0f407d78e52e63ded08c913468f1"
+                    "query": "urn:li:query:c4b3a21ef8c262ebbe99a5bdb6c29cb0be646392bb4af10b6f4a758af881470e"
                 },
                 {
                     "upstreamType": "FIELD_SET",
                     "upstreams": [
-                        "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:redshift,dev.public.upstream2,PROD),a)"
+                        "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:redshift,dev.public.upstream1,PROD),a)"
                     ],
                     "downstreamType": "FIELD",
                     "downstreams": [
                         "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:redshift,dev.public.downstream,PROD),a)"
                     ],
                     "confidenceScore": 0.2,
-                    "query": "urn:li:query:c4b3a21ef8c262ebbe99a5bdb6c29cb0be646392bb4af10b6f4a758af881470e"
+                    "query": "urn:li:query:377a73bbf094c8b176b15157c24242cdfc7a0f407d78e52e63ded08c913468f1"
                 },
                 {
                     "upstreamType": "FIELD_SET",
diff --git a/metadata-ingestion/tests/unit/sql_parsing/aggregator_goldens/test_overlapping_inserts_from_temp_tables.json b/metadata-ingestion/tests/unit/sql_parsing/aggregator_goldens/test_overlapping_inserts_from_temp_tables.json
index ed8009477bf4db..0f9870cda322c2 100644
--- a/metadata-ingestion/tests/unit/sql_parsing/aggregator_goldens/test_overlapping_inserts_from_temp_tables.json
+++ b/metadata-ingestion/tests/unit/sql_parsing/aggregator_goldens/test_overlapping_inserts_from_temp_tables.json
@@ -64,26 +64,26 @@
                 {
                     "upstreamType": "FIELD_SET",
                     "upstreams": [
-                        "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:redshift,dev.public.in_person_returns,PROD),customer_id)"
+                        "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:redshift,dev.public.online_returns,PROD),customer_id)"
                     ],
                     "downstreamType": "FIELD",
                     "downstreams": [
                         "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:redshift,dev.public.all_returns,PROD),customer_id)"
                     ],
                     "confidenceScore": 0.2,
-                    "query": "urn:li:query:composite_6b5a11e96e3d2b742e4e4ec3310bb538d0f5c0c6496b84e4bfe0e8014d5f5b45"
+                    "query": "urn:li:query:composite_adc1c41c0ad37c643776d9d93d524e6c435a7e70633da1ce7e3222dda4bb9fb8"
                 },
                 {
                     "upstreamType": "FIELD_SET",
                     "upstreams": [
-                        "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:redshift,dev.public.online_returns,PROD),customer_id)"
+                        "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:redshift,dev.public.in_person_returns,PROD),customer_id)"
                     ],
                     "downstreamType": "FIELD",
                     "downstreams": [
                         "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:redshift,dev.public.all_returns,PROD),customer_id)"
                     ],
                     "confidenceScore": 0.2,
-                    "query": "urn:li:query:composite_adc1c41c0ad37c643776d9d93d524e6c435a7e70633da1ce7e3222dda4bb9fb8"
+                    "query": "urn:li:query:composite_6b5a11e96e3d2b742e4e4ec3310bb538d0f5c0c6496b84e4bfe0e8014d5f5b45"
                 },
                 {
                     "upstreamType": "FIELD_SET",
@@ -100,26 +100,26 @@
                 {
                     "upstreamType": "FIELD_SET",
                     "upstreams": [
-                        "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:redshift,dev.public.in_person_returns,PROD),return_date)"
+                        "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:redshift,dev.public.online_returns,PROD),return_date)"
                     ],
                     "downstreamType": "FIELD",
                     "downstreams": [
                         "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:redshift,dev.public.all_returns,PROD),return_date)"
                     ],
                     "confidenceScore": 0.2,
-                    "query": "urn:li:query:composite_6b5a11e96e3d2b742e4e4ec3310bb538d0f5c0c6496b84e4bfe0e8014d5f5b45"
+                    "query": "urn:li:query:composite_adc1c41c0ad37c643776d9d93d524e6c435a7e70633da1ce7e3222dda4bb9fb8"
                 },
                 {
                     "upstreamType": "FIELD_SET",
                     "upstreams": [
-                        "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:redshift,dev.public.online_returns,PROD),return_date)"
+                        "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:redshift,dev.public.in_person_returns,PROD),return_date)"
                     ],
                     "downstreamType": "FIELD",
                     "downstreams": [
                         "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:redshift,dev.public.all_returns,PROD),return_date)"
                     ],
                     "confidenceScore": 0.2,
-                    "query": "urn:li:query:composite_adc1c41c0ad37c643776d9d93d524e6c435a7e70633da1ce7e3222dda4bb9fb8"
+                    "query": "urn:li:query:composite_6b5a11e96e3d2b742e4e4ec3310bb538d0f5c0c6496b84e4bfe0e8014d5f5b45"
                 },
                 {
                     "upstreamType": "FIELD_SET",
diff --git a/metadata-ingestion/tests/unit/stateful_ingestion/state/test_stateful_ingestion.py b/metadata-ingestion/tests/unit/stateful_ingestion/state/test_stateful_ingestion.py
index 96ab8f7a01a386..e69727f73b6bf4 100644
--- a/metadata-ingestion/tests/unit/stateful_ingestion/state/test_stateful_ingestion.py
+++ b/metadata-ingestion/tests/unit/stateful_ingestion/state/test_stateful_ingestion.py
@@ -226,9 +226,9 @@ def test_stateful_ingestion(pytestconfig, tmp_path, mock_time):
         pipeline_run1_config: Dict[str, Dict[str, Dict[str, Any]]] = dict(  # type: ignore
             base_pipeline_config  # type: ignore
         )
-        pipeline_run1_config["sink"]["config"][
-            "filename"
-        ] = f"{tmp_path}/{output_file_name}"
+        pipeline_run1_config["sink"]["config"]["filename"] = (
+            f"{tmp_path}/{output_file_name}"
+        )
         pipeline_run1 = Pipeline.create(pipeline_run1_config)
         pipeline_run1.run()
         pipeline_run1.raise_from_status()
@@ -254,16 +254,18 @@ def test_stateful_ingestion(pytestconfig, tmp_path, mock_time):
     ) as mock_state:
         mock_state.return_value = GenericCheckpointState(serde="utf-8")
         pipeline_run2 = None
-        pipeline_run2_config: Dict[str, Dict[str, Dict[str, Any]]] = dict(base_pipeline_config)  # type: ignore
+        pipeline_run2_config: Dict[str, Dict[str, Dict[str, Any]]] = dict(
+            base_pipeline_config  # type: ignore
+        )
         pipeline_run2_config["source"]["config"]["dataset_patterns"] = {
             "allow": ["dummy_dataset1", "dummy_dataset2"],
         }
         pipeline_run2_config["source"]["config"]["dpi_id_to_ingest"] = "job2"
         pipeline_run2_config["source"]["config"]["query_id_to_ingest"] = "query2"
 
-        pipeline_run2_config["sink"]["config"][
-            "filename"
-        ] = f"{tmp_path}/{output_file_name_after_deleted}"
+        pipeline_run2_config["sink"]["config"]["filename"] = (
+            f"{tmp_path}/{output_file_name_after_deleted}"
+        )
         pipeline_run2 = Pipeline.create(pipeline_run2_config)
         pipeline_run2.run()
         pipeline_run2.raise_from_status()
@@ -370,9 +372,9 @@ def test_stateful_ingestion_failure(pytestconfig, tmp_path, mock_time):
         pipeline_run1_config: Dict[str, Dict[str, Dict[str, Any]]] = dict(  # type: ignore
             base_pipeline_config  # type: ignore
         )
-        pipeline_run1_config["sink"]["config"][
-            "filename"
-        ] = f"{tmp_path}/{output_file_name}"
+        pipeline_run1_config["sink"]["config"]["filename"] = (
+            f"{tmp_path}/{output_file_name}"
+        )
         pipeline_run1 = Pipeline.create(pipeline_run1_config)
         pipeline_run1.run()
         pipeline_run1.raise_from_status()
@@ -398,14 +400,16 @@ def test_stateful_ingestion_failure(pytestconfig, tmp_path, mock_time):
     ) as mock_state:
         mock_state.return_value = GenericCheckpointState(serde="utf-8")
         pipeline_run2 = None
-        pipeline_run2_config: Dict[str, Dict[str, Dict[str, Any]]] = dict(base_pipeline_config)  # type: ignore
+        pipeline_run2_config: Dict[str, Dict[str, Dict[str, Any]]] = dict(
+            base_pipeline_config  # type: ignore
+        )
         pipeline_run2_config["source"]["config"]["dataset_patterns"] = {
             "allow": ["dummy_dataset1", "dummy_dataset2"],
         }
         pipeline_run2_config["source"]["config"]["report_failure"] = True
-        pipeline_run2_config["sink"]["config"][
-            "filename"
-        ] = f"{tmp_path}/{output_file_name_after_deleted}"
+        pipeline_run2_config["sink"]["config"]["filename"] = (
+            f"{tmp_path}/{output_file_name_after_deleted}"
+        )
         pipeline_run2 = Pipeline.create(pipeline_run2_config)
         pipeline_run2.run()
         pipeline_run2.pretty_print_summary()
diff --git a/metadata-ingestion/tests/unit/test_confluent_schema_registry.py b/metadata-ingestion/tests/unit/test_confluent_schema_registry.py
index 3500636f00eddf..effa6ba85acaeb 100644
--- a/metadata-ingestion/tests/unit/test_confluent_schema_registry.py
+++ b/metadata-ingestion/tests/unit/test_confluent_schema_registry.py
@@ -85,16 +85,18 @@ def new_get_latest_version(subject_name: str) -> RegisteredSchema:
             "get_latest_version",
             new_get_latest_version,
         ):
-            schema_str = confluent_schema_registry.get_schema_str_replace_confluent_ref_avro(
-                # The external reference would match by name.
-                schema=Schema(
-                    schema_str=schema_str_orig,
-                    schema_type="AVRO",
-                    references=[
-                        SchemaReference(
-                            name="TestTopic1", subject="schema_subject_1", version=1
-                        )
-                    ],
+            schema_str = (
+                confluent_schema_registry.get_schema_str_replace_confluent_ref_avro(
+                    # The external reference would match by name.
+                    schema=Schema(
+                        schema_str=schema_str_orig,
+                        schema_type="AVRO",
+                        references=[
+                            SchemaReference(
+                                name="TestTopic1", subject="schema_subject_1", version=1
+                            )
+                        ],
+                    )
                 )
             )
             assert schema_str == ConfluentSchemaRegistry._compact_schema(
@@ -106,16 +108,18 @@ def new_get_latest_version(subject_name: str) -> RegisteredSchema:
             "get_latest_version",
             new_get_latest_version,
         ):
-            schema_str = confluent_schema_registry.get_schema_str_replace_confluent_ref_avro(
-                # The external reference would match by subject.
-                schema=Schema(
-                    schema_str=schema_str_orig,
-                    schema_type="AVRO",
-                    references=[
-                        SchemaReference(
-                            name="schema_subject_1", subject="TestTopic1", version=1
-                        )
-                    ],
+            schema_str = (
+                confluent_schema_registry.get_schema_str_replace_confluent_ref_avro(
+                    # The external reference would match by subject.
+                    schema=Schema(
+                        schema_str=schema_str_orig,
+                        schema_type="AVRO",
+                        references=[
+                            SchemaReference(
+                                name="schema_subject_1", subject="TestTopic1", version=1
+                            )
+                        ],
+                    )
                 )
             )
             assert schema_str == ConfluentSchemaRegistry._compact_schema(
diff --git a/metadata-ingestion/tests/unit/test_dbt_source.py b/metadata-ingestion/tests/unit/test_dbt_source.py
index ff22ffedc6228f..d7899af69f8405 100644
--- a/metadata-ingestion/tests/unit/test_dbt_source.py
+++ b/metadata-ingestion/tests/unit/test_dbt_source.py
@@ -475,9 +475,9 @@ def test_get_column_type_redshift():
     # Test 'super' type which should not show any warnings/errors
     result_super = get_column_type(report, dataset_name, "super", "redshift")
     assert isinstance(result_super.type, NullTypeClass)
-    assert (
-        len(report.infos) == 0
-    ), "No warnings should be generated for known SUPER type"
+    assert len(report.infos) == 0, (
+        "No warnings should be generated for known SUPER type"
+    )
 
     # Test unknown type, which generates a warning but resolves to NullTypeClass
     unknown_type = "unknown_type"
diff --git a/metadata-ingestion/tests/unit/test_gcs_source.py b/metadata-ingestion/tests/unit/test_gcs_source.py
index 9d5f4e915b18cf..5b99862c92a04b 100644
--- a/metadata-ingestion/tests/unit/test_gcs_source.py
+++ b/metadata-ingestion/tests/unit/test_gcs_source.py
@@ -1,13 +1,17 @@
+from unittest import mock
+
 import pytest
 from pydantic import ValidationError
 
 from datahub.ingestion.api.common import PipelineContext
+from datahub.ingestion.graph.client import DataHubGraph
 from datahub.ingestion.source.data_lake_common.data_lake_utils import PLATFORM_GCS
 from datahub.ingestion.source.gcs.gcs_source import GCSSource
 
 
 def test_gcs_source_setup():
-    ctx = PipelineContext(run_id="test-gcs")
+    graph = mock.MagicMock(spec=DataHubGraph)
+    ctx = PipelineContext(run_id="test-gcs", graph=graph, pipeline_name="test-gcs")
 
     # Baseline: valid config
     source: dict = {
@@ -18,6 +22,7 @@ def test_gcs_source_setup():
             }
         ],
         "credential": {"hmac_access_id": "id", "hmac_access_secret": "secret"},
+        "stateful_ingestion": {"enabled": "true"},
     }
     gcs = GCSSource.create(source, ctx)
     assert gcs.s3_source.source_config.platform == PLATFORM_GCS
diff --git a/metadata-ingestion/tests/unit/test_iceberg.py b/metadata-ingestion/tests/unit/test_iceberg.py
index 3afa26b35dfe9f..48524450caf36e 100644
--- a/metadata-ingestion/tests/unit/test_iceberg.py
+++ b/metadata-ingestion/tests/unit/test_iceberg.py
@@ -88,15 +88,15 @@ def assert_field(
     expected_nullable: bool,
     expected_type: Any,
 ) -> None:
-    assert (
-        schema_field.description == expected_description
-    ), f"Field description '{schema_field.description}' is different from expected description '{expected_description}'"
-    assert (
-        schema_field.nullable == expected_nullable
-    ), f"Field nullable '{schema_field.nullable}' is different from expected nullable '{expected_nullable}'"
-    assert isinstance(
-        schema_field.type.type, expected_type
-    ), f"Field type {schema_field.type.type} is different from expected type {expected_type}"
+    assert schema_field.description == expected_description, (
+        f"Field description '{schema_field.description}' is different from expected description '{expected_description}'"
+    )
+    assert schema_field.nullable == expected_nullable, (
+        f"Field nullable '{schema_field.nullable}' is different from expected nullable '{expected_nullable}'"
+    )
+    assert isinstance(schema_field.type.type, expected_type), (
+        f"Field type {schema_field.type.type} is different from expected type {expected_type}"
+    )
 
 
 def test_config_no_catalog():
@@ -219,9 +219,9 @@ def test_iceberg_primitive_type_to_schema_field(
     ]:
         schema = Schema(column)
         schema_fields = iceberg_source_instance._get_schema_fields_for_schema(schema)
-        assert (
-            len(schema_fields) == 1
-        ), f"Expected 1 field, but got {len(schema_fields)}"
+        assert len(schema_fields) == 1, (
+            f"Expected 1 field, but got {len(schema_fields)}"
+        )
         assert_field(
             schema_fields[0],
             column.doc,
@@ -300,19 +300,19 @@ def test_iceberg_list_to_schema_field(
         iceberg_source_instance = with_iceberg_source()
         schema = Schema(list_column)
         schema_fields = iceberg_source_instance._get_schema_fields_for_schema(schema)
-        assert (
-            len(schema_fields) == 1
-        ), f"Expected 1 field, but got {len(schema_fields)}"
+        assert len(schema_fields) == 1, (
+            f"Expected 1 field, but got {len(schema_fields)}"
+        )
         assert_field(
             schema_fields[0], list_column.doc, list_column.optional, ArrayTypeClass
         )
-        assert isinstance(
-            schema_fields[0].type.type, ArrayType
-        ), f"Field type {schema_fields[0].type.type} was expected to be {ArrayType}"
+        assert isinstance(schema_fields[0].type.type, ArrayType), (
+            f"Field type {schema_fields[0].type.type} was expected to be {ArrayType}"
+        )
         arrayType: ArrayType = schema_fields[0].type.type
-        assert arrayType.nestedType == [
-            expected_array_nested_type
-        ], f"List Field nested type {arrayType.nestedType} was expected to be {expected_array_nested_type}"
+        assert arrayType.nestedType == [expected_array_nested_type], (
+            f"List Field nested type {arrayType.nestedType} was expected to be {expected_array_nested_type}"
+        )
 
 
 @pytest.mark.parametrize(
@@ -387,9 +387,9 @@ def test_iceberg_map_to_schema_field(
         schema_fields = iceberg_source_instance._get_schema_fields_for_schema(schema)
         # Converting an Iceberg Map type will be done by creating an array of struct(key, value) records.
         # The first field will be the array.
-        assert (
-            len(schema_fields) == 3
-        ), f"Expected 3 fields, but got {len(schema_fields)}"
+        assert len(schema_fields) == 3, (
+            f"Expected 3 fields, but got {len(schema_fields)}"
+        )
         assert_field(
             schema_fields[0], map_column.doc, map_column.optional, ArrayTypeClass
         )
diff --git a/metadata-ingestion/tests/unit/test_packages.py b/metadata-ingestion/tests/unit/test_packages.py
index f4045bac6e6ef4..ab538cf0c1ed09 100644
--- a/metadata-ingestion/tests/unit/test_packages.py
+++ b/metadata-ingestion/tests/unit/test_packages.py
@@ -1,7 +1,10 @@
 import pytest
 import setuptools
 
-from datahub.testing.check_imports import ensure_no_indirect_model_imports
+from datahub.testing.check_imports import (
+    ban_direct_datahub_imports,
+    ensure_no_indirect_model_imports,
+)
 from datahub.testing.check_str_enum import ensure_no_enum_mixin
 
 
@@ -16,6 +19,7 @@ def test_check_import_paths(pytestconfig: pytest.Config) -> None:
     root = pytestconfig.rootpath
 
     ensure_no_indirect_model_imports([root / "src", root / "tests"])
+    ban_direct_datahub_imports([root / "src", root / "tests"])
 
 
 def test_check_str_enum_usage(pytestconfig: pytest.Config) -> None:
diff --git a/metadata-ingestion/tests/unit/test_packaging.py b/metadata-ingestion/tests/unit/test_packaging.py
index 4b99be750a4da7..f9a3ae9562d3eb 100644
--- a/metadata-ingestion/tests/unit/test_packaging.py
+++ b/metadata-ingestion/tests/unit/test_packaging.py
@@ -1,6 +1,6 @@
 import pytest
 
-import datahub as datahub_metadata
+import datahub._version as datahub_version
 
 
 @pytest.mark.filterwarnings(
@@ -10,4 +10,4 @@ def test_datahub_version():
     # Simply importing pkg_resources checks for unsatisfied dependencies.
     import pkg_resources
 
-    assert pkg_resources.get_distribution(datahub_metadata.__package_name__).version
+    assert pkg_resources.get_distribution(datahub_version.__package_name__).version
diff --git a/metadata-ingestion/tests/unit/test_postgres_source.py b/metadata-ingestion/tests/unit/test_postgres_source.py
index 91a62b603bb584..25140cf1b997f8 100644
--- a/metadata-ingestion/tests/unit/test_postgres_source.py
+++ b/metadata-ingestion/tests/unit/test_postgres_source.py
@@ -21,9 +21,7 @@ def test_initial_database(create_engine_mock):
 
 @patch("datahub.ingestion.source.sql.postgres.create_engine")
 def test_get_inspectors_multiple_databases(create_engine_mock):
-    execute_mock = (
-        create_engine_mock.return_value.connect.return_value.__enter__.return_value.execute
-    )
+    execute_mock = create_engine_mock.return_value.connect.return_value.__enter__.return_value.execute
     execute_mock.return_value = [{"datname": "db1"}, {"datname": "db2"}]
 
     config = PostgresConfig.parse_obj({**_base_config(), "initial_database": "db0"})
@@ -37,9 +35,7 @@ def test_get_inspectors_multiple_databases(create_engine_mock):
 
 @patch("datahub.ingestion.source.sql.postgres.create_engine")
 def tests_get_inspectors_with_database_provided(create_engine_mock):
-    execute_mock = (
-        create_engine_mock.return_value.connect.return_value.__enter__.return_value.execute
-    )
+    execute_mock = create_engine_mock.return_value.connect.return_value.__enter__.return_value.execute
     execute_mock.return_value = [{"datname": "db1"}, {"datname": "db2"}]
 
     config = PostgresConfig.parse_obj({**_base_config(), "database": "custom_db"})
@@ -51,9 +47,7 @@ def tests_get_inspectors_with_database_provided(create_engine_mock):
 
 @patch("datahub.ingestion.source.sql.postgres.create_engine")
 def tests_get_inspectors_with_sqlalchemy_uri_provided(create_engine_mock):
-    execute_mock = (
-        create_engine_mock.return_value.connect.return_value.__enter__.return_value.execute
-    )
+    execute_mock = create_engine_mock.return_value.connect.return_value.__enter__.return_value.execute
     execute_mock.return_value = [{"datname": "db1"}, {"datname": "db2"}]
 
     config = PostgresConfig.parse_obj(
diff --git a/metadata-ingestion/tests/unit/test_rest_sink.py b/metadata-ingestion/tests/unit/test_rest_sink.py
index a76f96039c2c71..564cf613c04464 100644
--- a/metadata-ingestion/tests/unit/test_rest_sink.py
+++ b/metadata-ingestion/tests/unit/test_rest_sink.py
@@ -283,9 +283,9 @@
 def test_datahub_rest_emitter(requests_mock, record, path, snapshot):
     def match_request_text(request: requests.Request) -> bool:
         requested_snapshot = request.json()
-        assert (
-            requested_snapshot == snapshot
-        ), f"Expected snapshot to be {json.dumps(snapshot)}, got {json.dumps(requested_snapshot)}"
+        assert requested_snapshot == snapshot, (
+            f"Expected snapshot to be {json.dumps(snapshot)}, got {json.dumps(requested_snapshot)}"
+        )
         return True
 
     requests_mock.post(
diff --git a/metadata-ingestion/tests/unit/test_tableau_source.py b/metadata-ingestion/tests/unit/test_tableau_source.py
index 227519fdb464a8..ba5e5a9832a62f 100644
--- a/metadata-ingestion/tests/unit/test_tableau_source.py
+++ b/metadata-ingestion/tests/unit/test_tableau_source.py
@@ -3,8 +3,13 @@
 import pytest
 
 import datahub.ingestion.source.tableau.tableau_constant as c
-from datahub.ingestion.source.tableau.tableau import TableauSiteSource
+from datahub.ingestion.source.tableau.tableau import (
+    DEFAULT_PAGE_SIZE,
+    TableauPageSizeConfig,
+    TableauSiteSource,
+)
 from datahub.ingestion.source.tableau.tableau_common import (
+    TableauUpstreamReference,
     get_filter_pages,
     make_filter,
     optimize_query_filter,
@@ -247,3 +252,128 @@ def test_optimize_query_filter_handles_no_duplicates():
     assert len(result) == 2
     assert result[c.ID_WITH_IN] == ["id1", "id2"]
     assert result[c.PROJECT_NAME_WITH_IN] == ["project1", "project2"]
+
+
+def test_tableau_upstream_reference():
+    d = {
+        "id": "7127b695-3df5-4a3a-4837-eb0f4b572337",
+        "name": "TABLE1",
+        "database": None,
+        "schema": "SCHEMA1",
+        "fullName": "DB1.SCHEMA1.TABLE1",
+        "connectionType": "snowflake",
+        "description": "",
+        "columnsConnection": {"totalCount": 0},
+    }
+    ref = TableauUpstreamReference.create(d)
+    assert ref
+
+    assert ref.database == "DB1"
+    assert ref.schema == "SCHEMA1"
+    assert ref.table == "TABLE1"
+    assert ref.connection_type == "snowflake"
+
+    try:
+        ref = TableauUpstreamReference.create(None)  # type: ignore[arg-type]
+        raise AssertionError(
+            "TableauUpstreamReference.create with None should have raised exception"
+        )
+    except ValueError:
+        assert True
+
+
+class TestTableauPageSizeConfig:
+    def test_defaults(self):
+        config = TableauPageSizeConfig()
+        assert config.effective_database_server_page_size == DEFAULT_PAGE_SIZE
+        assert config.effective_workbook_page_size == 1
+        assert config.effective_sheet_page_size == DEFAULT_PAGE_SIZE
+        assert config.effective_dashboard_page_size == DEFAULT_PAGE_SIZE
+        assert config.effective_embedded_datasource_page_size == DEFAULT_PAGE_SIZE
+        assert (
+            config.effective_embedded_datasource_field_upstream_page_size
+            == DEFAULT_PAGE_SIZE * 10
+        )
+        assert config.effective_published_datasource_page_size == DEFAULT_PAGE_SIZE
+        assert (
+            config.effective_published_datasource_field_upstream_page_size
+            == DEFAULT_PAGE_SIZE * 10
+        )
+        assert config.effective_custom_sql_table_page_size == DEFAULT_PAGE_SIZE
+        assert config.effective_database_table_page_size == DEFAULT_PAGE_SIZE
+
+    def test_page_size_fallbacks(self):
+        page_size = 33
+        config = TableauPageSizeConfig(page_size=page_size)
+        assert config.effective_database_server_page_size == page_size
+        assert config.effective_workbook_page_size == 1
+        assert config.effective_sheet_page_size == page_size
+        assert config.effective_dashboard_page_size == page_size
+        assert config.effective_embedded_datasource_page_size == page_size
+        assert (
+            config.effective_embedded_datasource_field_upstream_page_size
+            == page_size * 10
+        )
+        assert config.effective_published_datasource_page_size == page_size
+        assert (
+            config.effective_published_datasource_field_upstream_page_size
+            == page_size * 10
+        )
+        assert config.effective_custom_sql_table_page_size == page_size
+        assert config.effective_database_table_page_size == page_size
+
+    def test_fine_grained(self):
+        any_page_size = 55
+        config = TableauPageSizeConfig(database_server_page_size=any_page_size)
+        assert config.page_size == DEFAULT_PAGE_SIZE
+        assert config.effective_database_server_page_size == any_page_size
+
+        config = TableauPageSizeConfig(workbook_page_size=any_page_size)
+        assert config.page_size == DEFAULT_PAGE_SIZE
+        assert config.effective_workbook_page_size == any_page_size
+
+        config = TableauPageSizeConfig(workbook_page_size=None)
+        assert config.page_size == DEFAULT_PAGE_SIZE
+        assert config.effective_workbook_page_size == DEFAULT_PAGE_SIZE
+
+        config = TableauPageSizeConfig(sheet_page_size=any_page_size)
+        assert config.page_size == DEFAULT_PAGE_SIZE
+        assert config.effective_sheet_page_size == any_page_size
+
+        config = TableauPageSizeConfig(dashboard_page_size=any_page_size)
+        assert config.page_size == DEFAULT_PAGE_SIZE
+        assert config.effective_dashboard_page_size == any_page_size
+
+        config = TableauPageSizeConfig(embedded_datasource_page_size=any_page_size)
+        assert config.page_size == DEFAULT_PAGE_SIZE
+        assert config.effective_embedded_datasource_page_size == any_page_size
+
+        config = TableauPageSizeConfig(
+            embedded_datasource_field_upstream_page_size=any_page_size
+        )
+        assert config.page_size == DEFAULT_PAGE_SIZE
+        assert (
+            config.effective_embedded_datasource_field_upstream_page_size
+            == any_page_size
+        )
+
+        config = TableauPageSizeConfig(published_datasource_page_size=any_page_size)
+        assert config.page_size == DEFAULT_PAGE_SIZE
+        assert config.effective_published_datasource_page_size == any_page_size
+
+        config = TableauPageSizeConfig(
+            published_datasource_field_upstream_page_size=any_page_size
+        )
+        assert config.page_size == DEFAULT_PAGE_SIZE
+        assert (
+            config.effective_published_datasource_field_upstream_page_size
+            == any_page_size
+        )
+
+        config = TableauPageSizeConfig(custom_sql_table_page_size=any_page_size)
+        assert config.page_size == DEFAULT_PAGE_SIZE
+        assert config.effective_custom_sql_table_page_size == any_page_size
+
+        config = TableauPageSizeConfig(database_table_page_size=any_page_size)
+        assert config.page_size == DEFAULT_PAGE_SIZE
+        assert config.effective_database_table_page_size == any_page_size
diff --git a/metadata-ingestion/tests/unit/urns/test_urn.py b/metadata-ingestion/tests/unit/urns/test_urn.py
index 8490364326d940..f9ed2b387f0789 100644
--- a/metadata-ingestion/tests/unit/urns/test_urn.py
+++ b/metadata-ingestion/tests/unit/urns/test_urn.py
@@ -10,6 +10,7 @@
     DataPlatformUrn,
     DatasetUrn,
     SchemaFieldUrn,
+    TagUrn,
     Urn,
 )
 from datahub.testing.doctest import assert_doctest
@@ -71,9 +72,18 @@ def test_urn_coercion() -> None:
 def test_urns_in_init() -> None:
     platform = DataPlatformUrn("abc")
     assert platform.urn() == "urn:li:dataPlatform:abc"
+    assert platform == DataPlatformUrn(platform)
+    assert platform == DataPlatformUrn(platform.urn())
 
     dataset_urn = DatasetUrn(platform, "def", "PROD")
     assert dataset_urn.urn() == "urn:li:dataset:(urn:li:dataPlatform:abc,def,PROD)"
+    assert dataset_urn == DatasetUrn(platform.urn(), "def", "PROD")
+    assert dataset_urn == DatasetUrn(platform.platform_name, "def", "PROD")
+
+    with pytest.raises(
+        InvalidUrnError, match="Expecting a DataPlatformUrn but got .*dataset.*"
+    ):
+        assert dataset_urn == DatasetUrn(dataset_urn, "def", "PROD")  # type: ignore
 
     schema_field = SchemaFieldUrn(dataset_urn, "foo")
     assert (
@@ -101,6 +111,11 @@ def test_urn_type_dispatch_2() -> None:
     with pytest.raises(InvalidUrnError, match="Passed an urn of type dataJob"):
         CorpUserUrn.from_string(urn)
 
+    with pytest.raises(
+        InvalidUrnError, match="Expecting a CorpUserUrn but got .*dataJob.*"
+    ):
+        CorpUserUrn(urn)  # type: ignore
+
 
 def test_urn_type_dispatch_3() -> None:
     # Creating a "generic" Urn.
@@ -133,6 +148,24 @@ def test_urn_type_dispatch_4() -> None:
     assert urn2.urn() == urn_str
 
 
+def test_urn_from_urn_simple() -> None:
+    # This capability is also tested by a bunch of other tests above.
+
+    tag_str = "urn:li:tag:legacy"
+    tag = TagUrn.from_string(tag_str)
+    assert tag_str == tag.urn()
+    assert tag.name == "legacy"
+    assert tag == TagUrn(tag)
+    assert tag == TagUrn(tag.urn())
+
+
+def test_urn_from_urn_tricky() -> None:
+    tag_str = "urn:li:tag:urn:li:tag:legacy"
+    tag = TagUrn(tag_str)
+    assert tag.urn() == tag_str
+    assert tag.name == "urn:li:tag:legacy"
+
+
 def test_urn_doctest() -> None:
     assert_doctest(datahub.utilities.urns._urn_base)
 
diff --git a/metadata-ingestion/tests/unit/utilities/test_cli_logging.py b/metadata-ingestion/tests/unit/utilities/test_cli_logging.py
index 5acb5f3e3608be..aa22a3b5e7ceba 100644
--- a/metadata-ingestion/tests/unit/utilities/test_cli_logging.py
+++ b/metadata-ingestion/tests/unit/utilities/test_cli_logging.py
@@ -36,7 +36,7 @@ def my_logging_fn():
     logger.warning("This is a warning message")
     logger.error("this is an error with no stack trace")
     try:
-        1 / 0
+        _ = 1 / 0
     except ZeroDivisionError:
         logger.exception("failed to divide by zero")
 
diff --git a/metadata-ingestion/tests/unit/utilities/test_lossy_collections.py b/metadata-ingestion/tests/unit/utilities/test_lossy_collections.py
index c5663f4a678198..e137d671e95d71 100644
--- a/metadata-ingestion/tests/unit/utilities/test_lossy_collections.py
+++ b/metadata-ingestion/tests/unit/utilities/test_lossy_collections.py
@@ -9,36 +9,38 @@
 
 @pytest.mark.parametrize("length, sampling", [(10, False), (100, True)])
 def test_lossylist_sampling(length, sampling):
-    l: LossyList[str] = LossyList()
+    l_dict: LossyList[str] = LossyList()
     for i in range(0, length):
-        l.append(f"{i} Hello World")
+        l_dict.append(f"{i} Hello World")
 
-    assert len(l) == length
-    assert l.sampled is sampling
+    assert len(l_dict) == length
+    assert l_dict.sampled is sampling
     if sampling:
-        assert f"... sampled of {length} total elements" in str(l)
+        assert f"... sampled of {length} total elements" in str(l_dict)
     else:
-        assert "sampled" not in str(l)
+        assert "sampled" not in str(l_dict)
 
-    list_version = [int(i.split(" ")[0]) for i in l]
+    list_version = [int(i.split(" ")[0]) for i in l_dict]
     print(list_version)
     assert sorted(list_version) == list_version
 
 
 @pytest.mark.parametrize("length, sampling", [(10, False), (100, True)])
 def test_lossyset_sampling(length, sampling):
-    l: LossySet[str] = LossySet()
+    lossy_set: LossySet[str] = LossySet()
     for i in range(0, length):
-        l.add(f"{i} Hello World")
+        lossy_set.add(f"{i} Hello World")
 
-    assert len(l) == min(10, length)
-    assert l.sampled is sampling
+    assert len(lossy_set) == min(10, length)
+    assert lossy_set.sampled is sampling
     if sampling:
-        assert f"... sampled with at most {length-10} elements missing" in str(l)
+        assert f"... sampled with at most {length - 10} elements missing" in str(
+            lossy_set
+        )
     else:
-        assert "sampled" not in str(l)
+        assert "sampled" not in str(lossy_set)
 
-    list_version = [int(i.split(" ")[0]) for i in l]
+    list_version = [int(i.split(" ")[0]) for i in lossy_set]
     set_version = set(list_version)
 
     assert len(list_version) == len(set_version)
@@ -49,35 +51,36 @@ def test_lossyset_sampling(length, sampling):
     "length, sampling, sub_length", [(4, False, 4), (10, False, 14), (100, True, 1000)]
 )
 def test_lossydict_sampling(length, sampling, sub_length):
-    l: LossyDict[int, LossyList[str]] = LossyDict()
+    lossy_dict: LossyDict[int, LossyList[str]] = LossyDict()
     elements_added = 0
     element_length_map = {}
     for i in range(0, length):
         list_length = random.choice(range(1, sub_length))
         element_length_map[i] = 0
         for _num_elements in range(0, list_length):
-            if not l.get(i):
+            if not lossy_dict.get(i):
                 elements_added += 1
                 # reset to 0 until we get it back
                 element_length_map[i] = 0
             else:
-                element_length_map[i] = len(l[i])
+                element_length_map[i] = len(lossy_dict[i])
 
-            current_list = l.get(i, LossyList())
-            current_list.append(f"{i}:{round(time.time(),2)} Hello World")
-            l[i] = current_list
+            current_list = lossy_dict.get(i, LossyList())
+            current_list.append(f"{i}:{round(time.time(), 2)} Hello World")
+            lossy_dict[i] = current_list
             element_length_map[i] += 1
 
-    assert len(l) == min(l.max_elements, length)
-    assert l.sampled is sampling
+    assert len(lossy_dict) == min(lossy_dict.max_elements, length)
+    assert lossy_dict.sampled is sampling
     if sampling:
-        assert re.search("sampled of at most .* entries.", str(l))
-        assert f"{l.max_elements} sampled of at most {elements_added} entries." in str(
-            l
+        assert re.search("sampled of at most .* entries.", str(lossy_dict))
+        assert (
+            f"{lossy_dict.max_elements} sampled of at most {elements_added} entries."
+            in str(lossy_dict)
         )
     else:
         # cheap way to determine that the dict isn't reporting sampled keys
-        assert not re.search("sampled of at most .* entries.", str(l))
+        assert not re.search("sampled of at most .* entries.", str(lossy_dict))
 
-    for k, v in l.items():
+    for k, v in lossy_dict.items():
         assert len(v) == element_length_map[k]
diff --git a/metadata-ingestion/tests/unit/utilities/test_partition_executor.py b/metadata-ingestion/tests/unit/utilities/test_partition_executor.py
index ce211c2d618062..89e95d185e8028 100644
--- a/metadata-ingestion/tests/unit/utilities/test_partition_executor.py
+++ b/metadata-ingestion/tests/unit/utilities/test_partition_executor.py
@@ -37,9 +37,9 @@ def task(key: str, id: str) -> None:
         saw_keys_in_parallel = False
         while executing_tasks or not done_tasks:
             keys_executing = [key for key, _ in executing_tasks]
-            assert list(sorted(keys_executing)) == list(
-                sorted(set(keys_executing))
-            ), "partitioning not working"
+            assert list(sorted(keys_executing)) == list(sorted(set(keys_executing))), (
+                "partitioning not working"
+            )
 
             if len(keys_executing) == 2:
                 saw_keys_in_parallel = True
diff --git a/metadata-ingestion/tests/unit/utilities/test_threaded_iterator_executor.py b/metadata-ingestion/tests/unit/utilities/test_threaded_iterator_executor.py
index 35c44c7b4a8479..fb7e2266e1c9d3 100644
--- a/metadata-ingestion/tests/unit/utilities/test_threaded_iterator_executor.py
+++ b/metadata-ingestion/tests/unit/utilities/test_threaded_iterator_executor.py
@@ -4,7 +4,7 @@
 def test_threaded_iterator_executor():
     def table_of(i):
         for j in range(1, 11):
-            yield f"{i}x{j}={i*j}"
+            yield f"{i}x{j}={i * j}"
 
     assert {
         res
diff --git a/metadata-ingestion/tests/unit/utilities/test_utilities.py b/metadata-ingestion/tests/unit/utilities/test_utilities.py
index c333ceb136fffc..52406d37de46a6 100644
--- a/metadata-ingestion/tests/unit/utilities/test_utilities.py
+++ b/metadata-ingestion/tests/unit/utilities/test_utilities.py
@@ -6,6 +6,7 @@
 from datahub.sql_parsing.sqlglot_lineage import sqlglot_lineage
 from datahub.testing.doctest import assert_doctest
 from datahub.utilities.delayed_iter import delayed_iter
+from datahub.utilities.groupby import groupby_unsorted
 from datahub.utilities.is_pytest import is_pytest_running
 from datahub.utilities.urns.dataset_urn import DatasetUrn
 
@@ -335,3 +336,13 @@ def test_logging_name_extraction() -> None:
 
 def test_is_pytest_running() -> None:
     assert is_pytest_running()
+
+
+def test_groupby_unsorted():
+    grouped = groupby_unsorted("ABCAC", key=lambda x: x)
+
+    assert list(grouped) == [
+        ("A", ["A", "A"]),
+        ("B", ["B"]),
+        ("C", ["C", "C"]),
+    ]
diff --git a/metadata-integration/java/acryl-spark-lineage/build.gradle b/metadata-integration/java/acryl-spark-lineage/build.gradle
index 8816264fbe50f7..c8b78f25d70992 100644
--- a/metadata-integration/java/acryl-spark-lineage/build.gradle
+++ b/metadata-integration/java/acryl-spark-lineage/build.gradle
@@ -57,7 +57,13 @@ dependencies {
   //implementation "io.acryl:datahub-client:0.10.2"
   implementation "io.openlineage:openlineage-spark_2.12:$openLineageVersion"
   compileOnly "org.apache.iceberg:iceberg-spark3-runtime:0.12.1"
-  compileOnly "org.apache.spark:spark-sql_2.12:3.1.3"
+  compileOnly("org.apache.spark:spark-sql_2.12:3.1.3") {
+    exclude group: 'org.eclipse.jetty', module: 'jetty-servlet'
+    exclude group: 'org.eclipse.jetty', module: 'jetty-server'
+    exclude group: 'org.eclipse.jetty', module: 'jetty-util'
+    exclude group: 'org.eclipse.jetty', module: 'jetty-webapp'
+    exclude group: 'org.eclipse.jetty', module: 'jetty-security'
+  }
   compileOnly "io.github.spark-redshift-community:spark-redshift_2.12:6.2.0-spark_3.5"
 
   testCompileOnly externalDependency.lombok
diff --git a/metadata-io/build.gradle b/metadata-io/build.gradle
index aab29101b30f71..b33f19bef95986 100644
--- a/metadata-io/build.gradle
+++ b/metadata-io/build.gradle
@@ -37,6 +37,7 @@ dependencies {
     exclude group: 'com.google.guava', module: 'guava'
     exclude group: 'io.grpc', module: 'grpc-protobuf'
   }
+  implementation externalDependency.dgraphNetty
   implementation externalDependency.slf4jApi
   runtimeOnly externalDependency.logbackClassic
   compileOnly externalDependency.lombok
diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/utils/ESUtils.java b/metadata-io/src/main/java/com/linkedin/metadata/search/utils/ESUtils.java
index 95fff81d13957c..34894a584ce69f 100644
--- a/metadata-io/src/main/java/com/linkedin/metadata/search/utils/ESUtils.java
+++ b/metadata-io/src/main/java/com/linkedin/metadata/search/utils/ESUtils.java
@@ -874,7 +874,7 @@ private static void filterSoftDeletedByDefault(
                                         || criterion.getField().equals(REMOVED + KEYWORD_SUFFIX)));
       }
       if (!removedInOrFilter) {
-        filterQuery.mustNot(QueryBuilders.matchQuery(REMOVED, true));
+        filterQuery.mustNot(QueryBuilders.termQuery(REMOVED, true));
       }
     }
   }
diff --git a/metadata-io/src/test/java/com/linkedin/metadata/graph/neo4j/Neo4jTestServerBuilder.java b/metadata-io/src/test/java/com/linkedin/metadata/graph/neo4j/Neo4jTestServerBuilder.java
index fa04de340e12f7..b34730f481c63b 100644
--- a/metadata-io/src/test/java/com/linkedin/metadata/graph/neo4j/Neo4jTestServerBuilder.java
+++ b/metadata-io/src/test/java/com/linkedin/metadata/graph/neo4j/Neo4jTestServerBuilder.java
@@ -18,11 +18,11 @@ private Neo4jTestServerBuilder(Neo4jBuilder builder) {
   }
 
   public Neo4jTestServerBuilder() {
-    this(new InProcessNeo4jBuilder().withProcedure(PathExplorer.class));
+    this(new InProcessNeo4jBuilder().withProcedure(PathExplorer.class).withDisabledServer());
   }
 
   public Neo4jTestServerBuilder(File workingDirectory) {
-    this(new InProcessNeo4jBuilder(workingDirectory.toPath()));
+    this(new InProcessNeo4jBuilder(workingDirectory.toPath()).withDisabledServer());
   }
 
   public Neo4j newServer() {
diff --git a/metadata-io/src/test/java/com/linkedin/metadata/search/query/request/AutocompleteRequestHandlerTest.java b/metadata-io/src/test/java/com/linkedin/metadata/search/query/request/AutocompleteRequestHandlerTest.java
index 23d493b7287f78..e91dc1e6a8b0b9 100644
--- a/metadata-io/src/test/java/com/linkedin/metadata/search/query/request/AutocompleteRequestHandlerTest.java
+++ b/metadata-io/src/test/java/com/linkedin/metadata/search/query/request/AutocompleteRequestHandlerTest.java
@@ -158,7 +158,7 @@ public void testDefaultAutocompleteRequest() {
     assertEquals("keyPart1.delimited", prefixQuery.fieldName());
 
     assertEquals(wrapper.mustNot().size(), 1);
-    MatchQueryBuilder removedFilter = (MatchQueryBuilder) wrapper.mustNot().get(0);
+    TermQueryBuilder removedFilter = (TermQueryBuilder) wrapper.mustNot().get(0);
     assertEquals(removedFilter.fieldName(), "removed");
     assertEquals(removedFilter.value(), true);
     HighlightBuilder highlightBuilder = sourceBuilder.highlighter();
@@ -203,7 +203,7 @@ public void testAutocompleteRequestWithField() {
         (MatchPhrasePrefixQueryBuilder) query.should().get(1);
     assertEquals("field.delimited", prefixQuery.fieldName());
 
-    MatchQueryBuilder removedFilter = (MatchQueryBuilder) wrapper.mustNot().get(0);
+    TermQueryBuilder removedFilter = (TermQueryBuilder) wrapper.mustNot().get(0);
     assertEquals(removedFilter.fieldName(), "removed");
     assertEquals(removedFilter.value(), true);
     HighlightBuilder highlightBuilder = sourceBuilder.highlighter();
diff --git a/metadata-io/src/test/java/com/linkedin/metadata/search/query/request/SearchRequestHandlerTest.java b/metadata-io/src/test/java/com/linkedin/metadata/search/query/request/SearchRequestHandlerTest.java
index 1fea4476d75abb..d4d5b03ea5fec9 100644
--- a/metadata-io/src/test/java/com/linkedin/metadata/search/query/request/SearchRequestHandlerTest.java
+++ b/metadata-io/src/test/java/com/linkedin/metadata/search/query/request/SearchRequestHandlerTest.java
@@ -436,7 +436,7 @@ private BoolQueryBuilder constructFilterQuery(
   }
 
   private void testFilterQuery(BoolQueryBuilder testQuery) {
-    Optional<MatchQueryBuilder> mustNotHaveRemovedCondition =
+    Optional<TermQueryBuilder> mustNotHaveRemovedCondition =
         testQuery.filter().stream()
             .filter(or -> or instanceof BoolQueryBuilder)
             .map(or -> (BoolQueryBuilder) or)
@@ -445,8 +445,8 @@ private void testFilterQuery(BoolQueryBuilder testQuery) {
                   System.out.println("processing: " + or.mustNot());
                   return or.mustNot().stream();
                 })
-            .filter(and -> and instanceof MatchQueryBuilder)
-            .map(and -> (MatchQueryBuilder) and)
+            .filter(and -> and instanceof TermQueryBuilder)
+            .map(and -> (TermQueryBuilder) and)
             .filter(match -> match.fieldName().equals("removed"))
             .findAny();
 
diff --git a/metadata-io/src/test/java/io/datahubproject/test/search/SearchTestContainer.java b/metadata-io/src/test/java/io/datahubproject/test/search/SearchTestContainer.java
index cda6a4c179f484..c82d1ca87b862c 100644
--- a/metadata-io/src/test/java/io/datahubproject/test/search/SearchTestContainer.java
+++ b/metadata-io/src/test/java/io/datahubproject/test/search/SearchTestContainer.java
@@ -5,7 +5,7 @@
 
 public interface SearchTestContainer {
 
-  String SEARCH_JAVA_OPTS = "-Xms446m -Xmx446m -XX:MaxDirectMemorySize=368435456";
+  String SEARCH_JAVA_OPTS = "-Xms512m -Xmx512m -XX:MaxDirectMemorySize=368435456";
 
   Duration STARTUP_TIMEOUT = Duration.ofMinutes(5); // usually < 1min
 
diff --git a/metadata-jobs/mae-consumer/src/main/java/com/linkedin/metadata/kafka/DataHubUsageEventsProcessor.java b/metadata-jobs/mae-consumer/src/main/java/com/linkedin/metadata/kafka/DataHubUsageEventsProcessor.java
index ce7376f1f8d662..d699f0bff68019 100644
--- a/metadata-jobs/mae-consumer/src/main/java/com/linkedin/metadata/kafka/DataHubUsageEventsProcessor.java
+++ b/metadata-jobs/mae-consumer/src/main/java/com/linkedin/metadata/kafka/DataHubUsageEventsProcessor.java
@@ -49,7 +49,8 @@ public DataHubUsageEventsProcessor(
   @KafkaListener(
       id = "${DATAHUB_USAGE_EVENT_KAFKA_CONSUMER_GROUP_ID:datahub-usage-event-consumer-job-client}",
       topics = "${DATAHUB_USAGE_EVENT_NAME:" + Topics.DATAHUB_USAGE_EVENT + "}",
-      containerFactory = "simpleKafkaConsumer")
+      containerFactory = "simpleKafkaConsumer",
+      autoStartup = "false")
   public void consume(final ConsumerRecord<String, String> consumerRecord) {
     try (Timer.Context i = MetricUtils.timer(this.getClass(), "consume").time()) {
       kafkaLagStats.update(System.currentTimeMillis() - consumerRecord.timestamp());
diff --git a/metadata-jobs/mae-consumer/src/main/java/com/linkedin/metadata/kafka/MCLKafkaListenerRegistrar.java b/metadata-jobs/mae-consumer/src/main/java/com/linkedin/metadata/kafka/MCLKafkaListenerRegistrar.java
index c909b0034a9125..20c044b42741e8 100644
--- a/metadata-jobs/mae-consumer/src/main/java/com/linkedin/metadata/kafka/MCLKafkaListenerRegistrar.java
+++ b/metadata-jobs/mae-consumer/src/main/java/com/linkedin/metadata/kafka/MCLKafkaListenerRegistrar.java
@@ -73,7 +73,7 @@ public void afterPropertiesSet() {
                   buildConsumerGroupName(key),
                   List.of(mclVersionedTopicName, mclTimeseriesTopicName),
                   hooks);
-          registerMCLKafkaListener(kafkaListenerEndpoint, true);
+          registerMCLKafkaListener(kafkaListenerEndpoint, false);
         });
   }
 
@@ -97,7 +97,7 @@ private KafkaListenerEndpoint createListenerEndpoint(
         new MethodKafkaListenerEndpoint<>();
     kafkaListenerEndpoint.setId(consumerGroupId);
     kafkaListenerEndpoint.setGroupId(consumerGroupId);
-    kafkaListenerEndpoint.setAutoStartup(true);
+    kafkaListenerEndpoint.setAutoStartup(false);
     kafkaListenerEndpoint.setTopics(topics.toArray(new String[topics.size()]));
     kafkaListenerEndpoint.setMessageHandlerMethodFactory(new DefaultMessageHandlerMethodFactory());
     kafkaListenerEndpoint.setBean(
diff --git a/metadata-jobs/mce-consumer/src/main/java/com/linkedin/metadata/kafka/MetadataChangeEventsProcessor.java b/metadata-jobs/mce-consumer/src/main/java/com/linkedin/metadata/kafka/MetadataChangeEventsProcessor.java
index 5d2f6452e69197..2152ed15cf0e93 100644
--- a/metadata-jobs/mce-consumer/src/main/java/com/linkedin/metadata/kafka/MetadataChangeEventsProcessor.java
+++ b/metadata-jobs/mce-consumer/src/main/java/com/linkedin/metadata/kafka/MetadataChangeEventsProcessor.java
@@ -62,7 +62,8 @@ public class MetadataChangeEventsProcessor {
           "${METADATA_CHANGE_EVENT_NAME:${KAFKA_MCE_TOPIC_NAME:"
               + Topics.METADATA_CHANGE_EVENT
               + "}}",
-      containerFactory = DEFAULT_EVENT_CONSUMER_NAME)
+      containerFactory = DEFAULT_EVENT_CONSUMER_NAME,
+      autoStartup = "false")
   @Deprecated
   public void consume(final ConsumerRecord<String, GenericRecord> consumerRecord) {
     try (Timer.Context i = MetricUtils.timer(this.getClass(), "consume").time()) {
diff --git a/metadata-jobs/mce-consumer/src/main/java/com/linkedin/metadata/kafka/MetadataChangeProposalsProcessor.java b/metadata-jobs/mce-consumer/src/main/java/com/linkedin/metadata/kafka/MetadataChangeProposalsProcessor.java
index 4e356f5fb3670a..d854a5517793ff 100644
--- a/metadata-jobs/mce-consumer/src/main/java/com/linkedin/metadata/kafka/MetadataChangeProposalsProcessor.java
+++ b/metadata-jobs/mce-consumer/src/main/java/com/linkedin/metadata/kafka/MetadataChangeProposalsProcessor.java
@@ -80,7 +80,8 @@ public void registerConsumerThrottle() {
   @KafkaListener(
       id = CONSUMER_GROUP_ID_VALUE,
       topics = "${METADATA_CHANGE_PROPOSAL_TOPIC_NAME:" + Topics.METADATA_CHANGE_PROPOSAL + "}",
-      containerFactory = MCP_EVENT_CONSUMER_NAME)
+      containerFactory = MCP_EVENT_CONSUMER_NAME,
+      autoStartup = "false")
   public void consume(final ConsumerRecord<String, GenericRecord> consumerRecord) {
     try (Timer.Context ignored = MetricUtils.timer(this.getClass(), "consume").time()) {
       kafkaLagStats.update(System.currentTimeMillis() - consumerRecord.timestamp());
diff --git a/metadata-jobs/mce-consumer/src/main/java/com/linkedin/metadata/kafka/batch/BatchMetadataChangeProposalsProcessor.java b/metadata-jobs/mce-consumer/src/main/java/com/linkedin/metadata/kafka/batch/BatchMetadataChangeProposalsProcessor.java
index fed93628fe4d79..5ee9cd6ba94d2f 100644
--- a/metadata-jobs/mce-consumer/src/main/java/com/linkedin/metadata/kafka/batch/BatchMetadataChangeProposalsProcessor.java
+++ b/metadata-jobs/mce-consumer/src/main/java/com/linkedin/metadata/kafka/batch/BatchMetadataChangeProposalsProcessor.java
@@ -73,7 +73,8 @@ public void registerConsumerThrottle() {
       id = CONSUMER_GROUP_ID_VALUE,
       topics = "${METADATA_CHANGE_PROPOSAL_TOPIC_NAME:" + Topics.METADATA_CHANGE_PROPOSAL + "}",
       containerFactory = "kafkaEventConsumer",
-      batch = "true")
+      batch = "true",
+      autoStartup = "false")
   public void consume(final List<ConsumerRecord<String, GenericRecord>> consumerRecords) {
     try (Timer.Context ignored = MetricUtils.timer(this.getClass(), "consume").time()) {
       List<MetadataChangeProposal> metadataChangeProposals =
diff --git a/metadata-jobs/pe-consumer/src/main/java/com/datahub/event/PlatformEventProcessor.java b/metadata-jobs/pe-consumer/src/main/java/com/datahub/event/PlatformEventProcessor.java
index 5d11697bed93d2..2befeccb951a38 100644
--- a/metadata-jobs/pe-consumer/src/main/java/com/datahub/event/PlatformEventProcessor.java
+++ b/metadata-jobs/pe-consumer/src/main/java/com/datahub/event/PlatformEventProcessor.java
@@ -58,7 +58,8 @@ public PlatformEventProcessor(
   @KafkaListener(
       id = "${PLATFORM_EVENT_KAFKA_CONSUMER_GROUP_ID:generic-platform-event-job-client}",
       topics = {"${PLATFORM_EVENT_TOPIC_NAME:" + Topics.PLATFORM_EVENT + "}"},
-      containerFactory = PE_EVENT_CONSUMER_NAME)
+      containerFactory = PE_EVENT_CONSUMER_NAME,
+      autoStartup = "false")
   public void consume(final ConsumerRecord<String, GenericRecord> consumerRecord) {
     try (Timer.Context i = MetricUtils.timer(this.getClass(), "consume").time()) {
 
diff --git a/metadata-models/src/main/pegasus/com/linkedin/dataprocess/DataProcessInstanceInput.pdl b/metadata-models/src/main/pegasus/com/linkedin/dataprocess/DataProcessInstanceInput.pdl
index d005cd557cf77e..3e5192267c8c2e 100644
--- a/metadata-models/src/main/pegasus/com/linkedin/dataprocess/DataProcessInstanceInput.pdl
+++ b/metadata-models/src/main/pegasus/com/linkedin/dataprocess/DataProcessInstanceInput.pdl
@@ -10,12 +10,13 @@ import com.linkedin.common.Urn
 }
 record DataProcessInstanceInput {
   /**
-   * Input datasets to be consumed
+   * Input assets consumed
    */
   @Relationship = {
     "/*": {
       "name": "Consumes",
-      "entityTypes": [ "dataset" ]
+      "entityTypes": [ "dataset", "mlModel"],
+      "isLineage": true
     }
   }
   @Searchable = {
diff --git a/metadata-models/src/main/pegasus/com/linkedin/dataprocess/DataProcessInstanceOutput.pdl b/metadata-models/src/main/pegasus/com/linkedin/dataprocess/DataProcessInstanceOutput.pdl
index fe782dbe01ca9b..6b41cf9ba63417 100644
--- a/metadata-models/src/main/pegasus/com/linkedin/dataprocess/DataProcessInstanceOutput.pdl
+++ b/metadata-models/src/main/pegasus/com/linkedin/dataprocess/DataProcessInstanceOutput.pdl
@@ -15,7 +15,9 @@ record DataProcessInstanceOutput {
   @Relationship = {
     "/*": {
       "name": "Produces",
-      "entityTypes": [ "dataset", "mlModel" ]
+      "entityTypes": [ "dataset", "mlModel" ],
+      "isLineage": true,
+      "isUpstream": false
     }
   }
   @Searchable = {
diff --git a/metadata-service/auth-config/src/main/java/com/datahub/authentication/AuthenticationConfiguration.java b/metadata-service/auth-config/src/main/java/com/datahub/authentication/AuthenticationConfiguration.java
index 81cc5e60552a77..d258fcfeb65750 100644
--- a/metadata-service/auth-config/src/main/java/com/datahub/authentication/AuthenticationConfiguration.java
+++ b/metadata-service/auth-config/src/main/java/com/datahub/authentication/AuthenticationConfiguration.java
@@ -12,6 +12,9 @@ public class AuthenticationConfiguration {
   /** Whether user existence is enforced */
   private boolean enforceExistenceEnabled;
 
+  /** Paths to be excluded from filtering * */
+  private String excludedPaths;
+
   /**
    * List of configurations for {@link com.datahub.plugins.auth.authentication.Authenticator}s to be
    * registered
diff --git a/metadata-service/auth-filter/src/main/java/com/datahub/auth/authentication/filter/AuthenticationFilter.java b/metadata-service/auth-filter/src/main/java/com/datahub/auth/authentication/filter/AuthenticationFilter.java
index 30f98180f80180..492e165c0781a0 100644
--- a/metadata-service/auth-filter/src/main/java/com/datahub/auth/authentication/filter/AuthenticationFilter.java
+++ b/metadata-service/auth-filter/src/main/java/com/datahub/auth/authentication/filter/AuthenticationFilter.java
@@ -25,38 +25,43 @@
 import com.datahub.plugins.factory.PluginConfigFactory;
 import com.datahub.plugins.loader.IsolatedClassLoader;
 import com.datahub.plugins.loader.PluginPermissionManagerImpl;
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableMap;
 import com.linkedin.gms.factory.config.ConfigurationProvider;
 import com.linkedin.metadata.entity.EntityService;
 import jakarta.inject.Named;
 import jakarta.servlet.Filter;
 import jakarta.servlet.FilterChain;
-import jakarta.servlet.FilterConfig;
 import jakarta.servlet.ServletException;
-import jakarta.servlet.ServletRequest;
-import jakarta.servlet.ServletResponse;
 import jakarta.servlet.http.HttpServletRequest;
 import jakarta.servlet.http.HttpServletResponse;
 import java.io.IOException;
 import java.nio.file.Path;
 import java.nio.file.Paths;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.Set;
 import java.util.stream.Collectors;
+import javax.annotation.PostConstruct;
 import lombok.extern.slf4j.Slf4j;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.beans.factory.annotation.Value;
-import org.springframework.web.context.support.SpringBeanAutowiringSupport;
+import org.springframework.stereotype.Component;
+import org.springframework.util.StringUtils;
+import org.springframework.web.filter.OncePerRequestFilter;
 
 /**
  * A servlet {@link Filter} for authenticating requests inbound to the Metadata Service. This filter
  * is applied to the GraphQL Servlet, the Rest.li Servlet, and the Auth (token) Servlet.
  */
+@Component
 @Slf4j
-public class AuthenticationFilter implements Filter {
+public class AuthenticationFilter extends OncePerRequestFilter {
 
   @Autowired private ConfigurationProvider configurationProvider;
 
@@ -72,18 +77,32 @@ public class AuthenticationFilter implements Filter {
   private boolean _logAuthenticatorExceptions;
 
   private AuthenticatorChain authenticatorChain;
+  private Set<String> excludedPathPatterns;
 
-  @Override
-  public void init(FilterConfig filterConfig) throws ServletException {
-    SpringBeanAutowiringSupport.processInjectionBasedOnCurrentContext(this);
+  @PostConstruct
+  public void init() {
     buildAuthenticatorChain();
+    initializeExcludedPaths();
     log.info("AuthenticationFilter initialized.");
   }
 
+  private void initializeExcludedPaths() {
+    excludedPathPatterns = new HashSet<>();
+    String excludedPaths = configurationProvider.getAuthentication().getExcludedPaths();
+    if (StringUtils.hasText(excludedPaths)) {
+      excludedPathPatterns.addAll(
+          Arrays.stream(excludedPaths.split(","))
+              .map(String::trim)
+              .filter(path -> !path.isBlank())
+              .toList());
+    }
+  }
+
   @Override
-  public void doFilter(ServletRequest request, ServletResponse response, FilterChain chain)
-      throws IOException, ServletException {
-    AuthenticationRequest context = buildAuthContext((HttpServletRequest) request);
+  protected void doFilterInternal(
+      HttpServletRequest request, HttpServletResponse response, FilterChain chain)
+      throws ServletException, IOException {
+    AuthenticationRequest context = buildAuthContext(request);
     Authentication authentication = null;
     try {
       authentication = this.authenticatorChain.authenticate(context, _logAuthenticatorExceptions);
@@ -119,6 +138,34 @@ public void doFilter(ServletRequest request, ServletResponse response, FilterCha
     AuthenticationContext.remove();
   }
 
+  @VisibleForTesting
+  @Override
+  public boolean shouldNotFilter(HttpServletRequest request) {
+    String path = request.getServletPath();
+    if (path == null) {
+      return false;
+    }
+
+    // Check if the path matches any of the excluded patterns
+    boolean shouldExclude =
+        excludedPathPatterns.stream()
+            .anyMatch(
+                pattern -> {
+                  if (pattern.endsWith("/*")) {
+                    // Handle wildcard patterns
+                    String basePattern = pattern.substring(0, pattern.length() - 2);
+                    return path.startsWith(basePattern);
+                  }
+                  return path.equals(pattern);
+                });
+
+    if (shouldExclude) {
+      log.debug("Skipping authentication for excluded path: {}", path);
+    }
+
+    return shouldExclude;
+  }
+
   @Override
   public void destroy() {
     // Nothing
diff --git a/metadata-service/auth-filter/src/test/java/com/datahub/auth/authentication/AuthenticationFilterTest.java b/metadata-service/auth-filter/src/test/java/com/datahub/auth/authentication/AuthenticationFilterTest.java
index 382e881542b0b8..0ffabb9db305dd 100644
--- a/metadata-service/auth-filter/src/test/java/com/datahub/auth/authentication/AuthenticationFilterTest.java
+++ b/metadata-service/auth-filter/src/test/java/com/datahub/auth/authentication/AuthenticationFilterTest.java
@@ -2,6 +2,8 @@
 
 import static com.datahub.authentication.AuthenticationConstants.*;
 import static org.mockito.Mockito.*;
+import static org.testng.Assert.assertFalse;
+import static org.testng.Assert.assertTrue;
 
 import com.datahub.auth.authentication.filter.AuthenticationFilter;
 import com.datahub.authentication.Actor;
@@ -9,27 +11,34 @@
 import com.datahub.authentication.token.StatefulTokenService;
 import com.datahub.authentication.token.TokenException;
 import jakarta.servlet.FilterChain;
+import jakarta.servlet.FilterConfig;
 import jakarta.servlet.ServletException;
 import jakarta.servlet.http.HttpServletRequest;
 import jakarta.servlet.http.HttpServletResponse;
 import java.io.IOException;
+import java.util.Arrays;
 import java.util.Collections;
+import java.util.HashSet;
 import java.util.List;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.test.context.ContextConfiguration;
 import org.springframework.test.context.testng.AbstractTestNGSpringContextTests;
+import org.springframework.test.util.ReflectionTestUtils;
 import org.testng.annotations.Test;
 
 @ContextConfiguration(classes = {AuthTestConfiguration.class})
 public class AuthenticationFilterTest extends AbstractTestNGSpringContextTests {
 
-  @Autowired AuthenticationFilter _authenticationFilter;
+  @Autowired AuthenticationFilter authenticationFilter;
 
-  @Autowired StatefulTokenService _statefulTokenService;
+  @Autowired StatefulTokenService statefulTokenService;
 
   @Test
   public void testExpiredToken() throws ServletException, IOException, TokenException {
-    _authenticationFilter.init(null);
+    FilterConfig mockFilterConfig = mock(FilterConfig.class);
+    when(mockFilterConfig.getInitParameterNames()).thenReturn(Collections.emptyEnumeration());
+
+    authenticationFilter.init(mockFilterConfig);
     HttpServletRequest servletRequest = mock(HttpServletRequest.class);
     HttpServletResponse servletResponse = mock(HttpServletResponse.class);
     FilterChain filterChain = mock(FilterChain.class);
@@ -46,8 +55,47 @@ public void testExpiredToken() throws ServletException, IOException, TokenExcept
         .thenReturn(Collections.enumeration(List.of(AUTHORIZATION_HEADER_NAME)));
     when(servletRequest.getHeader(AUTHORIZATION_HEADER_NAME)).thenReturn("Bearer " + token);
 
-    _authenticationFilter.doFilter(servletRequest, servletResponse, filterChain);
+    authenticationFilter.doFilter(servletRequest, servletResponse, filterChain);
     verify(servletResponse, times(1))
         .sendError(eq(HttpServletResponse.SC_UNAUTHORIZED), anyString());
   }
+
+  @Test
+  public void testExcludedPaths() throws ServletException {
+    // Mock configuration setup
+    FilterConfig mockFilterConfig = mock(FilterConfig.class);
+    when(mockFilterConfig.getInitParameterNames()).thenReturn(Collections.emptyEnumeration());
+    authenticationFilter.init(mockFilterConfig);
+
+    // Test cases for different path patterns
+    HttpServletRequest exactPathRequest = mock(HttpServletRequest.class);
+    when(exactPathRequest.getServletPath()).thenReturn("/health");
+
+    HttpServletRequest wildcardPathRequest = mock(HttpServletRequest.class);
+    when(wildcardPathRequest.getServletPath()).thenReturn("/schema-registry/api/config");
+
+    HttpServletRequest nonExcludedRequest = mock(HttpServletRequest.class);
+    when(nonExcludedRequest.getServletPath()).thenReturn("/protected/resource");
+
+    // Set excluded paths in the filter
+    ReflectionTestUtils.setField(
+        authenticationFilter,
+        "excludedPathPatterns",
+        new HashSet<>(Arrays.asList("/health", "/schema-registry/*")));
+
+    // Verify exact path match
+    assertTrue(
+        authenticationFilter.shouldNotFilter(exactPathRequest),
+        "Exact path match should be excluded from filtering");
+
+    // Verify wildcard path match
+    assertTrue(
+        authenticationFilter.shouldNotFilter(wildcardPathRequest),
+        "Path matching wildcard pattern should be excluded from filtering");
+
+    // Verify non-excluded path
+    assertFalse(
+        authenticationFilter.shouldNotFilter(nonExcludedRequest),
+        "Non-excluded path should not be excluded from filtering");
+  }
 }
diff --git a/metadata-service/auth-servlet-impl/src/main/java/com/datahub/auth/authentication/AuthServiceController.java b/metadata-service/auth-servlet-impl/src/main/java/com/datahub/auth/authentication/AuthServiceController.java
index 5d4542cf0826e8..1f86c56049190d 100644
--- a/metadata-service/auth-servlet-impl/src/main/java/com/datahub/auth/authentication/AuthServiceController.java
+++ b/metadata-service/auth-servlet-impl/src/main/java/com/datahub/auth/authentication/AuthServiceController.java
@@ -35,11 +35,13 @@
 import org.springframework.http.HttpStatus;
 import org.springframework.http.ResponseEntity;
 import org.springframework.web.bind.annotation.PostMapping;
+import org.springframework.web.bind.annotation.RequestMapping;
 import org.springframework.web.bind.annotation.RestController;
 import org.springframework.web.client.HttpClientErrorException;
 
 @Slf4j
 @RestController
+@RequestMapping("/auth")
 public class AuthServiceController {
 
   private static final String USER_ID_FIELD_NAME = "userId";
diff --git a/metadata-service/configuration/src/main/resources/application.yaml b/metadata-service/configuration/src/main/resources/application.yaml
index 69b86962442b91..04f7409d5c39a3 100644
--- a/metadata-service/configuration/src/main/resources/application.yaml
+++ b/metadata-service/configuration/src/main/resources/application.yaml
@@ -5,6 +5,7 @@ baseUrl: ${DATAHUB_BASE_URL:http://localhost:9002}
 authentication:
   # Enable if you want all requests to the Metadata Service to be authenticated.
   enabled: ${METADATA_SERVICE_AUTH_ENABLED:true}
+  excludedPaths: /schema-registry/*,/health,/config,/config/search/export
 
   # Disable if you want to skip validation of deleted user's tokens
   enforceExistenceEnabled: ${METADATA_SERVICE_AUTH_ENFORCE_EXISTENCE_ENABLED:true}
@@ -330,16 +331,34 @@ neo4j:
   connectionLivenessCheckTimeout: ${NEO4J_CONNECTION_LIVENESS_CHECK_TIMEOUT_IN_SECONDS:-1}
 
 spring:
+  error:
+    include-message: never
+    include-stacktrace: never
+    include-exception: false
+    whitelabel:
+      enabled: false
+  jmx:
+    enabled: true
+  web:
+    resources:
+      add-mappings: false # do not serve static files
   mvc:
     servlet:
-      path: /openapi
+      path: /
+    throw-exception-if-no-handler-found: true # throw exception on 404 to be handled
   kafka:
     security:
       protocol: ${KAFKA_PROPERTIES_SECURITY_PROTOCOL:PLAINTEXT}
 
 springdoc:
   cache:
-    disabled: true
+    disabled: false
+  swagger-ui:
+    path: /openapi/swagger-ui/index.html
+  api-docs:
+    path: /openapi/v3/api-docs
+  groups:
+    enabled: true
 
 metadataTests:
   enabled: ${METADATA_TESTS_ENABLED:false}
@@ -563,8 +582,6 @@ graphQL:
     depthLimit: ${GRAPHQL_QUERY_DEPTH_LIMIT:50}
     introspectionEnabled: ${GRAPHQL_QUERY_INTROSPECTION_ENABLED:true}
 
-springdoc.api-docs.groups.enabled: true
-
 forms:
   hook:
     enabled: ${FORMS_HOOK_ENABLED:true}
diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/kafka/KafkaEventConsumerFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/kafka/KafkaEventConsumerFactory.java
index a1ee4df360b7ec..625623d008e127 100644
--- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/kafka/KafkaEventConsumerFactory.java
+++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/kafka/KafkaEventConsumerFactory.java
@@ -172,6 +172,7 @@ private KafkaListenerContainerFactory<?> buildDefaultKafkaListenerContainerFacto
     factory.setConsumerFactory(factoryWithOverrides);
     factory.setContainerCustomizer(new ThreadPoolContainerCustomizer());
     factory.setConcurrency(kafkaEventConsumerConcurrency);
+    factory.setAutoStartup(false);
 
     /* Sets up a delegating error handler for Deserialization errors, if disabled will
      use DefaultErrorHandler (does back-off retry and then logs) rather than stopping the container. Stopping the container
@@ -202,6 +203,7 @@ protected KafkaListenerContainerFactory<?> duheKafkaEventConsumer(
     factory.setConsumerFactory(kafkaConsumerFactory);
     factory.setContainerCustomizer(new ThreadPoolContainerCustomizer());
     factory.setConcurrency(1);
+    factory.setAutoStartup(false);
 
     log.info(
         "Event-based DUHE KafkaListenerContainerFactory built successfully. Consumer concurrency = 1");
diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/kafka/SimpleKafkaConsumerFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/kafka/SimpleKafkaConsumerFactory.java
index 0193ded97f81b5..acab78dcfd5f52 100644
--- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/kafka/SimpleKafkaConsumerFactory.java
+++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/kafka/SimpleKafkaConsumerFactory.java
@@ -53,6 +53,7 @@ protected KafkaListenerContainerFactory<?> createInstance(
         new ConcurrentKafkaListenerContainerFactory<>();
     factory.setContainerCustomizer(new ThreadPoolContainerCustomizer());
     factory.setConsumerFactory(new DefaultKafkaConsumerFactory<>(customizedProperties));
+    factory.setAutoStartup(false);
 
     log.info("Simple KafkaListenerContainerFactory built successfully");
 
diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/kafka/common/KafkaInitializationManager.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/kafka/common/KafkaInitializationManager.java
new file mode 100644
index 00000000000000..dcc818d82ee560
--- /dev/null
+++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/kafka/common/KafkaInitializationManager.java
@@ -0,0 +1,52 @@
+package com.linkedin.gms.factory.kafka.common;
+
+import java.util.concurrent.atomic.AtomicBoolean;
+import lombok.extern.slf4j.Slf4j;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.boot.availability.AvailabilityChangeEvent;
+import org.springframework.boot.availability.ReadinessState;
+import org.springframework.context.event.EventListener;
+import org.springframework.kafka.config.KafkaListenerEndpointRegistry;
+import org.springframework.stereotype.Component;
+
+@Slf4j
+@Component
+public class KafkaInitializationManager {
+  private final AtomicBoolean isInitialized = new AtomicBoolean(false);
+  private final KafkaListenerEndpointRegistry registry;
+
+  @Autowired
+  public KafkaInitializationManager(KafkaListenerEndpointRegistry registry) {
+    this.registry = registry;
+    log.info(
+        "Created KafkaInitializationManager. Waiting for application to be ready to enable kafka consumers.");
+  }
+
+  @EventListener
+  public void onStateChange(AvailabilityChangeEvent<ReadinessState> event) {
+    if (event.getState() == ReadinessState.ACCEPTING_TRAFFIC) {
+      initialize(this.getClass().getSimpleName());
+    }
+  }
+
+  public void initialize(String initializerName) {
+    if (isInitialized.compareAndSet(false, true)) {
+      int containerCount = registry.getAllListenerContainers().size();
+      log.info("Starting {} kafka consumers. Initialized by {}", containerCount, initializerName);
+      registry
+          .getAllListenerContainers()
+          .forEach(
+              container -> {
+                if (!container.isRunning()) {
+                  container.start();
+                  log.info("Started container: {}", container.getListenerId());
+                }
+              });
+      log.info("All {} kafka containers started.", containerCount);
+    }
+  }
+
+  public boolean isInitialized() {
+    return isInitialized.get();
+  }
+}
diff --git a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/OnBootApplicationListener.java b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/OnBootApplicationListener.java
index 921246fa98f7a2..2d599f340f7588 100644
--- a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/OnBootApplicationListener.java
+++ b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/OnBootApplicationListener.java
@@ -1,22 +1,11 @@
 package com.linkedin.metadata.boot;
 
-import com.linkedin.gms.factory.config.ConfigurationProvider;
-import com.linkedin.gms.factory.kafka.schemaregistry.InternalSchemaRegistryFactory;
+import com.linkedin.gms.factory.kafka.common.KafkaInitializationManager;
 import io.datahubproject.metadata.context.OperationContext;
-import java.io.IOException;
-import java.util.Set;
-import java.util.concurrent.ExecutorService;
-import java.util.concurrent.Executors;
 import javax.annotation.Nonnull;
 import lombok.extern.slf4j.Slf4j;
-import org.apache.http.HttpStatus;
-import org.apache.http.StatusLine;
-import org.apache.http.client.methods.HttpGet;
-import org.apache.http.impl.client.CloseableHttpClient;
-import org.apache.http.impl.client.HttpClients;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.beans.factory.annotation.Qualifier;
-import org.springframework.beans.factory.annotation.Value;
 import org.springframework.context.event.ContextRefreshedEvent;
 import org.springframework.context.event.EventListener;
 import org.springframework.stereotype.Component;
@@ -27,98 +16,39 @@
 @Component
 public class OnBootApplicationListener {
 
-  public static final String SCHEMA_REGISTRY_SERVLET_NAME = "dispatcher-schema-registry";
-
-  private static final Set<Integer> ACCEPTED_HTTP_CODES =
-      Set.of(
-          HttpStatus.SC_OK,
-          HttpStatus.SC_MOVED_PERMANENTLY,
-          HttpStatus.SC_MOVED_TEMPORARILY,
-          HttpStatus.SC_FORBIDDEN,
-          HttpStatus.SC_UNAUTHORIZED);
-
-  private static final String ROOT_WEB_APPLICATION_CONTEXT_ID =
-      String.format("%s:", WebApplicationContext.class.getName());
-
-  private final CloseableHttpClient httpClient = HttpClients.createDefault();
-
-  private final ExecutorService executorService = Executors.newSingleThreadExecutor();
-
   @Autowired
   @Qualifier("bootstrapManager")
   private BootstrapManager _bootstrapManager;
 
-  @Autowired
-  @Qualifier("configurationProvider")
-  private ConfigurationProvider provider;
-
-  @Value("${bootstrap.servlets.waitTimeout}")
-  private int _servletsWaitTimeout;
-
   @Autowired
   @Qualifier("systemOperationContext")
   private OperationContext systemOperationContext;
 
+  @Autowired private KafkaInitializationManager kafkaInitializationManager;
+
   @EventListener(ContextRefreshedEvent.class)
   public void onApplicationEvent(@Nonnull ContextRefreshedEvent event) {
+    String contextId = event.getApplicationContext().getId();
+    log.info("Context refreshed for ID: {}", contextId);
 
-    if (SCHEMA_REGISTRY_SERVLET_NAME.equals(event.getApplicationContext().getId())) {
-      log.info("Loading servlet {} without interruption.", SCHEMA_REGISTRY_SERVLET_NAME);
-      return;
-    }
-
-    log.warn(
-        "OnBootApplicationListener context refreshed! {} event: {}",
-        ROOT_WEB_APPLICATION_CONTEXT_ID.equals(event.getApplicationContext().getId()),
-        event);
-    String schemaRegistryType = provider.getKafka().getSchemaRegistry().getType();
-    if (ROOT_WEB_APPLICATION_CONTEXT_ID.equals(event.getApplicationContext().getId())) {
+    // For the root application context
+    if (event.getApplicationContext() instanceof WebApplicationContext) {
+      log.info("Root WebApplicationContext initialized, starting bootstrap process");
 
-      // Handle race condition, if ebean code is executed while waiting/bootstrapping (i.e.
-      // AuthenticationFilter)
+      // Initialize Ebean first
       try {
         Class.forName("io.ebean.XServiceProvider");
       } catch (ClassNotFoundException e) {
-        log.error(
-            "Failure to initialize required class `io.ebean.XServiceProvider` during initialization.");
+        log.error("Failed to initialize io.ebean.XServiceProvider", e);
         throw new RuntimeException(e);
       }
 
-      if (InternalSchemaRegistryFactory.TYPE.equals(schemaRegistryType)) {
-        executorService.submit(isSchemaRegistryAPIServletReady());
-      } else {
-        _bootstrapManager.start(systemOperationContext);
-      }
-    }
-  }
+      // Initialize consumers
+      kafkaInitializationManager.initialize(this.getClass().getSimpleName());
 
-  public Runnable isSchemaRegistryAPIServletReady() {
-    return () -> {
-      final HttpGet request = new HttpGet(provider.getKafka().getSchemaRegistry().getUrl());
-      int timeouts = _servletsWaitTimeout;
-      boolean openAPIServletReady = false;
-      while (!openAPIServletReady && timeouts > 0) {
-        try {
-          log.info("Sleeping for 1 second");
-          Thread.sleep(1000);
-          StatusLine statusLine = httpClient.execute(request).getStatusLine();
-          if (ACCEPTED_HTTP_CODES.contains(statusLine.getStatusCode())) {
-            log.info("Connected! Authentication not tested.");
-            openAPIServletReady = true;
-          }
-        } catch (IOException | InterruptedException e) {
-          log.info("Failed to connect to open servlet: {}", e.getMessage());
-        }
-        timeouts--;
-      }
-      if (!openAPIServletReady) {
-        log.error(
-            "Failed to bootstrap DataHub, OpenAPI servlet was not ready after {} seconds",
-            timeouts);
-        System.exit(1);
-      } else {
-        _bootstrapManager.start(systemOperationContext);
-      }
-    };
+      _bootstrapManager.start(systemOperationContext);
+    } else {
+      log.debug("Ignoring non-web application context refresh");
+    }
   }
 }
diff --git a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/factories/IngestRetentionPoliciesStepFactory.java b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/factories/IngestRetentionPoliciesStepFactory.java
index 56a857ee45feb8..388f48650a20f0 100644
--- a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/factories/IngestRetentionPoliciesStepFactory.java
+++ b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/factories/IngestRetentionPoliciesStepFactory.java
@@ -12,6 +12,7 @@
 import org.springframework.context.annotation.Configuration;
 import org.springframework.context.annotation.Import;
 import org.springframework.context.annotation.Scope;
+import org.springframework.core.io.support.PathMatchingResourcePatternResolver;
 
 @Configuration
 @Import({RetentionServiceFactory.class})
@@ -43,6 +44,7 @@ protected IngestRetentionPoliciesStep createInstance() {
         _entityService,
         _enableRetention,
         _applyOnBootstrap,
-        _pluginRegistryPath);
+        _pluginRegistryPath,
+        new PathMatchingResourcePatternResolver());
   }
 }
diff --git a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/kafka/DataHubUpgradeKafkaListener.java b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/kafka/DataHubUpgradeKafkaListener.java
index e69ab342740e43..50be0149ce2d4e 100644
--- a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/kafka/DataHubUpgradeKafkaListener.java
+++ b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/kafka/DataHubUpgradeKafkaListener.java
@@ -97,7 +97,8 @@ public void onPartitionsAssigned(
       id = CONSUMER_GROUP,
       topics = {TOPIC_NAME},
       containerFactory = "duheKafkaEventConsumer",
-      concurrency = "1")
+      concurrency = "1",
+      autoStartup = "false")
   public void checkSystemVersion(final ConsumerRecord<String, GenericRecord> consumerRecord) {
     try (Timer.Context i = MetricUtils.timer(this.getClass(), "checkSystemVersion").time()) {
       final GenericRecord record = consumerRecord.value();
diff --git a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/IngestDefaultGlobalSettingsStep.java b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/IngestDefaultGlobalSettingsStep.java
index c63d71475c2fc8..a41553379b1faa 100644
--- a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/IngestDefaultGlobalSettingsStep.java
+++ b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/IngestDefaultGlobalSettingsStep.java
@@ -78,7 +78,7 @@ public void execute(@Nonnull OperationContext systemOperationContext)
     // 1. Read from the file into JSON.
     JsonNode defaultSettingsObj;
     try {
-      defaultSettingsObj = mapper.readTree(new ClassPathResource(_resourcePath).getFile());
+      defaultSettingsObj = mapper.readTree(new ClassPathResource(_resourcePath).getInputStream());
     } catch (Exception e) {
       throw new RuntimeException(
           String.format(
diff --git a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/IngestPoliciesStep.java b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/IngestPoliciesStep.java
index dac2879487469c..04d73895802a8d 100644
--- a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/IngestPoliciesStep.java
+++ b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/IngestPoliciesStep.java
@@ -74,7 +74,7 @@ public void execute(@Nonnull OperationContext systemOperationContext)
     log.info("Ingesting default access policies from: {}...", _policiesResource);
 
     // 1. Read from the file into JSON.
-    final JsonNode policiesObj = mapper.readTree(_policiesResource.getFile());
+    final JsonNode policiesObj = mapper.readTree(_policiesResource.getInputStream());
 
     if (!policiesObj.isArray()) {
       throw new RuntimeException(
diff --git a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/IngestRetentionPoliciesStep.java b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/IngestRetentionPoliciesStep.java
index b5ca0dee142df3..01e785d06cc7d8 100644
--- a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/IngestRetentionPoliciesStep.java
+++ b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/IngestRetentionPoliciesStep.java
@@ -15,7 +15,6 @@
 import com.linkedin.retention.DataHubRetentionConfig;
 import io.datahubproject.metadata.context.OperationContext;
 import jakarta.annotation.Nonnull;
-import java.io.File;
 import java.io.IOException;
 import java.net.URISyntaxException;
 import java.util.Collections;
@@ -23,7 +22,8 @@
 import java.util.Map;
 import lombok.RequiredArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
-import org.springframework.core.io.ClassPathResource;
+import org.springframework.core.io.Resource;
+import org.springframework.core.io.support.ResourcePatternResolver;
 
 @Slf4j
 @RequiredArgsConstructor
@@ -34,6 +34,7 @@ public class IngestRetentionPoliciesStep implements BootstrapStep {
   private final boolean _enableRetention;
   private final boolean _applyOnBootstrap;
   private final String pluginPath;
+  private final ResourcePatternResolver resolver;
 
   private static final ObjectMapper YAML_MAPPER = new ObjectMapper(new YAMLFactory());
 
@@ -77,12 +78,19 @@ public void execute(@Nonnull OperationContext systemOperationContext)
       return;
     }
 
-    // 1. Read default retention config
-    final Map<DataHubRetentionKey, DataHubRetentionConfig> retentionPolicyMap =
-        parseFileOrDir(new ClassPathResource("./boot/retention.yaml").getFile());
-
-    // 2. Read plugin retention config files from input path and overlay
-    retentionPolicyMap.putAll(parseFileOrDir(new File(pluginPath)));
+    // 1. Read default retention config from classpath
+    Resource defaultResource = resolver.getResource("classpath:boot/retention.yaml");
+    Map<DataHubRetentionKey, DataHubRetentionConfig> retentionPolicyMap =
+        parseYamlRetentionConfig(defaultResource);
+
+    // 2. Read plugin retention config files from filesystem path
+    if (!pluginPath.isEmpty()) {
+      String pattern = "file:" + pluginPath + "/**/*.{yaml,yml}";
+      Resource[] resources = resolver.getResources(pattern);
+      for (Resource resource : resources) {
+        retentionPolicyMap.putAll(parseYamlRetentionConfig(resource));
+      }
+    }
 
     // 4. Set the specified retention policies
     log.info("Setting {} policies", retentionPolicyMap.size());
@@ -106,39 +114,6 @@ public void execute(@Nonnull OperationContext systemOperationContext)
     BootstrapStep.setUpgradeResult(systemOperationContext, UPGRADE_ID_URN, _entityService);
   }
 
-  // Parse input yaml file or yaml files in the input directory to generate a retention policy map
-  private Map<DataHubRetentionKey, DataHubRetentionConfig> parseFileOrDir(File retentionFileOrDir)
-      throws IOException {
-    // If path does not exist return empty
-    if (!retentionFileOrDir.exists()) {
-      return Collections.emptyMap();
-    }
-
-    // If directory, parse the yaml files under the directory
-    if (retentionFileOrDir.isDirectory()) {
-      Map<DataHubRetentionKey, DataHubRetentionConfig> result = new HashMap<>();
-
-      for (File retentionFile : retentionFileOrDir.listFiles()) {
-        if (!retentionFile.isFile()) {
-          log.info(
-              "Element {} in plugin directory {} is not a file. Skipping",
-              retentionFile.getPath(),
-              retentionFileOrDir.getPath());
-          continue;
-        }
-        result.putAll(parseFileOrDir(retentionFile));
-      }
-      return result;
-    }
-    // If file, parse the yaml file and return result;
-    if (!retentionFileOrDir.getPath().endsWith(".yaml")
-        && retentionFileOrDir.getPath().endsWith(".yml")) {
-      log.info("File {} is not a YAML file. Skipping", retentionFileOrDir.getPath());
-      return Collections.emptyMap();
-    }
-    return parseYamlRetentionConfig(retentionFileOrDir);
-  }
-
   /**
    * Parse yaml retention config
    *
@@ -147,8 +122,11 @@ private Map<DataHubRetentionKey, DataHubRetentionConfig> parseFileOrDir(File ret
    * converted into the {@link com.linkedin.retention.DataHubRetentionConfig} class.
    */
   private Map<DataHubRetentionKey, DataHubRetentionConfig> parseYamlRetentionConfig(
-      File retentionConfigFile) throws IOException {
-    final JsonNode retentionPolicies = YAML_MAPPER.readTree(retentionConfigFile);
+      Resource resource) throws IOException {
+    if (!resource.exists()) {
+      return Collections.emptyMap();
+    }
+    final JsonNode retentionPolicies = YAML_MAPPER.readTree(resource.getInputStream());
     if (!retentionPolicies.isArray()) {
       throw new IllegalArgumentException(
           "Retention config file must contain an array of retention policies");
diff --git a/metadata-service/factories/src/main/java/com/linkedin/r2/transport/http/server/AbstractJakartaR2Servlet.java b/metadata-service/factories/src/main/java/com/linkedin/r2/transport/http/server/AbstractJakartaR2Servlet.java
index 8d589c4ab2408b..9c4117202ed570 100644
--- a/metadata-service/factories/src/main/java/com/linkedin/r2/transport/http/server/AbstractJakartaR2Servlet.java
+++ b/metadata-service/factories/src/main/java/com/linkedin/r2/transport/http/server/AbstractJakartaR2Servlet.java
@@ -1,35 +1,9 @@
-/*
-   Copyright (c) 2012 LinkedIn Corp.
-   Copyright (c) 2023 Acryl Data, Inc.
-
-   Licensed under the Apache License, Version 2.0 (the "License");
-   you may not use this file except in compliance with the License.
-   You may obtain a copy of the License at
-
-       http://www.apache.org/licenses/LICENSE-2.0
-
-   Unless required by applicable law or agreed to in writing, software
-   distributed under the License is distributed on an "AS IS" BASIS,
-   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-   See the License for the specific language governing permissions and
-   limitations under the License.
-
-   MODIFICATIONS:
-   Changed javax packages to jakarta for Spring Boot 3 support
-*/
-
-/* $Id$ */
 package com.linkedin.r2.transport.http.server;
 
 import com.linkedin.data.ByteString;
 import com.linkedin.r2.message.RequestContext;
-import com.linkedin.r2.message.rest.RestException;
-import com.linkedin.r2.message.rest.RestRequest;
-import com.linkedin.r2.message.rest.RestRequestBuilder;
-import com.linkedin.r2.message.rest.RestResponse;
-import com.linkedin.r2.message.rest.RestStatus;
+import com.linkedin.r2.message.rest.*;
 import com.linkedin.r2.transport.common.WireAttributeHelper;
-import com.linkedin.r2.transport.common.bridge.common.TransportCallback;
 import com.linkedin.r2.transport.common.bridge.common.TransportResponse;
 import com.linkedin.r2.transport.common.bridge.common.TransportResponseImpl;
 import com.linkedin.r2.transport.http.common.HttpConstants;
@@ -40,210 +14,177 @@
 import java.io.IOException;
 import java.net.URI;
 import java.net.URISyntaxException;
-import java.util.Enumeration;
-import java.util.Map;
-import java.util.concurrent.CountDownLatch;
+import java.time.Duration;
+import java.util.Collections;
+import java.util.Optional;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.TimeUnit;
-import java.util.concurrent.atomic.AtomicReference;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-/**
- * @author Steven Ihde
- * @author Chris Pettitt
- * @author Fatih Emekci
- * @version $Revision$
- */
+import java.util.concurrent.TimeoutException;
+import lombok.extern.slf4j.Slf4j;
+
+@Slf4j
 public abstract class AbstractJakartaR2Servlet extends HttpServlet {
-  private static final Logger _log = LoggerFactory.getLogger(AbstractJakartaR2Servlet.class);
-  private static final long serialVersionUID = 0L;
+  private static final long serialVersionUID = 1L;
 
-  // servlet timeout in ms.
-  protected final long _timeout;
+  private final Duration timeout;
 
   protected abstract HttpDispatcher getDispatcher();
 
-  public AbstractJakartaR2Servlet(long timeout) {
-    _timeout = timeout;
+  protected AbstractJakartaR2Servlet(Duration timeout) {
+    this.timeout = timeout;
   }
 
   @Override
-  protected void service(final HttpServletRequest req, final HttpServletResponse resp)
+  protected void service(HttpServletRequest req, HttpServletResponse resp)
       throws ServletException, IOException {
-    RequestContext requestContext = JakartaServletHelper.readRequestContext(req);
-
-    RestRequest restRequest;
-
     try {
-      restRequest = readFromServletRequest(req);
-    } catch (URISyntaxException e) {
-      writeToServletError(resp, RestStatus.BAD_REQUEST, e.toString());
-      return;
-    }
+      RequestContext requestContext = JakartaServletHelper.readRequestContext(req);
+      RestRequest restRequest = createRestRequest(req);
 
-    final AtomicReference<TransportResponse<RestResponse>> result = new AtomicReference<>();
-    final CountDownLatch latch = new CountDownLatch(1);
+      CompletableFuture<TransportResponse<RestResponse>> future = new CompletableFuture<>();
 
-    TransportCallback<RestResponse> callback =
-        new TransportCallback<RestResponse>() {
-          @Override
-          public void onResponse(TransportResponse<RestResponse> response) {
-            result.set(response);
-            latch.countDown();
-          }
-        };
+      getDispatcher().handleRequest(restRequest, requestContext, future::complete);
 
-    getDispatcher().handleRequest(restRequest, requestContext, callback);
+      TransportResponse<RestResponse> result =
+          future
+              .orTimeout(timeout.toMillis(), TimeUnit.MILLISECONDS)
+              .exceptionally(this::handleException)
+              .join();
 
-    try {
-      if (latch.await(_timeout, TimeUnit.MILLISECONDS)) {
-        writeToServletResponse(result.get(), resp);
-      } else {
-        writeToServletError(
-            resp, RestStatus.INTERNAL_SERVER_ERROR, "Server Timeout after " + _timeout + "ms.");
-      }
-    } catch (InterruptedException e) {
-      throw new ServletException("Interrupted!", e);
-    }
-  }
+      writeResponse(result, resp);
 
-  protected void writeToServletResponse(
-      TransportResponse<RestResponse> response, HttpServletResponse resp) throws IOException {
-    Map<String, String> wireAttrs = response.getWireAttributes();
-    for (Map.Entry<String, String> e : WireAttributeHelper.toWireAttributes(wireAttrs).entrySet()) {
-      resp.setHeader(e.getKey(), e.getValue());
-    }
-
-    RestResponse restResponse = null;
-    if (response.hasError()) {
-      Throwable e = response.getError();
-      if (e instanceof RestException) {
-        restResponse = ((RestException) e).getResponse();
-      }
-      if (restResponse == null) {
-        restResponse = RestStatus.responseForError(RestStatus.INTERNAL_SERVER_ERROR, e);
-      }
-    } else {
-      restResponse = response.getResponse();
-    }
-
-    resp.setStatus(restResponse.getStatus());
-    Map<String, String> headers = restResponse.getHeaders();
-    for (Map.Entry<String, String> e : headers.entrySet()) {
-      // TODO multi-valued headers
-      resp.setHeader(e.getKey(), e.getValue());
+    } catch (URISyntaxException e) {
+      writeError(resp, RestStatus.BAD_REQUEST, "Invalid URI: " + e.getMessage());
+    } catch (Exception e) {
+      log.error("Unexpected error processing request", e);
+      writeError(resp, RestStatus.INTERNAL_SERVER_ERROR, "Internal server error");
     }
+  }
 
-    for (String cookie : restResponse.getCookies()) {
-      resp.addHeader(HttpConstants.RESPONSE_COOKIE_HEADER_NAME, cookie);
+  private TransportResponse<RestResponse> handleException(Throwable ex) {
+    if (ex instanceof TimeoutException) {
+      RestResponse errorResponse =
+          RestStatus.responseForError(
+              RestStatus.INTERNAL_SERVER_ERROR,
+              new RuntimeException("Server timeout after " + timeout.toSeconds() + " seconds"));
+      return TransportResponseImpl.error(new RestException(errorResponse));
     }
+    return TransportResponseImpl.error(ex);
+  }
 
-    final ByteString entity = restResponse.getEntity();
-    entity.write(resp.getOutputStream());
-
-    resp.getOutputStream().close();
+  private RestRequest createRestRequest(HttpServletRequest req)
+      throws IOException, ServletException, URISyntaxException {
+    String pathInfo = extractPathInfo(req);
+    String queryString = Optional.ofNullable(req.getQueryString()).map(q -> "?" + q).orElse("");
+
+    URI uri = new URI(pathInfo + queryString);
+
+    RestRequestBuilder builder = new RestRequestBuilder(uri).setMethod(req.getMethod());
+
+    // Handle headers
+    Collections.list(req.getHeaderNames())
+        .forEach(
+            headerName -> {
+              if (headerName.equalsIgnoreCase(HttpConstants.REQUEST_COOKIE_HEADER_NAME)) {
+                Collections.list(req.getHeaders(headerName)).forEach(builder::addCookie);
+              } else {
+                Collections.list(req.getHeaders(headerName))
+                    .forEach(value -> builder.addHeaderValue(headerName, value));
+              }
+            });
+
+    // Handle request body
+    int contentLength = req.getContentLength();
+    ByteString entity =
+        (contentLength > 0)
+            ? ByteString.read(req.getInputStream(), contentLength)
+            : ByteString.read(req.getInputStream());
+
+    builder.setEntity(entity);
+
+    return builder.build();
   }
 
-  protected void writeToServletError(HttpServletResponse resp, int statusCode, String message)
+  private void writeResponse(TransportResponse<RestResponse> response, HttpServletResponse resp)
       throws IOException {
-    RestResponse restResponse = RestStatus.responseForStatus(statusCode, message);
-    writeToServletResponse(TransportResponseImpl.success(restResponse), resp);
-  }
+    // Write wire attributes
+    WireAttributeHelper.toWireAttributes(response.getWireAttributes()).forEach(resp::setHeader);
+
+    // Get response or create error response
+    RestResponse restResponse =
+        Optional.of(response)
+            .filter(TransportResponse::hasError)
+            .map(
+                r -> {
+                  Throwable error = r.getError();
+                  if (error instanceof RestException) {
+                    return ((RestException) error).getResponse();
+                  }
+                  return RestStatus.responseForError(RestStatus.INTERNAL_SERVER_ERROR, error);
+                })
+            .orElseGet(response::getResponse);
+
+    // Write status and headers
+    resp.setStatus(restResponse.getStatus());
+    restResponse.getHeaders().forEach(resp::setHeader);
 
-  protected RestRequest readFromServletRequest(HttpServletRequest req)
-      throws IOException, ServletException, URISyntaxException {
-    StringBuilder sb = new StringBuilder();
-    sb.append(extractPathInfo(req));
-    String query = req.getQueryString();
-    if (query != null) {
-      sb.append('?');
-      sb.append(query);
-    }
+    // Write cookies
+    restResponse
+        .getCookies()
+        .forEach(cookie -> resp.addHeader(HttpConstants.RESPONSE_COOKIE_HEADER_NAME, cookie));
 
-    URI uri = new URI(sb.toString());
-
-    RestRequestBuilder rb = new RestRequestBuilder(uri);
-    rb.setMethod(req.getMethod());
-
-    for (Enumeration<String> headerNames = req.getHeaderNames(); headerNames.hasMoreElements(); ) {
-      String headerName = headerNames.nextElement();
-      if (headerName.equalsIgnoreCase(HttpConstants.REQUEST_COOKIE_HEADER_NAME)) {
-        for (Enumeration<String> cookies = req.getHeaders(headerName);
-            cookies.hasMoreElements(); ) {
-          rb.addCookie(cookies.nextElement());
-        }
-      } else {
-        for (Enumeration<String> headerValues = req.getHeaders(headerName);
-            headerValues.hasMoreElements(); ) {
-          rb.addHeaderValue(headerName, headerValues.nextElement());
-        }
-      }
+    // Write response body
+    try (var outputStream = resp.getOutputStream()) {
+      restResponse.getEntity().write(outputStream);
     }
+  }
 
-    int length = req.getContentLength();
-    if (length > 0) {
-      rb.setEntity(ByteString.read(req.getInputStream(), length));
-    } else {
-      // Known cases for not sending a content-length header in a request
-      // 1. Chunked transfer encoding
-      // 2. HTTP/2
-      rb.setEntity(ByteString.read(req.getInputStream()));
-    }
-    return rb.build();
+  private void writeError(HttpServletResponse resp, int statusCode, String message)
+      throws IOException {
+    RestResponse errorResponse = RestStatus.responseForStatus(statusCode, message);
+    writeResponse(TransportResponseImpl.success(errorResponse), resp);
   }
 
-  /**
-   * Attempts to return a "non decoded" pathInfo by stripping off the contextPath and servletPath
-   * parts of the requestURI. As a defensive measure, this method will return the "decoded" pathInfo
-   * directly by calling req.getPathInfo() if it is unable to strip off the contextPath or
-   * servletPath.
-   *
-   * @throws javax.servlet.ServletException if resulting pathInfo is empty
-   */
   protected static String extractPathInfo(HttpServletRequest req) throws ServletException {
-    // For "http:hostname:8080/contextPath/servletPath/pathInfo" the RequestURI is
-    // "/contextPath/servletPath/pathInfo"
-    // where the contextPath, servletPath and pathInfo parts all contain their leading slash.
-
-    // stripping contextPath and servletPath this way is not fully compatible with the HTTP spec.
-    // If a
-    // request for, say "/%75scp-proxy/reso%75rces" is made (where %75 decodes to 'u')
-    // the stripping off of contextPath and servletPath will fail because the requestUri string will
-    // include the encoded char but the contextPath and servletPath strings will not.
     String requestUri = req.getRequestURI();
-    String contextPath = req.getContextPath();
-    StringBuilder builder = new StringBuilder();
-    if (contextPath != null) {
-      builder.append(contextPath);
-    }
+    String contextPath = Optional.ofNullable(req.getContextPath()).orElse("");
+    String servletPath = Optional.ofNullable(req.getServletPath()).orElse("");
 
-    String servletPath = req.getServletPath();
-    if (servletPath != null) {
-      builder.append(servletPath);
-    }
-    String prefix = builder.toString();
-    String pathInfo;
-    if (prefix.length() == 0) {
+    String prefix = contextPath + servletPath;
+    String pathInfo = null;
+
+    if (prefix.isEmpty()) {
       pathInfo = requestUri;
-    } else if (requestUri.startsWith(prefix)) {
+    } else if (servletPath.startsWith("/gms") && requestUri.startsWith(prefix)) {
       pathInfo = requestUri.substring(prefix.length());
-    } else {
-      _log.warn(
-          "Unable to extract 'non decoded' pathInfo, returning 'decoded' pathInfo instead.  This may cause issues processing request URIs containing special characters. requestUri="
-              + requestUri);
-      return req.getPathInfo();
     }
 
-    if (pathInfo.length() == 0) {
-      // We prefer to keep servlet mapping trivial with R2 and have R2
-      // TransportDispatchers make most of the routing decisions based on the 'pathInfo'
-      // and query parameters in the URI.
-      // If pathInfo is null, it's highly likely that the servlet was mapped to an exact
-      // path or to a file extension, making such R2-based services too reliant on the
-      // servlet container for routing
-      throw new ServletException(
-          "R2 servlet should only be mapped via wildcard path mapping e.g. /r2/*. "
-              + "Exact path matching (/r2) and file extension mappings (*.r2) are currently not supported");
+    if (pathInfo == null || pathInfo.isEmpty()) {
+      log.debug(
+          "Previously invalid servlet mapping detected. Request details: method='{}', requestUri='{}', contextPath='{}', "
+              + "servletPath='{}', prefix='{}', pathInfo='{}', queryString='{}', protocol='{}', remoteAddr='{}', "
+              + "serverName='{}', serverPort={}, contentType='{}', characterEncoding='{}'",
+          req.getMethod(),
+          requestUri,
+          contextPath,
+          servletPath,
+          prefix,
+          pathInfo,
+          req.getQueryString(),
+          req.getProtocol(),
+          req.getRemoteAddr(),
+          req.getServerName(),
+          req.getServerPort(),
+          req.getContentType(),
+          req.getCharacterEncoding());
+
+      /*  NOTE: Working around what was previously considered an error.
+       *       throw new ServletException(
+       *       "R2 servlet must be mapped using wildcard path mapping (e.g., /r2/*). " +
+       *               "Exact path matching (/r2) and file extension mappings (*.r2) are not supported.");
+       **/
+
+      pathInfo = requestUri;
     }
 
     return pathInfo;
diff --git a/metadata-service/factories/src/main/java/com/linkedin/r2/transport/http/server/RAPJakartaServlet.java b/metadata-service/factories/src/main/java/com/linkedin/r2/transport/http/server/RAPJakartaServlet.java
index 75e808d5245c83..f0db3cd01ffcf9 100644
--- a/metadata-service/factories/src/main/java/com/linkedin/r2/transport/http/server/RAPJakartaServlet.java
+++ b/metadata-service/factories/src/main/java/com/linkedin/r2/transport/http/server/RAPJakartaServlet.java
@@ -1,44 +1,21 @@
-/*
-   Copyright (c) 2012 LinkedIn Corp.
-   Copyright (c) 2023 Acryl Data, Inc.
-
-   Licensed under the Apache License, Version 2.0 (the "License");
-   you may not use this file except in compliance with the License.
-   You may obtain a copy of the License at
-
-       http://www.apache.org/licenses/LICENSE-2.0
-
-   Unless required by applicable law or agreed to in writing, software
-   distributed under the License is distributed on an "AS IS" BASIS,
-   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-   See the License for the specific language governing permissions and
-   limitations under the License.
-
-   MODIFICATIONS:
-   Changed javax dependencies to jakarta for Spring Boot 3 support
-*/
-
-/** $Id: $ */
 package com.linkedin.r2.transport.http.server;
 
 import com.linkedin.r2.transport.common.bridge.server.TransportDispatcher;
+import java.time.Duration;
+import lombok.Getter;
 
-/**
- * @author Steven Ihde
- * @version $Revision: $
- */
 public class RAPJakartaServlet extends AbstractJakartaR2Servlet {
   private static final long serialVersionUID = 0L;
 
-  private final HttpDispatcher _dispatcher;
+  @Getter private final HttpDispatcher dispatcher;
 
-  public RAPJakartaServlet(HttpDispatcher dispatcher) {
-    super(Long.MAX_VALUE);
-    _dispatcher = dispatcher;
+  public RAPJakartaServlet(HttpDispatcher dispatcher, int timeoutSeconds) {
+    super(Duration.ofSeconds(timeoutSeconds));
+    this.dispatcher = dispatcher;
   }
 
-  public RAPJakartaServlet(TransportDispatcher dispatcher) {
-    this(HttpDispatcherFactory.create((dispatcher)));
+  public RAPJakartaServlet(TransportDispatcher dispatcher, int timeoutSeconds) {
+    this(HttpDispatcherFactory.create((dispatcher)), timeoutSeconds);
   }
 
   /**
@@ -48,22 +25,7 @@ public RAPJakartaServlet(TransportDispatcher dispatcher) {
    */
   public RAPJakartaServlet(
       HttpDispatcher dispatcher, boolean useContinuations, int timeOut, int timeOutDelta) {
-    super(timeOut);
-    _dispatcher = dispatcher;
-  }
-
-  /**
-   * Initialize the RAPJakartaServlet.
-   *
-   * @see #AbstractJakartaR2Servlet
-   */
-  public RAPJakartaServlet(
-      TransportDispatcher dispatcher, boolean useContinuations, int timeOut, int timeOutDelta) {
-    this(HttpDispatcherFactory.create((dispatcher)), useContinuations, timeOut, timeOutDelta);
-  }
-
-  @Override
-  protected HttpDispatcher getDispatcher() {
-    return _dispatcher;
+    super(Duration.ofSeconds(timeOut));
+    this.dispatcher = dispatcher;
   }
 }
diff --git a/metadata-service/factories/src/main/java/com/linkedin/restli/server/RAPServletFactory.java b/metadata-service/factories/src/main/java/com/linkedin/restli/server/RAPServletFactory.java
index be060477aeb1f6..73db7ec73f4d7e 100644
--- a/metadata-service/factories/src/main/java/com/linkedin/restli/server/RAPServletFactory.java
+++ b/metadata-service/factories/src/main/java/com/linkedin/restli/server/RAPServletFactory.java
@@ -26,6 +26,9 @@ public class RAPServletFactory {
   @Value("#{systemEnvironment['RESTLI_SERVLET_THREADS']}")
   private Integer environmentThreads;
 
+  @Value("${RESTLI_TIMEOUT_SECONDS:60}")
+  private int restliTimeoutSeconds;
+
   @Value("${" + INGESTION_MAX_SERIALIZED_STRING_LENGTH + ":16000000}")
   private int maxSerializedStringLength;
 
@@ -71,6 +74,7 @@ public RAPJakartaServlet rapServlet(
     RestLiServer restLiServer = new RestLiServer(config, springInjectResourceFactory, parseqEngine);
     return new RAPJakartaServlet(
         new FilterChainDispatcher(
-            new DelegatingTransportDispatcher(restLiServer, restLiServer), FilterChains.empty()));
+            new DelegatingTransportDispatcher(restLiServer, restLiServer), FilterChains.empty()),
+        restliTimeoutSeconds);
   }
 }
diff --git a/metadata-service/factories/src/main/java/com/linkedin/restli/server/RestliHandlerServlet.java b/metadata-service/factories/src/main/java/com/linkedin/restli/server/RestliHandlerServlet.java
index bfc25b7ddaef50..4a38b331d95fed 100644
--- a/metadata-service/factories/src/main/java/com/linkedin/restli/server/RestliHandlerServlet.java
+++ b/metadata-service/factories/src/main/java/com/linkedin/restli/server/RestliHandlerServlet.java
@@ -8,26 +8,25 @@
 import java.io.IOException;
 import lombok.AllArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
-import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.web.HttpRequestHandler;
 import org.springframework.web.context.support.HttpRequestHandlerServlet;
 
 @Slf4j
 @AllArgsConstructor
 public class RestliHandlerServlet extends HttpRequestHandlerServlet implements HttpRequestHandler {
-  @Autowired private RAPJakartaServlet _r2Servlet;
+  private final RAPJakartaServlet r2Servlet;
 
   @Override
   public void init(ServletConfig config) throws ServletException {
     log.info("Initializing RestliHandlerServlet");
-    this._r2Servlet.init(config);
+    this.r2Servlet.init(config);
     log.info("Initialized RestliHandlerServlet");
   }
 
   @Override
   public void service(HttpServletRequest req, HttpServletResponse res)
       throws ServletException, IOException {
-    _r2Servlet.service(req, res);
+    r2Servlet.service(req, res);
   }
 
   @Override
diff --git a/metadata-service/factories/src/test/java/com/linkedin/gms/factory/kafka/common/KafkaInitializationManagerTest.java b/metadata-service/factories/src/test/java/com/linkedin/gms/factory/kafka/common/KafkaInitializationManagerTest.java
new file mode 100644
index 00000000000000..6404f238b5cf93
--- /dev/null
+++ b/metadata-service/factories/src/test/java/com/linkedin/gms/factory/kafka/common/KafkaInitializationManagerTest.java
@@ -0,0 +1,165 @@
+package com.linkedin.gms.factory.kafka.common;
+
+import static org.mockito.Mockito.*;
+import static org.testng.Assert.*;
+
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.boot.availability.AvailabilityChangeEvent;
+import org.springframework.boot.availability.ReadinessState;
+import org.springframework.boot.test.context.SpringBootTest;
+import org.springframework.boot.test.mock.mockito.MockBean;
+import org.springframework.kafka.annotation.EnableKafka;
+import org.springframework.kafka.config.KafkaListenerEndpointRegistry;
+import org.springframework.kafka.listener.MessageListenerContainer;
+import org.springframework.test.annotation.DirtiesContext;
+import org.springframework.test.context.ContextConfiguration;
+import org.springframework.test.context.testng.AbstractTestNGSpringContextTests;
+import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.Test;
+
+@SpringBootTest
+@ContextConfiguration(classes = {KafkaInitializationManager.class})
+@EnableKafka
+@DirtiesContext
+public class KafkaInitializationManagerTest extends AbstractTestNGSpringContextTests {
+
+  @Autowired private KafkaInitializationManager manager;
+
+  @MockBean private KafkaListenerEndpointRegistry registry;
+
+  private MessageListenerContainer container1;
+  private MessageListenerContainer container2;
+  private List<MessageListenerContainer> containers;
+
+  @BeforeMethod
+  public void setUp() {
+    // Reset the initialization state before each test
+    // This is necessary because Spring maintains the same instance across tests
+    try {
+      java.lang.reflect.Field field =
+          KafkaInitializationManager.class.getDeclaredField("isInitialized");
+      field.setAccessible(true);
+      ((java.util.concurrent.atomic.AtomicBoolean) field.get(manager)).set(false);
+    } catch (Exception e) {
+      throw new RuntimeException("Failed to reset initialization state", e);
+    }
+
+    // Create fresh mocks for each test
+    container1 = mock(MessageListenerContainer.class);
+    container2 = mock(MessageListenerContainer.class);
+    containers = Arrays.asList(container1, container2);
+
+    // Clear any previous mock interactions
+    reset(registry);
+  }
+
+  @Test
+  public void testInitialize_WhenNotInitialized() {
+    // Arrange
+    when(registry.getAllListenerContainers()).thenReturn(containers);
+    when(container1.isRunning()).thenReturn(false);
+    when(container2.isRunning()).thenReturn(false);
+    when(container1.getListenerId()).thenReturn("container1");
+    when(container2.getListenerId()).thenReturn("container2");
+
+    // Act
+    manager.initialize("TestInitializer");
+
+    // Assert
+    verify(container1, times(1)).start();
+    verify(container2, times(1)).start();
+    assertTrue(manager.isInitialized());
+  }
+
+  @Test
+  public void testInitialize_WhenAlreadyInitialized() {
+    // Arrange
+    when(registry.getAllListenerContainers()).thenReturn(containers);
+
+    // First initialization
+    manager.initialize("FirstInitializer");
+
+    // Clear previous interactions
+    clearInvocations(container1, container2);
+
+    // Act
+    manager.initialize("SecondInitializer");
+
+    // Assert
+    verify(container1, never()).start();
+    verify(container2, never()).start();
+    assertTrue(manager.isInitialized());
+  }
+
+  @Test
+  public void testInitialize_WithNoContainers() {
+    // Arrange
+    when(registry.getAllListenerContainers()).thenReturn(Collections.emptyList());
+
+    // Act
+    manager.initialize("TestInitializer");
+
+    // Assert
+    assertTrue(manager.isInitialized());
+  }
+
+  @Test
+  public void testInitialize_WithAlreadyRunningContainer() {
+    // Arrange
+    when(registry.getAllListenerContainers()).thenReturn(Collections.singletonList(container1));
+    when(container1.isRunning()).thenReturn(true);
+    when(container1.getListenerId()).thenReturn("container1");
+
+    // Act
+    manager.initialize("TestInitializer");
+
+    // Assert
+    verify(container1, never()).start();
+    assertTrue(manager.isInitialized());
+  }
+
+  @Test
+  public void testOnStateChange_WhenAcceptingTraffic() {
+    // Arrange
+    AvailabilityChangeEvent<ReadinessState> event =
+        new AvailabilityChangeEvent<>(this, ReadinessState.ACCEPTING_TRAFFIC);
+    when(registry.getAllListenerContainers()).thenReturn(containers);
+    when(container1.isRunning()).thenReturn(false);
+    when(container2.isRunning()).thenReturn(false);
+    when(container1.getListenerId()).thenReturn("container1");
+    when(container2.getListenerId()).thenReturn("container2");
+
+    // Act
+    manager.onStateChange(event);
+
+    // Assert
+    assertTrue(manager.isInitialized());
+    // The method is called twice in the implementation - once for size() and once for iteration
+    verify(registry, times(2)).getAllListenerContainers();
+    verify(container1).start();
+    verify(container2).start();
+  }
+
+  @Test
+  public void testOnStateChange_WhenNotAcceptingTraffic() {
+    // Arrange
+    AvailabilityChangeEvent<ReadinessState> event =
+        new AvailabilityChangeEvent<>(this, ReadinessState.REFUSING_TRAFFIC);
+
+    // Act
+    manager.onStateChange(event);
+
+    // Assert
+    assertFalse(manager.isInitialized());
+    verify(registry, never()).getAllListenerContainers();
+  }
+
+  @Test
+  public void testIsInitialized_DefaultState() {
+    // Assert
+    assertFalse(manager.isInitialized());
+  }
+}
diff --git a/metadata-service/factories/src/test/java/com/linkedin/metadata/boot/steps/IngestRetentionPoliciesStepTest.java b/metadata-service/factories/src/test/java/com/linkedin/metadata/boot/steps/IngestRetentionPoliciesStepTest.java
new file mode 100644
index 00000000000000..4f4c4b8bb68ae5
--- /dev/null
+++ b/metadata-service/factories/src/test/java/com/linkedin/metadata/boot/steps/IngestRetentionPoliciesStepTest.java
@@ -0,0 +1,257 @@
+package com.linkedin.metadata.boot.steps;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+import com.linkedin.common.urn.Urn;
+import com.linkedin.metadata.boot.BootstrapStep;
+import com.linkedin.metadata.entity.EntityService;
+import com.linkedin.metadata.entity.RetentionService;
+import com.linkedin.retention.DataHubRetentionConfig;
+import io.datahubproject.metadata.context.OperationContext;
+import java.io.ByteArrayInputStream;
+import org.mockito.Mockito;
+import org.springframework.core.io.Resource;
+import org.springframework.core.io.support.PathMatchingResourcePatternResolver;
+import org.testng.annotations.Test;
+
+public class IngestRetentionPoliciesStepTest {
+
+  private static final String UPGRADE_ID = "ingest-retention-policies";
+  private static final Urn UPGRADE_ID_URN = BootstrapStep.getUpgradeUrn(UPGRADE_ID);
+
+  private static final String DEFAULT_RETENTION_YAML =
+      "- entity: \"dataHubExecutionRequest\"\n"
+          + "  aspect: \"dataHubExecutionRequestResult\"\n"
+          + "  config:\n"
+          + "    retention:\n"
+          + "      version:\n"
+          + "        maxVersions: 1\n"
+          + "- entity: \"*\"\n"
+          + "  aspect: \"*\"\n"
+          + "  config:\n"
+          + "    retention:\n"
+          + "      version:\n"
+          + "        maxVersions: 20\n";
+
+  @Test
+  public void testExecuteWhenRetentionIsDisabled() throws Exception {
+    final EntityService<?> mockEntityService = Mockito.mock(EntityService.class);
+    final RetentionService<?> mockRetentionService = Mockito.mock(RetentionService.class);
+    final OperationContext mockContext = mock(OperationContext.class);
+    final PathMatchingResourcePatternResolver mockResolver =
+        mock(PathMatchingResourcePatternResolver.class);
+
+    final IngestRetentionPoliciesStep retentionStep =
+        new IngestRetentionPoliciesStep(
+            mockRetentionService, mockEntityService, false, false, "", mockResolver);
+
+    retentionStep.execute(mockContext);
+
+    // Verify that no interactions occur with retention service when retention is disabled
+    Mockito.verify(mockRetentionService, Mockito.times(0)).setRetention(any(), any(), any(), any());
+    Mockito.verify(mockRetentionService, Mockito.times(0)).batchApplyRetention(any(), any());
+  }
+
+  @Test
+  public void testExecuteWhenAlreadyApplied() throws Exception {
+    final EntityService<?> mockEntityService = Mockito.mock(EntityService.class);
+    final RetentionService<?> mockRetentionService = Mockito.mock(RetentionService.class);
+    final OperationContext mockContext = mock(OperationContext.class);
+    final PathMatchingResourcePatternResolver mockResolver =
+        mock(PathMatchingResourcePatternResolver.class);
+
+    when(mockEntityService.exists(mockContext, UPGRADE_ID_URN, true)).thenReturn(true);
+
+    final IngestRetentionPoliciesStep retentionStep =
+        new IngestRetentionPoliciesStep(
+            mockRetentionService, mockEntityService, true, true, "", mockResolver);
+
+    retentionStep.execute(mockContext);
+
+    Mockito.verify(mockRetentionService, Mockito.times(0)).setRetention(any(), any(), any(), any());
+    Mockito.verify(mockRetentionService, Mockito.times(0)).batchApplyRetention(any(), any());
+  }
+
+  @Test
+  public void testExecuteWithDefaultRetentionOnly() throws Exception {
+    final EntityService<?> mockEntityService = Mockito.mock(EntityService.class);
+    final RetentionService<?> mockRetentionService = Mockito.mock(RetentionService.class);
+    final OperationContext mockContext = mock(OperationContext.class);
+    final PathMatchingResourcePatternResolver mockResolver =
+        mock(PathMatchingResourcePatternResolver.class);
+    final Resource mockResource = mock(Resource.class);
+
+    when(mockEntityService.exists(mockContext, UPGRADE_ID_URN, true)).thenReturn(false);
+
+    // Mock resource loading
+    when(mockResolver.getResource("classpath:boot/retention.yaml")).thenReturn(mockResource);
+    when(mockResource.exists()).thenReturn(true);
+    when(mockResource.getInputStream())
+        .thenReturn(new ByteArrayInputStream(DEFAULT_RETENTION_YAML.getBytes()));
+
+    when(mockRetentionService.setRetention(
+            any(OperationContext.class),
+            eq("dataHubExecutionRequest"),
+            eq("dataHubExecutionRequestResult"),
+            any(DataHubRetentionConfig.class)))
+        .thenReturn(true);
+
+    when(mockRetentionService.setRetention(
+            any(OperationContext.class), eq("*"), eq("*"), any(DataHubRetentionConfig.class)))
+        .thenReturn(true);
+
+    final IngestRetentionPoliciesStep retentionStep =
+        new IngestRetentionPoliciesStep(
+            mockRetentionService, mockEntityService, true, true, "", mockResolver);
+
+    retentionStep.execute(mockContext);
+
+    Mockito.verify(mockRetentionService, Mockito.times(1))
+        .setRetention(
+            any(OperationContext.class),
+            eq("dataHubExecutionRequest"),
+            eq("dataHubExecutionRequestResult"),
+            any(DataHubRetentionConfig.class));
+
+    Mockito.verify(mockRetentionService, Mockito.times(1))
+        .setRetention(
+            any(OperationContext.class), eq("*"), eq("*"), any(DataHubRetentionConfig.class));
+
+    Mockito.verify(mockRetentionService, Mockito.times(1)).batchApplyRetention(null, null);
+
+    Mockito.verify(mockEntityService, Mockito.times(1))
+        .ingestProposal(any(), any(), any(), Mockito.eq(false));
+  }
+
+  @Test
+  public void testExecuteWithoutBatchApply() throws Exception {
+    final EntityService<?> mockEntityService = Mockito.mock(EntityService.class);
+    final RetentionService<?> mockRetentionService = Mockito.mock(RetentionService.class);
+    final OperationContext mockContext = mock(OperationContext.class);
+    final PathMatchingResourcePatternResolver mockResolver =
+        mock(PathMatchingResourcePatternResolver.class);
+    final Resource mockResource = mock(Resource.class);
+
+    when(mockEntityService.exists(mockContext, UPGRADE_ID_URN, true)).thenReturn(false);
+
+    // Mock resource loading
+    when(mockResolver.getResource("classpath:boot/retention.yaml")).thenReturn(mockResource);
+    when(mockResource.exists()).thenReturn(true);
+    when(mockResource.getInputStream())
+        .thenReturn(new ByteArrayInputStream(DEFAULT_RETENTION_YAML.getBytes()));
+
+    when(mockRetentionService.setRetention(
+            any(OperationContext.class),
+            eq("dataHubExecutionRequest"),
+            eq("dataHubExecutionRequestResult"),
+            any(DataHubRetentionConfig.class)))
+        .thenReturn(true);
+
+    when(mockRetentionService.setRetention(
+            any(OperationContext.class), eq("*"), eq("*"), any(DataHubRetentionConfig.class)))
+        .thenReturn(true);
+
+    final IngestRetentionPoliciesStep retentionStep =
+        new IngestRetentionPoliciesStep(
+            mockRetentionService, mockEntityService, true, false, "", mockResolver);
+
+    retentionStep.execute(mockContext);
+
+    Mockito.verify(mockRetentionService, Mockito.times(1))
+        .setRetention(
+            any(OperationContext.class),
+            eq("dataHubExecutionRequest"),
+            eq("dataHubExecutionRequestResult"),
+            any(DataHubRetentionConfig.class));
+
+    Mockito.verify(mockRetentionService, Mockito.times(1))
+        .setRetention(
+            any(OperationContext.class), eq("*"), eq("*"), any(DataHubRetentionConfig.class));
+
+    Mockito.verify(mockRetentionService, Mockito.times(0)).batchApplyRetention(any(), any());
+  }
+
+  @Test
+  public void testExecuteWithPluginPath() throws Exception {
+    final EntityService<?> mockEntityService = Mockito.mock(EntityService.class);
+    final RetentionService<?> mockRetentionService = Mockito.mock(RetentionService.class);
+    final OperationContext mockContext = mock(OperationContext.class);
+    final PathMatchingResourcePatternResolver mockResolver =
+        mock(PathMatchingResourcePatternResolver.class);
+    final Resource mockDefaultResource = mock(Resource.class);
+    final Resource mockPluginResource = mock(Resource.class);
+
+    when(mockEntityService.exists(mockContext, UPGRADE_ID_URN, true)).thenReturn(false);
+
+    // Mock resource loading
+    when(mockResolver.getResource("classpath:boot/retention.yaml")).thenReturn(mockDefaultResource);
+    when(mockResolver.getResources("file:/test/plugins/**/*.{yaml,yml}"))
+        .thenReturn(new Resource[] {mockPluginResource});
+
+    // Mock default resource loading
+    when(mockDefaultResource.exists()).thenReturn(true);
+    when(mockDefaultResource.getInputStream())
+        .thenReturn(new ByteArrayInputStream(DEFAULT_RETENTION_YAML.getBytes()));
+
+    // Mock plugin resource loading
+    String pluginYaml =
+        "- entity: \"container\"\n"
+            + "  aspect: \"containerProperties\"\n"
+            + "  config:\n"
+            + "    retention:\n"
+            + "      version:\n"
+            + "        maxVersions: 5\n";
+
+    when(mockPluginResource.exists()).thenReturn(true);
+    when(mockPluginResource.getInputStream())
+        .thenReturn(new ByteArrayInputStream(pluginYaml.getBytes()));
+
+    when(mockRetentionService.setRetention(
+            any(OperationContext.class),
+            eq("dataHubExecutionRequest"),
+            eq("dataHubExecutionRequestResult"),
+            any(DataHubRetentionConfig.class)))
+        .thenReturn(true);
+
+    when(mockRetentionService.setRetention(
+            any(OperationContext.class), eq("*"), eq("*"), any(DataHubRetentionConfig.class)))
+        .thenReturn(true);
+
+    when(mockRetentionService.setRetention(
+            any(OperationContext.class),
+            eq("container"),
+            eq("containerProperties"),
+            any(DataHubRetentionConfig.class)))
+        .thenReturn(true);
+
+    final IngestRetentionPoliciesStep retentionStep =
+        new IngestRetentionPoliciesStep(
+            mockRetentionService, mockEntityService, true, true, "/test/plugins", mockResolver);
+
+    retentionStep.execute(mockContext);
+
+    // Verify retention operations occur for both default and plugin configs
+    Mockito.verify(mockRetentionService, Mockito.times(1))
+        .setRetention(
+            any(OperationContext.class),
+            eq("dataHubExecutionRequest"),
+            eq("dataHubExecutionRequestResult"),
+            any(DataHubRetentionConfig.class));
+    Mockito.verify(mockRetentionService, Mockito.times(1))
+        .setRetention(
+            any(OperationContext.class), eq("*"), eq("*"), any(DataHubRetentionConfig.class));
+    Mockito.verify(mockRetentionService, Mockito.times(1))
+        .setRetention(
+            any(OperationContext.class),
+            eq("container"),
+            eq("containerProperties"),
+            any(DataHubRetentionConfig.class));
+
+    Mockito.verify(mockRetentionService, Mockito.times(1)).batchApplyRetention(null, null);
+    Mockito.verify(mockEntityService, Mockito.times(1))
+        .ingestProposal(any(), any(), any(), Mockito.eq(false));
+  }
+}
diff --git a/metadata-service/factories/src/test/java/com/linkedin/r2/transport/http/server/RAPJakartaServletTest.java b/metadata-service/factories/src/test/java/com/linkedin/r2/transport/http/server/RAPJakartaServletTest.java
new file mode 100644
index 00000000000000..09b582b7c7c179
--- /dev/null
+++ b/metadata-service/factories/src/test/java/com/linkedin/r2/transport/http/server/RAPJakartaServletTest.java
@@ -0,0 +1,303 @@
+package com.linkedin.r2.transport.http.server;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertNotNull;
+import static org.testng.Assert.assertTrue;
+
+import com.linkedin.r2.filter.R2Constants;
+import com.linkedin.r2.message.RequestContext;
+import com.linkedin.r2.message.rest.RestRequest;
+import com.linkedin.r2.message.rest.RestResponse;
+import com.linkedin.r2.message.rest.RestStatus;
+import com.linkedin.r2.transport.common.bridge.common.TransportCallback;
+import com.linkedin.r2.transport.common.bridge.common.TransportResponse;
+import com.linkedin.r2.transport.common.bridge.common.TransportResponseImpl;
+import com.linkedin.r2.transport.common.bridge.server.TransportDispatcher;
+import jakarta.servlet.ServletException;
+import jakarta.servlet.ServletOutputStream;
+import jakarta.servlet.http.HttpServletRequest;
+import jakarta.servlet.http.HttpServletResponse;
+import java.io.ByteArrayInputStream;
+import java.io.IOException;
+import java.security.cert.X509Certificate;
+import java.util.Collections;
+import org.mockito.ArgumentCaptor;
+import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.DataProvider;
+import org.testng.annotations.Test;
+
+public class RAPJakartaServletTest {
+  private static final int TIMEOUT_SECONDS = 5;
+  private static final String REMOTE_ADDRESS = "192.168.1.1";
+  private static final int REMOTE_PORT = 12345;
+
+  private HttpDispatcher mockHttpDispatcher;
+  private TransportDispatcher mockTransportDispatcher;
+  private HttpServletRequest mockRequest;
+  private HttpServletResponse mockResponse;
+  private ServletOutputStream mockOutputStream;
+
+  @BeforeMethod
+  public void setUp() throws IOException {
+    mockHttpDispatcher = mock(HttpDispatcher.class);
+    mockTransportDispatcher = mock(TransportDispatcher.class);
+    mockRequest = mock(HttpServletRequest.class);
+    mockResponse = mock(HttpServletResponse.class);
+    mockOutputStream = mock(ServletOutputStream.class);
+
+    // Basic request setup
+    when(mockRequest.getMethod()).thenReturn("GET");
+    when(mockRequest.getRequestURI()).thenReturn("/test");
+    when(mockRequest.getContextPath()).thenReturn("");
+    when(mockRequest.getServletPath()).thenReturn("");
+    when(mockRequest.getPathInfo()).thenReturn("/test");
+
+    // Protocol and server info
+    when(mockRequest.getProtocol()).thenReturn("HTTP/1.1");
+    when(mockRequest.getScheme()).thenReturn("http");
+    when(mockRequest.getServerName()).thenReturn("localhost");
+    when(mockRequest.getServerPort()).thenReturn(8080);
+
+    // Remote client info
+    when(mockRequest.getRemoteAddr()).thenReturn(REMOTE_ADDRESS);
+    when(mockRequest.getRemotePort()).thenReturn(REMOTE_PORT);
+
+    // Security defaults (non-secure)
+    when(mockRequest.isSecure()).thenReturn(false);
+    when(mockRequest.getAttribute("javax.servlet.request.X509Certificate")).thenReturn(null);
+    when(mockRequest.getAttribute("javax.servlet.request.cipher_suite")).thenReturn(null);
+
+    // Headers and content
+    when(mockRequest.getHeaderNames()).thenReturn(Collections.enumeration(Collections.emptyList()));
+    when(mockRequest.getInputStream())
+        .thenReturn(new TestServletInputStream(new ByteArrayInputStream(new byte[0])));
+    when(mockResponse.getOutputStream()).thenReturn(mockOutputStream);
+  }
+
+  @Test
+  public void testConstructorWithHttpDispatcher() {
+    RAPJakartaServlet servlet = new RAPJakartaServlet(mockHttpDispatcher, TIMEOUT_SECONDS);
+    assertEquals(servlet.getDispatcher(), mockHttpDispatcher);
+  }
+
+  @Test
+  public void testConstructorWithTransportDispatcher() {
+    RAPJakartaServlet servlet = new RAPJakartaServlet(mockTransportDispatcher, TIMEOUT_SECONDS);
+    assertNotNull(servlet.getDispatcher());
+  }
+
+  @Test
+  public void testSuccessfulRequest() throws ServletException, IOException {
+    // Setup
+    RAPJakartaServlet servlet = new RAPJakartaServlet(mockHttpDispatcher, TIMEOUT_SECONDS);
+    RestResponse mockRestResponse = RestStatus.responseForStatus(200, "OK");
+    TransportResponse<RestResponse> transportResponse =
+        TransportResponseImpl.success(mockRestResponse);
+
+    doAnswer(
+            invocation -> {
+              TransportCallback<RestResponse> callback = invocation.getArgument(2);
+              callback.onResponse(transportResponse);
+              return null;
+            })
+        .when(mockHttpDispatcher)
+        .handleRequest(
+            any(RestRequest.class), any(RequestContext.class), any(TransportCallback.class));
+
+    // Execute
+    servlet.service(mockRequest, mockResponse);
+
+    // Verify status code
+    verify(mockResponse).setStatus(200);
+
+    // Capture and verify the output stream write
+    ArgumentCaptor<byte[]> dataCaptor = ArgumentCaptor.forClass(byte[].class);
+    ArgumentCaptor<Integer> offsetCaptor = ArgumentCaptor.forClass(Integer.class);
+    ArgumentCaptor<Integer> lengthCaptor = ArgumentCaptor.forClass(Integer.class);
+
+    verify(mockOutputStream)
+        .write(dataCaptor.capture(), offsetCaptor.capture(), lengthCaptor.capture());
+
+    // Convert captured bytes to string for readable verification
+    byte[] capturedData = dataCaptor.getValue();
+    int offset = offsetCaptor.getValue();
+    int length = lengthCaptor.getValue();
+
+    String responseBody = new String(capturedData, offset, length);
+    System.out.println("Actual success response body: " + responseBody); // For debugging
+
+    assertEquals(responseBody, "OK", "Expected response body to be 'OK' but got: " + responseBody);
+  }
+
+  @Test
+  public void testSecureRequest() throws ServletException, IOException {
+    // Setup secure request
+    X509Certificate[] certs = new X509Certificate[1];
+    certs[0] = mock(X509Certificate.class);
+    String cipherSuite = "TLS_AES_256_GCM_SHA384";
+
+    when(mockRequest.isSecure()).thenReturn(true);
+    when(mockRequest.getAttribute("javax.servlet.request.X509Certificate")).thenReturn(certs);
+    when(mockRequest.getAttribute("javax.servlet.request.cipher_suite")).thenReturn(cipherSuite);
+
+    RAPJakartaServlet servlet = new RAPJakartaServlet(mockHttpDispatcher, TIMEOUT_SECONDS);
+
+    ArgumentCaptor<RequestContext> contextCaptor = ArgumentCaptor.forClass(RequestContext.class);
+
+    // Execute
+    servlet.service(mockRequest, mockResponse);
+
+    // Verify
+    verify(mockHttpDispatcher)
+        .handleRequest(
+            any(RestRequest.class), contextCaptor.capture(), any(TransportCallback.class));
+
+    RequestContext capturedContext = contextCaptor.getValue();
+    assertTrue((Boolean) capturedContext.getLocalAttr(R2Constants.IS_SECURE));
+    assertEquals(capturedContext.getLocalAttr(R2Constants.CIPHER_SUITE), cipherSuite);
+    assertEquals(capturedContext.getLocalAttr(R2Constants.CLIENT_CERT), certs[0]);
+  }
+
+  @Test
+  public void testBadRequest() throws ServletException, IOException {
+    // Setup with invalid URI components that will cause URISyntaxException
+    when(mockRequest.getRequestURI()).thenReturn("/test");
+    when(mockRequest.getQueryString()).thenReturn("invalid=%%test"); // Invalid percent encoding
+    when(mockRequest.getContextPath()).thenReturn("");
+    when(mockRequest.getServletPath()).thenReturn("");
+
+    RAPJakartaServlet servlet = new RAPJakartaServlet(mockHttpDispatcher, TIMEOUT_SECONDS);
+
+    // Execute
+    servlet.service(mockRequest, mockResponse);
+
+    // Verify
+    verify(mockResponse).setStatus(400);
+
+    // Capture all write invocations to the output stream
+    ArgumentCaptor<byte[]> dataCaptor = ArgumentCaptor.forClass(byte[].class);
+    ArgumentCaptor<Integer> offsetCaptor = ArgumentCaptor.forClass(Integer.class);
+    ArgumentCaptor<Integer> lengthCaptor = ArgumentCaptor.forClass(Integer.class);
+
+    verify(mockOutputStream)
+        .write(dataCaptor.capture(), offsetCaptor.capture(), lengthCaptor.capture());
+
+    // Convert captured bytes to string for readable verification
+    byte[] capturedData = dataCaptor.getValue();
+    int offset = offsetCaptor.getValue();
+    int length = lengthCaptor.getValue();
+
+    String responseBody = new String(capturedData, offset, length);
+    System.out.println("Actual response body: " + responseBody); // For debugging
+
+    assertTrue(
+        responseBody.contains("Invalid URI"),
+        "Expected response to contain 'Invalid URI' but got: " + responseBody);
+  }
+
+  @Test
+  public void testRequestTimeout() throws ServletException, IOException {
+    // Setup with very short timeout
+    RAPJakartaServlet servlet = new RAPJakartaServlet(mockHttpDispatcher, 1);
+
+    doAnswer(
+            invocation -> {
+              Thread.sleep(2000);
+              return null;
+            })
+        .when(mockHttpDispatcher)
+        .handleRequest(
+            any(RestRequest.class), any(RequestContext.class), any(TransportCallback.class));
+
+    // Execute
+    servlet.service(mockRequest, mockResponse);
+
+    // Verify status code
+    verify(mockResponse).setStatus(500);
+
+    // Capture and verify the output stream write
+    ArgumentCaptor<byte[]> dataCaptor = ArgumentCaptor.forClass(byte[].class);
+    ArgumentCaptor<Integer> offsetCaptor = ArgumentCaptor.forClass(Integer.class);
+    ArgumentCaptor<Integer> lengthCaptor = ArgumentCaptor.forClass(Integer.class);
+
+    verify(mockOutputStream)
+        .write(dataCaptor.capture(), offsetCaptor.capture(), lengthCaptor.capture());
+
+    // Convert captured bytes to string for readable verification
+    byte[] capturedData = dataCaptor.getValue();
+    int offset = offsetCaptor.getValue();
+    int length = lengthCaptor.getValue();
+
+    String responseBody = new String(capturedData, offset, length);
+    System.out.println("Actual timeout response body: " + responseBody); // For debugging
+
+    assertTrue(
+        responseBody.contains("Server timeout after 1 seconds"),
+        "Expected response to contain timeout message but got: " + responseBody);
+  }
+
+  @DataProvider(name = "headerTestCases")
+  public Object[][] headerTestCases() {
+    return new Object[][] {
+      {"Content-Type", "application/json"},
+      {"Accept", "text/plain"},
+      {"X-Custom-Header", "custom-value"}
+    };
+  }
+
+  @Test(dataProvider = "headerTestCases")
+  public void testHeaderHandling(String headerName, String headerValue)
+      throws ServletException, IOException {
+    // Setup
+    RAPJakartaServlet servlet = new RAPJakartaServlet(mockHttpDispatcher, TIMEOUT_SECONDS);
+    when(mockRequest.getHeaderNames())
+        .thenReturn(Collections.enumeration(Collections.singletonList(headerName)));
+    when(mockRequest.getHeaders(headerName))
+        .thenReturn(Collections.enumeration(Collections.singletonList(headerValue)));
+
+    ArgumentCaptor<RestRequest> requestCaptor = ArgumentCaptor.forClass(RestRequest.class);
+
+    // Execute
+    servlet.service(mockRequest, mockResponse);
+
+    // Verify
+    verify(mockHttpDispatcher)
+        .handleRequest(
+            requestCaptor.capture(), any(RequestContext.class), any(TransportCallback.class));
+    RestRequest capturedRequest = requestCaptor.getValue();
+    assertEquals(capturedRequest.getHeader(headerName), headerValue);
+  }
+
+  private static class TestServletInputStream extends jakarta.servlet.ServletInputStream {
+    private final ByteArrayInputStream basis;
+
+    public TestServletInputStream(ByteArrayInputStream basis) {
+      this.basis = basis;
+    }
+
+    @Override
+    public int read() throws IOException {
+      return basis.read();
+    }
+
+    @Override
+    public boolean isFinished() {
+      return basis.available() == 0;
+    }
+
+    @Override
+    public boolean isReady() {
+      return true;
+    }
+
+    @Override
+    public void setReadListener(jakarta.servlet.ReadListener readListener) {
+      // Not implemented for test
+    }
+  }
+}
diff --git a/metadata-service/graphql-servlet-impl/src/main/java/com/datahub/graphql/GraphQLController.java b/metadata-service/graphql-servlet-impl/src/main/java/com/datahub/graphql/GraphQLController.java
index 2f66b30f55844c..6e02db6eb65239 100644
--- a/metadata-service/graphql-servlet-impl/src/main/java/com/datahub/graphql/GraphQLController.java
+++ b/metadata-service/graphql-servlet-impl/src/main/java/com/datahub/graphql/GraphQLController.java
@@ -37,10 +37,12 @@
 import org.springframework.web.HttpRequestMethodNotSupportedException;
 import org.springframework.web.bind.annotation.GetMapping;
 import org.springframework.web.bind.annotation.PostMapping;
+import org.springframework.web.bind.annotation.RequestMapping;
 import org.springframework.web.bind.annotation.RestController;
 
 @Slf4j
 @RestController
+@RequestMapping("/api")
 public class GraphQLController {
 
   public GraphQLController() {
diff --git a/metadata-service/graphql-servlet-impl/src/main/java/com/datahub/graphql/GraphiQLController.java b/metadata-service/graphql-servlet-impl/src/main/java/com/datahub/graphql/GraphiQLController.java
index 61f2720c6cfca4..0461af4db11cc8 100644
--- a/metadata-service/graphql-servlet-impl/src/main/java/com/datahub/graphql/GraphiQLController.java
+++ b/metadata-service/graphql-servlet-impl/src/main/java/com/datahub/graphql/GraphiQLController.java
@@ -32,7 +32,7 @@ public GraphiQLController() {
     }
   }
 
-  @GetMapping(value = "/graphiql", produces = MediaType.TEXT_HTML_VALUE)
+  @GetMapping(value = "/api/graphiql", produces = MediaType.TEXT_HTML_VALUE)
   @ResponseBody
   CompletableFuture<String> graphiQL() {
     return GraphQLConcurrencyUtils.supplyAsync(
diff --git a/metadata-service/openapi-analytics-servlet/src/main/resources/JavaSpring/apiController.mustache b/metadata-service/openapi-analytics-servlet/src/main/resources/JavaSpring/apiController.mustache
index ddf1c130d7be69..fbeb0ba7454a8b 100644
--- a/metadata-service/openapi-analytics-servlet/src/main/resources/JavaSpring/apiController.mustache
+++ b/metadata-service/openapi-analytics-servlet/src/main/resources/JavaSpring/apiController.mustache
@@ -62,7 +62,7 @@ import java.util.concurrent.Callable;
 {{/isJava8or11}}
 {{>generatedAnnotation}}
 @Controller
-@RequestMapping("/v2/analytics")
+@RequestMapping("/openapi/v2/analytics")
 {{#operations}}
 public class {{classname}}Controller implements {{classname}} {
 
diff --git a/metadata-service/openapi-entity-servlet/src/main/resources/JavaSpring/apiController.mustache b/metadata-service/openapi-entity-servlet/src/main/resources/JavaSpring/apiController.mustache
index fbf354ff91688f..78bb061b5d50a4 100644
--- a/metadata-service/openapi-entity-servlet/src/main/resources/JavaSpring/apiController.mustache
+++ b/metadata-service/openapi-entity-servlet/src/main/resources/JavaSpring/apiController.mustache
@@ -80,7 +80,7 @@ import java.util.concurrent.Callable;
 {{/useOas2}}
 {{#operations}}
 @RestController
-@RequestMapping("/v2/entity")
+@RequestMapping("/openapi/v2/entity")
 public class {{classname}}Controller implements {{classname}} {
 
     private static final Logger log = LoggerFactory.getLogger({{classname}}Controller.class);
diff --git a/metadata-service/openapi-entity-servlet/src/test/java/io/datahubproject/openapi/v2/delegates/EntityApiDelegateImplTest.java b/metadata-service/openapi-entity-servlet/src/test/java/io/datahubproject/openapi/v2/delegates/EntityApiDelegateImplTest.java
index 31b35b65ea1a8c..fc91b0990a333a 100644
--- a/metadata-service/openapi-entity-servlet/src/test/java/io/datahubproject/openapi/v2/delegates/EntityApiDelegateImplTest.java
+++ b/metadata-service/openapi-entity-servlet/src/test/java/io/datahubproject/openapi/v2/delegates/EntityApiDelegateImplTest.java
@@ -301,7 +301,7 @@ public void customModelTest() throws Exception {
 
     mockMvc
         .perform(
-            MockMvcRequestBuilders.post("/v2/entity/dataset")
+            MockMvcRequestBuilders.post("/openapi/v2/entity/dataset")
                 .content(body)
                 .contentType(MediaType.APPLICATION_JSON)
                 .accept(MediaType.APPLICATION_JSON))
diff --git a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/config/GlobalControllerExceptionHandler.java b/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/config/GlobalControllerExceptionHandler.java
index 01493d71643481..b78121573811f6 100644
--- a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/config/GlobalControllerExceptionHandler.java
+++ b/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/config/GlobalControllerExceptionHandler.java
@@ -18,6 +18,7 @@
 import org.springframework.http.ResponseEntity;
 import org.springframework.web.bind.annotation.ControllerAdvice;
 import org.springframework.web.bind.annotation.ExceptionHandler;
+import org.springframework.web.servlet.NoHandlerFoundException;
 import org.springframework.web.servlet.mvc.support.DefaultHandlerExceptionResolver;
 
 @Slf4j
@@ -88,4 +89,13 @@ public ResponseEntity<Map<String, String>> handleGenericException(
     return new ResponseEntity<>(
         Map.of("error", "Internal server error occurred"), HttpStatus.INTERNAL_SERVER_ERROR);
   }
+
+  @ExceptionHandler(NoHandlerFoundException.class)
+  public static ResponseEntity<Map<String, String>> handleNoHandlerFoundException(
+      NoHandlerFoundException ex, HttpServletRequest request) {
+    String message = String.format("No endpoint %s %s.", ex.getHttpMethod(), ex.getRequestURL());
+
+    log.error("No handler found for request: " + request.getRequestURI());
+    return new ResponseEntity<>(Map.of("error", message), HttpStatus.NOT_FOUND);
+  }
 }
diff --git a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/config/SpringWebConfig.java b/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/config/SpringWebConfig.java
index c4b4431e77c4ef..c6d5f1452fea91 100644
--- a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/config/SpringWebConfig.java
+++ b/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/config/SpringWebConfig.java
@@ -1,15 +1,7 @@
 package io.datahubproject.openapi.config;
 
-import static com.linkedin.metadata.Constants.INGESTION_MAX_SERIALIZED_STRING_LENGTH;
-import static com.linkedin.metadata.Constants.MAX_JACKSON_STRING_SIZE;
-
-import com.fasterxml.jackson.annotation.JsonInclude;
-import com.fasterxml.jackson.core.StreamReadConstraints;
-import com.fasterxml.jackson.databind.DeserializationFeature;
-import com.fasterxml.jackson.databind.ObjectMapper;
 import com.linkedin.gms.factory.config.ConfigurationProvider;
 import com.linkedin.metadata.models.registry.EntityRegistry;
-import io.datahubproject.openapi.converter.StringToChangeCategoryConverter;
 import io.datahubproject.openapi.v3.OpenAPIV3Generator;
 import io.swagger.v3.oas.annotations.OpenAPIDefinition;
 import io.swagger.v3.oas.annotations.info.Info;
@@ -18,29 +10,24 @@
 import io.swagger.v3.oas.models.OpenAPI;
 import java.util.Collections;
 import java.util.LinkedHashMap;
-import java.util.List;
 import java.util.Map;
 import java.util.Set;
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
-import javax.annotation.Nonnull;
 import org.springdoc.core.models.GroupedOpenApi;
-import org.springframework.beans.factory.annotation.Value;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
-import org.springframework.format.FormatterRegistry;
-import org.springframework.http.converter.ByteArrayHttpMessageConverter;
-import org.springframework.http.converter.FormHttpMessageConverter;
-import org.springframework.http.converter.HttpMessageConverter;
-import org.springframework.http.converter.StringHttpMessageConverter;
-import org.springframework.http.converter.json.MappingJackson2HttpMessageConverter;
-import org.springframework.web.servlet.config.annotation.AsyncSupportConfigurer;
+import org.springframework.core.annotation.Order;
+import org.springframework.web.servlet.config.annotation.EnableWebMvc;
+import org.springframework.web.servlet.config.annotation.ResourceHandlerRegistry;
 import org.springframework.web.servlet.config.annotation.WebMvcConfigurer;
 
+@EnableWebMvc
 @OpenAPIDefinition(
     info = @Info(title = "DataHub OpenAPI", version = "2.0.0"),
     servers = {@Server(url = "/openapi/", description = "Default Server URL")})
+@Order(2)
 @Configuration
 public class SpringWebConfig implements WebMvcConfigurer {
   private static final Set<String> OPERATIONS_PACKAGES =
@@ -52,35 +39,6 @@ public class SpringWebConfig implements WebMvcConfigurer {
   private static final Set<String> OPENLINEAGE_PACKAGES =
       Set.of("io.datahubproject.openapi.openlineage");
 
-  @Value("${datahub.gms.async.request-timeout-ms}")
-  private long asyncTimeoutMilliseconds;
-
-  @Override
-  public void configureMessageConverters(List<HttpMessageConverter<?>> messageConverters) {
-    messageConverters.add(new StringHttpMessageConverter());
-    messageConverters.add(new ByteArrayHttpMessageConverter());
-    messageConverters.add(new FormHttpMessageConverter());
-
-    ObjectMapper objectMapper = new ObjectMapper();
-    int maxSize =
-        Integer.parseInt(
-            System.getenv()
-                .getOrDefault(INGESTION_MAX_SERIALIZED_STRING_LENGTH, MAX_JACKSON_STRING_SIZE));
-    objectMapper
-        .getFactory()
-        .setStreamReadConstraints(StreamReadConstraints.builder().maxStringLength(maxSize).build());
-    objectMapper.setSerializationInclusion(JsonInclude.Include.NON_NULL);
-    objectMapper.configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false);
-    MappingJackson2HttpMessageConverter jsonConverter =
-        new MappingJackson2HttpMessageConverter(objectMapper);
-    messageConverters.add(jsonConverter);
-  }
-
-  @Override
-  public void addFormatters(FormatterRegistry registry) {
-    registry.addConverter(new StringToChangeCategoryConverter());
-  }
-
   @Bean
   public GroupedOpenApi v3OpenApiGroup(
       final EntityRegistry entityRegistry, final ConfigurationProvider configurationProvider) {
@@ -153,6 +111,13 @@ public GroupedOpenApi openlineageOpenApiGroup() {
         .build();
   }
 
+  @Override
+  public void addResourceHandlers(ResourceHandlerRegistry registry) {
+    registry
+        .addResourceHandler("/swagger-ui/**")
+        .addResourceLocations("classpath:/META-INF/resources/webjars/swagger-ui/");
+  }
+
   /** Concatenates two maps. */
   private <K, V> Map<K, V> concat(Supplier<Map<K, V>> a, Supplier<Map<K, V>> b) {
     return a.get() == null
@@ -167,10 +132,4 @@ private <K, V> Map<K, V> concat(Supplier<Map<K, V>> a, Supplier<Map<K, V>> b) {
                         (v1, v2) -> v2,
                         LinkedHashMap::new));
   }
-
-  @Override
-  public void configureAsyncSupport(@Nonnull AsyncSupportConfigurer configurer) {
-    WebMvcConfigurer.super.configureAsyncSupport(configurer);
-    configurer.setDefaultTimeout(asyncTimeoutMilliseconds);
-  }
 }
diff --git a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/operations/elastic/OperationsController.java b/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/operations/elastic/OperationsController.java
index 6151b866e5208b..64333009dda7a9 100644
--- a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/operations/elastic/OperationsController.java
+++ b/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/operations/elastic/OperationsController.java
@@ -60,7 +60,7 @@
 import org.springframework.web.bind.annotation.RestController;
 
 @RestController
-@RequestMapping("/operations/elasticSearch")
+@RequestMapping("/openapi/operations/elasticSearch")
 @Slf4j
 @Tag(
     name = "ElasticSearchOperations",
diff --git a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/operations/test/IdController.java b/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/operations/test/IdController.java
index 99d3879ab9a320..9be9526afac340 100644
--- a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/operations/test/IdController.java
+++ b/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/operations/test/IdController.java
@@ -19,7 +19,7 @@
 import org.springframework.web.bind.annotation.RestController;
 
 @RestController
-@RequestMapping("/operations/identity")
+@RequestMapping("/openapi/operations/identity")
 @Slf4j
 @Tag(name = "Identity", description = "An API for checking identity")
 public class IdController {
diff --git a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/operations/throttle/ThrottleController.java b/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/operations/throttle/ThrottleController.java
index 3c44d94428f428..5db2183d3acd62 100644
--- a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/operations/throttle/ThrottleController.java
+++ b/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/operations/throttle/ThrottleController.java
@@ -28,7 +28,7 @@
 import org.springframework.web.bind.annotation.RestController;
 
 @RestController
-@RequestMapping("/operations/throttle")
+@RequestMapping("/openapi/operations/throttle")
 @Slf4j
 @Tag(name = "GMS Throttle Control", description = "An API for GMS throttle control.")
 public class ThrottleController {
diff --git a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v1/entities/EntitiesController.java b/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v1/entities/EntitiesController.java
index 03050868efdcab..dd359cbc464319 100644
--- a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v1/entities/EntitiesController.java
+++ b/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v1/entities/EntitiesController.java
@@ -61,7 +61,7 @@
 */
 @Deprecated
 @RestController
-@RequestMapping("/entities/v1")
+@RequestMapping("/openapi/entities/v1")
 @Slf4j
 @Tag(
     name = "Entities",
diff --git a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v1/relationships/RelationshipsController.java b/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v1/relationships/RelationshipsController.java
index b2b7eb557ca32a..f778bf54aaeac7 100644
--- a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v1/relationships/RelationshipsController.java
+++ b/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v1/relationships/RelationshipsController.java
@@ -50,7 +50,7 @@
 @Deprecated
 @RestController
 @RequiredArgsConstructor
-@RequestMapping("/relationships/v1")
+@RequestMapping("/openapi/relationships/v1")
 @Slf4j
 @Tag(name = "Relationships", description = "APIs for accessing relationships of entities")
 public class RelationshipsController {
diff --git a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v1/timeline/TimelineControllerV1.java b/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v1/timeline/TimelineControllerV1.java
index 30cdb632d54773..47b558f6a9b3a3 100644
--- a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v1/timeline/TimelineControllerV1.java
+++ b/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v1/timeline/TimelineControllerV1.java
@@ -38,7 +38,7 @@
 @Deprecated
 @RestController
 @AllArgsConstructor
-@RequestMapping("/timeline/v1")
+@RequestMapping("/openapi/timeline/v1")
 @Tag(
     name = "Timeline",
     description =
diff --git a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v2/controller/EntityController.java b/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v2/controller/EntityController.java
index b1c5709ef01470..bafa8c43f60f55 100644
--- a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v2/controller/EntityController.java
+++ b/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v2/controller/EntityController.java
@@ -73,7 +73,7 @@
 
 @RestController
 @RequiredArgsConstructor
-@RequestMapping("/v2/entity")
+@RequestMapping("/openapi/v2/entity")
 @Slf4j
 public class EntityController
     extends GenericEntitiesController<
diff --git a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v2/controller/PlatformEntitiesController.java b/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v2/controller/PlatformEntitiesController.java
index 6c99d972dde03f..bfd1f952f95f74 100644
--- a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v2/controller/PlatformEntitiesController.java
+++ b/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v2/controller/PlatformEntitiesController.java
@@ -39,7 +39,7 @@
 
 @RestController
 @RequiredArgsConstructor
-@RequestMapping("/v2/platform/entities/v1")
+@RequestMapping("/openapiv2/platform/entities/v1")
 @Slf4j
 @Tag(
     name = "Platform Entities",
diff --git a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v2/controller/RelationshipController.java b/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v2/controller/RelationshipController.java
index a0412676b5cbce..1e431bd291ca17 100644
--- a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v2/controller/RelationshipController.java
+++ b/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v2/controller/RelationshipController.java
@@ -9,7 +9,7 @@
 
 @RestController
 @RequiredArgsConstructor
-@RequestMapping("/v2/relationship")
+@RequestMapping("/openapi/v2/relationship")
 @Slf4j
 @Tag(
     name = "Generic Relationships",
diff --git a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v2/controller/TimelineControllerV2.java b/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v2/controller/TimelineControllerV2.java
index f3d0d5188b1e7e..0245ae289f07ee 100644
--- a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v2/controller/TimelineControllerV2.java
+++ b/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v2/controller/TimelineControllerV2.java
@@ -34,7 +34,7 @@
 
 @RestController
 @AllArgsConstructor
-@RequestMapping("/v2/timeline/v1")
+@RequestMapping("/openapi/v2/timeline/v1")
 @Tag(
     name = "Timeline",
     description =
diff --git a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v2/controller/TimeseriesController.java b/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v2/controller/TimeseriesController.java
index 3b896dc5000822..212084f25ddd5d 100644
--- a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v2/controller/TimeseriesController.java
+++ b/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v2/controller/TimeseriesController.java
@@ -38,7 +38,7 @@
 import org.springframework.web.bind.annotation.RestController;
 
 @RestController
-@RequestMapping("/v2/timeseries")
+@RequestMapping("/openapi/v2/timeseries")
 @Slf4j
 @Tag(
     name = "Generic Timeseries Aspects",
diff --git a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v3/controller/EntityController.java b/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v3/controller/EntityController.java
index a4583082d57c7f..573feec64d2283 100644
--- a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v3/controller/EntityController.java
+++ b/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v3/controller/EntityController.java
@@ -93,7 +93,7 @@
 
 @RestController("EntityControllerV3")
 @RequiredArgsConstructor
-@RequestMapping("/v3/entity")
+@RequestMapping("/openapi/v3/entity")
 @Slf4j
 @Hidden
 public class EntityController
diff --git a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v3/controller/RelationshipController.java b/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v3/controller/RelationshipController.java
index 8f317e86227239..3b7b897b55d1c5 100644
--- a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v3/controller/RelationshipController.java
+++ b/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v3/controller/RelationshipController.java
@@ -9,7 +9,7 @@
 
 @RestController("RelationshipControllerV3")
 @RequiredArgsConstructor
-@RequestMapping("/v3/relationship")
+@RequestMapping("/openapi/v3/relationship")
 @Slf4j
 @Tag(
     name = "Generic Relationships",
diff --git a/metadata-service/openapi-servlet/src/test/java/io/datahubproject/openapi/config/SpringWebConfigTest.java b/metadata-service/openapi-servlet/src/test/java/io/datahubproject/openapi/config/SpringWebConfigTest.java
new file mode 100644
index 00000000000000..c449ed20e47ba4
--- /dev/null
+++ b/metadata-service/openapi-servlet/src/test/java/io/datahubproject/openapi/config/SpringWebConfigTest.java
@@ -0,0 +1,77 @@
+package io.datahubproject.openapi.config;
+
+import static org.mockito.Mockito.mock;
+import static org.testng.AssertJUnit.assertEquals;
+
+import com.linkedin.gms.factory.config.ConfigurationProvider;
+import com.linkedin.metadata.models.registry.EntityRegistry;
+import com.linkedin.metadata.spring.YamlPropertySourceFactory;
+import io.datahubproject.openapi.v3.OpenAPIV3Generator;
+import io.swagger.v3.oas.models.Components;
+import io.swagger.v3.oas.models.OpenAPI;
+import io.swagger.v3.oas.models.Paths;
+import io.swagger.v3.oas.models.media.Schema;
+import java.util.HashMap;
+import java.util.Map;
+import org.mockito.MockedStatic;
+import org.mockito.Mockito;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.context.annotation.Import;
+import org.springframework.context.annotation.PropertySource;
+import org.springframework.test.context.testng.AbstractTestNGSpringContextTests;
+import org.testng.annotations.Test;
+
+@Import(ConfigurationProvider.class)
+@PropertySource(value = "classpath:/application.yaml", factory = YamlPropertySourceFactory.class)
+public class SpringWebConfigTest extends AbstractTestNGSpringContextTests {
+  @Autowired private ConfigurationProvider configurationProvider;
+
+  @Test
+  void testComponentsMergeWithDuplicateKeys() {
+    // Setup
+    SpringWebConfig config = new SpringWebConfig();
+    EntityRegistry entityRegistry = mock(EntityRegistry.class);
+
+    // Create test schemas with duplicate keys
+    Map<String, Schema> schemas1 = new HashMap<>();
+    schemas1.put("TestSchema", new Schema().type("string").description("First schema"));
+
+    Map<String, Schema> schemas2 = new HashMap<>();
+    schemas2.put("TestSchema", new Schema().type("object").description("Second schema"));
+
+    // Create OpenAPI objects with proper initialization
+    OpenAPI openApi1 =
+        new OpenAPI().components(new Components().schemas(schemas1)).paths(new Paths());
+
+    OpenAPI openApi2 =
+        new OpenAPI().components(new Components().schemas(schemas2)).paths(new Paths());
+
+    // Mock OpenAPIV3Generator
+    try (MockedStatic<OpenAPIV3Generator> mockedGenerator =
+        Mockito.mockStatic(OpenAPIV3Generator.class)) {
+      mockedGenerator
+          .when(
+              () ->
+                  OpenAPIV3Generator.generateOpenApiSpec(
+                      Mockito.any(EntityRegistry.class), Mockito.any(ConfigurationProvider.class)))
+          .thenReturn(openApi2);
+
+      // Get the GroupedOpenApi
+      var groupedApi = config.v3OpenApiGroup(entityRegistry, configurationProvider);
+
+      // Execute the customizer
+      groupedApi.getOpenApiCustomizers().get(0).customise(openApi1);
+
+      // Verify the merged components
+      Map<String, Schema> mergedSchemas = openApi1.getComponents().getSchemas();
+
+      // Assert that we have the expected number of schemas
+      assertEquals(1, mergedSchemas.size());
+
+      // Assert that the duplicate key contains the second schema (v2 value)
+      Schema resultSchema = mergedSchemas.get("TestSchema");
+      assertEquals("object", resultSchema.getType());
+      assertEquals("Second schema", resultSchema.getDescription());
+    }
+  }
+}
diff --git a/metadata-service/openapi-servlet/src/test/java/io/datahubproject/openapi/v3/controller/EntityControllerTest.java b/metadata-service/openapi-servlet/src/test/java/io/datahubproject/openapi/v3/controller/EntityControllerTest.java
index e82ab50a0defeb..5080f6a12cdd28 100644
--- a/metadata-service/openapi-servlet/src/test/java/io/datahubproject/openapi/v3/controller/EntityControllerTest.java
+++ b/metadata-service/openapi-servlet/src/test/java/io/datahubproject/openapi/v3/controller/EntityControllerTest.java
@@ -178,7 +178,7 @@ public void testSearchOrderPreserved() throws Exception {
     // test ASCENDING
     mockMvc
         .perform(
-            MockMvcRequestBuilders.get("/v3/entity/dataset")
+            MockMvcRequestBuilders.get("/openapi/v3/entity/dataset")
                 .param("sortOrder", "ASCENDING")
                 .accept(MediaType.APPLICATION_JSON))
         .andExpect(status().is2xxSuccessful())
@@ -192,7 +192,7 @@ public void testSearchOrderPreserved() throws Exception {
     // test DESCENDING
     mockMvc
         .perform(
-            MockMvcRequestBuilders.get("/v3/entity/dataset")
+            MockMvcRequestBuilders.get("/openapi/v3/entity/dataset")
                 .accept(MediaType.APPLICATION_JSON)
                 .param("sortOrder", "DESCENDING"))
         .andExpect(status().is2xxSuccessful())
@@ -211,14 +211,14 @@ public void testDeleteEntity() throws Exception {
     // test delete entity
     mockMvc
         .perform(
-            MockMvcRequestBuilders.delete(String.format("/v3/entity/dataset/%s", TEST_URN))
+            MockMvcRequestBuilders.delete(String.format("/openapi/v3/entity/dataset/%s", TEST_URN))
                 .accept(MediaType.APPLICATION_JSON))
         .andExpect(status().is2xxSuccessful());
 
     // test delete entity by aspect key
     mockMvc
         .perform(
-            MockMvcRequestBuilders.delete(String.format("/v3/entity/dataset/%s", TEST_URN))
+            MockMvcRequestBuilders.delete(String.format("/openapi/v3/entity/dataset/%s", TEST_URN))
                 .param("aspects", "datasetKey")
                 .accept(MediaType.APPLICATION_JSON))
         .andExpect(status().is2xxSuccessful());
@@ -229,7 +229,7 @@ public void testDeleteEntity() throws Exception {
     reset(mockEntityService);
     mockMvc
         .perform(
-            MockMvcRequestBuilders.delete(String.format("/v3/entity/dataset/%s", TEST_URN))
+            MockMvcRequestBuilders.delete(String.format("/openapi/v3/entity/dataset/%s", TEST_URN))
                 .param("aspects", "status")
                 .accept(MediaType.APPLICATION_JSON))
         .andExpect(status().is2xxSuccessful());
@@ -240,7 +240,7 @@ public void testDeleteEntity() throws Exception {
     reset(mockEntityService);
     mockMvc
         .perform(
-            MockMvcRequestBuilders.delete(String.format("/v3/entity/dataset/%s", TEST_URN))
+            MockMvcRequestBuilders.delete(String.format("/openapi/v3/entity/dataset/%s", TEST_URN))
                 .param("clear", "true")
                 .accept(MediaType.APPLICATION_JSON))
         .andExpect(status().is2xxSuccessful());
@@ -374,7 +374,7 @@ public void testTimeseriesAspect() throws Exception {
     // test timeseries latest aspect
     mockMvc
         .perform(
-            MockMvcRequestBuilders.get("/v3/entity/dataset/{urn}/datasetprofile", TEST_URN)
+            MockMvcRequestBuilders.get("/openapi/v3/entity/dataset/{urn}/datasetprofile", TEST_URN)
                 .accept(MediaType.APPLICATION_JSON))
         .andExpect(status().is2xxSuccessful())
         .andExpect(MockMvcResultMatchers.jsonPath("$.value.rowCount").value(10))
@@ -384,7 +384,7 @@ public void testTimeseriesAspect() throws Exception {
     // test oldd aspect
     mockMvc
         .perform(
-            MockMvcRequestBuilders.get("/v3/entity/dataset/{urn}/datasetprofile", TEST_URN)
+            MockMvcRequestBuilders.get("/openapi/v3/entity/dataset/{urn}/datasetprofile", TEST_URN)
                 .param("version", "150")
                 .accept(MediaType.APPLICATION_JSON))
         .andExpect(status().is2xxSuccessful())
@@ -471,7 +471,7 @@ public void testGetEntityBatchWithMultipleEntities() throws Exception {
 
     mockMvc
         .perform(
-            MockMvcRequestBuilders.post("/v3/entity/dataset/batchGet")
+            MockMvcRequestBuilders.post("/openapi/v3/entity/dataset/batchGet")
                 .content(requestBody)
                 .contentType(MediaType.APPLICATION_JSON)
                 .accept(MediaType.APPLICATION_JSON))
@@ -486,7 +486,7 @@ public void testGetEntityBatchWithInvalidUrn() throws Exception {
 
     mockMvc
         .perform(
-            MockMvcRequestBuilders.post("/v3/entity/dataset/batchGet")
+            MockMvcRequestBuilders.post("/openapi/v3/entity/dataset/batchGet")
                 .content(requestBody)
                 .contentType(MediaType.APPLICATION_JSON)
                 .accept(MediaType.APPLICATION_JSON))
@@ -531,7 +531,7 @@ public void testScrollEntitiesWithMultipleSortFields() throws Exception {
 
     mockMvc
         .perform(
-            MockMvcRequestBuilders.post("/v3/entity/scroll")
+            MockMvcRequestBuilders.post("/openapi/v3/entity/scroll")
                 .content("{\"entities\":[\"dataset\"]}")
                 .param("sortCriteria", "name", "urn")
                 .contentType(MediaType.APPLICATION_JSON)
@@ -575,7 +575,7 @@ public void testScrollEntitiesWithPitKeepAlive() throws Exception {
 
     mockMvc
         .perform(
-            MockMvcRequestBuilders.post("/v3/entity/scroll")
+            MockMvcRequestBuilders.post("/openapi/v3/entity/scroll")
                 .content("{\"entities\":[\"dataset\"]}")
                 .param("pitKeepAlive", "10m")
                 .contentType(MediaType.APPLICATION_JSON)
@@ -598,7 +598,7 @@ public void testEntityVersioningFeatureFlagDisabled() throws Exception {
         .perform(
             MockMvcRequestBuilders.post(
                     String.format(
-                        "/v3/entity/versioning/%s/relationship/versionOf/%s",
+                        "/openapi/v3/entity/versioning/%s/relationship/versionOf/%s",
                         VERSION_SET_URN, TEST_URN))
                 .content("{}")
                 .contentType(MediaType.APPLICATION_JSON)
@@ -610,7 +610,7 @@ public void testEntityVersioningFeatureFlagDisabled() throws Exception {
         .perform(
             MockMvcRequestBuilders.delete(
                     String.format(
-                        "/v3/entity/versioning/%s/relationship/versionOf/%s",
+                        "/openapi/v3/entity/versioning/%s/relationship/versionOf/%s",
                         VERSION_SET_URN, TEST_URN))
                 .accept(MediaType.APPLICATION_JSON))
         .andExpect(status().is4xxClientError());
@@ -630,7 +630,7 @@ public void testInvalidVersionSetUrn() throws Exception {
         .perform(
             MockMvcRequestBuilders.post(
                     String.format(
-                        "/v3/entity/versioning/%s/relationship/versionOf/%s",
+                        "/openapi/v3/entity/versioning/%s/relationship/versionOf/%s",
                         INVALID_VERSION_SET_URN, TEST_URN))
                 .content("{}")
                 .contentType(MediaType.APPLICATION_JSON)
@@ -642,7 +642,7 @@ public void testInvalidVersionSetUrn() throws Exception {
         .perform(
             MockMvcRequestBuilders.delete(
                     String.format(
-                        "/v3/entity/versioning/%s/relationship/versionOf/%s",
+                        "/openapi/v3/entity/versioning/%s/relationship/versionOf/%s",
                         INVALID_VERSION_SET_URN, TEST_URN))
                 .accept(MediaType.APPLICATION_JSON))
         .andExpect(status().is4xxClientError());
diff --git a/metadata-service/schema-registry-servlet/src/main/java/io/datahubproject/openapi/schema/registry/SchemaRegistryController.java b/metadata-service/schema-registry-servlet/src/main/java/io/datahubproject/openapi/schema/registry/SchemaRegistryController.java
index 09043c6dd5e87e..145ecae92ff0b6 100644
--- a/metadata-service/schema-registry-servlet/src/main/java/io/datahubproject/openapi/schema/registry/SchemaRegistryController.java
+++ b/metadata-service/schema-registry-servlet/src/main/java/io/datahubproject/openapi/schema/registry/SchemaRegistryController.java
@@ -2,7 +2,6 @@
 
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.common.collect.ImmutableSet;
-import com.linkedin.gms.factory.kafka.schemaregistry.InternalSchemaRegistryFactory;
 import com.linkedin.metadata.registry.SchemaRegistryService;
 import io.datahubproject.schema_registry.openapi.generated.CompatibilityCheckResponse;
 import io.datahubproject.schema_registry.openapi.generated.Config;
@@ -29,10 +28,12 @@
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
+import javax.annotation.PostConstruct;
 import lombok.RequiredArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
 import org.springframework.beans.factory.annotation.Qualifier;
-import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
+import org.springframework.core.Ordered;
+import org.springframework.core.annotation.Order;
 import org.springframework.http.HttpStatus;
 import org.springframework.http.ResponseEntity;
 import org.springframework.web.bind.annotation.RequestMapping;
@@ -42,11 +43,9 @@
 /** DataHub Rest Controller implementation for Confluent's Schema Registry OpenAPI spec. */
 @Slf4j
 @RestController
-@RequestMapping("/api")
+@RequestMapping("/schema-registry/api")
+@Order(Ordered.HIGHEST_PRECEDENCE)
 @RequiredArgsConstructor
-@ConditionalOnProperty(
-    name = "kafka.schemaRegistry.type",
-    havingValue = InternalSchemaRegistryFactory.TYPE)
 public class SchemaRegistryController
     implements CompatibilityApi,
         ConfigApi,
@@ -66,6 +65,11 @@ public class SchemaRegistryController
   @Qualifier("schemaRegistryService")
   private final SchemaRegistryService _schemaRegistryService;
 
+  @PostConstruct
+  public void init() {
+    log.info("SchemaRegistryController initialized with base path: /schema-registry/api");
+  }
+
   @Override
   public Optional<ObjectMapper> getObjectMapper() {
     return Optional.ofNullable(objectMapper);
diff --git a/metadata-service/schema-registry-servlet/src/main/java/io/datahubproject/openapi/schema/registry/config/SpringWebSchemaRegistryConfig.java b/metadata-service/schema-registry-servlet/src/main/java/io/datahubproject/openapi/schema/registry/config/SpringWebSchemaRegistryConfig.java
index 5f7d8d9604de6e..c2200eef6e60f5 100644
--- a/metadata-service/schema-registry-servlet/src/main/java/io/datahubproject/openapi/schema/registry/config/SpringWebSchemaRegistryConfig.java
+++ b/metadata-service/schema-registry-servlet/src/main/java/io/datahubproject/openapi/schema/registry/config/SpringWebSchemaRegistryConfig.java
@@ -1,23 +1,50 @@
 package io.datahubproject.openapi.schema.registry.config;
 
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.linkedin.gms.factory.kafka.schemaregistry.InternalSchemaRegistryFactory;
+import com.linkedin.metadata.registry.SchemaRegistryService;
+import io.datahubproject.openapi.schema.registry.SchemaRegistryController;
 import io.swagger.v3.oas.annotations.OpenAPIDefinition;
 import io.swagger.v3.oas.annotations.info.Info;
 import io.swagger.v3.oas.annotations.servers.Server;
+import jakarta.servlet.http.HttpServletRequest;
 import java.util.List;
+import lombok.extern.slf4j.Slf4j;
+import org.springframework.beans.factory.annotation.Qualifier;
+import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.ComponentScan;
 import org.springframework.context.annotation.Configuration;
+import org.springframework.core.annotation.Order;
 import org.springframework.http.converter.ByteArrayHttpMessageConverter;
 import org.springframework.http.converter.FormHttpMessageConverter;
 import org.springframework.http.converter.HttpMessageConverter;
 import org.springframework.http.converter.StringHttpMessageConverter;
 import org.springframework.http.converter.json.MappingJackson2HttpMessageConverter;
+import org.springframework.web.servlet.config.annotation.EnableWebMvc;
 import org.springframework.web.servlet.config.annotation.WebMvcConfigurer;
 
+@Slf4j
+@EnableWebMvc
 @OpenAPIDefinition(
     info = @Info(title = "DataHub OpenAPI", version = "1.0.0"),
     servers = {@Server(url = "/schema-registry/", description = "Schema Registry Server URL")})
+@Order(3)
+@ConditionalOnProperty(
+    name = "kafka.schemaRegistry.type",
+    havingValue = InternalSchemaRegistryFactory.TYPE)
 @Configuration
+@ComponentScan(basePackages = {"io.datahubproject.openapi.schema.registry"})
 public class SpringWebSchemaRegistryConfig implements WebMvcConfigurer {
 
+  @Bean
+  public SchemaRegistryController schemaRegistryController(
+      ObjectMapper objectMapper,
+      HttpServletRequest request,
+      @Qualifier("schemaRegistryService") SchemaRegistryService schemaRegistryService) {
+    return new SchemaRegistryController(objectMapper, request, schemaRegistryService);
+  }
+
   @Override
   public void configureMessageConverters(List<HttpMessageConverter<?>> messageConverters) {
     messageConverters.add(new StringHttpMessageConverter());
diff --git a/metadata-service/schema-registry-servlet/src/test/java/io/datahubproject/openapi/test/SchemaRegistryControllerTest.java b/metadata-service/schema-registry-servlet/src/test/java/io/datahubproject/openapi/test/SchemaRegistryControllerTest.java
index e8deed00672da7..fd862d65fb2ef4 100644
--- a/metadata-service/schema-registry-servlet/src/test/java/io/datahubproject/openapi/test/SchemaRegistryControllerTest.java
+++ b/metadata-service/schema-registry-servlet/src/test/java/io/datahubproject/openapi/test/SchemaRegistryControllerTest.java
@@ -74,7 +74,7 @@ static void kafkaProperties(DynamicPropertyRegistry registry) {
     kafka.start();
     registry.add("kafka.bootstrapServers", kafka::getBootstrapServers);
     registry.add("kafka.schemaRegistry.type", () -> "INTERNAL");
-    registry.add("kafka.schemaRegistry.url", () -> "http://localhost:53222/api/");
+    registry.add("kafka.schemaRegistry.url", () -> "http://localhost:53222/schema-registry/api/");
   }
 
   @Autowired EventProducer _producer;
diff --git a/metadata-service/war/build.gradle b/metadata-service/war/build.gradle
index 5e00207c1726ee..d83b29c4ad82bc 100644
--- a/metadata-service/war/build.gradle
+++ b/metadata-service/war/build.gradle
@@ -1,6 +1,7 @@
 plugins {
-  id 'war'
+  id 'org.springframework.boot'
   id 'com.palantir.docker'
+  id 'java'
 }
 
 apply from: '../../gradle/coverage/java-coverage.gradle'
@@ -31,6 +32,12 @@ dependencies {
   runtimeOnly externalDependency.mariadbConnector
   runtimeOnly externalDependency.mysqlConnector
   runtimeOnly externalDependency.postgresql
+
+  implementation(externalDependency.springBootStarterWeb) {
+    exclude module: "spring-boot-starter-tomcat"
+  }
+  implementation externalDependency.springBootStarterJetty
+  implementation externalDependency.jettyJmx
   implementation externalDependency.springWebMVC
   implementation externalDependency.springBootAutoconfigure
   implementation externalDependency.servletApi
@@ -38,6 +45,8 @@ dependencies {
   implementation spec.product.pegasus.restliDocgen
   implementation spec.product.pegasus.restliSpringBridge
 
+  compileOnly externalDependency.lombok
+  annotationProcessor externalDependency.lombok
   runtimeOnly externalDependency.log4jCore
   runtimeOnly externalDependency.log4j2Api
   runtimeOnly externalDependency.logbackClassic
@@ -51,31 +60,30 @@ dependencies {
 }
 configurations.all{
   exclude group: "com.charleskorn.kaml", module:"kaml"
-  
-}
-configurations {
-  jetty11
 }
 
-dependencies {
-  jetty11 "org.eclipse.jetty:jetty-runner:11.0.19"
+tasks.register('run') {
+  group = 'application'  // Add a group for better organization
+  description = 'Runs the application with Gretty'
+  dependsOn 'bootRun' // spring task
 }
 
-task run(type: JavaExec, dependsOn: build) {
-  main = "org.eclipse.jetty.runner.Runner"
-  args = ["${projectDir}/build/libs/war.war"]
-  classpath configurations.jetty11
+bootJar {
+  // backwards compatible with old war archive name
+  archiveFileName = 'war.war'
+  mainClass = 'com.linkedin.gms.GMSApplication'
 }
 
-war {
-  archiveFileName = "war.war"
+bootRun {
+  environment "SCHEMA_REGISTRY_TYPE", "INTERNAL"
+  environment "KAFKA_SCHEMAREGISTRY_URL", "http://localhost:8080/schema-registry/api"
 }
 
 docker {
   name "${docker_registry}/${docker_repo}:v${version}"
   version "v${version}"
   dockerfile file("${rootProject.projectDir}/docker/${docker_repo}/Dockerfile")
-  files war.outputs.files
+  files bootJar.outputs.files
   files fileTree(rootProject.projectDir) {
     include '.dockerignore'
     include 'docker/monitoring/*'
@@ -107,7 +115,7 @@ docker {
     buildArgs(dockerBuildArgs)
   }
 }
-tasks.getByPath(":metadata-service:war:docker").dependsOn([build, war])
+tasks.getByPath(":metadata-service:war:docker").dependsOn([bootJar])
 
 task cleanLocalDockerImages {
   doLast {
@@ -115,3 +123,12 @@ task cleanLocalDockerImages {
   }
 }
 dockerClean.finalizedBy(cleanLocalDockerImages)
+
+test {
+  jacoco {
+    // This contains quickstart tested code for jetty startup
+    excludes = ["com.linkedin.gms.CommonApplicationConfig",
+                "com.linkedin.gms.ServletConfig",
+                "com.linkedin.gms.GMSApplication"]
+  }
+}
\ No newline at end of file
diff --git a/metadata-service/war/src/main/java/com/linkedin/gms/CommonApplicationConfig.java b/metadata-service/war/src/main/java/com/linkedin/gms/CommonApplicationConfig.java
index 69fb9df2b04dbb..3f3ed9f75cceb3 100644
--- a/metadata-service/war/src/main/java/com/linkedin/gms/CommonApplicationConfig.java
+++ b/metadata-service/war/src/main/java/com/linkedin/gms/CommonApplicationConfig.java
@@ -1,11 +1,22 @@
 package com.linkedin.gms;
 
-import com.datahub.auth.authentication.filter.AuthenticationFilter;
 import com.linkedin.metadata.spring.YamlPropertySourceFactory;
+import java.lang.management.ManagementFactory;
+import javax.management.MBeanServer;
+import org.eclipse.jetty.jmx.MBeanContainer;
+import org.eclipse.jetty.server.Connector;
+import org.eclipse.jetty.server.HttpConfiguration;
+import org.eclipse.jetty.server.HttpConnectionFactory;
+import org.eclipse.jetty.server.ServerConnector;
+import org.slf4j.LoggerFactory;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.boot.web.embedded.jetty.JettyServletWebServerFactory;
+import org.springframework.boot.web.server.WebServerFactoryCustomizer;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.ComponentScan;
 import org.springframework.context.annotation.Configuration;
 import org.springframework.context.annotation.PropertySource;
+import org.springframework.core.env.Environment;
 
 /**
  * Common configuration for all servlets. Generally this list also includes dependencies of the
@@ -43,12 +54,39 @@
       "com.linkedin.gms.factory.notifications",
       "com.linkedin.gms.factory.telemetry"
     })
-@PropertySource(value = "classpath:/application.yaml", factory = YamlPropertySourceFactory.class)
 @Configuration
+@PropertySource(value = "classpath:/application.yaml", factory = YamlPropertySourceFactory.class)
 public class CommonApplicationConfig {
 
-  @Bean("authenticationFilter")
-  public AuthenticationFilter authenticationFilter() {
-    return new AuthenticationFilter();
+  @Autowired private Environment environment;
+
+  @Bean
+  public WebServerFactoryCustomizer<JettyServletWebServerFactory> jettyCustomizer() {
+    return factory -> {
+      // Configure HTTP
+      factory.addServerCustomizers(
+          server -> {
+            // HTTP Configuration
+            HttpConfiguration httpConfig = new HttpConfiguration();
+            httpConfig.setRequestHeaderSize(32768);
+
+            // HTTP Connector
+            ServerConnector connector =
+                new ServerConnector(server, new HttpConnectionFactory(httpConfig));
+
+            // Get port from environment directly
+            int port = environment.getProperty("server.port", Integer.class, 8080);
+            connector.setPort(port);
+
+            // Set connectors
+            server.setConnectors(new Connector[] {connector});
+
+            // JMX Configuration
+            MBeanServer mBeanServer = ManagementFactory.getPlatformMBeanServer();
+            MBeanContainer mBeanContainer = new MBeanContainer(mBeanServer);
+            mBeanContainer.beanAdded(null, LoggerFactory.getILoggerFactory());
+            server.addBean(mBeanContainer);
+          });
+    };
   }
 }
diff --git a/metadata-service/war/src/main/java/com/linkedin/gms/GMSApplication.java b/metadata-service/war/src/main/java/com/linkedin/gms/GMSApplication.java
new file mode 100644
index 00000000000000..ec8aaeffbba1d0
--- /dev/null
+++ b/metadata-service/war/src/main/java/com/linkedin/gms/GMSApplication.java
@@ -0,0 +1,21 @@
+package com.linkedin.gms;
+
+import org.springframework.boot.SpringApplication;
+import org.springframework.boot.autoconfigure.SpringBootApplication;
+import org.springframework.boot.builder.SpringApplicationBuilder;
+import org.springframework.boot.web.servlet.support.SpringBootServletInitializer;
+import org.springframework.context.annotation.Import;
+
+@SpringBootApplication
+@Import({CommonApplicationConfig.class, ServletConfig.class})
+public class GMSApplication extends SpringBootServletInitializer {
+
+  @Override
+  protected SpringApplicationBuilder configure(SpringApplicationBuilder application) {
+    return application.sources(GMSApplication.class);
+  }
+
+  public static void main(String[] args) {
+    SpringApplication.run(GMSApplication.class, args);
+  }
+}
diff --git a/metadata-service/war/src/main/java/com/linkedin/gms/ServletConfig.java b/metadata-service/war/src/main/java/com/linkedin/gms/ServletConfig.java
new file mode 100644
index 00000000000000..fbd8d72b76c8f2
--- /dev/null
+++ b/metadata-service/war/src/main/java/com/linkedin/gms/ServletConfig.java
@@ -0,0 +1,158 @@
+package com.linkedin.gms;
+
+import static com.linkedin.metadata.Constants.INGESTION_MAX_SERIALIZED_STRING_LENGTH;
+import static com.linkedin.metadata.Constants.MAX_JACKSON_STRING_SIZE;
+
+import com.datahub.auth.authentication.filter.AuthenticationFilter;
+import com.datahub.gms.servlet.Config;
+import com.datahub.gms.servlet.ConfigSearchExport;
+import com.datahub.gms.servlet.HealthCheck;
+import com.fasterxml.jackson.annotation.JsonInclude;
+import com.fasterxml.jackson.core.StreamReadConstraints;
+import com.fasterxml.jackson.databind.DeserializationFeature;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.linkedin.r2.transport.http.server.RAPJakartaServlet;
+import com.linkedin.restli.server.RestliHandlerServlet;
+import io.datahubproject.openapi.converter.StringToChangeCategoryConverter;
+import java.util.List;
+import javax.annotation.Nonnull;
+import lombok.extern.slf4j.Slf4j;
+import org.springframework.beans.factory.annotation.Value;
+import org.springframework.boot.web.servlet.FilterRegistrationBean;
+import org.springframework.boot.web.servlet.ServletRegistrationBean;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.ComponentScan;
+import org.springframework.context.annotation.Configuration;
+import org.springframework.core.Ordered;
+import org.springframework.core.annotation.Order;
+import org.springframework.format.FormatterRegistry;
+import org.springframework.http.converter.ByteArrayHttpMessageConverter;
+import org.springframework.http.converter.FormHttpMessageConverter;
+import org.springframework.http.converter.HttpMessageConverter;
+import org.springframework.http.converter.StringHttpMessageConverter;
+import org.springframework.http.converter.json.MappingJackson2HttpMessageConverter;
+import org.springframework.web.servlet.config.annotation.AsyncSupportConfigurer;
+import org.springframework.web.servlet.config.annotation.EnableWebMvc;
+import org.springframework.web.servlet.config.annotation.WebMvcConfigurer;
+
+/**
+ * Common configuration for all servlets. Generally this list also includes dependencies of the
+ * embedded MAE/MCE consumers.
+ */
+@Slf4j
+@Configuration
+@Order(1)
+@EnableWebMvc
+@ComponentScan(
+    basePackages = {"io.datahubproject.openapi.schema.registry.config", "com.linkedin.gms.servlet"})
+public class ServletConfig implements WebMvcConfigurer {
+
+  @Value("${datahub.gms.async.request-timeout-ms}")
+  private long asyncTimeoutMilliseconds;
+
+  @Bean
+  public FilterRegistrationBean<AuthenticationFilter> authFilter(AuthenticationFilter filter) {
+    FilterRegistrationBean<AuthenticationFilter> registration = new FilterRegistrationBean<>();
+    registration.setFilter(filter);
+    registration.setOrder(Ordered.HIGHEST_PRECEDENCE);
+    registration.setAsyncSupported(true);
+
+    // Register filter for all paths - exclusions are handled by shouldNotFilter()
+    registration.addUrlPatterns("/*");
+
+    return registration;
+  }
+
+  @Bean
+  public ServletRegistrationBean<HealthCheck> healthCheckServlet() {
+    ServletRegistrationBean<HealthCheck> registration =
+        new ServletRegistrationBean<>(new HealthCheck());
+    registration.setName("healthCheck");
+    registration.addUrlMappings("/health");
+    registration.setLoadOnStartup(15);
+    registration.setAsyncSupported(true);
+    return registration;
+  }
+
+  @Bean
+  public ServletRegistrationBean<Config> configServlet() {
+    ServletRegistrationBean<Config> registration = new ServletRegistrationBean<>(new Config());
+    registration.setName("config");
+    registration.addUrlMappings("/config");
+    registration.setLoadOnStartup(15);
+    registration.setAsyncSupported(true);
+    return registration;
+  }
+
+  @Bean
+  public ServletRegistrationBean<ConfigSearchExport> configSearchExportServlet() {
+    ServletRegistrationBean<ConfigSearchExport> registration =
+        new ServletRegistrationBean<>(new ConfigSearchExport());
+    registration.setName("config-search-export");
+    registration.addUrlMappings("/config/search/export");
+    registration.setLoadOnStartup(15);
+    registration.setAsyncSupported(true);
+    return registration;
+  }
+
+  /**
+   * SpringBoot is now the default, explicitly map these to legacy rest.li servlet. Additions are
+   * more likely to be built on the Spring side so we're preventing unexpected behavior for the most
+   * likely changes. Rest.li API is intended to be deprecated and removed.
+   *
+   * @param r2Servlet the restli servlet
+   * @return registration
+   */
+  @Bean
+  public ServletRegistrationBean<RestliHandlerServlet> restliServletRegistration(
+      RAPJakartaServlet r2Servlet) {
+    ServletRegistrationBean<RestliHandlerServlet> registration =
+        new ServletRegistrationBean<>(new RestliHandlerServlet(r2Servlet));
+    registration.addUrlMappings(
+        "/aspects/*",
+        "/entities/*",
+        "/entitiesV2/*",
+        "/entitiesVersionedV2/*",
+        "/usageStats/*",
+        "/platform/*",
+        "/relationships/*",
+        "/analytics/*",
+        "/operations/*",
+        "/runs/*");
+    registration.setLoadOnStartup(2);
+    registration.setOrder(Integer.MAX_VALUE); // lowest priority
+    return registration;
+  }
+
+  @Override
+  public void configureMessageConverters(List<HttpMessageConverter<?>> messageConverters) {
+    messageConverters.add(new StringHttpMessageConverter());
+    messageConverters.add(new ByteArrayHttpMessageConverter());
+    messageConverters.add(new FormHttpMessageConverter());
+
+    ObjectMapper objectMapper = new ObjectMapper();
+    int maxSize =
+        Integer.parseInt(
+            System.getenv()
+                .getOrDefault(INGESTION_MAX_SERIALIZED_STRING_LENGTH, MAX_JACKSON_STRING_SIZE));
+    objectMapper
+        .getFactory()
+        .setStreamReadConstraints(StreamReadConstraints.builder().maxStringLength(maxSize).build());
+    objectMapper.setSerializationInclusion(JsonInclude.Include.NON_NULL);
+    objectMapper.configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false);
+    MappingJackson2HttpMessageConverter jsonConverter =
+        new MappingJackson2HttpMessageConverter(objectMapper);
+    messageConverters.add(jsonConverter);
+  }
+
+  @Override
+  public void addFormatters(FormatterRegistry registry) {
+    registry.addConverter(new StringToChangeCategoryConverter());
+  }
+
+  @Override
+  public void configureAsyncSupport(@Nonnull AsyncSupportConfigurer configurer) {
+    WebMvcConfigurer.super.configureAsyncSupport(configurer);
+    configurer.setDefaultTimeout(asyncTimeoutMilliseconds);
+  }
+}
diff --git a/metadata-service/war/src/main/java/com/linkedin/gms/SpringApplicationInitializer.java b/metadata-service/war/src/main/java/com/linkedin/gms/SpringApplicationInitializer.java
deleted file mode 100644
index 821d1bf0ad6f8c..00000000000000
--- a/metadata-service/war/src/main/java/com/linkedin/gms/SpringApplicationInitializer.java
+++ /dev/null
@@ -1,11 +0,0 @@
-package com.linkedin.gms;
-
-import javax.annotation.Nonnull;
-import org.springframework.context.ApplicationContextInitializer;
-import org.springframework.context.ConfigurableApplicationContext;
-
-public class SpringApplicationInitializer
-    implements ApplicationContextInitializer<ConfigurableApplicationContext> {
-  @Override
-  public void initialize(@Nonnull ConfigurableApplicationContext applicationContext) {}
-}
diff --git a/metadata-service/war/src/main/java/com/linkedin/gms/WebApplicationInitializer.java b/metadata-service/war/src/main/java/com/linkedin/gms/WebApplicationInitializer.java
deleted file mode 100644
index 4ed84e48a5049f..00000000000000
--- a/metadata-service/war/src/main/java/com/linkedin/gms/WebApplicationInitializer.java
+++ /dev/null
@@ -1,166 +0,0 @@
-package com.linkedin.gms;
-
-import static com.linkedin.metadata.boot.OnBootApplicationListener.SCHEMA_REGISTRY_SERVLET_NAME;
-
-import com.datahub.auth.authentication.filter.AuthenticationFilter;
-import com.datahub.gms.servlet.Config;
-import com.datahub.gms.servlet.ConfigSearchExport;
-import com.datahub.gms.servlet.HealthCheck;
-import com.linkedin.gms.servlet.AuthServletConfig;
-import com.linkedin.gms.servlet.GraphQLServletConfig;
-import com.linkedin.gms.servlet.OpenAPIServletConfig;
-import com.linkedin.gms.servlet.RestliServletConfig;
-import com.linkedin.gms.servlet.SchemaRegistryServletConfig;
-import jakarta.servlet.DispatcherType;
-import jakarta.servlet.FilterRegistration;
-import jakarta.servlet.ServletContext;
-import jakarta.servlet.ServletRegistration;
-import java.util.ArrayList;
-import java.util.EnumSet;
-import java.util.List;
-import org.springframework.web.context.ContextLoaderListener;
-import org.springframework.web.context.support.AnnotationConfigWebApplicationContext;
-import org.springframework.web.context.support.HttpRequestHandlerServlet;
-import org.springframework.web.servlet.DispatcherServlet;
-
-/** This class is before Spring Context is loaded, previously web.xml based */
-public class WebApplicationInitializer
-    implements org.springframework.web.WebApplicationInitializer {
-  @Override
-  public void onStartup(ServletContext container) {
-    AnnotationConfigWebApplicationContext rootContext = new AnnotationConfigWebApplicationContext();
-    rootContext.register(CommonApplicationConfig.class);
-
-    ContextLoaderListener contextLoaderListener = new ContextLoaderListener(rootContext);
-    container.addListener(contextLoaderListener);
-    container.setInitParameter(
-        "contextInitializerClasses", "com.linkedin.gms.SpringApplicationInitializer");
-
-    // Auth filter
-    List<String> servletNames = new ArrayList<>();
-
-    // Independent dispatcher
-    schemaRegistryServlet(container);
-
-    // Spring Dispatcher servlets
-    DispatcherServlet dispatcherServlet = new DispatcherServlet(rootContext);
-    servletNames.add(authServlet(rootContext, dispatcherServlet, container));
-    servletNames.add(graphQLServlet(rootContext, dispatcherServlet, container));
-    servletNames.add(openAPIServlet(rootContext, dispatcherServlet, container));
-    // Restli non-Dispatcher default
-    servletNames.add(restliServlet(rootContext, container));
-
-    FilterRegistration.Dynamic filterRegistration =
-        container.addFilter("authenticationFilter", AuthenticationFilter.class);
-    filterRegistration.setAsyncSupported(true);
-    filterRegistration.addMappingForServletNames(
-        EnumSet.of(DispatcherType.ASYNC, DispatcherType.REQUEST),
-        false,
-        servletNames.toArray(String[]::new));
-
-    // Non-Spring servlets
-    healthCheckServlet(container);
-    configServlet(container);
-  }
-
-  /*
-   * This is a servlet exclusive to DataHub's implementation of the Confluent OpenAPI spec which is not behind
-   * DataHub's authentication layer as it is not compatible with confluent consumers & producers.
-   */
-  private void schemaRegistryServlet(ServletContext container) {
-    AnnotationConfigWebApplicationContext webContext = new AnnotationConfigWebApplicationContext();
-    webContext.setId(SCHEMA_REGISTRY_SERVLET_NAME);
-    webContext.register(SchemaRegistryServletConfig.class);
-
-    DispatcherServlet dispatcherServlet = new DispatcherServlet(webContext);
-    ServletRegistration.Dynamic registration =
-        container.addServlet(SCHEMA_REGISTRY_SERVLET_NAME, dispatcherServlet);
-    registration.addMapping("/schema-registry/*");
-    registration.setLoadOnStartup(1);
-    registration.setAsyncSupported(true);
-  }
-
-  private String authServlet(
-      AnnotationConfigWebApplicationContext rootContext,
-      DispatcherServlet dispatcherServlet,
-      ServletContext container) {
-    final String servletName = "dispatcher-auth";
-    rootContext.register(AuthServletConfig.class);
-
-    ServletRegistration.Dynamic registration = container.addServlet(servletName, dispatcherServlet);
-    registration.addMapping("/auth/*");
-    registration.setLoadOnStartup(5);
-    registration.setAsyncSupported(true);
-
-    return servletName;
-  }
-
-  private String graphQLServlet(
-      AnnotationConfigWebApplicationContext rootContext,
-      DispatcherServlet dispatcherServlet,
-      ServletContext container) {
-    final String servletName = "dispatcher-graphql";
-    rootContext.register(GraphQLServletConfig.class);
-
-    ServletRegistration.Dynamic registration = container.addServlet(servletName, dispatcherServlet);
-    registration.addMapping("/api/*");
-    registration.setLoadOnStartup(5);
-    registration.setAsyncSupported(true);
-
-    return servletName;
-  }
-
-  private String openAPIServlet(
-      AnnotationConfigWebApplicationContext rootContext,
-      DispatcherServlet dispatcherServlet,
-      ServletContext container) {
-    final String servletName = "dispatcher-openapi";
-    rootContext.register(OpenAPIServletConfig.class);
-
-    ServletRegistration.Dynamic registration = container.addServlet(servletName, dispatcherServlet);
-    registration.addMapping("/openapi/*");
-    registration.setLoadOnStartup(5);
-    registration.setAsyncSupported(true);
-
-    return servletName;
-  }
-
-  private String restliServlet(
-      AnnotationConfigWebApplicationContext rootContext, ServletContext container) {
-    final String servletName = "restliRequestHandler";
-
-    rootContext.register(RestliServletConfig.class);
-
-    ServletRegistration.Dynamic registration =
-        container.addServlet(servletName, HttpRequestHandlerServlet.class);
-    registration.addMapping("/*");
-    registration.setLoadOnStartup(10);
-    registration.setAsyncSupported(true);
-    registration.setInitParameter(
-        "org.springframework.web.servlet.FrameworkServlet.ORDER",
-        String.valueOf(Integer.MAX_VALUE - 1));
-
-    return servletName;
-  }
-
-  private void healthCheckServlet(ServletContext container) {
-    ServletRegistration.Dynamic registration =
-        container.addServlet("healthCheck", new HealthCheck());
-    registration.addMapping("/health");
-    registration.setLoadOnStartup(15);
-    registration.setAsyncSupported(true);
-  }
-
-  private void configServlet(ServletContext container) {
-    ServletRegistration.Dynamic registration = container.addServlet("config", new Config());
-    registration.addMapping("/config");
-    registration.setLoadOnStartup(15);
-    registration.setAsyncSupported(true);
-
-    ServletRegistration.Dynamic registration2 =
-        container.addServlet("config-search-export", new ConfigSearchExport());
-    registration2.addMapping("/config/search/export");
-    registration2.setLoadOnStartup(15);
-    registration2.setAsyncSupported(true);
-  }
-}
diff --git a/metadata-service/war/src/main/java/com/linkedin/gms/servlet/AuthServletConfig.java b/metadata-service/war/src/main/java/com/linkedin/gms/servlet/AuthServletConfig.java
index 437c8722a747f5..08074c26d44c55 100644
--- a/metadata-service/war/src/main/java/com/linkedin/gms/servlet/AuthServletConfig.java
+++ b/metadata-service/war/src/main/java/com/linkedin/gms/servlet/AuthServletConfig.java
@@ -2,9 +2,7 @@
 
 import org.springframework.context.annotation.ComponentScan;
 import org.springframework.context.annotation.Configuration;
-import org.springframework.web.servlet.config.annotation.EnableWebMvc;
 
-@EnableWebMvc
 @ComponentScan(basePackages = {"com.datahub.auth.authentication"})
 @Configuration
 public class AuthServletConfig {}
diff --git a/metadata-service/war/src/main/java/com/linkedin/gms/servlet/GraphQLServletConfig.java b/metadata-service/war/src/main/java/com/linkedin/gms/servlet/GraphQLServletConfig.java
index 42413df0757e64..9f08306518482a 100644
--- a/metadata-service/war/src/main/java/com/linkedin/gms/servlet/GraphQLServletConfig.java
+++ b/metadata-service/war/src/main/java/com/linkedin/gms/servlet/GraphQLServletConfig.java
@@ -2,9 +2,7 @@
 
 import org.springframework.context.annotation.ComponentScan;
 import org.springframework.context.annotation.Configuration;
-import org.springframework.web.servlet.config.annotation.EnableWebMvc;
 
-@EnableWebMvc
 @ComponentScan(
     basePackages = {
       "com.datahub.graphql",
diff --git a/metadata-service/war/src/main/java/com/linkedin/gms/servlet/OpenAPIServletConfig.java b/metadata-service/war/src/main/java/com/linkedin/gms/servlet/OpenAPIServletConfig.java
index d259f763d3d440..ebab789b058799 100644
--- a/metadata-service/war/src/main/java/com/linkedin/gms/servlet/OpenAPIServletConfig.java
+++ b/metadata-service/war/src/main/java/com/linkedin/gms/servlet/OpenAPIServletConfig.java
@@ -2,9 +2,7 @@
 
 import org.springframework.context.annotation.ComponentScan;
 import org.springframework.context.annotation.Configuration;
-import org.springframework.web.servlet.config.annotation.EnableWebMvc;
 
-@EnableWebMvc
 @ComponentScan(
     basePackages = {
       "io.datahubproject.openapi.config",
diff --git a/metadata-service/war/src/main/java/com/linkedin/gms/servlet/RestliServletConfig.java b/metadata-service/war/src/main/java/com/linkedin/gms/servlet/RestliServletConfig.java
index 222e2356dfb1ca..b079efac4a8c5b 100644
--- a/metadata-service/war/src/main/java/com/linkedin/gms/servlet/RestliServletConfig.java
+++ b/metadata-service/war/src/main/java/com/linkedin/gms/servlet/RestliServletConfig.java
@@ -1,20 +1,8 @@
 package com.linkedin.gms.servlet;
 
-import com.linkedin.metadata.spring.YamlPropertySourceFactory;
-import com.linkedin.r2.transport.http.server.RAPJakartaServlet;
-import com.linkedin.restli.server.RestliHandlerServlet;
-import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.ComponentScan;
 import org.springframework.context.annotation.Configuration;
-import org.springframework.context.annotation.PropertySource;
-import org.springframework.web.context.support.HttpRequestHandlerServlet;
 
 @ComponentScan(basePackages = {"com.linkedin.restli.server"})
-@PropertySource(value = "classpath:/application.yaml", factory = YamlPropertySourceFactory.class)
 @Configuration
-public class RestliServletConfig {
-  @Bean("restliRequestHandler")
-  public HttpRequestHandlerServlet restliHandlerServlet(final RAPJakartaServlet r2Servlet) {
-    return new RestliHandlerServlet(r2Servlet);
-  }
-}
+public class RestliServletConfig {}
diff --git a/metadata-service/war/src/main/java/com/linkedin/gms/servlet/SchemaRegistryServletConfig.java b/metadata-service/war/src/main/java/com/linkedin/gms/servlet/SchemaRegistryServletConfig.java
deleted file mode 100644
index beb633888bd349..00000000000000
--- a/metadata-service/war/src/main/java/com/linkedin/gms/servlet/SchemaRegistryServletConfig.java
+++ /dev/null
@@ -1,10 +0,0 @@
-package com.linkedin.gms.servlet;
-
-import org.springframework.context.annotation.ComponentScan;
-import org.springframework.context.annotation.Configuration;
-import org.springframework.web.servlet.config.annotation.EnableWebMvc;
-
-@EnableWebMvc
-@ComponentScan(basePackages = {"io.datahubproject.openapi.schema.registry"})
-@Configuration
-public class SchemaRegistryServletConfig {}
diff --git a/metadata-service/war/src/test/java/com/linkedin/gms/SpringTest.java b/metadata-service/war/src/test/java/com/linkedin/gms/SpringTest.java
index 27dabf1e33ddd8..74d2f2c8107148 100644
--- a/metadata-service/war/src/test/java/com/linkedin/gms/SpringTest.java
+++ b/metadata-service/war/src/test/java/com/linkedin/gms/SpringTest.java
@@ -3,30 +3,51 @@
 import static org.testng.AssertJUnit.assertNotNull;
 
 import com.linkedin.gms.factory.telemetry.DailyReport;
-import com.linkedin.metadata.models.registry.ConfigEntityRegistry;
+import com.linkedin.metadata.boot.BootstrapManager;
 import com.linkedin.metadata.models.registry.EntityRegistry;
+import io.datahubproject.metadata.context.OperationContext;
+import io.datahubproject.test.metadata.context.TestOperationContexts;
 import io.ebean.Database;
 import org.springframework.boot.test.context.SpringBootTest;
+import org.springframework.boot.test.context.TestConfiguration;
 import org.springframework.boot.test.mock.mockito.MockBean;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Primary;
 import org.springframework.test.context.ContextConfiguration;
 import org.springframework.test.context.testng.AbstractTestNGSpringContextTests;
 import org.testng.annotations.Test;
 
 @SpringBootTest(
     webEnvironment = SpringBootTest.WebEnvironment.MOCK,
-    properties = {"telemetry.enabledServer=true"})
-@ContextConfiguration(classes = CommonApplicationConfig.class)
+    properties = {
+      "telemetry.enabledServer=true",
+      "spring.main.allow-bean-definition-overriding=true"
+    })
+@ContextConfiguration(classes = {CommonApplicationConfig.class, SpringTest.TestBeans.class})
 public class SpringTest extends AbstractTestNGSpringContextTests {
 
   // Mock Beans take precedence, we add these to avoid needing to configure data sources etc. while
   // still testing prod config
   @MockBean private Database database;
-  @MockBean private ConfigEntityRegistry configEntityRegistry;
-  @MockBean private EntityRegistry entityRegistry;
+  @MockBean private BootstrapManager bootstrapManager;
 
   @Test
   public void testTelemetry() {
     DailyReport dailyReport = this.applicationContext.getBean(DailyReport.class);
     assertNotNull(dailyReport);
   }
+
+  @TestConfiguration
+  public static class TestBeans {
+    @Bean
+    public OperationContext systemOperationContext() {
+      return TestOperationContexts.systemContextNoSearchAuthorization();
+    }
+
+    @Primary
+    @Bean
+    public EntityRegistry entityRegistry(OperationContext systemOperationContext) {
+      return systemOperationContext.getEntityRegistry();
+    }
+  }
 }
diff --git a/python-build/generate_release_scripts.py b/python-build/generate_release_scripts.py
new file mode 100644
index 00000000000000..36253a24cfa3b0
--- /dev/null
+++ b/python-build/generate_release_scripts.py
@@ -0,0 +1,85 @@
+import dataclasses
+import pathlib
+
+REPO_ROOT = pathlib.Path(__file__).parent.parent
+
+
+@dataclasses.dataclass
+class Package:
+    # TODO: This doesn't have the actual package names.
+    directory: str
+    main_module_name: str
+
+    def root_from_directory(self) -> str:
+        ups = self.directory.count("/") + 1
+
+        return "/".join([".."] * ups)
+
+
+packages = [
+    Package(directory="metadata-ingestion", main_module_name="datahub"),
+    Package(
+        directory="metadata-ingestion-modules/airflow-plugin",
+        main_module_name="datahub_airflow_plugin",
+    ),
+    Package(
+        directory="metadata-ingestion-modules/dagster-plugin",
+        main_module_name="datahub_dagster_plugin",
+    ),
+    Package(
+        directory="metadata-ingestion-modules/gx-plugin",
+        main_module_name="datahub_gx_plugin",
+    ),
+    Package(
+        directory="metadata-ingestion-modules/prefect-plugin",
+        main_module_name="prefect_datahub",
+    ),
+]
+
+generation_header = f"# Auto-generated by {pathlib.Path(__file__).relative_to(REPO_ROOT)}. Do not edit manually."
+
+template = """\
+#!/bin/bash
+%s
+
+set -euxo pipefail
+
+ROOT=%s
+MODULE=%s
+
+if [[ ! ${RELEASE_SKIP_TEST:-} ]] && [[ ! ${RELEASE_SKIP_INSTALL:-} ]]; then
+    ${ROOT}/gradlew build  # also runs tests
+elif [[ ! ${RELEASE_SKIP_INSTALL:-} ]]; then
+    ${ROOT}/gradlew install
+fi
+
+# Check packaging constraint.
+python -c 'import setuptools; where="./src"; assert setuptools.find_packages(where) == setuptools.find_namespace_packages(where), "you seem to be missing or have extra __init__.py files"'
+
+# Update the release version.
+if [[ ! ${RELEASE_VERSION:-} ]]; then
+    echo "RELEASE_VERSION is not set"
+    exit 1
+fi
+sed -i.bak "s/__version__ = .*$/__version__ = \\"$(echo $RELEASE_VERSION|sed s/-/+/)\\"/" src/${MODULE}/_version.py
+
+# Build and upload the release.
+rm -rf build dist || true
+python -m build
+if [[ ! ${RELEASE_SKIP_UPLOAD:-} ]]; then
+    python -m twine upload 'dist/*'
+fi
+mv src/${MODULE}/_version.py.bak src/${MODULE}/_version.py
+"""
+
+for package in packages:
+    script_path = REPO_ROOT / package.directory / "scripts/release.sh"
+
+    script_path.write_text(
+        template
+        % (
+            generation_header,
+            package.root_from_directory(),
+            package.main_module_name,
+        )
+    )
diff --git a/smoke-test/build.gradle b/smoke-test/build.gradle
index 60d08e0206cdab..c38468ca8cd8b0 100644
--- a/smoke-test/build.gradle
+++ b/smoke-test/build.gradle
@@ -73,17 +73,16 @@ task installDev(type: Exec) {
 task pythonLint(type: Exec, dependsOn: installDev) {
     commandLine 'bash', '-c',
         "source ${venv_name}/bin/activate && set -x && " +
-        "black --check --diff tests/ && " +
-        "isort --check --diff tests/ && " +
-        "ruff --statistics tests/ && " +
+        "ruff check tests/ && " +
+        "ruff format --check tests/ && " +
         "mypy tests/"
 }
+
 task pythonLintFix(type: Exec, dependsOn: installDev) {
     commandLine 'bash', '-c',
         "source ${venv_name}/bin/activate && set -x && " +
-        "black tests/ && " +
-        "isort tests/ && " +
-        "ruff --fix tests/ && " +
+        "ruff check --fix tests/ && " +
+        "ruff format tests/ && " +
         "mypy tests/"
 }
 
@@ -155,3 +154,19 @@ task lint {
 task lintFix {
     dependsOn pythonLintFix
 }
+
+task cleanPythonCache(type: Exec) {
+  commandLine 'bash', '-c',
+    "find . -type f -name '*.py[co]' -delete -o -type d -name __pycache__ -delete -o -type d -empty -delete"
+}
+
+
+clean {
+  delete venv_name
+  delete 'build'
+  delete 'dist'
+  delete '.ruff_cache'
+  delete '.mypy_cache'
+  delete '.pytest_cache'
+}
+clean.dependsOn cleanPythonCache
\ No newline at end of file
diff --git a/smoke-test/pyproject.toml b/smoke-test/pyproject.toml
index c7745d0e9a3640..55e286c73c01b9 100644
--- a/smoke-test/pyproject.toml
+++ b/smoke-test/pyproject.toml
@@ -7,30 +7,65 @@ name = "smoke-test"
 version = "0.0.0"
 description = ""
 authors = [
-	{ name="Acryl Data", email="eng@acryl.io" },
+    { name="Acryl Data", email="eng@acryl.io" },
 ]
 requires-python = ">=3.9"
 
+[tool.ruff]
+# Enable ruff format
+target-version = "py310"
+line-length = 88
+extend-exclude = ["tmp", "venv"]
 
-[tool.black]
-extend-exclude = '''
-# A regex preceded with ^/ will apply only to files and directories
-# in the root of the project.
-tmp
-venv
-'''
-include = '\.pyi?$'
-target-version = ['py310']
+[tool.ruff.format]
+quote-style = "double"
+indent-style = "space"
+skip-magic-trailing-comma = false
+line-ending = "auto"
 
-[tool.isort]
-profile = 'black'
+[tool.ruff.lint.isort]
+combine-as-imports = true
+known-first-party = ["datahub"]
+extra-standard-library = ["__future__", "datahub.utilities._markupsafe_compat", "datahub.sql_parsing._sqlglot_patch"]
+section-order = ["future", "standard-library", "third-party", "first-party", "local-folder"]
+force-sort-within-sections = false
+force-wrap-aliases = false
+split-on-trailing-comma = false
+order-by-type = true
+relative-imports-order = "closest-to-furthest"
+force-single-line = false
+single-line-exclusions = ["typing"]
+length-sort = false
+from-first = false
+required-imports = []
+classes = ["typing"]
 
-[tool.ruff]
+[tool.ruff.lint]
+extend-select = [
+    "B",    # flake8-bugbear
+    "C90",  # mccabe complexity
+    "E",    # pycodestyle errors
+    "F",    # pyflakes
+    "G010", # logging.warn -> logging.warning
+    "I",    # isort
+    "TID",  # flake8-tidy-imports
+]
 ignore = [
-	'E501', # Ignore line length, since black handles that.
-	'D203', # Ignore 1 blank line required before class docstring.
+    "E501",  # Line length violations (handled by formatter)
+    "B904",  # exception with `raise ... from err` or `raise ... from None` to distinguish
+    "TID252",# Prefer absolute imports over relative imports
 ]
 
+[tool.ruff.lint.mccabe]
+max-complexity = 15
+
+[tool.ruff.lint.flake8-tidy-imports]
+# Disallow all relative imports.
+ban-relative-imports = "all"
+
+[tool.ruff.lint.per-file-ignores]
+"__init__.py" = ["F401"]
+
 [tool.mypy]
 exclude = "^(venv/|build/|dist/)"
 ignore_missing_imports = true
@@ -40,7 +75,4 @@ disallow_untyped_decorators = true
 warn_unused_configs = true
 # eventually we'd like to enable these
 disallow_incomplete_defs = false
-disallow_untyped_defs = false
-
-[tool.pyright]
-extraPaths = ['tests']
+disallow_untyped_defs = false
\ No newline at end of file
diff --git a/smoke-test/requirements.txt b/smoke-test/requirements.txt
index c5294d95177cb9..fadc3dbec1f2b5 100644
--- a/smoke-test/requirements.txt
+++ b/smoke-test/requirements.txt
@@ -9,10 +9,8 @@ joblib
 pytest-xdist
 networkx
 # libaries for linting below this
-black==23.7.0
-isort==5.12.0
 mypy==1.5.1
-ruff==0.0.287
+ruff==0.9.2
 # stub version are copied from metadata-ingestion/setup.py and that should be the source of truth
 types-requests>=2.28.11.6,<=2.31.0.3
 types-PyYAML
diff --git a/smoke-test/tests/aspect_generators/timeseries/dataset_profile_gen.py b/smoke-test/tests/aspect_generators/timeseries/dataset_profile_gen.py
index f808e7a58a3292..d97f119ea55a0c 100644
--- a/smoke-test/tests/aspect_generators/timeseries/dataset_profile_gen.py
+++ b/smoke-test/tests/aspect_generators/timeseries/dataset_profile_gen.py
@@ -5,7 +5,6 @@
     DatasetProfileClass,
     TimeWindowSizeClass,
 )
-
 from tests.utils import get_timestampmillis_at_start_of_day
 
 
diff --git a/smoke-test/tests/assertions/assertions_test.py b/smoke-test/tests/assertions/assertions_test.py
index 8615a55d7289f0..53d1362fbeef91 100644
--- a/smoke-test/tests/assertions/assertions_test.py
+++ b/smoke-test/tests/assertions/assertions_test.py
@@ -3,6 +3,7 @@
 
 import pytest
 import tenacity
+
 from datahub.emitter.mce_builder import make_dataset_urn, make_schema_field_urn
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.ingestion.api.common import PipelineContext, RecordEnvelope
@@ -22,7 +23,6 @@
     PartitionSpecClass,
     PartitionTypeClass,
 )
-
 from tests.utils import delete_urns_from_file, get_sleep_info, ingest_file_via_rest
 
 restli_default_headers = {
diff --git a/smoke-test/tests/assertions/custom_assertions_test.py b/smoke-test/tests/assertions/custom_assertions_test.py
index 25336185185645..78e06dfe1be2fa 100644
--- a/smoke-test/tests/assertions/custom_assertions_test.py
+++ b/smoke-test/tests/assertions/custom_assertions_test.py
@@ -2,11 +2,11 @@
 from typing import Any
 
 import pytest
+
 from datahub.emitter.mce_builder import make_dataset_urn
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.ingestion.graph.client import DataHubGraph
 from datahub.metadata.schema_classes import StatusClass
-
 from tests.consistency_utils import wait_for_writes_to_sync
 from tests.utils import delete_urn
 
diff --git a/smoke-test/tests/cli/datahub_cli.py b/smoke-test/tests/cli/datahub_cli.py
index 099002195aa51a..51a43f47fd9e2f 100644
--- a/smoke-test/tests/cli/datahub_cli.py
+++ b/smoke-test/tests/cli/datahub_cli.py
@@ -1,11 +1,11 @@
 import json
 
 import pytest
+
 from datahub.metadata.schema_classes import (
     BrowsePathsV2Class,
     EditableDatasetPropertiesClass,
 )
-
 from tests.utils import ingest_file_via_rest, wait_for_writes_to_sync
 
 ingested_dataset_run_id = ""
diff --git a/smoke-test/tests/cli/datahub_graph_test.py b/smoke-test/tests/cli/datahub_graph_test.py
index e29492c0d3f9d9..0285f7d83b8272 100644
--- a/smoke-test/tests/cli/datahub_graph_test.py
+++ b/smoke-test/tests/cli/datahub_graph_test.py
@@ -2,9 +2,9 @@
 
 import pytest
 import tenacity
+
 from datahub.ingestion.graph.client import DataHubGraph
 from datahub.metadata.schema_classes import KafkaSchemaClass, SchemaMetadataClass
-
 from tests.utils import delete_urns_from_file, get_sleep_info, ingest_file_via_rest
 
 sleep_sec, sleep_times = get_sleep_info()
diff --git a/smoke-test/tests/cli/delete_cmd/test_timeseries_delete.py b/smoke-test/tests/cli/delete_cmd/test_timeseries_delete.py
index 8c526e32c70536..745224a099fc53 100644
--- a/smoke-test/tests/cli/delete_cmd/test_timeseries_delete.py
+++ b/smoke-test/tests/cli/delete_cmd/test_timeseries_delete.py
@@ -5,12 +5,12 @@
 from json import JSONDecodeError
 from typing import Any, Dict, List, Optional
 
-import datahub.emitter.mce_builder as builder
 from click.testing import CliRunner, Result
+
+import datahub.emitter.mce_builder as builder
 from datahub.emitter.serialization_helper import pre_json_transform
 from datahub.entrypoints import datahub
 from datahub.metadata.schema_classes import DatasetProfileClass
-
 from tests.aspect_generators.timeseries.dataset_profile_gen import gen_dataset_profiles
 from tests.utils import get_strftime_from_timestamp_millis, wait_for_writes_to_sync
 
@@ -97,6 +97,7 @@ def datahub_delete(auth_session, params: List[str]) -> None:
     args: List[str] = ["delete"]
     args.extend(params)
     args.append("--hard")
+    logger.info(f"Running delete command with args: {args}")
     delete_result: Result = runner.invoke(
         datahub,
         args,
diff --git a/smoke-test/tests/cli/ingest_cmd/test_timeseries_rollback.py b/smoke-test/tests/cli/ingest_cmd/test_timeseries_rollback.py
index 4d148041e1a0be..f78ae89d7feca2 100644
--- a/smoke-test/tests/cli/ingest_cmd/test_timeseries_rollback.py
+++ b/smoke-test/tests/cli/ingest_cmd/test_timeseries_rollback.py
@@ -1,12 +1,12 @@
 import json
 from typing import Dict, List, Optional
 
-import datahub.emitter.mce_builder as builder
 from click.testing import CliRunner, Result
+
+import datahub.emitter.mce_builder as builder
 from datahub.emitter.serialization_helper import post_json_transform
 from datahub.entrypoints import datahub
 from datahub.metadata.schema_classes import DatasetProfileClass
-
 from tests.utils import ingest_file_via_rest, wait_for_writes_to_sync
 
 runner = CliRunner(mix_stderr=False)
diff --git a/smoke-test/tests/cli/user_groups_cmd/test_group_cmd.py b/smoke-test/tests/cli/user_groups_cmd/test_group_cmd.py
index 502848a722902e..5c5e0c62f6b10e 100644
--- a/smoke-test/tests/cli/user_groups_cmd/test_group_cmd.py
+++ b/smoke-test/tests/cli/user_groups_cmd/test_group_cmd.py
@@ -5,10 +5,10 @@
 
 import yaml
 from click.testing import CliRunner, Result
+
 from datahub.api.entities.corpgroup.corpgroup import CorpGroup
 from datahub.entrypoints import datahub
 from datahub.ingestion.graph.client import DataHubGraph
-
 from tests.utils import wait_for_writes_to_sync
 
 runner = CliRunner(mix_stderr=False)
diff --git a/smoke-test/tests/cli/user_groups_cmd/test_user_cmd.py b/smoke-test/tests/cli/user_groups_cmd/test_user_cmd.py
index 3d0902451ebfcd..12f349bed38055 100644
--- a/smoke-test/tests/cli/user_groups_cmd/test_user_cmd.py
+++ b/smoke-test/tests/cli/user_groups_cmd/test_user_cmd.py
@@ -5,6 +5,7 @@
 
 import yaml
 from click.testing import CliRunner, Result
+
 from datahub.api.entities.corpuser.corpuser import CorpUser
 from datahub.entrypoints import datahub
 
diff --git a/smoke-test/tests/cypress/integration_test.py b/smoke-test/tests/cypress/integration_test.py
index fa7274158f9e23..36fdd53a30c361 100644
--- a/smoke-test/tests/cypress/integration_test.py
+++ b/smoke-test/tests/cypress/integration_test.py
@@ -172,7 +172,7 @@ def ingest_cleanup_data(auth_session, graph_client):
 
 def _get_js_files(base_path: str):
     file_paths = []
-    for root, dirs, files in os.walk(base_path):
+    for root, _, files in os.walk(base_path):
         for file in files:
             if file.endswith(".js"):
                 file_paths.append(os.path.relpath(os.path.join(root, file), base_path))
diff --git a/smoke-test/tests/data_process_instance/test_data_process_instance.py b/smoke-test/tests/data_process_instance/test_data_process_instance.py
index a8aca6034d5be1..c0d2c62ba404a9 100644
--- a/smoke-test/tests/data_process_instance/test_data_process_instance.py
+++ b/smoke-test/tests/data_process_instance/test_data_process_instance.py
@@ -4,6 +4,8 @@
 from random import randint
 
 import pytest
+
+import datahub.metadata.schema_classes as models
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.ingestion.api.common import PipelineContext, RecordEnvelope
 from datahub.ingestion.api.sink import NoopWriteCallback
@@ -14,7 +16,6 @@
     ContainerPropertiesClass,
     DataPlatformInstanceClass,
     DataPlatformInstancePropertiesClass,
-    DataProcessInstanceKeyClass,
     DataProcessInstancePropertiesClass,
     DataProcessInstanceRunEventClass,
     MLHyperParamClass,
@@ -23,7 +24,6 @@
     SubTypesClass,
     TimeWindowSizeClass,
 )
-
 from tests.utils import (
     delete_urns_from_file,
     ingest_file_via_rest,
@@ -54,109 +54,114 @@ def close(self):
         self.sink.close()
 
 
+def create_status_mcp(entity_urn: str):
+    return MetadataChangeProposalWrapper(
+        entityUrn=entity_urn,
+        aspect=models.StatusClass(removed=False),
+    )
+
+
 def create_test_data(filename: str):
+    input_dataset_urn = (
+        "urn:li:dataset:(urn:li:dataPlatform:kafka,my_input_dataset,PROD)"
+    )
+    input_model_urn = "urn:li:mlModel:(urn:li:dataPlatform:mlflow,my_input_model,PROD)"
+    output_dataset_urn = (
+        "urn:li:dataset:(urn:li:dataPlatform:kafka,my_output_dataset,PROD)"
+    )
+    output_model_urn = (
+        "urn:li:mlModel:(urn:li:dataPlatform:mlflow,my_output_model,PROD)"
+    )
+    data_platform_instance_urn = (
+        "urn:li:dataPlatformInstance:(urn:li:dataPlatform:airflow,1234567890)"
+    )
+    container_urn = "urn:li:container:testGroup1"
     mcps = [
-        # Key aspect
-        MetadataChangeProposalWrapper(
-            entityType="dataProcessInstance",
-            entityUrn=dpi_urn,
-            aspectName="dataProcessInstanceKey",
-            aspect=DataProcessInstanceKeyClass(id=dpi_id),
-        ),
-        # Properties aspect
-        MetadataChangeProposalWrapper(
-            entityType="dataProcessInstance",
-            entityUrn=dpi_urn,
-            aspectName="dataProcessInstanceProperties",
-            aspect=DataProcessInstancePropertiesClass(
-                name="Test Pipeline Run",
-                type="BATCH_SCHEDULED",
-                created=AuditStampClass(
-                    time=1640692800000, actor="urn:li:corpuser:datahub"
-                ),
-            ),
-        ),
-        # Run Event aspect
-        MetadataChangeProposalWrapper(
-            entityType="dataProcessInstance",
-            entityUrn=dpi_urn,
-            aspectName="dataProcessInstanceRunEvent",
-            aspect=DataProcessInstanceRunEventClass(
-                timestampMillis=1704067200000,
-                eventGranularity=TimeWindowSizeClass(unit="WEEK", multiple=1),
-                status="COMPLETE",
-            ),
-        ),
-        # Platform Instance aspect
-        MetadataChangeProposalWrapper(
-            entityType="dataProcessInstance",
-            entityUrn=dpi_urn,
-            aspectName="dataPlatformInstance",
-            aspect=DataPlatformInstanceClass(
-                platform="urn:li:dataPlatform:airflow",
-                instance="urn:li:dataPlatformInstance:(urn:li:dataPlatform:airflow,1234567890)",
-            ),
-        ),
-        MetadataChangeProposalWrapper(
-            entityType="dataPlatformInstance",
-            entityUrn="urn:li:dataPlatformInstance:(urn:li:dataPlatform:airflow,1234567890)",
-            aspectName="dataPlatformInstanceProperties",
-            aspect=DataPlatformInstancePropertiesClass(
-                name="my process instance",
-            ),
-        ),
-        # SubTypes aspect
-        MetadataChangeProposalWrapper(
-            entityType="dataProcessInstance",
-            entityUrn=dpi_urn,
-            aspectName="subTypes",
-            aspect=SubTypesClass(typeNames=["TEST", "BATCH_JOB"]),
-        ),
-        # Container aspect
-        MetadataChangeProposalWrapper(
-            entityType="dataProcessInstance",
-            entityUrn=dpi_urn,
-            aspectName="container",
-            aspect=ContainerClass(container="urn:li:container:testGroup1"),
-        ),
+        create_status_mcp(urn)
+        for urn in [
+            input_dataset_urn,
+            input_model_urn,
+            output_dataset_urn,
+            output_model_urn,
+            data_platform_instance_urn,
+        ]
+    ]
+    mcps += [
         MetadataChangeProposalWrapper(
-            entityType="container",
-            entityUrn="urn:li:container:testGroup1",
-            aspectName="containerProperties",
+            entityUrn=container_urn,
             aspect=ContainerPropertiesClass(name="testGroup1"),
-        ),
-        # ML Training Run Properties aspect
+        )
+    ]
+    mcps += [
         MetadataChangeProposalWrapper(
-            entityType="dataProcessInstance",
+            entityUrn=data_platform_instance_urn,
+            aspect=DataPlatformInstancePropertiesClass(name="my process instance"),
+        )
+    ]
+    mcps += [
+        e
+        for e in MetadataChangeProposalWrapper.construct_many(
             entityUrn=dpi_urn,
-            aspectName="mlTrainingRunProperties",
-            aspect=MLTrainingRunPropertiesClass(
-                id="test-training-run-123",
-                trainingMetrics=[
-                    MLMetricClass(
-                        name="accuracy",
-                        description="accuracy of the model",
-                        value="0.95",
+            aspects=[
+                # Properties aspect
+                DataProcessInstancePropertiesClass(
+                    name="Test Pipeline Run",
+                    type="BATCH_SCHEDULED",
+                    created=AuditStampClass(
+                        time=1640692800000, actor="urn:li:corpuser:datahub"
                     ),
-                    MLMetricClass(
-                        name="loss",
-                        description="accuracy loss of the model",
-                        value="0.05",
-                    ),
-                ],
-                hyperParams=[
-                    MLHyperParamClass(
-                        name="learningRate",
-                        description="rate of learning",
-                        value="0.001",
-                    ),
-                    MLHyperParamClass(
-                        name="batchSize", description="size of the batch", value="32"
-                    ),
-                ],
-                outputUrls=["s3://my-bucket/ml/output"],
-            ),
-        ),
+                ),
+                # # Run Event aspect
+                DataProcessInstanceRunEventClass(
+                    timestampMillis=1704067200000,
+                    eventGranularity=TimeWindowSizeClass(unit="WEEK", multiple=1),
+                    status="COMPLETE",
+                ),
+                # Platform Instance aspect
+                DataPlatformInstanceClass(
+                    platform="urn:li:dataPlatform:airflow",
+                    instance="urn:li:dataPlatformInstance:(urn:li:dataPlatform:airflow,1234567890)",
+                ),
+                # SubTypes aspect
+                SubTypesClass(typeNames=["TEST", "BATCH_JOB"]),
+                ContainerClass(container="urn:li:container:testGroup1"),
+                # ML Training Run Properties aspect
+                MLTrainingRunPropertiesClass(
+                    id="test-training-run-123",
+                    trainingMetrics=[
+                        MLMetricClass(
+                            name="accuracy",
+                            description="accuracy of the model",
+                            value="0.95",
+                        ),
+                        MLMetricClass(
+                            name="loss",
+                            description="accuracy loss of the model",
+                            value="0.05",
+                        ),
+                    ],
+                    hyperParams=[
+                        MLHyperParamClass(
+                            name="learningRate",
+                            description="rate of learning",
+                            value="0.001",
+                        ),
+                        MLHyperParamClass(
+                            name="batchSize",
+                            description="size of the batch",
+                            value="32",
+                        ),
+                    ],
+                    outputUrls=["s3://my-bucket/ml/output"],
+                ),
+                models.DataProcessInstanceInputClass(
+                    inputs=[input_dataset_urn, input_model_urn]
+                ),
+                models.DataProcessInstanceOutputClass(
+                    outputs=[output_dataset_urn, output_model_urn]
+                ),
+            ],
+        )
     ]
 
     file_emitter = FileEmitter(filename)
@@ -181,7 +186,7 @@ def ingest_cleanup_data(auth_session, graph_client, request):
         os.remove(filename)
 
 
-@pytest.mark.integration
+# @pytest.mark.integration
 def test_search_dpi(auth_session, ingest_cleanup_data):
     """Test DPI search and validation of returned fields using GraphQL."""
 
@@ -211,13 +216,6 @@ def test_search_dpi(auth_session, ingest_cleanup_data):
                             container {
                                 urn
                             }
-                            platform {
-                                urn
-                                name
-                                properties {
-                                    type
-                                }
-                            }
                             mlTrainingRunProperties {
                                 id
                                 trainingMetrics {
@@ -250,14 +248,14 @@ def test_search_dpi(auth_session, ingest_cleanup_data):
     assert res_data, "Response should not be empty"
     assert "data" in res_data, "Response should contain 'data' field"
     print("RESPONSE DATA:" + str(res_data))
-    assert (
-        "scrollAcrossEntities" in res_data["data"]
-    ), "Response should contain 'scrollAcrossEntities' field"
+    assert "scrollAcrossEntities" in res_data["data"], (
+        "Response should contain 'scrollAcrossEntities' field"
+    )
 
     search_results = res_data["data"]["scrollAcrossEntities"]
-    assert (
-        "searchResults" in search_results
-    ), "Response should contain 'searchResults' field"
+    assert "searchResults" in search_results, (
+        "Response should contain 'searchResults' field"
+    )
 
     results = search_results["searchResults"]
     assert len(results) > 0, "Should find at least one result"
diff --git a/smoke-test/tests/database/test_database.py b/smoke-test/tests/database/test_database.py
index 656c9e98887bf0..df7dc797cfdbad 100644
--- a/smoke-test/tests/database/test_database.py
+++ b/smoke-test/tests/database/test_database.py
@@ -1,8 +1,8 @@
 import logging
 
 import pytest
-from datahub.emitter.mce_builder import make_dataset_urn
 
+from datahub.emitter.mce_builder import make_dataset_urn
 from tests.utilities.concurrent_openapi import run_tests
 from tests.utils import delete_urns, wait_for_writes_to_sync
 
diff --git a/smoke-test/tests/dataproduct/test_dataproduct.py b/smoke-test/tests/dataproduct/test_dataproduct.py
index c26b443dd34a72..8d484820d1ed45 100644
--- a/smoke-test/tests/dataproduct/test_dataproduct.py
+++ b/smoke-test/tests/dataproduct/test_dataproduct.py
@@ -6,6 +6,7 @@
 
 import pytest
 import tenacity
+
 from datahub.emitter.mce_builder import datahub_guid, make_dataset_urn
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.ingestion.api.common import PipelineContext, RecordEnvelope
@@ -19,7 +20,6 @@
     DomainsClass,
 )
 from datahub.utilities.urns.urn import Urn
-
 from tests.utils import (
     delete_urns_from_file,
     get_sleep_info,
@@ -135,9 +135,9 @@ def validate_relationships(
                 urn_match[dataset_urn] = True
 
     urns_missing = [k for k in urn_match if urn_match[k] is False]
-    assert (
-        urns_missing == []
-    ), "All dataset urns should have a DataProductContains relationship to the data product"
+    assert urns_missing == [], (
+        "All dataset urns should have a DataProductContains relationship to the data product"
+    )
 
     dataset_urns_matched = set()
     for e in graph_client.get_related_entities(
@@ -147,9 +147,9 @@ def validate_relationships(
     ):
         dataset_urns_matched.add(e.urn)
 
-    assert (
-        set(dataset_urns) == dataset_urns_matched
-    ), "All dataset urns should be navigable from the data product"
+    assert set(dataset_urns) == dataset_urns_matched, (
+        "All dataset urns should be navigable from the data product"
+    )
 
 
 @tenacity.retry(
@@ -247,6 +247,6 @@ def test_create_data_product(graph_client, ingest_cleanup_data):
                 urn_match[dataset_urn] = True
 
     urns_missing = [k for k in urn_match if urn_match[k] is False]
-    assert set(urns_missing) == set(
-        dataset_urns
-    ), f"All dataset urns should no longer have a DataProductContains relationship to the data product {data_product_urn}"
+    assert set(urns_missing) == set(dataset_urns), (
+        f"All dataset urns should no longer have a DataProductContains relationship to the data product {data_product_urn}"
+    )
diff --git a/smoke-test/tests/delete/delete_test.py b/smoke-test/tests/delete/delete_test.py
index a13fc883b4055d..e7fc5c793dc7d8 100644
--- a/smoke-test/tests/delete/delete_test.py
+++ b/smoke-test/tests/delete/delete_test.py
@@ -2,8 +2,8 @@
 import os
 
 import pytest
-from datahub.cli.cli_utils import get_aspects_for_entity
 
+from datahub.cli.cli_utils import get_aspects_for_entity
 from tests.utils import (
     delete_urns_from_file,
     ingest_file_via_rest,
diff --git a/smoke-test/tests/lineage/test_lineage.py b/smoke-test/tests/lineage/test_lineage.py
index 4a43fd591ae2e1..dd5309667806c5 100644
--- a/smoke-test/tests/lineage/test_lineage.py
+++ b/smoke-test/tests/lineage/test_lineage.py
@@ -3,9 +3,11 @@
 from enum import Enum
 from typing import Any, Dict, Iterable, List, Optional, Set, Tuple, Union
 
-import datahub.emitter.mce_builder as builder
 import networkx as nx
 import pytest
+from pydantic import BaseModel, validator
+
+import datahub.emitter.mce_builder as builder
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.ingestion.graph.client import DataHubGraph
 from datahub.metadata.schema_classes import (
@@ -18,17 +20,9 @@
     DatasetLineageTypeClass,
     DatasetPropertiesClass,
     EdgeClass,
-)
-from datahub.metadata.schema_classes import (
     FineGrainedLineageClass as FineGrainedLineage,
-)
-from datahub.metadata.schema_classes import (
     FineGrainedLineageDownstreamTypeClass as FineGrainedLineageDownstreamType,
-)
-from datahub.metadata.schema_classes import (
     FineGrainedLineageUpstreamTypeClass as FineGrainedLineageUpstreamType,
-)
-from datahub.metadata.schema_classes import (
     OtherSchemaClass,
     QueryLanguageClass,
     QueryPropertiesClass,
@@ -43,8 +37,6 @@
 )
 from datahub.utilities.urns.dataset_urn import DatasetUrn
 from datahub.utilities.urns.urn import Urn
-from pydantic import BaseModel, validator
-
 from tests.utils import ingest_file_via_rest, wait_for_writes_to_sync
 
 logger = logging.getLogger(__name__)
@@ -100,7 +92,7 @@ def _explain_sal_result(result: dict) -> str:
         explain += "Entities: "
         try:
             for e in entities:
-                explain += f"\t{e.replace('urn:li:','')}\n"
+                explain += f"\t{e.replace('urn:li:', '')}\n"
             for entity in entities:
                 paths = [
                     x["paths"][0]["path"]
@@ -357,9 +349,9 @@ def get_expectation_for_query(self, query: ImpactQuery) -> LineageExpectation:
                             lineage_expectation.impacted_entities[impacted_entity]
                         )
                     else:
-                        entries_to_add[
-                            impacted_dataset_entity
-                        ] = lineage_expectation.impacted_entities[impacted_entity]
+                        entries_to_add[impacted_dataset_entity] = (
+                            lineage_expectation.impacted_entities[impacted_entity]
+                        )
                     entries_to_remove.append(impacted_entity)
             for impacted_entity in entries_to_remove:
                 del lineage_expectation.impacted_entities[impacted_entity]
@@ -764,9 +756,9 @@ def test_expectation(self, graph: DataHubGraph) -> bool:
                     ]
                 )
                 try:
-                    assert (
-                        impacted_entities == impacted_entities_expectation
-                    ), f"Expected impacted entities to be {impacted_entities_expectation}, found {impacted_entities}"
+                    assert impacted_entities == impacted_entities_expectation, (
+                        f"Expected impacted entities to be {impacted_entities_expectation}, found {impacted_entities}"
+                    )
                 except Exception:
                     # breakpoint()
                     raise
@@ -791,10 +783,14 @@ def test_expectation(self, graph: DataHubGraph) -> bool:
                     try:
                         assert len(impacted_entity_paths) == len(
                             expectation.impacted_entities[impacted_entity]
-                        ), f"Expected length of impacted entity paths to be {len(expectation.impacted_entities[impacted_entity])}, found {len(impacted_entity_paths)}"
+                        ), (
+                            f"Expected length of impacted entity paths to be {len(expectation.impacted_entities[impacted_entity])}, found {len(impacted_entity_paths)}"
+                        )
                         assert set(impacted_entity_paths) == set(
                             expectation.impacted_entities[impacted_entity]
-                        ), f"Expected impacted entity paths to be {expectation.impacted_entities[impacted_entity]}, found {impacted_entity_paths}"
+                        ), (
+                            f"Expected impacted entity paths to be {expectation.impacted_entities[impacted_entity]}, found {impacted_entity_paths}"
+                        )
                     except Exception:
                         # breakpoint()
                         raise
diff --git a/smoke-test/tests/managed_ingestion/managed_ingestion_test.py b/smoke-test/tests/managed_ingestion/managed_ingestion_test.py
index 5d6179de6be644..7fe9421af85b1a 100644
--- a/smoke-test/tests/managed_ingestion/managed_ingestion_test.py
+++ b/smoke-test/tests/managed_ingestion/managed_ingestion_test.py
@@ -489,9 +489,9 @@ def test_create_list_get_ingestion_execution_request(auth_session):
 
     assert res_data
     assert res_data["data"]
-    assert (
-        res_data["data"]["createIngestionExecutionRequest"] is not None
-    ), f"res_data was {res_data}"
+    assert res_data["data"]["createIngestionExecutionRequest"] is not None, (
+        f"res_data was {res_data}"
+    )
     assert "errors" not in res_data
 
     execution_request_urn = res_data["data"]["createIngestionExecutionRequest"]
diff --git a/smoke-test/tests/ml_models/test_ml_models.py b/smoke-test/tests/ml_models/test_ml_models.py
index 59821ab3e3cc41..534f40035083f2 100644
--- a/smoke-test/tests/ml_models/test_ml_models.py
+++ b/smoke-test/tests/ml_models/test_ml_models.py
@@ -4,6 +4,7 @@
 from random import randint
 
 import pytest
+
 from datahub.emitter.mce_builder import make_ml_model_group_urn, make_ml_model_urn
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.ingestion.api.common import PipelineContext, RecordEnvelope
@@ -14,7 +15,6 @@
     MLModelGroupPropertiesClass,
     MLModelPropertiesClass,
 )
-
 from tests.utils import (
     delete_urns_from_file,
     get_sleep_info,
diff --git a/smoke-test/tests/patch/test_datajob_patches.py b/smoke-test/tests/patch/test_datajob_patches.py
index e025c8a2aebeb5..68b08f2b1ade72 100644
--- a/smoke-test/tests/patch/test_datajob_patches.py
+++ b/smoke-test/tests/patch/test_datajob_patches.py
@@ -10,7 +10,6 @@
     EdgeClass,
 )
 from datahub.specific.datajob import DataJobPatchBuilder
-
 from tests.patch.common_patch_tests import (
     helper_test_custom_properties_patch,
     helper_test_dataset_tags_patch,
diff --git a/smoke-test/tests/patch/test_dataset_patches.py b/smoke-test/tests/patch/test_dataset_patches.py
index bafacbc8d2b0b4..cb131f4356f429 100644
--- a/smoke-test/tests/patch/test_dataset_patches.py
+++ b/smoke-test/tests/patch/test_dataset_patches.py
@@ -15,7 +15,6 @@
     UpstreamLineageClass,
 )
 from datahub.specific.dataset import DatasetPatchBuilder
-
 from tests.patch.common_patch_tests import (
     helper_test_custom_properties_patch,
     helper_test_dataset_tags_patch,
diff --git a/smoke-test/tests/platform_resources/test_platform_resource.py b/smoke-test/tests/platform_resources/test_platform_resource.py
index 23a147bda695b3..191f8c44694985 100644
--- a/smoke-test/tests/platform_resources/test_platform_resource.py
+++ b/smoke-test/tests/platform_resources/test_platform_resource.py
@@ -4,13 +4,13 @@
 from typing import List
 
 import pytest
+
 from datahub.api.entities.platformresource.platform_resource import (
     ElasticPlatformResourceQuery,
     PlatformResource,
     PlatformResourceKey,
     PlatformResourceSearchFields,
 )
-
 from tests.utils import wait_for_writes_to_sync
 
 logger = logging.getLogger(__name__)
diff --git a/smoke-test/tests/read_only/test_search.py b/smoke-test/tests/read_only/test_search.py
index 36ecf68395f919..66bbeb408d0529 100644
--- a/smoke-test/tests/read_only/test_search.py
+++ b/smoke-test/tests/read_only/test_search.py
@@ -153,6 +153,6 @@ def test_openapi_v3_entity(auth_session, entity_type):
 
     expected_data = {"urn": first_urn}
 
-    assert (
-        actual_data["urn"] == expected_data["urn"]
-    ), f"Mismatch: expected {expected_data}, got {actual_data}"
+    assert actual_data["urn"] == expected_data["urn"], (
+        f"Mismatch: expected {expected_data}, got {actual_data}"
+    )
diff --git a/smoke-test/tests/read_only/test_services_up.py b/smoke-test/tests/read_only/test_services_up.py
index 12ff04965548f0..79812b46476fa8 100644
--- a/smoke-test/tests/read_only/test_services_up.py
+++ b/smoke-test/tests/read_only/test_services_up.py
@@ -27,6 +27,6 @@ def test_gms_config_accessible(auth_session) -> None:
     default_cli_version: str = gms_config["managedIngestion"]["defaultCliVersion"]
     print(f"Default CLI version: {default_cli_version}")
     assert not default_cli_version.startswith("@")
-    assert "." in default_cli_version or looks_like_a_short_sha(
-        default_cli_version
-    ), "Default CLI version does not look like a version string"
+    assert "." in default_cli_version or looks_like_a_short_sha(default_cli_version), (
+        "Default CLI version does not look like a version string"
+    )
diff --git a/smoke-test/tests/restli/restli_test.py b/smoke-test/tests/restli/restli_test.py
index 7e9cf6aa76235c..c10ae3fe584f08 100644
--- a/smoke-test/tests/restli/restli_test.py
+++ b/smoke-test/tests/restli/restli_test.py
@@ -4,6 +4,7 @@
 from typing import List
 
 import pytest
+
 from datahub.emitter.aspect import JSON_CONTENT_TYPE
 from datahub.emitter.mce_builder import make_dashboard_urn
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
@@ -16,7 +17,6 @@
     MetadataChangeProposalClass,
 )
 from datahub.utilities.urns.urn import guess_entity_type
-
 from tests.utils import delete_urns
 
 generated_urns: List[str] = []
@@ -51,7 +51,7 @@ def make_mcp(self) -> MetadataChangeProposalClass:
         return mcp
 
 
-@pytest.fixture(scope="module")
+@pytest.fixture(scope="module", autouse=True)
 def ingest_cleanup_data(auth_session, graph_client, request):
     yield
     delete_urns(graph_client, generated_urns)
diff --git a/smoke-test/tests/restli/test_restli_batch_ingestion.py b/smoke-test/tests/restli/test_restli_batch_ingestion.py
new file mode 100644
index 00000000000000..ab33a2b26605db
--- /dev/null
+++ b/smoke-test/tests/restli/test_restli_batch_ingestion.py
@@ -0,0 +1,187 @@
+import time
+from typing import List
+
+import pytest
+
+import datahub.metadata.schema_classes as models
+from datahub.emitter.mce_builder import make_dashboard_urn
+from datahub.emitter.mcp import MetadataChangeProposalWrapper
+from datahub.emitter.serialization_helper import pre_json_transform
+from datahub.metadata._schema_classes import MetadataChangeProposalClass
+from datahub.metadata.schema_classes import (
+    AuditStampClass,
+    ChangeAuditStampsClass,
+    DashboardInfoClass,
+)
+from datahub.metadata.urns import MlModelUrn
+from tests.consistency_utils import wait_for_writes_to_sync
+from tests.restli.restli_test import MetadataChangeProposalInvalidWrapper
+from tests.utils import delete_urns
+
+generated_urns: List[str] = []
+
+
+@pytest.fixture(scope="module", autouse=True)
+def ingest_cleanup_data(auth_session, graph_client, request):
+    yield
+    delete_urns(graph_client, generated_urns)
+
+
+def _create_valid_dashboard_mcps() -> List[MetadataChangeProposalClass]:
+    mcps = []
+    num_valid_mcp = 5
+
+    audit_stamp = pre_json_transform(
+        ChangeAuditStampsClass(
+            created=AuditStampClass(
+                time=int(time.time() * 1000),
+                actor="urn:li:corpuser:datahub",
+            )
+        ).to_obj()
+    )
+
+    valid_dashboard_info = DashboardInfoClass(
+        title="Dummy Title For Testing",
+        description="Dummy Description For Testing",
+        lastModified=audit_stamp,
+    )
+
+    for i in range(num_valid_mcp):
+        mcp_valid = MetadataChangeProposalWrapper(
+            entityUrn=make_dashboard_urn(
+                platform="looker", name=f"dummy-test-invalid-{i}"
+            ),
+            aspectName="dashboardInfo",
+            aspect=valid_dashboard_info,
+        )
+        mcps.append(mcp_valid.make_mcp())
+    generated_urns.extend([mcp.entityUrn for mcp in mcps if mcp.entityUrn])
+
+    return mcps
+
+
+def _create_invalid_dashboard_mcp() -> MetadataChangeProposalClass:
+    audit_stamp = pre_json_transform(
+        ChangeAuditStampsClass(
+            created=AuditStampClass(
+                time=int(time.time() * 1000),
+                actor="urn:li:corpuser:datahub",
+            )
+        ).to_obj()
+    )
+
+    invalid_dashboard_info = {
+        "title": "Dummy Title For Testing",
+        "description": "Dummy Description For Testing",
+        "lastModified": audit_stamp,
+        "notValidField": "invalid field value",
+    }
+
+    mcp_invalid = MetadataChangeProposalInvalidWrapper(
+        entityUrn=make_dashboard_urn(platform="looker", name="dummy-test-valid"),
+        aspectName="dashboardInfo",
+        aspect=invalid_dashboard_info,
+    )
+    generated_urns.append(mcp_invalid.entityUrn) if mcp_invalid.entityUrn else None
+    return mcp_invalid.make_mcp()
+
+
+def _create_invalid_dataset_mcps() -> List[MetadataChangeProposalWrapper]:
+    dataset_urn = "urn:li:dataset:(urn:li:dataPlatform:kafka,my_dataset,PROD)"
+    model_urn = MlModelUrn("mlflow", "my_model", "PROD").urn()
+    bad_mcps = [
+        MetadataChangeProposalWrapper(
+            entityUrn=dataset_urn,
+            aspect=models.StatusClass(removed=False),
+        ),
+        MetadataChangeProposalWrapper(
+            entityUrn=dataset_urn,
+            aspect=models.UpstreamLineageClass(
+                upstreams=[
+                    models.UpstreamClass(
+                        dataset=model_urn,
+                        type=models.DatasetLineageTypeClass.TRANSFORMED,
+                    )
+                ]
+            ),
+        ),
+    ]
+    return bad_mcps
+
+
+def test_restli_batch_ingestion_sync(graph_client):
+    # Positive Test (all valid MetadataChangeProposal)
+    mcps = _create_valid_dashboard_mcps()
+    ret = graph_client.emit_mcps(mcps, async_flag=False)
+    assert ret >= 0
+
+    # Negative Test (contains invalid MetadataChangeProposal)
+    invalid_mcp = _create_invalid_dashboard_mcp()
+    mcps.append(invalid_mcp)
+    ret = graph_client.emit_mcps(mcps, async_flag=False)
+    assert ret >= 0
+
+    # Expected that invalid field of MetadataChangeProposal is ignored,
+    # Rest Fields are persistd into DB
+    aspect = graph_client.get_aspect(
+        entity_urn=invalid_mcp.entityUrn, aspect_type=DashboardInfoClass
+    )
+
+    assert aspect is not None
+    assert isinstance(aspect, DashboardInfoClass)
+    assert aspect.title == "Dummy Title For Testing"
+    assert aspect.description == "Dummy Description For Testing"
+    assert aspect.lastModified is not None
+
+
+def test_restli_batch_ingestion_async(graph_client):
+    # Positive Test (all valid MetadataChangeProposal)
+    mcps = _create_valid_dashboard_mcps()
+    ret = graph_client.emit_mcps(mcps, async_flag=True)
+    assert ret >= 0
+
+    # Negative Test (contains invalid MetadataChangeProposal)
+    invalid_mcp = _create_invalid_dashboard_mcp()
+    mcps.append(invalid_mcp)
+    ret = graph_client.emit_mcps(mcps, async_flag=True)
+    assert ret >= 0
+
+    # Expected that invalid field of MetadataChangeProposal is ignored,
+    # Rest Fields are persistd into DB
+    wait_for_writes_to_sync()
+    aspect = graph_client.get_aspect(
+        entity_urn=invalid_mcp.entityUrn, aspect_type=DashboardInfoClass
+    )
+
+    assert aspect is not None
+    assert isinstance(aspect, DashboardInfoClass)
+    assert aspect.title == "Dummy Title For Testing"
+    assert aspect.description == "Dummy Description For Testing"
+    assert aspect.lastModified is not None
+
+
+def test_restli_batch_ingestion_exception_sync(graph_client):
+    """
+    Test Batch ingestion when an exception occurs in sync mode
+    """
+    bad_mcps = _create_invalid_dataset_mcps()
+    generated_urns.extend([mcp.entityUrn for mcp in bad_mcps if mcp.entityUrn])
+
+    try:
+        graph_client.emit_mcps(bad_mcps, async_flag=False)
+        raise AssertionError("should have thrown an exception")
+    except Exception as e:
+        if isinstance(e, AssertionError):
+            raise e
+        print(f"Error emitting MCPs due to {e}")
+
+
+def test_restli_batch_ingestion_exception_async(graph_client):
+    """
+    Test Batch ingestion when an exception occurs in async mode
+    """
+    bad_mcps = _create_invalid_dataset_mcps()
+    generated_urns.extend([mcp.entityUrn for mcp in bad_mcps if mcp.entityUrn])
+    # TODO expectation is that it throws exception, but it doesn't currently.this test case need to change after fix.
+    ret = graph_client.emit_mcps(bad_mcps, async_flag=True)
+    assert ret >= 0
diff --git a/smoke-test/tests/schema_fields/schema_evolution.py b/smoke-test/tests/schema_fields/schema_evolution.py
index 5481742969f5d1..962b71058150fd 100644
--- a/smoke-test/tests/schema_fields/schema_evolution.py
+++ b/smoke-test/tests/schema_fields/schema_evolution.py
@@ -1,14 +1,14 @@
 import time
 from enum import Enum
 
-import datahub.metadata.schema_classes as models
 import pytest
+from tenacity import retry, stop_after_delay, wait_fixed
+
+import datahub.metadata.schema_classes as models
 from datahub.cli.cli_utils import get_aspects_for_entity
 from datahub.emitter.mce_builder import make_dataset_urn, make_schema_field_urn
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.ingestion.graph.client import DataHubGraph
-from tenacity import retry, stop_after_delay, wait_fixed
-
 from tests.utils import ingest_file_via_rest, wait_for_writes_to_sync
 
 _MAX_DELAY_UNTIL_WRITES_VISIBLE_SECS = 30
@@ -70,7 +70,7 @@ def test_setup(auth_session, graph_client):
 
     ingest_file_via_rest(
         auth_session, "tests/schema_fields/schema_field_side_effect_data.json"
-    ).config.run_id
+    )
 
     assert "schemaMetadata" in get_aspects_for_entity(
         session,
diff --git a/smoke-test/tests/schema_fields/test_schemafields.py b/smoke-test/tests/schema_fields/test_schemafields.py
index 30d705ef2b4e70..329a120a2f8587 100644
--- a/smoke-test/tests/schema_fields/test_schemafields.py
+++ b/smoke-test/tests/schema_fields/test_schemafields.py
@@ -4,15 +4,15 @@
 import time
 from random import randint
 
-import datahub.metadata.schema_classes as models
 import pytest
+
+import datahub.metadata.schema_classes as models
 from datahub.emitter.mce_builder import make_dataset_urn, make_schema_field_urn
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.ingestion.api.common import PipelineContext, RecordEnvelope
 from datahub.ingestion.api.sink import NoopWriteCallback
 from datahub.ingestion.graph.client import DataHubGraph
 from datahub.ingestion.sink.file import FileSink, FileSinkConfig
-
 from tests.utils import (
     delete_urns_from_file,
     get_sleep_info,
diff --git a/smoke-test/tests/setup/lineage/ingest_data_job_change.py b/smoke-test/tests/setup/lineage/ingest_data_job_change.py
index 85924569ba9f0b..8143bd0b519df5 100644
--- a/smoke-test/tests/setup/lineage/ingest_data_job_change.py
+++ b/smoke-test/tests/setup/lineage/ingest_data_job_change.py
@@ -12,7 +12,6 @@
     SchemaFieldDataTypeClass,
     StringTypeClass,
 )
-
 from tests.setup.lineage.constants import (
     AIRFLOW_DATA_PLATFORM,
     SNOWFLAKE_DATA_PLATFORM,
diff --git a/smoke-test/tests/setup/lineage/ingest_dataset_join_change.py b/smoke-test/tests/setup/lineage/ingest_dataset_join_change.py
index 36fe3d19eff41d..aea5318eeba4ff 100644
--- a/smoke-test/tests/setup/lineage/ingest_dataset_join_change.py
+++ b/smoke-test/tests/setup/lineage/ingest_dataset_join_change.py
@@ -8,7 +8,6 @@
     StringTypeClass,
     UpstreamClass,
 )
-
 from tests.setup.lineage.constants import (
     DATASET_ENTITY_TYPE,
     SNOWFLAKE_DATA_PLATFORM,
diff --git a/smoke-test/tests/setup/lineage/ingest_input_datasets_change.py b/smoke-test/tests/setup/lineage/ingest_input_datasets_change.py
index 9ff0cfcc1182fa..623cb9665a3528 100644
--- a/smoke-test/tests/setup/lineage/ingest_input_datasets_change.py
+++ b/smoke-test/tests/setup/lineage/ingest_input_datasets_change.py
@@ -11,7 +11,6 @@
     SchemaFieldDataTypeClass,
     StringTypeClass,
 )
-
 from tests.setup.lineage.constants import (
     AIRFLOW_DATA_PLATFORM,
     BQ_DATA_PLATFORM,
diff --git a/smoke-test/tests/setup/lineage/ingest_time_lineage.py b/smoke-test/tests/setup/lineage/ingest_time_lineage.py
index 3d34ea5426c58c..3597ac5ff738bd 100644
--- a/smoke-test/tests/setup/lineage/ingest_time_lineage.py
+++ b/smoke-test/tests/setup/lineage/ingest_time_lineage.py
@@ -1,7 +1,6 @@
 from typing import List
 
 from datahub.ingestion.graph.client import DataHubGraph
-
 from tests.setup.lineage.ingest_data_job_change import (
     get_data_job_change_urns,
     ingest_data_job_change,
diff --git a/smoke-test/tests/setup/lineage/utils.py b/smoke-test/tests/setup/lineage/utils.py
index a58e73cb927987..6d21337c669f4b 100644
--- a/smoke-test/tests/setup/lineage/utils.py
+++ b/smoke-test/tests/setup/lineage/utils.py
@@ -24,7 +24,6 @@
     SchemaMetadataClass,
     UpstreamClass,
 )
-
 from tests.setup.lineage.constants import (
     DATA_FLOW_ENTITY_TYPE,
     DATA_FLOW_INFO_ASPECT_NAME,
diff --git a/smoke-test/tests/structured_properties/test_structured_properties.py b/smoke-test/tests/structured_properties/test_structured_properties.py
index e3c33aa406efc4..8b0f61d8a5ea26 100644
--- a/smoke-test/tests/structured_properties/test_structured_properties.py
+++ b/smoke-test/tests/structured_properties/test_structured_properties.py
@@ -1,16 +1,18 @@
 import logging
 import os
+import re
 import tempfile
 from random import randint
 from typing import Iterable, List, Optional, Union
 
+import pydantic
 import pytest
 
-# import tenacity
 from datahub.api.entities.dataset.dataset import Dataset
 from datahub.api.entities.structuredproperties.structuredproperties import (
     StructuredProperties,
 )
+from datahub.configuration.common import GraphError, OperationalError
 from datahub.emitter.mce_builder import make_dataset_urn, make_schema_field_urn
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.ingestion.graph.client import DataHubGraph
@@ -24,7 +26,6 @@
 from datahub.specific.dataset import DatasetPatchBuilder
 from datahub.utilities.urns.structured_properties_urn import StructuredPropertyUrn
 from datahub.utilities.urns.urn import Urn
-
 from tests.consistency_utils import wait_for_writes_to_sync
 from tests.utilities.file_emitter import FileEmitter
 from tests.utils import (
@@ -174,10 +175,6 @@ def to_es_filter_name(
         return f"structuredProperties.{qualified_name}"
 
 
-# @tenacity.retry(
-#     stop=tenacity.stop_after_attempt(sleep_times),
-#     wait=tenacity.wait_fixed(sleep_sec),
-# )
 def test_structured_property_string(ingest_cleanup_data, graph_client):
     property_name = f"retention{randint(10, 10000)}Policy"
 
@@ -187,24 +184,13 @@ def test_structured_property_string(ingest_cleanup_data, graph_client):
         dataset_urns[0], property_name, ["30d"], graph=graph_client
     )
 
-    try:
+    with pytest.raises(OperationalError):
+        # Cannot add a number to a string property.
         attach_property_to_entity(
             dataset_urns[0], property_name, 200030, graph=graph_client
         )
-        raise AssertionError(
-            "Should not be able to attach a number to a string property"
-        )
-    except Exception as e:
-        if not isinstance(e, AssertionError):
-            pass
-        else:
-            raise e
 
 
-# @tenacity.retry(
-#     stop=tenacity.stop_after_attempt(sleep_times),
-#     wait=tenacity.wait_fixed(sleep_sec),
-# )
 def test_structured_property_double(ingest_cleanup_data, graph_client):
     property_name = f"expiryTime{randint(10, 10000)}"
 
@@ -214,35 +200,19 @@ def test_structured_property_double(ingest_cleanup_data, graph_client):
         dataset_urns[0], property_name, 2000034, graph=graph_client
     )
 
-    try:
+    with pytest.raises(OperationalError):
+        # Cannot add a string to a number property.
         attach_property_to_entity(
             dataset_urns[0], property_name, "30 days", graph=graph_client
         )
-        raise AssertionError(
-            "Should not be able to attach a string to a number property"
-        )
-    except Exception as e:
-        if not isinstance(e, AssertionError):
-            pass
-        else:
-            raise e
 
-    try:
+    with pytest.raises(OperationalError):
+        # Cannot add a list to a number property.
         attach_property_to_entity(
             dataset_urns[0], property_name, [2000034, 2000035], graph=graph_client
         )
-        raise AssertionError("Should not be able to attach a list to a number property")
-    except Exception as e:
-        if not isinstance(e, AssertionError):
-            pass
-        else:
-            raise e
-
-
-# @tenacity.retry(
-#     stop=tenacity.stop_after_attempt(sleep_times),
-#     wait=tenacity.wait_fixed(sleep_sec),
-# )
+
+
 def test_structured_property_double_multiple(ingest_cleanup_data, graph_client):
     property_name = f"versions{randint(10, 10000)}"
 
@@ -255,10 +225,6 @@ def test_structured_property_double_multiple(ingest_cleanup_data, graph_client):
     )
 
 
-# @tenacity.retry(
-#     stop=tenacity.stop_after_attempt(sleep_times),
-#     wait=tenacity.wait_fixed(sleep_sec),
-# )
 def test_structured_property_string_allowed_values(ingest_cleanup_data, graph_client):
     property_name = f"enumProperty{randint(10, 10000)}"
 
@@ -277,18 +243,13 @@ def test_structured_property_string_allowed_values(ingest_cleanup_data, graph_cl
         dataset_urns[0], property_name, ["foo", "bar"], graph=graph_client
     )
 
-    try:
+    with pytest.raises(
+        OperationalError, match=re.escape("value: {string=baz} should be one of [")
+    ):
+        # Cannot add a value that isn't in the allowed values list.
         attach_property_to_entity(
             dataset_urns[0], property_name, ["foo", "baz"], graph=graph_client
         )
-        raise AssertionError(
-            "Should not be able to attach a value not in allowed values"
-        )
-    except Exception as e:
-        if "value: {string=baz} should be one of [" in str(e):
-            pass
-        else:
-            raise e
 
 
 def test_structured_property_definition_evolution(ingest_cleanup_data, graph_client):
@@ -305,7 +266,8 @@ def test_structured_property_definition_evolution(ingest_cleanup_data, graph_cli
         ],
     )
 
-    try:
+    with pytest.raises(OperationalError):
+        # Cannot change cardinality from MULTIPLE to SINGLE.
         create_property_definition(
             property_name,
             graph_client,
@@ -316,20 +278,8 @@ def test_structured_property_definition_evolution(ingest_cleanup_data, graph_cli
                 PropertyValueClass(value="bar"),
             ],
         )
-        raise AssertionError(
-            "Should not be able to change cardinality from MULTIPLE to SINGLE"
-        )
-    except Exception as e:
-        if isinstance(e, AssertionError):
-            raise e
-        else:
-            pass
 
 
-# @tenacity.retry(
-#     stop=tenacity.stop_after_attempt(sleep_times),
-#     wait=tenacity.wait_fixed(sleep_sec),
-# )
 def test_structured_property_schema_field(ingest_cleanup_data, graph_client):
     property_name = f"deprecationDate{randint(10, 10000)}"
 
@@ -355,7 +305,8 @@ def test_structured_property_schema_field(ingest_cleanup_data, graph_client):
         graph=graph_client,
     ) == ["2020-10-01"]
 
-    try:
+    with pytest.raises(OperationalError):
+        # Cannot add a number to a date property.
         attach_property_to_entity(
             schema_field_urns[0],
             property_name,
@@ -363,30 +314,19 @@ def test_structured_property_schema_field(ingest_cleanup_data, graph_client):
             graph=graph_client,
             namespace="io.datahubproject.test",
         )
-        raise AssertionError("Should not be able to attach a number to a DATE property")
-    except Exception as e:
-        if not isinstance(e, AssertionError):
-            pass
-        else:
-            raise e
 
 
 def test_structured_properties_yaml_load_with_bad_entity_type(
     ingest_cleanup_data, graph_client
 ):
-    try:
+    with pytest.raises(
+        pydantic.ValidationError,
+        match="urn:li:entityType:dataset is not a valid entity type urn",
+    ):
         StructuredProperties.create(
             "tests/structured_properties/bad_entity_type.yaml",
             graph=graph_client,
         )
-        raise AssertionError(
-            "Should not be able to create structured properties with bad entity type"
-        )
-    except Exception as e:
-        if "urn:li:entityType:dataset is not a valid entity type urn" in str(e):
-            pass
-        else:
-            raise e
 
 
 def test_dataset_yaml_loader(ingest_cleanup_data, graph_client):
@@ -618,35 +558,24 @@ def test_dataset_structured_property_soft_delete_validation(
     graph_client.soft_delete_entity(urn=property_urn)
 
     # Attempt to modify soft deleted definition
-    try:
+    with pytest.raises(
+        OperationalError,
+        match="Cannot mutate a soft deleted Structured Property Definition",
+    ):
         create_property_definition(
             property_name=property_name,
             graph=graph_client,
             value_type=value_type,
             cardinality="SINGLE",
         )
-        raise AssertionError(
-            "Should not be able to modify soft deleted structured property"
-        )
-    except Exception as e:
-        if "Cannot mutate a soft deleted Structured Property Definition" in str(e):
-            pass
-        else:
-            raise e
 
     # Attempt to add soft deleted structured property to entity
-    try:
+    with pytest.raises(
+        OperationalError, match="Cannot apply a soft deleted Structured Property value"
+    ):
         attach_property_to_entity(
             dataset_urns[0], property_name, "test string", graph=graph_client
         )
-        raise AssertionError(
-            "Should not be able to apply a soft deleted structured property to another entity"
-        )
-    except Exception as e:
-        if "Cannot apply a soft deleted Structured Property value" in str(e):
-            pass
-        else:
-            raise e
 
 
 def test_dataset_structured_property_soft_delete_read_mutation(
@@ -686,7 +615,7 @@ def test_dataset_structured_property_soft_delete_read_mutation(
 
 
 def test_dataset_structured_property_soft_delete_search_filter_validation(
-    ingest_cleanup_data, graph_client, caplog
+    ingest_cleanup_data, graph_client: DataHubGraph, caplog: pytest.LogCaptureFixture
 ):
     # Create a test structured property
     dataset_property_name = f"softDeleteSearchFilter{randint(10, 10000)}"
@@ -723,7 +652,9 @@ def test_dataset_structured_property_soft_delete_search_filter_validation(
     wait_for_writes_to_sync()
 
     # Perform search, make sure it validates filter and rejects as invalid request
-    try:
+    with pytest.raises(
+        GraphError, match="Cannot filter on deleted Structured Property"
+    ):
         list(
             graph_client.get_urns_by_filter(
                 extraFilters=[
@@ -735,14 +666,6 @@ def test_dataset_structured_property_soft_delete_search_filter_validation(
                 ]
             )
         )
-        raise AssertionError(
-            "Should not be able to filter by soft deleted structured property"
-        )
-    except Exception as e:
-        if "Cannot filter on deleted Structured Property" in str(e):
-            pass
-        else:
-            raise e
 
 
 def test_dataset_structured_property_delete(ingest_cleanup_data, graph_client, caplog):
diff --git a/smoke-test/tests/test_stateful_ingestion.py b/smoke-test/tests/test_stateful_ingestion.py
index 4436cf26c2fd78..819715013ab48f 100644
--- a/smoke-test/tests/test_stateful_ingestion.py
+++ b/smoke-test/tests/test_stateful_ingestion.py
@@ -1,5 +1,8 @@
 from typing import Any, Dict, Optional, cast
 
+from sqlalchemy import create_engine
+from sqlalchemy.sql import text
+
 from datahub.ingestion.api.committable import StatefulCommittable
 from datahub.ingestion.run.pipeline import Pipeline
 from datahub.ingestion.source.sql.mysql import MySQLConfig, MySQLSource
@@ -8,9 +11,6 @@
 from datahub.ingestion.source.state.stale_entity_removal_handler import (
     StaleEntityRemovalHandler,
 )
-from sqlalchemy import create_engine
-from sqlalchemy.sql import text
-
 from tests.utils import get_mysql_password, get_mysql_url, get_mysql_username
 
 
@@ -36,7 +36,7 @@ def run_and_get_pipeline(pipeline_config_dict: Dict[str, Any]) -> Pipeline:
 
     def validate_all_providers_have_committed_successfully(pipeline: Pipeline) -> None:
         provider_count: int = 0
-        for name, provider in pipeline.ctx.get_committables():
+        for _, provider in pipeline.ctx.get_committables():
             provider_count += 1
             assert isinstance(provider, StatefulCommittable)
             stateful_committable = cast(StatefulCommittable, provider)
diff --git a/smoke-test/tests/timeline/timeline_test.py b/smoke-test/tests/timeline/timeline_test.py
index d33ad3b24ceb5c..ddb5f8faf4a9b2 100644
--- a/smoke-test/tests/timeline/timeline_test.py
+++ b/smoke-test/tests/timeline/timeline_test.py
@@ -1,9 +1,9 @@
 import json
 
 import pytest
+
 from datahub.cli import timeline_cli
 from datahub.cli.cli_utils import guess_entity_type, post_entity
-
 from tests.utils import ingest_file_via_rest, wait_for_writes_to_sync
 
 pytestmark = pytest.mark.no_cypress_suite1
diff --git a/smoke-test/tests/tokens/revokable_access_token_test.py b/smoke-test/tests/tokens/revokable_access_token_test.py
index 006daae39333ed..30022345f15b45 100644
--- a/smoke-test/tests/tokens/revokable_access_token_test.py
+++ b/smoke-test/tests/tokens/revokable_access_token_test.py
@@ -432,7 +432,9 @@ def generateAccessToken_v2(session, actorUrn):
     return response.json()
 
 
-def listAccessTokens(session, filters=[]):
+def listAccessTokens(session, filters):
+    if filters is None:
+        filters = []
     # Get count of existing tokens
     input = {"start": 0, "count": 20}
 
diff --git a/smoke-test/tests/tokens/session_access_token_test.py b/smoke-test/tests/tokens/session_access_token_test.py
index a16abc44453036..5328fd4237f688 100644
--- a/smoke-test/tests/tokens/session_access_token_test.py
+++ b/smoke-test/tests/tokens/session_access_token_test.py
@@ -2,10 +2,10 @@
 import time
 
 import pytest
-from datahub.emitter.mcp import MetadataChangeProposalWrapper
-from datahub.metadata.schema_classes import AuditStampClass, CorpUserStatusClass
 from requests.exceptions import HTTPError
 
+from datahub.emitter.mcp import MetadataChangeProposalWrapper
+from datahub.metadata.schema_classes import AuditStampClass, CorpUserStatusClass
 from tests.utils import (
     get_admin_credentials,
     get_frontend_url,
diff --git a/smoke-test/tests/utilities/file_emitter.py b/smoke-test/tests/utilities/file_emitter.py
index ddbcff8db31d8b..d5539d143af737 100644
--- a/smoke-test/tests/utilities/file_emitter.py
+++ b/smoke-test/tests/utilities/file_emitter.py
@@ -7,7 +7,7 @@
 
 class FileEmitter:
     def __init__(
-        self, filename: str, run_id: str = f"test_{int(time.time()*1000.0)}"
+        self, filename: str, run_id: str = f"test_{int(time.time() * 1000.0)}"
     ) -> None:
         self.sink: FileSink = FileSink(
             ctx=PipelineContext(run_id=run_id),
diff --git a/smoke-test/tests/utils.py b/smoke-test/tests/utils.py
index 2ddf233f5029ac..64121b9fa1ef5e 100644
--- a/smoke-test/tests/utils.py
+++ b/smoke-test/tests/utils.py
@@ -5,11 +5,11 @@
 from typing import Any, Dict, List, Tuple
 
 import requests
-from datahub.cli import cli_utils, env_utils
-from datahub.ingestion.run.pipeline import Pipeline
 from joblib import Parallel, delayed
 from requests.structures import CaseInsensitiveDict
 
+from datahub.cli import cli_utils, env_utils
+from datahub.ingestion.run.pipeline import Pipeline
 from tests.consistency_utils import wait_for_writes_to_sync
 
 TIME: int = 1581407189000