databricks · vikrantpuppala · Apr 13, 2026 · Apr 13, 2026 · Apr 13, 2026 · Apr 13, 2026
diff --git a/.github/actions/setup-jfrog/action.yml b/.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"
diff --git a/.github/actions/setup-poetry/action.yml b/.github/actions/setup-poetry/action.yml
@@ -0,0 +1,63 @@
+name: Setup Poetry with JFrog
+description: Install Poetry, configure JFrog as primary PyPI source, and install project dependencies
+
+inputs:
+  python-version:
+    description: Python version to set up
+    required: true
+  install-args:
+    description: Extra arguments for poetry install (e.g. --all-extras)
+    required: false
+    default: ""
+  cache-path:
+    description: Path to the virtualenv for caching (e.g. .venv or .venv-pyarrow)
+    required: false
+    default: ".venv"
+  cache-suffix:
+    description: Extra suffix for the cache key to avoid collisions across job variants
+    required: false
+    default: ""
+
+runs:
+  using: composite
+  steps:
+    - name: Setup JFrog
+      uses: ./.github/actions/setup-jfrog
+
+    - name: Set up python ${{ inputs.python-version }}
+      id: setup-python
+      uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5
+      with:
+        python-version: ${{ inputs.python-version }}
+
+    - name: Install Poetry
+      shell: bash
+      run: |
+        pip install poetry==2.2.1
+        poetry config virtualenvs.create true
+        poetry config virtualenvs.in-project true
+        poetry config installer.parallel true
+
+    - name: Configure Poetry JFrog source
+      shell: bash
+      run: |
+        poetry config repositories.jfrog https://databricks.jfrog.io/artifactory/api/pypi/db-pypi/simple
+        poetry config http-basic.jfrog gha-service-account "${JFROG_ACCESS_TOKEN}"
+        poetry source add --priority=primary jfrog https://databricks.jfrog.io/artifactory/api/pypi/db-pypi/simple
+        poetry lock
+
+    - name: Load cached venv
+      id: cached-poetry-dependencies
+      uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4
+      with:
+        path: ${{ inputs.cache-path }}
+        key: venv-${{ runner.os }}-${{ steps.setup-python.outputs.python-version }}-${{ inputs.cache-suffix }}${{ github.event.repository.name }}-${{ hashFiles('**/poetry.lock') }}
+
+    - name: Install dependencies
+      if: steps.cached-poetry-dependencies.outputs.cache-hit != 'true'
+      shell: bash
+      run: poetry install --no-interaction --no-root
+
+    - name: Install library
+      shell: bash
+      run: poetry install --no-interaction ${{ inputs.install-args }}
diff --git a/.github/workflows/code-coverage.yml b/.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 }}
@@ -16,63 +19,19 @@ jobs:
      DATABRICKS_CATALOG: peco
      DATABRICKS_USER: ${{ secrets.TEST_PECO_SP_ID }}
    steps:
-      #----------------------------------------------
-      #       check-out repo and set-up python
-      #----------------------------------------------
      - name: Check out repository
        uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4
        with:
          fetch-depth: 0
-      - name: Set up python
-        id: setup-python
-        uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5
-        with:
-          python-version: "3.10"
-      #----------------------------------------------
-      #  -----  install system dependencies  -----
-      #----------------------------------------------
      - name: Install system dependencies
        run: |
          sudo apt-get update
          sudo apt-get install -y libkrb5-dev
-      #----------------------------------------------
-      #  -----  install & configure poetry  -----
-      #----------------------------------------------
-      - name: Install Poetry
-        uses: snok/install-poetry@76e04a911780d5b312d89783f7b1cd627778900a # v1
+      - name: Setup Poetry
+        uses: ./.github/actions/setup-poetry
        with:
-          version: "2.2.1"
-          virtualenvs-create: true
-          virtualenvs-in-project: true
-          installer-parallel: true
-
-      #----------------------------------------------
-      #       load cached venv if cache exists
-      #----------------------------------------------
-      - name: Load cached venv
-        id: cached-poetry-dependencies
-        uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4
-        with:
-          path: .venv
-          key: venv-${{ runner.os }}-${{ steps.setup-python.outputs.python-version }}-${{ github.event.repository.name }}-${{ hashFiles('**/poetry.lock') }}
-      #----------------------------------------------
-      # install dependencies if cache does not exist
-      #----------------------------------------------
-      - name: Install dependencies
-        if: steps.cached-poetry-dependencies.outputs.cache-hit != 'true'
-        run: poetry install --no-interaction --no-root
-      #----------------------------------------------
-      # install your root project, if required
-      #----------------------------------------------
-      - name: Install Kerberos system dependencies
-        run: |
-          sudo apt-get update
-          sudo apt-get install -y libkrb5-dev
-      - name: Install library
-        run: poetry install --no-interaction --all-extras
-      #----------------------------------------------
-      #    run parallel tests with coverage
-      #----------------------------------------------
+          python-version: "3.10"
+          install-args: "--all-extras"
      - name: Run parallel tests with coverage
        continue-on-error: false
        run: |
@@ -83,24 +42,15 @@ jobs:
            --cov-report=xml \
            --cov-report=term \
            -v
-
-      #----------------------------------------------
-      #    run telemetry tests with coverage (isolated)
-      #----------------------------------------------
      - name: Run telemetry tests with coverage (isolated)
        continue-on-error: false
        run: |
-          # Run test_concurrent_telemetry.py separately for isolation
          poetry run pytest tests/e2e/test_concurrent_telemetry.py \
            --cov=src \
            --cov-append \
            --cov-report=xml \
            --cov-report=term \
            -v
-
-      #----------------------------------------------
-      #         check for coverage override
-      #----------------------------------------------
      - name: Check for coverage override
        id: override
        env:
@@ -116,9 +66,6 @@ jobs:
            echo "override=false" >> $GITHUB_OUTPUT
            echo "No coverage override found"
          fi
-      #----------------------------------------------
-      #         check coverage percentage
-      #----------------------------------------------
      - name: Check coverage percentage
        if: steps.override.outputs.override == 'false'
        run: |
@@ -127,40 +74,29 @@ jobs:
            echo "ERROR: Coverage file not found at $COVERAGE_FILE"
            exit 1
          fi
-
-          # Install xmllint if not available
          if ! command -v xmllint &> /dev/null; then
            sudo apt-get update && sudo apt-get install -y libxml2-utils
          fi
-
          COVERED=$(xmllint --xpath "string(//coverage/@lines-covered)" "$COVERAGE_FILE")
          TOTAL=$(xmllint --xpath "string(//coverage/@lines-valid)" "$COVERAGE_FILE")
          PERCENTAGE=$(python3 -c "covered=${COVERED}; total=${TOTAL}; print(round((covered/total)*100, 2))")
-
          echo "Branch Coverage: $PERCENTAGE%"
          echo "Required Coverage: 85%"
-
-          # Use Python to compare the coverage with 85
          python3 -c "import sys; sys.exit(0 if float('$PERCENTAGE') >= 85 else 1)"
          if [ $? -eq 1 ]; then
            echo "ERROR: Coverage is $PERCENTAGE%, which is less than the required 85%"
            exit 1
          else
            echo "SUCCESS: Coverage is $PERCENTAGE%, which meets the required 85%"
          fi
-
-      #----------------------------------------------
-      #         coverage enforcement summary
-      #----------------------------------------------
      - name: Coverage enforcement summary
        env:
          OVERRIDE: ${{ steps.override.outputs.override }}
          REASON: ${{ steps.override.outputs.reason }}
        run: |
          if [ "$OVERRIDE" == "true" ]; then
-            echo "⚠️ Coverage checks bypassed: $REASON"
+            echo "Coverage checks bypassed: $REASON"
            echo "Please ensure this override is justified and temporary"
          else
-            echo "✅ Coverage checks enforced - minimum 85% required"
+            echo "Coverage checks enforced - minimum 85% required"
          fi
-