Improve QueryId/TaskId handling

core/trino-main/src/main/java/io/trino/Session.java

@@ -541,7 +541,7 @@ public ConnectorSession toConnectorSession(CatalogHandle catalogHandle)
     public SessionRepresentation toSessionRepresentation()
     {
         return new SessionRepresentation(
-                queryId.toString(),
+                queryId.id(),
                 querySpan,
                 transactionId,
                 clientTransactionSupport,

core/trino-main/src/main/java/io/trino/dispatcher/QueuedStatementResource.java

@@ -263,7 +263,7 @@ private static URI getQueuedUri(QueryId queryId, Slug slug, long token, External
     {
         return externalUriInfo.baseUriBuilder()
                 .path("/v1/statement/queued/")
-                .path(queryId.toString())
+                .path(queryId.id())
                 .path(slug.makeSlug(QUEUED_QUERY, token))
                 .path(String.valueOf(token))
                 .build();
@@ -280,7 +280,7 @@ private static QueryResults createQueryResults(
     {
         QueryState state = queryError.map(error -> FAILED).orElse(QUEUED);
         return new QueryResults(
-                queryId.toString(),
+                queryId.id(),
                 getQueryInfoUri(queryInfoUrl, queryId, externalUriInfo),
                 null,
                 nextUri,
@@ -325,7 +325,7 @@ public Query(String query, SessionContext sessionContext, DispatchManager dispat
             this.queryInfoUrl = queryInfoUrlFactory.getQueryInfoUrl(queryId);
             requireNonNull(tracer, "tracer is null");
             this.querySpan = tracer.spanBuilder("query")
-                    .setAttribute(TrinoAttributes.QUERY_ID, queryId.toString())
+                    .setAttribute(TrinoAttributes.QUERY_ID, queryId.id())
                     .startSpan();
         }
 
@@ -446,7 +446,7 @@ private URI getRedirectUri(CoordinatorLocation coordinatorLocation, ExternalUriI
         {
             return coordinatorLocation.getUri(externalUriInfo)
                     .path("/v1/statement/executing")
-                    .path(queryId.toString())
+                    .path(queryId.id())
                     .path(slug.makeSlug(EXECUTING_QUERY, 0))
                     .path("0")
                     .build();

core/trino-main/src/main/java/io/trino/execution/QueryTracker.java

@@ -146,7 +146,7 @@ public T getQuery(QueryId queryId)
             throws NoSuchElementException
     {
         return tryGetQuery(queryId)
-                .orElseThrow(() -> new NoSuchElementException(queryId.toString()));
+                .orElseThrow(() -> new NoSuchElementException(queryId.id()));
     }
 
     public boolean hasQuery(QueryId queryId)

core/trino-main/src/main/java/io/trino/execution/TaskId.java

@@ -39,12 +39,17 @@ public static TaskId valueOf(String taskId)
     }
 
     private final String fullId;
+    private final StageId stageId;
+    private final int partitionId;
+    private final int attemptId;
 
     public TaskId(StageId stageId, int partitionId, int attemptId)
     {
-        requireNonNull(stageId, "stageId is null");
+        this.stageId = requireNonNull(stageId, "stageId is null");
         checkArgument(partitionId >= 0, "partitionId is negative: %s", partitionId);
         checkArgument(attemptId >= 0, "attemptId is negative: %s", attemptId);
+        this.partitionId = partitionId;
+        this.attemptId = attemptId;
 
         // There is a strange JDK bug related to the CompactStrings implementation in JDK20+ which causes some fullId values
         // to get corrupted when this particular line is JIT-optimized. Changing implicit concatenation to a String.join call
@@ -55,27 +60,32 @@ public TaskId(StageId stageId, int partitionId, int attemptId)
     private TaskId(String fullId)
     {
         this.fullId = requireNonNull(fullId, "fullId is null");
+        List<String> parts = parseDottedId(fullId, 4, "taskId");
+        this.stageId = new StageId(new QueryId(parts.get(0)), parseInt(parts.get(1)));
+        this.partitionId = parseInt(parts.get(2));
+        this.attemptId = parseInt(parts.get(3));
+        checkArgument(partitionId >= 0, "partitionId is negative: %s", partitionId);
+        checkArgument(attemptId >= 0, "attemptId is negative: %s", attemptId);
     }
 
     public QueryId getQueryId()
     {
-        return new QueryId(parseDottedId(fullId, 4, "taskId").get(0));
+        return stageId.getQueryId();
     }
 
     public StageId getStageId()
     {
-        List<String> ids = parseDottedId(fullId, 4, "taskId");
-        return StageId.valueOf(ids.subList(0, 2));
+        return stageId;
     }
 
     public int getPartitionId()
     {
-        return parseInt(parseDottedId(fullId, 4, "taskId").get(2));
+        return partitionId;
     }
 
     public int getAttemptId()
     {
-        return parseInt(parseDottedId(fullId, 4, "taskId").get(3));
+        return attemptId;
     }
 
     @Override

core/trino-main/src/main/java/io/trino/server/protocol/Query.java

@@ -523,7 +523,7 @@ private synchronized QueryResultsResponse getNextResult(long token, ExternalUriI
 
         // first time through, self is null
         QueryResults queryResults = new QueryResults(
-                queryId.toString(),
+                queryId.id(),
                 getQueryInfoUri(queryInfoUrl, queryId, externalUriInfo),
                 partialCancelUri,
                 nextResultsUri,
@@ -714,7 +714,7 @@ private URI createNextResultsUri(ExternalUriInfo externalUriInfo, long nextToken
     {
         return externalUriInfo.baseUriBuilder()
                 .path("/v1/statement/executing")
-                .path(queryId.toString())
+                .path(queryId.id())
                 .path(slug.makeSlug(EXECUTING_QUERY, nextToken))
                 .path(String.valueOf(nextToken))
                 .build();
@@ -724,7 +724,7 @@ private URI createPartialCancelUri(int stage, ExternalUriInfo externalUriInfo, l
     {
         return externalUriInfo.baseUriBuilder()
                 .path("/v1/statement/executing/partialCancel")
-                .path(queryId.toString())
+                .path(queryId.id())
                 .path(String.valueOf(stage))
                 .path(slug.makeSlug(EXECUTING_QUERY, nextToken))
                 .path(String.valueOf(nextToken))

core/trino-main/src/main/java/io/trino/server/protocol/QueryInfoUrlFactory.java

@@ -45,7 +45,7 @@ public QueryInfoUrlFactory(ServerConfig serverConfig)
     public Optional<URI> getQueryInfoUrl(QueryId queryId)
     {
         return queryInfoUrlTemplate
-                .map(template -> template.replace("${QUERY_ID}", queryId.toString()))
+                .map(template -> template.replace("${QUERY_ID}", queryId.id()))
                 .map(URI::create);
     }
 
@@ -54,7 +54,7 @@ public static URI getQueryInfoUri(Optional<URI> queryInfoUrl, QueryId queryId, E
         return queryInfoUrl.orElseGet(() ->
                 externalUriInfo.baseUriBuilder()
                         .path("ui/query.html")
-                        .replaceQuery(queryId.toString())
+                        .replaceQuery(queryId.id())
                         .build());
     }
 }

core/trino-main/src/main/java/io/trino/server/remotetask/HttpLocationFactory.java

@@ -50,7 +50,7 @@ public URI createQueryLocation(QueryId queryId)
         requireNonNull(queryId, "queryId is null");
         return uriBuilderFrom(baseUri)
                 .appendPath("/v1/query")
-                .appendPath(queryId.toString())
+                .appendPath(queryId.id())
                 .build();
     }
 

core/trino-main/src/test/java/io/trino/execution/scheduler/TestScaledWriterScheduler.java

@@ -34,6 +34,7 @@
 import io.trino.metadata.Split;
 import io.trino.node.InternalNode;
 import io.trino.node.TestingInternalNodeManager;
+import io.trino.spi.QueryId;
 import io.trino.spi.metrics.Metrics;
 import io.trino.spi.predicate.TupleDomain;
 import io.trino.sql.planner.Partitioning;
@@ -229,7 +230,7 @@ private static TaskStatus buildTaskStatus(int maxWriterCount, DataSize writerInp
     private static TaskStatus buildTaskStatus(boolean isOutputBufferOverUtilized, long outputDataSize, Optional<Integer> maxWriterCount, DataSize writerInputDataSize)
     {
         return new TaskStatus(
-                TaskId.valueOf("taskId"),
+                new TaskId(new StageId(new QueryId("query_id"), 0), 0, 0),
                 "task-instance-id",
                 0,
                 TaskState.RUNNING,
@@ -258,10 +259,12 @@ private static class TestingStageExecution
             implements StageExecution
     {
         private final PlanFragment fragment;
+        private final StageId stageId;
 
         public TestingStageExecution(PlanFragment fragment)
         {
             this.fragment = requireNonNull(fragment, "fragment is null");
+            this.stageId = new StageId(new QueryId("query_id"), 0);
         }
 
         @Override
@@ -357,7 +360,7 @@ public void recordSplitSourceMetrics(PlanNodeId nodeId, Metrics metrics, long st
         @Override
         public Optional<RemoteTask> scheduleTask(InternalNode node, int partition, Multimap<PlanNodeId, Split> initialSplits)
         {
-            return Optional.of(new TestingRemoteTask(TaskId.valueOf("taskId"), "nodeId", fragment));
+            return Optional.of(new TestingRemoteTask(new TaskId(stageId, partition, 0), "nodeId", fragment));
         }
 
         @Override

core/trino-spi/pom.xml

@@ -230,6 +230,29 @@
                                     <new>method void io.trino.spi.connector.Connector::shutdown()</new>
                                     <justification>Require connector to implement shutdown to prevent leaks</justification>
                                 </item>
+                                <item>
+                                    <ignore>true</ignore>
+                                    <code>java.method.visibilityReduced</code>
+                                    <old>method java.lang.String io.trino.spi.QueryId::validateId(java.lang.String)</old>
+                                    <new>method java.lang.String io.trino.spi.QueryId::validateId(java.lang.String, java.lang.String)</new>
+                                    <oldVisibility>public</oldVisibility>
+                                    <newVisibility>package</newVisibility>
+                                    <justification>Hidden method that is not used outside of QueryId</justification>
+                                </item>
+                                <item>
+                                    <code>java.annotation.removed</code>
+                                    <old>method java.lang.String io.trino.spi.QueryId::toString()</old>
+                                    <new>method java.lang.String io.trino.spi.QueryId::toString()</new>
+                                    <annotation>@com.fasterxml.jackson.annotation.JsonValue</annotation>
+                                    <justification>QueryId converted to a record</justification>
+                                </item>
+                                <item>
+                                    <ignore>true</ignore>
+                                    <code>java.class.kindChanged</code>
+                                    <old>class io.trino.spi.QueryId</old>
+                                    <new>class io.trino.spi.QueryId</new>
+                                    <justification>QueryId converted to a record</justification>
+                                </item>
                             </differences>
                         </revapi.differences>
                     </analysisConfiguration>

core/trino-spi/src/main/java/io/trino/spi/QueryId.java

@@ -15,15 +15,12 @@
 
 import com.fasterxml.jackson.annotation.JsonCreator;
 import com.fasterxml.jackson.annotation.JsonValue;
-import com.google.errorprone.annotations.FormatMethod;
 
 import java.util.List;
-import java.util.Objects;
 
-import static java.lang.String.format;
 import static java.util.Objects.requireNonNull;
 
-public final class QueryId
+public record QueryId(String id)
 {
     @JsonCreator
     public static QueryId valueOf(String queryId)
@@ -32,48 +29,30 @@ public static QueryId valueOf(String queryId)
         return new QueryId(queryId);
     }
 
-    private final String id;
-
-    public QueryId(String id)
+    public QueryId
     {
-        this.id = validateId(id);
+        requireNonNull(id, "id is null");
+        checkArgument(!id.isEmpty(), "id is empty");
+        validateId("queryId", id);
     }
 
+    // For backward compatibility
+    @JsonValue
+    @Deprecated // Use id() instead
     public String getId()
     {
         return id;
     }
 
     @Override
-    @JsonValue
     public String toString()
     {
         return id;
     }
 
-    @Override
-    public int hashCode()
-    {
-        return Objects.hash(id);
-    }
-
-    @Override
-    public boolean equals(Object obj)
-    {
-        if (this == obj) {
-            return true;
-        }
-        if (obj == null || getClass() != obj.getClass()) {
-            return false;
-        }
-        QueryId other = (QueryId) obj;
-        return Objects.equals(this.id, other.id);
-    }
-
     //
     // Id helper methods
     //
-
     // Check if the string matches [_a-z0-9]+ , but without the overhead of regex
     private static boolean isValidId(String id)
     {
@@ -86,34 +65,57 @@ private static boolean isValidId(String id)
         return true;
     }
 
-    public static String validateId(String id)
+    private static boolean isValidDottedId(char[] chars)
     {
-        requireNonNull(id, "id is null");
-        checkArgument(!id.isEmpty(), "id is empty");
-        checkArgument(isValidId(id), "Invalid id %s", id);
+        for (int i = 0; i < chars.length; i++) {
+            if (!(chars[i] == '_' || chars[i] == '.' || chars[i] >= 'a' && chars[i] <= 'z' || chars[i] >= '0' && chars[i] <= '9')) {
+                return false;
+            }
+        }
+        return true;
+    }
+
+    static String validateId(String name, String id)
+    {
+        if (!isValidId(id)) {
+            throw new IllegalArgumentException("Invalid " + name + " " + id);
+        }
         return id;
     }
 
     public static List<String> parseDottedId(String id, int expectedParts, String name)
     {
         requireNonNull(id, "id is null");
-        checkArgument(expectedParts > 0, "expectedParts must be at least 1");
+        checkArgument(expectedParts > 1, "expectedParts must be at least 2");
         requireNonNull(name, "name is null");
 
-        List<String> ids = List.of(id.split("\\."));
-        checkArgument(ids.size() == expectedParts, "Invalid %s %s", name, id);
-
-        for (String part : ids) {
-            validateId(part);
+        char[] chars = id.toCharArray();
+        if (!isValidDottedId(chars)) {
+            throw new IllegalArgumentException("Invalid " + name + " " + id);
+        }
+        String[] parts = new String[expectedParts];
+        int startOffset = 0;
+        int partIndex = 0;
+        for (int i = 0, length = chars.length; i < length; i++) {
+            if (chars[i] == '.') {
+                if (i <= startOffset || i == length - 1) {
+                    throw new IllegalArgumentException("Invalid " + name + " " + id);
+                }
+                parts[partIndex++] = new String(chars, startOffset, i - startOffset);
+                startOffset = i + 1;
+            }
+        }
+        parts[partIndex++] = new String(chars, startOffset, chars.length - startOffset);
+        if (partIndex != expectedParts) {
+            throw new IllegalArgumentException("Invalid " + name + " " + id);
         }
-        return ids;
+        return List.of(parts);
     }
 
-    @FormatMethod
-    private static void checkArgument(boolean condition, String message, Object... messageArgs)
+    private static void checkArgument(boolean condition, String message)
     {
         if (!condition) {
-            throw new IllegalArgumentException(format(message, messageArgs));
+            throw new IllegalArgumentException(message);
         }
     }
 }

plugin/trino-bigquery/src/test/java/io/trino/plugin/bigquery/BaseBigQueryConnectorTest.java

@@ -882,7 +882,7 @@ public void testQueryLabeling()
 
     private void assertLabelForTable(String expectedView, QueryId queryId, String traceToken)
     {
-        String expectedLabel = "q_" + queryId.toString() + "__t_" + traceToken;
+        String expectedLabel = "q_" + queryId.id() + "__t_" + traceToken;
 
         @Language("SQL")
         String checkForLabelQuery =

plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/BaseIcebergConnectorTest.java

@@ -9449,7 +9449,7 @@ private QueryId executeWithQueryId(String sql)
     private void assertQueryIdAndUserStored(String tableName, QueryId queryId)
     {
         assertThat(getFieldFromLatestSnapshotSummary(tableName, TRINO_QUERY_ID_NAME))
-                .isEqualTo(queryId.toString());
+                .isEqualTo(queryId.id());
         assertThat(getFieldFromLatestSnapshotSummary(tableName, TRINO_USER_NAME))
                 .isEqualTo("user");
     }