Resolve credentials lazily in SqlStage

JohnFitzpatrick44 · JohnFitzpatrick44 · commit 3c08065b74ec · 2026-05-12T12:37:27.000-04:00
Use a Supplier to resolve creds lazily in SqlStage.
This allows the use of creds that are created
post-stage creation, e.g. per-split vended credentials.
diff --git a/core/trino-main/src/main/java/io/trino/execution/SqlStage.java b/core/trino-main/src/main/java/io/trino/execution/SqlStage.java
@@ -56,9 +56,11 @@
 import java.util.concurrent.Executor;
 import java.util.concurrent.TimeUnit;
 import java.util.function.Function;
+import java.util.function.Supplier;
 
 import static com.google.common.base.MoreObjects.toStringHelper;
 import static com.google.common.base.Preconditions.checkArgument;
+import static com.google.common.base.Suppliers.memoize;
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static io.trino.server.DynamicFilterService.getOutboundDynamicFilters;
 import static java.util.Objects.requireNonNull;
@@ -77,7 +79,7 @@ public final class SqlStage
 {
     private final Session session;
     private final StageStateMachine stateMachine;
-    private final Map<PlanNodeId, ConnectorTableCredentials> tableCredentials;
+    private final Supplier<Map<PlanNodeId, ConnectorTableCredentials>> tableCredentialsProvider;
     private final RemoteTaskFactory remoteTaskFactory;
     private final NodeTaskMap nodeTaskMap;
     private final boolean summarizeTaskInfo;
@@ -135,7 +137,7 @@ public static SqlStage createSqlStage(
                 nodeTaskMap,
                 summarizeTaskInfo,
                 bucketCountProvider,
-                extractTableCredentials(session, metadata, fragment));
+                memoize(() -> extractTableCredentials(session, metadata, fragment)));
         sqlStage.initialize();
         return sqlStage;
     }
@@ -147,7 +149,7 @@ private SqlStage(
             NodeTaskMap nodeTaskMap,
             boolean summarizeTaskInfo,
             LocalExchangeBucketCountProvider bucketCountProvider,
-            Map<PlanNodeId, ConnectorTableCredentials> tableCredentials)
+            Supplier<Map<PlanNodeId, ConnectorTableCredentials>> tableCredentialsProvider)
     {
         this.session = requireNonNull(session, "session is null");
         this.stateMachine = stateMachine;
@@ -157,7 +159,7 @@ private SqlStage(
         this.bucketCountProvider = requireNonNull(bucketCountProvider, "bucketCountProvider is null");
 
         this.outboundDynamicFilterIds = getOutboundDynamicFilters(stateMachine.getFragment());
-        this.tableCredentials = ImmutableMap.copyOf(tableCredentials);
+        this.tableCredentialsProvider = requireNonNull(tableCredentialsProvider, "tableCredentialsProvider is null");
     }
 
     private static Map<PlanNodeId, ConnectorTableCredentials> extractTableCredentials(Session session, Metadata metadata, PlanFragment fragment)
@@ -308,7 +310,7 @@ public synchronized Optional<RemoteTask> createTask(
                 node,
                 speculative,
                 fragment,
-                tableCredentials,
+                tableCredentialsProvider.get(),
                 splits,
                 outputBuffers,
                 nodeTaskMap.createPartitionedSplitCountTracker(node, taskId),
diff --git a/core/trino-main/src/test/java/io/trino/execution/TestSqlStage.java b/core/trino-main/src/test/java/io/trino/execution/TestSqlStage.java
@@ -17,20 +17,31 @@
 import com.google.common.collect.ImmutableListMultimap;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.Multimap;
 import com.google.common.util.concurrent.SettableFuture;
+import io.airlift.units.DataSize;
 import io.opentelemetry.api.trace.Span;
+import io.trino.Session;
+import io.trino.connector.CatalogHandle;
 import io.trino.cost.StatsAndCosts;
+import io.trino.execution.NodeTaskMap.PartitionedSplitCountTracker;
+import io.trino.execution.buffer.OutputBuffers;
 import io.trino.execution.buffer.PipelinedOutputBuffers;
 import io.trino.execution.scheduler.SplitSchedulerStats;
+import io.trino.metadata.AbstractMockMetadata;
+import io.trino.metadata.Metadata;
 import io.trino.metadata.Split;
 import io.trino.node.InternalNode;
 import io.trino.operator.RetryPolicy;
 import io.trino.spi.NodeVersion;
 import io.trino.spi.QueryId;
+import io.trino.spi.connector.ConnectorTableCredentials;
+import io.trino.spi.connector.ConnectorTableHandle;
 import io.trino.sql.planner.Partitioning;
 import io.trino.sql.planner.PartitioningScheme;
 import io.trino.sql.planner.PlanFragment;
 import io.trino.sql.planner.Symbol;
+import io.trino.sql.planner.plan.DynamicFilterId;
 import io.trino.sql.planner.plan.PlanFragmentId;
 import io.trino.sql.planner.plan.PlanNode;
 import io.trino.sql.planner.plan.PlanNodeId;
@@ -48,18 +59,25 @@
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 import java.util.OptionalInt;
+import java.util.Set;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Future;
 import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
 
 import static io.airlift.concurrent.Threads.daemonThreadsNamed;
 import static io.airlift.tracing.Tracing.noopTracer;
 import static io.trino.SessionTestUtils.TEST_SESSION;
 import static io.trino.execution.SqlStage.createSqlStage;
+import static io.trino.execution.TaskTestUtils.PLAN_FRAGMENT;
+import static io.trino.execution.TaskTestUtils.SPLIT;
+import static io.trino.execution.TaskTestUtils.TABLE_SCAN_NODE_ID;
 import static io.trino.execution.buffer.PipelinedOutputBuffers.BufferType.ARBITRARY;
 import static io.trino.metadata.AbstractMockMetadata.dummyMetadata;
 import static io.trino.sql.planner.SystemPartitioningHandle.SINGLE_DISTRIBUTION;
@@ -110,6 +128,131 @@ public void testFinalStageInfo()
         }
     }
 
