Add retryExecutionStrategy in client tags

prestodb · Nov 2, 2022 · cb986f6 · cb986f6
1 parent a62173e
commit cb986f6
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/...-spark-base/src/main/java/com/facebook/presto/spark/PrestoSparkQueryExecutionFactory.java b/...-spark-base/src/main/java/com/facebook/presto/spark/PrestoSparkQueryExecutionFactory.java
@@ -108,6 +108,7 @@
 import java.security.MessageDigest;
 import java.security.NoSuchAlgorithmException;
 import java.util.Collection;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
@@ -579,6 +580,11 @@ public IPrestoSparkQueryExecution create(
             // Update presto settings in Session, if present
             Session.SessionBuilder sessionBuilder = Session.builder(session);
             prestoSparkRetryExecutionSettings.getPrestoSettings().forEach(sessionBuilder::setSystemProperty);
+
+            Set<String> clientTags = new HashSet<>(session.getClientTags());
+            clientTags.add(retryExecutionStrategy.get().name());
+            sessionBuilder.setClientTags(clientTags);
+
             session = sessionBuilder.build();
         }