From dad888fee98b389e0a55fa425d3d3b0dbd052e62 Mon Sep 17 00:00:00 2001
From: James Petty <pettja@amazon.com>
Date: Wed, 29 Mar 2023 11:05:29 -0400
Subject: [PATCH] Include QueryId in assertDistrubutedQuery failure message

---
 .../main/java/io/trino/testing/QueryAssertions.java    | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/testing/trino-testing/src/main/java/io/trino/testing/QueryAssertions.java b/testing/trino-testing/src/main/java/io/trino/testing/QueryAssertions.java
index b1c8b31558df..5593921163f8 100644
--- a/testing/trino-testing/src/main/java/io/trino/testing/QueryAssertions.java
+++ b/testing/trino-testing/src/main/java/io/trino/testing/QueryAssertions.java
@@ -279,7 +279,15 @@ private static void assertDistributedQuery(
             actualResults = resultWithQueryId.getResult().toTestTypes();
         }
         catch (RuntimeException ex) {
-            fail("Execution of 'actual' query failed: " + actual, ex);
+            if (queryId == null && ex instanceof QueryFailedException queryFailedException) {
+                queryId = queryFailedException.getQueryId();
+            }
+            if (queryId != null) {
+                fail("Execution of 'actual' query " + queryId + " failed: " + actual, ex);
+            }
+            else {
+                fail("Execution of 'actual' query failed: " + actual, ex);
+            }
         }
         if (planAssertion.isPresent()) {
             try {