Migrate CI to databricks-protected runners and route PyPI through JFrog

Protected runners are required for Databricks OSS repos. Add a
setup-jfrog composite action (OIDC-based, matching databricks-odbc) that
sets PIP_INDEX_URL so all pip/poetry installs go through the JFrog PyPI
proxy. Every workflow now runs on the databricks-protected-runner-group
with id-token: write for the OIDC exchange.

Co-authored-by: Isaac
Signed-off-by: Vikrant Puppala <vikrant.puppala@databricks.com>

Author: vikrantpuppala

.github/actions/setup-jfrog/action.yml

@@ -0,0 +1,32 @@
+name: Setup JFrog OIDC
+description: Obtain a JFrog access token via GitHub OIDC and configure pip to use JFrog PyPI proxy
+
+runs:
+  using: composite
+  steps:
+    - name: Get JFrog OIDC token
+      shell: bash
+      run: |
+        set -euo pipefail
+        ID_TOKEN=$(curl -sLS \
+          -H "User-Agent: actions/oidc-client" \
+          -H "Authorization: Bearer $ACTIONS_ID_TOKEN_REQUEST_TOKEN" \
+          "${ACTIONS_ID_TOKEN_REQUEST_URL}&audience=jfrog-github" | jq .value | tr -d '"')
+        echo "::add-mask::${ID_TOKEN}"
+        ACCESS_TOKEN=$(curl -sLS -XPOST -H "Content-Type: application/json" \
+          "https://databricks.jfrog.io/access/api/v1/oidc/token" \
+          -d "{\"grant_type\": \"urn:ietf:params:oauth:grant-type:token-exchange\", \"subject_token_type\":\"urn:ietf:params:oauth:token-type:id_token\", \"subject_token\": \"${ID_TOKEN}\", \"provider_name\": \"github-actions\"}" | jq .access_token | tr -d '"')
+        echo "::add-mask::${ACCESS_TOKEN}"
+        if [ -z "$ACCESS_TOKEN" ] || [ "$ACCESS_TOKEN" = "null" ]; then
+          echo "FAIL: Could not extract JFrog access token"
+          exit 1
+        fi
+        echo "JFROG_ACCESS_TOKEN=${ACCESS_TOKEN}" >> "$GITHUB_ENV"
+        echo "JFrog OIDC token obtained successfully"
+
+    - name: Configure pip
+      shell: bash
+      run: |
+        set -euo pipefail
+        echo "PIP_INDEX_URL=https://gha-service-account:${JFROG_ACCESS_TOKEN}@databricks.jfrog.io/artifactory/api/pypi/db-pypi/simple" >> "$GITHUB_ENV"
+        echo "pip configured to use JFrog registry"

.github/workflows/code-coverage.yml

@@ -2,12 +2,15 @@ name: Code Coverage
 
 permissions:
   contents: read
+  id-token: write
 
 on: [pull_request, workflow_dispatch]
 
 jobs:
   test-with-coverage:
-    runs-on: ubuntu-latest
+    runs-on:
+      group: databricks-protected-runner-group
+      labels: linux-ubuntu-latest
     environment: azure-prod
     env:
       DATABRICKS_SERVER_HOSTNAME: ${{ secrets.DATABRICKS_HOST }}
@@ -23,6 +26,8 @@ jobs:
         uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4
         with:
           fetch-depth: 0
+      - name: Setup JFrog
+        uses: ./.github/actions/setup-jfrog
       - name: Set up python
         id: setup-python
         uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5

.github/workflows/code-quality-checks.yml

@@ -4,10 +4,13 @@ on: [pull_request]
 
 permissions:
   contents: read
+  id-token: write
 
 jobs:
   run-unit-tests:
-    runs-on: ubuntu-latest
+    runs-on:
+      group: databricks-protected-runner-group
+      labels: linux-ubuntu-latest
     strategy:
       matrix:
         python-version: ["3.9", "3.10", "3.11", "3.12", "3.13", "3.14"]
@@ -27,6 +30,8 @@ jobs:
       #----------------------------------------------
       - name: Check out repository
         uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4
+      - name: Setup JFrog
+        uses: ./.github/actions/setup-jfrog
       - name: Set up python ${{ matrix.python-version }}
         id: setup-python
         uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5
@@ -92,7 +97,9 @@ jobs:
       - name: Run tests
         run: poetry run python -m pytest tests/unit
   run-unit-tests-with-arrow:
-    runs-on: ubuntu-latest
+    runs-on:
+      group: databricks-protected-runner-group
+      labels: linux-ubuntu-latest
     strategy:
       matrix:
         python-version: ["3.9", "3.10", "3.11", "3.12", "3.13", "3.14"]
@@ -111,6 +118,8 @@ jobs:
       #----------------------------------------------
       -   name: Check out repository
           uses: actions/checkout@ee0669bd1cc54295c223e0bb666b733df41de1c5 # v2
+      -   name: Setup JFrog
+          uses: ./.github/actions/setup-jfrog
       -   name: Set up python ${{ matrix.python-version }}
           id: setup-python
           uses: actions/setup-python@e9aba2c848f5ebd159c070c61ea2c4e2b122355e # v2
@@ -179,7 +188,9 @@ jobs:
       -   name: Run tests
           run: poetry run python -m pytest tests/unit
   check-linting:
-    runs-on: ubuntu-latest
+    runs-on:
+      group: databricks-protected-runner-group
+      labels: linux-ubuntu-latest
     strategy:
       matrix:
         python-version: ["3.9", "3.10", "3.11", "3.12", "3.13", "3.14"]
@@ -189,6 +200,8 @@ jobs:
       #----------------------------------------------
       - name: Check out repository
         uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4
+      - name: Setup JFrog
+        uses: ./.github/actions/setup-jfrog
       - name: Set up python ${{ matrix.python-version }}
         id: setup-python
         uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5
@@ -232,7 +245,9 @@ jobs:
         run: poetry run black --check src
 
   check-types:
-    runs-on: ubuntu-latest
+    runs-on:
+      group: databricks-protected-runner-group
+      labels: linux-ubuntu-latest
     strategy:
       matrix:
         python-version: ["3.9", "3.10", "3.11", "3.12", "3.13", "3.14"]
@@ -242,6 +257,8 @@ jobs:
       #----------------------------------------------
       - name: Check out repository
         uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4
+      - name: Setup JFrog
+        uses: ./.github/actions/setup-jfrog
       - name: Set up python ${{ matrix.python-version }}
         id: setup-python
         uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5

.github/workflows/daily-telemetry-e2e.yml

@@ -14,10 +14,13 @@ on:
 
 permissions:
   contents: read
+  id-token: write
 
 jobs:
   telemetry-e2e-tests:
-    runs-on: ubuntu-latest
+    runs-on:
+      group: databricks-protected-runner-group
+      labels: linux-ubuntu-latest
     environment: azure-prod
 
     env:
@@ -33,7 +36,10 @@ jobs:
       #----------------------------------------------
       - name: Check out repository
         uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4
-      
+
+      - name: Setup JFrog
+        uses: ./.github/actions/setup-jfrog
+
       - name: Set up python
         id: setup-python
         uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5

.github/workflows/dco-check.yml

@@ -8,7 +8,9 @@ permissions:
 
 jobs:
   check:
-    runs-on: ubuntu-latest
+    runs-on:
+      group: databricks-protected-runner-group
+      labels: linux-ubuntu-latest
     steps:
       - name: Check for DCO
         id: dco-check

.github/workflows/integration.yml

@@ -8,10 +8,13 @@ on:
 
 permissions:
   contents: read
+  id-token: write
 
 jobs:
   run-non-telemetry-tests:
-    runs-on: ubuntu-latest
+    runs-on:
+      group: databricks-protected-runner-group
+      labels: linux-ubuntu-latest
     environment: azure-prod
     env:
       DATABRICKS_SERVER_HOSTNAME: ${{ secrets.DATABRICKS_HOST }}
@@ -25,6 +28,8 @@ jobs:
       #----------------------------------------------
       - name: Check out repository
         uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4
+      - name: Setup JFrog
+        uses: ./.github/actions/setup-jfrog
       - name: Set up python
         id: setup-python
         uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5
@@ -71,7 +76,9 @@ jobs:
             -n auto
 
   run-telemetry-tests:
-    runs-on: ubuntu-latest
+    runs-on:
+      group: databricks-protected-runner-group
+      labels: linux-ubuntu-latest
     needs: run-non-telemetry-tests  # Run after non-telemetry tests complete
     environment: azure-prod
     env:
@@ -83,6 +90,8 @@ jobs:
     steps:
       - name: Check out repository
         uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4
+      - name: Setup JFrog
+        uses: ./.github/actions/setup-jfrog
       - name: Set up python
         id: setup-python
         uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5