+    @Test
+    public void testTableCredentialsAreLazilyResolved()
+    {
+        AtomicInteger callCount = new AtomicInteger();
+        Metadata metadata = new AbstractMockMetadata()
+        {
+            @Override
+            public Optional<ConnectorTableCredentials> getTableCredentials(Session s, CatalogHandle catalogHandle, ConnectorTableHandle tableHandle)
+            {
+                callCount.incrementAndGet();
+                return Optional.empty();
+            }
+        };
+
+        SqlStage stage = createSqlStage(
+                metadata,
+                new StageId(new QueryId("query"), 0),
+                createScanPlanFragment(),
+                ImmutableMap.of(),
+                new MockRemoteTaskFactory(executor, scheduledExecutor),
+                TEST_SESSION,
+                true,
+                new NodeTaskMap(new FinalizerService()),
+                executor,
+                noopTracer(),
+                Span.getInvalid(),
+                new SplitSchedulerStats(),
+                (_, _) -> OptionalInt.empty());
+
+        // Extraction does not run at stage construction
+        assertThat(callCount).hasValue(0);
+
+        // First createTask triggers extraction
+        createTestTask(stage, 0);
+        assertThat(callCount).hasValue(1);
+
+        // Subsequent createTasks reuse the memoized result
+        createTestTask(stage, 1);
+        createTestTask(stage, 2);
+        assertThat(callCount).hasValue(1);
+
+        stage.finish();
+    }
+
+    @Test
+    public void testTableCredentialsSeeLateUpdate()
+    {
+        AtomicReference<Optional<ConnectorTableCredentials>> response = new AtomicReference<>(Optional.empty());
+        Metadata metadata = new AbstractMockMetadata()
+        {
+            @Override
+            public Optional<ConnectorTableCredentials> getTableCredentials(Session s, CatalogHandle catalogHandle, ConnectorTableHandle tableHandle)
+            {
+                return response.get();
+            }
+        };
+
+        AtomicReference<Map<PlanNodeId, ConnectorTableCredentials>> capturedCredentials = new AtomicReference<>();
+        RemoteTaskFactory capturingFactory = new RemoteTaskFactory() {
+            private final RemoteTaskFactory delegate = new MockRemoteTaskFactory(executor, scheduledExecutor);
+
+            @Override
+            public RemoteTask createRemoteTask(
+                    Session session,
+                    Span stageSpan,
+                    TaskId taskId,
+                    InternalNode node,
+                    boolean speculative,
+                    PlanFragment fragment,
+                    Map<PlanNodeId, ConnectorTableCredentials> tableCredentials,
+                    Multimap<PlanNodeId, Split> initialSplits,
+                    OutputBuffers outputBuffers,
+                    PartitionedSplitCountTracker partitionedSplitCountTracker,
+                    Set<DynamicFilterId> outboundDynamicFilterIds,
+                    Optional<DataSize> estimatedMemory,
+                    boolean summarizeTaskInfo)
+            {
+                capturedCredentials.set(tableCredentials);
+                return delegate.createRemoteTask(session, stageSpan, taskId, node, speculative, fragment, tableCredentials, initialSplits, outputBuffers, partitionedSplitCountTracker, outboundDynamicFilterIds, estimatedMemory, summarizeTaskInfo);
+            }
+        };
+
+        SqlStage stage = createSqlStage(
+                metadata,
+                new StageId(new QueryId("query"), 0),
+                createScanPlanFragment(),
+                ImmutableMap.of(),
+                capturingFactory,
+                TEST_SESSION,
+                true,
+                new NodeTaskMap(new FinalizerService()),
+                executor,
+                noopTracer(),
+                Span.getInvalid(),
+                new SplitSchedulerStats(),
+                (_, _) -> OptionalInt.empty());
+
+        // Simulate a split source populating credentials after stage construction but before first task creation.
+        ConnectorTableCredentials credentials = new ConnectorTableCredentials() {};
+        response.set(Optional.of(credentials));
+
+        createTestTask(stage, 0);
+
+        assertThat(capturedCredentials.get())
+                .containsExactly(Map.entry(TABLE_SCAN_NODE_ID, credentials));
+
+        stage.finish();
+    }
+
+    private void createTestTask(SqlStage stage, int partition)
+    {
+        InternalNode node = new InternalNode("node" + partition, URI.create("http://node/" + partition), NodeVersion.UNKNOWN, false);
+        stage.createTask(
+                node,
+                partition,
+                0,
+                Optional.empty(),
+                OptionalInt.empty(),
+                PipelinedOutputBuffers.createInitial(ARBITRARY),
+                ImmutableListMultimap.of(TABLE_SCAN_NODE_ID, SPLIT.split()),
+                ImmutableSet.of(),
+                Optional.empty(),
+                false);
+    }
+
     private void testFinalStageInfoInternal()
             throws Exception
     {
@@ -255,4 +398,9 @@ private static PlanFragment createExchangePlanFragment()
                 ImmutableMap.of(),
                 Optional.empty());
     }
+
+    private static PlanFragment createScanPlanFragment()
+    {
+        return PLAN_FRAGMENT.withOutputPartitioning(Optional.empty(), OptionalInt.empty());
+    }
 }
