From 5525ac88e78841f2fd0693542528c03ebc826b59 Mon Sep 17 00:00:00 2001
From: Danila Baklazhenko <danilabaklazhenko@gmail.com>
Date: Sat, 8 Mar 2025 13:18:37 +0200
Subject: [PATCH] Add check updates

---
 .github/workflows/extended.yml            | 57 ++++++++++++++++++++++-
 .github/workflows/pr_comment_commands.yml | 26 ++++++-----
 2 files changed, 71 insertions(+), 12 deletions(-)

diff --git a/.github/workflows/extended.yml b/.github/workflows/extended.yml
index b8e05e0b4d9c..5232d92a9349 100644
--- a/.github/workflows/extended.yml
+++ b/.github/workflows/extended.yml
@@ -33,7 +33,23 @@ on:
     branches:
       - main
   workflow_dispatch:
-    
+    inputs:
+      pr_number:
+        description: 'Pull request number'
+        required: false
+        type: string
+      check_run_id:
+        description: 'Check run ID for status updates'
+        required: false
+        type: string
+      pr_head_sha:
+        description: 'PR head SHA'
+        required: false
+        type: string
+
+permissions:
+  contents: read
+  checks: write
       
 jobs:
 
@@ -132,5 +148,44 @@ jobs:
           cargo test --profile release-nonlto --test sqllogictests -- --include-sqlite
           cargo clean
 
+  # If the workflow was triggered by the PR comment we need to manually update check status to display in UI
+  update-check-status:
+    needs: [linux-build-lib, linux-test-extended, hash-collisions, sqllogictest-sqlite]
+    runs-on: ubuntu-latest
+    if: ${{ always() && github.event_name == 'workflow_dispatch' }}
+    steps:
+      - name: Determine workflow status
+        id: status
+        run: |
+          if [[ "${{ contains(needs.*.result, 'failure') || contains(needs.*.result, 'cancelled') }}" == "true" ]]; then
+            echo "workflow_status=failure" >> $GITHUB_OUTPUT
+            echo "conclusion=failure" >> $GITHUB_OUTPUT
+          else
+            echo "workflow_status=completed" >> $GITHUB_OUTPUT
+            echo "conclusion=success" >> $GITHUB_OUTPUT
+          fi
+      
+      - name: Update check run
+        uses: actions/github-script@v7
+        with:
+          github-token: ${{ secrets.GITHUB_TOKEN }}
+          script: |
+            const workflowRunUrl = `https://github.com/${context.repo.owner}/${context.repo.repo}/actions/runs/${context.runId}`;
+            
+            await github.rest.checks.update({
+              owner: context.repo.owner,
+              repo: context.repo.repo,
+              check_run_id: ${{ github.event.inputs.check_run_id }},
+              status: 'completed',
+              conclusion: '${{ steps.status.outputs.conclusion }}',
+              output: {
+                title: '${{ steps.status.outputs.conclusion == 'success' && 'Extended Tests Passed' || 'Extended Tests Failed' }}',
+                summary: `Extended tests have completed with status: ${{ steps.status.outputs.conclusion }}.\n\n[View workflow run](${workflowRunUrl})`
+              },
+              details_url: workflowRunUrl
+            });
+
+
+
 
 
diff --git a/.github/workflows/pr_comment_commands.yml b/.github/workflows/pr_comment_commands.yml
index 45f1cc4e3072..b46b9a30d9b5 100644
--- a/.github/workflows/pr_comment_commands.yml
+++ b/.github/workflows/pr_comment_commands.yml
@@ -47,20 +47,11 @@ jobs:
             
             // Extract the branch name
             const branchName = pullRequest.head.ref;
-            
-            // Dispatch the workflow with the PR branch name
-            await github.rest.actions.createWorkflowDispatch({
-              owner: context.repo.owner,
-              repo: context.repo.repo,
-              workflow_id: 'extended.yml',
-              ref: branchName,
-            });
-
-            // Create a check run that links to the Actions tab so the run will be visible in GitHub UI
             const headSha = pullRequest.head.sha;
             const workflowRunsUrl = `https://github.com/${context.repo.owner}/${context.repo.repo}/actions?query=workflow%3A%22Datafusion+extended+tests%22+branch%3A${branchName}`;
             
-            await github.rest.checks.create({
+            // Create a check run that links to the Actions tab so the run will be visible in GitHub UI
+            const check = await github.rest.checks.create({
               owner: context.repo.owner,
               repo: context.repo.repo,
               name: 'Extended Tests',
@@ -73,6 +64,19 @@ jobs:
               details_url: workflowRunsUrl
             });
 
+            // Dispatch the workflow with the PR branch name
+            await github.rest.actions.createWorkflowDispatch({
+              owner: context.repo.owner,
+              repo: context.repo.repo,
+              workflow_id: 'extended.yml',
+              ref: branchName,
+              inputs: {
+                pr_number: context.payload.issue.number.toString(),
+                check_run_id: check.data.id.toString(),
+                pr_head_sha: headSha
+              }
+            });
+
       - name: Add reaction to comment 
         uses: actions/github-script@v7
         with: