Allow coordinator only procedures to output metrics in result

raunaqmorarka · raunaqmorarka · commit 1ee97a9457d7 · 2025-09-18T18:47:34.000+05:30
diff --git a/core/trino-main/src/main/java/io/trino/metadata/Metadata.java b/core/trino-main/src/main/java/io/trino/metadata/Metadata.java
@@ -121,7 +121,7 @@ Optional<TableExecuteHandle> getTableHandleForExecute(
 
     void finishTableExecute(Session session, TableExecuteHandle handle, Collection<Slice> fragments, List<Object> tableExecuteState);
 
-    void executeTableExecute(Session session, TableExecuteHandle handle);
+    Map<String, Long> executeTableExecute(Session session, TableExecuteHandle handle);
 
     TableProperties getTableProperties(Session session, TableHandle handle);
 
diff --git a/core/trino-main/src/main/java/io/trino/metadata/MetadataManager.java b/core/trino-main/src/main/java/io/trino/metadata/MetadataManager.java
@@ -362,11 +362,11 @@ public void finishTableExecute(Session session, TableExecuteHandle tableExecuteH
     }
 
     @Override
-    public void executeTableExecute(Session session, TableExecuteHandle tableExecuteHandle)
+    public Map<String, Long> executeTableExecute(Session session, TableExecuteHandle tableExecuteHandle)
     {
         CatalogHandle catalogHandle = tableExecuteHandle.catalogHandle();
         ConnectorMetadata metadata = getMetadata(session, catalogHandle);
-        metadata.executeTableExecute(session.toConnectorSession(catalogHandle), tableExecuteHandle.connectorHandle());
+        return metadata.executeTableExecute(session.toConnectorSession(catalogHandle), tableExecuteHandle.connectorHandle());
     }
 
     @Override
diff --git a/core/trino-main/src/main/java/io/trino/operator/SimpleTableExecuteOperator.java b/core/trino-main/src/main/java/io/trino/operator/SimpleTableExecuteOperator.java
@@ -13,20 +13,26 @@
  */
 package io.trino.operator;
 
+import com.google.common.collect.ImmutableList;
+import io.airlift.slice.Slices;
 import io.trino.Session;
 import io.trino.metadata.Metadata;
 import io.trino.metadata.TableExecuteHandle;
 import io.trino.spi.Page;
+import io.trino.spi.PageBuilder;
+import io.trino.spi.block.BlockBuilder;
+import io.trino.spi.type.Type;
 import io.trino.sql.planner.plan.PlanNodeId;
 
+import java.util.List;
+import java.util.Map;
+
 import static com.google.common.base.Preconditions.checkState;
 import static java.util.Objects.requireNonNull;
 
 public class SimpleTableExecuteOperator
         implements Operator
 {
-    private static final Page PAGE = new Page(0);
-
     public static class SimpleTableExecuteOperatorOperatorFactory
             implements OperatorFactory
     {
@@ -35,20 +41,23 @@ public static class SimpleTableExecuteOperatorOperatorFactory
         private final Metadata metadata;
         private final Session session;
         private final TableExecuteHandle executeHandle;
+        private final List<Type> types;
         private boolean closed;
 
         public SimpleTableExecuteOperatorOperatorFactory(
                 int operatorId,
                 PlanNodeId planNodeId,
                 Metadata metadata,
                 Session session,
-                TableExecuteHandle executeHandle)
+                TableExecuteHandle executeHandle,
+                List<Type> types)
         {
             this.operatorId = operatorId;
             this.planNodeId = requireNonNull(planNodeId, "planNodeId is null");
             this.metadata = requireNonNull(metadata, "planNodeId is null");
             this.session = requireNonNull(session, "planNodeId is null");
             this.executeHandle = requireNonNull(executeHandle, "executeHandle is null");
+            this.types = ImmutableList.copyOf(requireNonNull(types, "types is null"));
         }
 
         @Override
@@ -60,7 +69,8 @@ public Operator createOperator(DriverContext driverContext)
                     context,
                     metadata,
                     session,
-                    executeHandle);
+                    executeHandle,
+                    types);
         }
 
         @Override
@@ -77,27 +87,31 @@ public OperatorFactory duplicate()
                     planNodeId,
                     metadata,
                     session,
-                    executeHandle);
+                    executeHandle,
+                    types);
         }
     }
 
     private final OperatorContext operatorContext;
     private final Metadata metadata;
     private final Session session;
     private final TableExecuteHandle executeHandle;
+    private final List<Type> types;
 
     private boolean finished;
 
     public SimpleTableExecuteOperator(
             OperatorContext operatorContext,
             Metadata metadata,
             Session session,
-            TableExecuteHandle executeHandle)
+            TableExecuteHandle executeHandle,
+            List<Type> types)
     {
         this.operatorContext = requireNonNull(operatorContext, "operatorContext is null");
         this.metadata = requireNonNull(metadata, "metadata is null");
         this.session = requireNonNull(session, "session is null");
         this.executeHandle = requireNonNull(executeHandle, "executeHandle is null");
+        this.types = ImmutableList.copyOf(requireNonNull(types, "types is null"));
     }
 
     @Override
@@ -125,9 +139,17 @@ public Page getOutput()
             return null;
         }
 
-        metadata.executeTableExecute(session, executeHandle);
+        Map<String, Long> metrics = metadata.executeTableExecute(session, executeHandle);
         finished = true;
-        return PAGE;
+        PageBuilder pageBuilder = new PageBuilder(types);
+        BlockBuilder metricNameBuilder = pageBuilder.getBlockBuilder(0);
+        BlockBuilder metricValueBuilder = pageBuilder.getBlockBuilder(1);
+        for (Map.Entry<String, Long> entry : metrics.entrySet()) {
+            types.get(0).writeSlice(metricNameBuilder, Slices.utf8Slice(entry.getKey()));
+            types.get(1).writeLong(metricValueBuilder, entry.getValue());
+            pageBuilder.declarePosition();
+        }
+        return pageBuilder.build();
     }
 
     @Override
diff --git a/core/trino-main/src/main/java/io/trino/sql/analyzer/StatementAnalyzer.java b/core/trino-main/src/main/java/io/trino/sql/analyzer/StatementAnalyzer.java
@@ -1323,6 +1323,9 @@ protected Scope visitTableExecute(TableExecute node, Optional<Scope> scope)
             analysis.setUpdateType("ALTER TABLE EXECUTE");
             analysis.setUpdateTarget(executeHandle.catalogHandle().getVersion(), tableName, Optional.of(table), Optional.empty());
 
+            if (!procedureMetadata.getExecutionMode().isReadsData()) {
+                return createAndAssignScope(node, scope, Field.newUnqualified("metric_name", VARCHAR), Field.newUnqualified("metric_value", BIGINT));
+            }
             return createAndAssignScope(node, scope, Field.newUnqualified("rows", BIGINT));
         }
 
diff --git a/core/trino-main/src/main/java/io/trino/sql/planner/LocalExecutionPlanner.java b/core/trino-main/src/main/java/io/trino/sql/planner/LocalExecutionPlanner.java
@@ -3476,7 +3476,8 @@ public PhysicalOperation visitSimpleTableExecuteNode(SimpleTableExecuteNode node
                             node.getId(),
                             metadata,
                             session,
-                            node.getExecuteHandle());
+                            node.getExecuteHandle(),
+                            getSymbolTypes(node.getOutputSymbols()));
 
             return new PhysicalOperation(operatorFactory, makeLayout(node));
         }
diff --git a/core/trino-main/src/main/java/io/trino/sql/planner/LogicalPlanner.java b/core/trino-main/src/main/java/io/trino/sql/planner/LogicalPlanner.java
@@ -966,7 +966,9 @@ private RelationPlan createTableExecutePlan(Analysis analysis, TableExecute stat
         if (!analysis.isTableExecuteReadsData()) {
             SimpleTableExecuteNode node = new SimpleTableExecuteNode(
                     idAllocator.getNextId(),
-                    symbolAllocator.newSymbol("rows", BIGINT),
+                    ImmutableList.of(
+                            symbolAllocator.newSymbol("metricName", VARCHAR),
+                            symbolAllocator.newSymbol("metricValue", BIGINT)),
                     executeHandle);
             return new RelationPlan(node, analysis.getRootScope(), node.getOutputSymbols(), Optional.empty());
         }
diff --git a/core/trino-main/src/main/java/io/trino/sql/planner/optimizations/UnaliasSymbolReferences.java b/core/trino-main/src/main/java/io/trino/sql/planner/optimizations/UnaliasSymbolReferences.java
@@ -687,12 +687,11 @@ public PlanAndMappings visitSimpleTableExecuteNode(SimpleTableExecuteNode node,
         {
             Map<Symbol, Symbol> mapping = new HashMap<>(context.getCorrelationMapping());
             SymbolMapper mapper = symbolMapper(mapping);
-            Symbol newOutput = mapper.map(node.getOutput());
 
             return new PlanAndMappings(
                     new SimpleTableExecuteNode(
                             node.getId(),
-                            newOutput,
+                            mapper.map(node.getOutputSymbols()),
                             node.getExecuteHandle()),
                     mapping);
         }
diff --git a/core/trino-main/src/main/java/io/trino/sql/planner/plan/SimpleTableExecuteNode.java b/core/trino-main/src/main/java/io/trino/sql/planner/plan/SimpleTableExecuteNode.java
@@ -27,17 +27,17 @@
 public class SimpleTableExecuteNode
         extends PlanNode
 {
-    private final Symbol output;
+    private final List<Symbol> outputs;
     private final TableExecuteHandle executeHandle;
 
     @JsonCreator
     public SimpleTableExecuteNode(
             @JsonProperty("id") PlanNodeId id,
-            @JsonProperty("output") Symbol output,
+            @JsonProperty("outputs") List<Symbol> outputs,
             @JsonProperty("executeHandle") TableExecuteHandle executeHandle)
     {
         super(id);
-        this.output = requireNonNull(output, "output is null");
+        this.outputs = ImmutableList.copyOf(requireNonNull(outputs, "outputs is null"));
         this.executeHandle = requireNonNull(executeHandle, "executeHandle is null");
     }
 
@@ -48,16 +48,11 @@ public List<PlanNode> getSources()
         return ImmutableList.of();
     }
 
+    @JsonProperty("outputs")
     @Override
     public List<Symbol> getOutputSymbols()
     {
-        return ImmutableList.of(output);
-    }
-
-    @JsonProperty
-    public Symbol getOutput()
-    {
-        return output;
+        return outputs;
     }
 
     @Override
diff --git a/core/trino-main/src/main/java/io/trino/tracing/TracingConnectorMetadata.java b/core/trino-main/src/main/java/io/trino/tracing/TracingConnectorMetadata.java
@@ -181,11 +181,11 @@ public void finishTableExecute(ConnectorSession session, ConnectorTableExecuteHa
     }
 
     @Override
-    public void executeTableExecute(ConnectorSession session, ConnectorTableExecuteHandle tableExecuteHandle)
+    public Map<String, Long> executeTableExecute(ConnectorSession session, ConnectorTableExecuteHandle tableExecuteHandle)
     {
         Span span = startSpan("executeTableExecute", tableExecuteHandle);
         try (var _ = scopedSpan(span)) {
-            delegate.executeTableExecute(session, tableExecuteHandle);
+            return delegate.executeTableExecute(session, tableExecuteHandle);
         }
     }
 
diff --git a/core/trino-main/src/main/java/io/trino/tracing/TracingMetadata.java b/core/trino-main/src/main/java/io/trino/tracing/TracingMetadata.java
@@ -236,11 +236,11 @@ public void finishTableExecute(Session session, TableExecuteHandle handle, Colle
     }
 
     @Override
-    public void executeTableExecute(Session session, TableExecuteHandle handle)
+    public Map<String, Long> executeTableExecute(Session session, TableExecuteHandle handle)
     {
         Span span = startSpan("executeTableExecute", handle);
         try (var _ = scopedSpan(span)) {
-            delegate.executeTableExecute(session, handle);
+            return delegate.executeTableExecute(session, handle);
         }
     }
 
diff --git a/core/trino-main/src/test/java/io/trino/connector/MockConnector.java b/core/trino-main/src/test/java/io/trino/connector/MockConnector.java
@@ -879,7 +879,10 @@ public Optional<ConnectorTableExecuteHandle> getTableHandleForExecute(ConnectorS
         }
 
         @Override
-        public void executeTableExecute(ConnectorSession session, ConnectorTableExecuteHandle tableExecuteHandle) {}
+        public Map<String, Long> executeTableExecute(ConnectorSession session, ConnectorTableExecuteHandle tableExecuteHandle)
+        {
+            return ImmutableMap.of();
+        }
 
         @Override
         public void finishTableExecute(ConnectorSession session, ConnectorTableExecuteHandle tableExecuteHandle, Collection<Slice> fragments, List<Object> tableExecuteState) {}
diff --git a/core/trino-main/src/test/java/io/trino/metadata/AbstractMockMetadata.java b/core/trino-main/src/test/java/io/trino/metadata/AbstractMockMetadata.java
@@ -170,7 +170,7 @@ public void finishTableExecute(Session session, TableExecuteHandle handle, Colle
     }
 
     @Override
-    public void executeTableExecute(Session session, TableExecuteHandle handle)
+    public Map<String, Long> executeTableExecute(Session session, TableExecuteHandle handle)
     {
         throw new UnsupportedOperationException();
     }
diff --git a/core/trino-spi/pom.xml b/core/trino-spi/pom.xml
@@ -818,6 +818,11 @@
                                     <code>java.method.removed</code>
                                     <old>method io.trino.spi.block.Block io.trino.spi.block.RowBlock::getFieldBlock(int)</old>
                                 </item>
+                                <item>
+                                    <code>java.method.returnTypeChanged</code>
+                                    <old>method void io.trino.spi.connector.ConnectorMetadata::executeTableExecute(io.trino.spi.connector.ConnectorSession, io.trino.spi.connector.ConnectorTableExecuteHandle)</old>
+                                    <new>method java.util.Map&lt;java.lang.String, java.lang.Long&gt; io.trino.spi.connector.ConnectorMetadata::executeTableExecute(io.trino.spi.connector.ConnectorSession, io.trino.spi.connector.ConnectorTableExecuteHandle)</new>
+                                </item>
                             </differences>
                         </revapi.differences>
                     </analysisConfiguration>
diff --git a/core/trino-spi/src/main/java/io/trino/spi/connector/ConnectorMetadata.java b/core/trino-spi/src/main/java/io/trino/spi/connector/ConnectorMetadata.java
@@ -160,9 +160,10 @@ default void finishTableExecute(ConnectorSession session, ConnectorTableExecuteH
     }
 
     /**
-     * Execute a {@link TableProcedureExecutionMode#coordinatorOnly() coordinator-only} table procedure.
+     * Execute a {@link TableProcedureExecutionMode#coordinatorOnly() coordinator-only} table procedure
+     * and return procedure execution metrics that will be populated in the query output.
      */
-    default void executeTableExecute(ConnectorSession session, ConnectorTableExecuteHandle tableExecuteHandle)
+    default Map<String, Long> executeTableExecute(ConnectorSession session, ConnectorTableExecuteHandle tableExecuteHandle)
     {
         throw new TrinoException(GENERIC_INTERNAL_ERROR, "ConnectorMetadata executeTableExecute() is not implemented");
     }
diff --git a/lib/trino-plugin-toolkit/src/main/java/io/trino/plugin/base/classloader/ClassLoaderSafeConnectorMetadata.java b/lib/trino-plugin-toolkit/src/main/java/io/trino/plugin/base/classloader/ClassLoaderSafeConnectorMetadata.java
@@ -211,10 +211,10 @@ public Optional<ConnectorTableExecuteHandle> getTableHandleForExecute(ConnectorS
     }
 
     @Override
-    public void executeTableExecute(ConnectorSession session, ConnectorTableExecuteHandle tableExecuteHandle)
+    public Map<String, Long> executeTableExecute(ConnectorSession session, ConnectorTableExecuteHandle tableExecuteHandle)
     {
         try (ThreadContextClassLoader _ = new ThreadContextClassLoader(classLoader)) {
-            delegate.executeTableExecute(session, tableExecuteHandle);
+            return delegate.executeTableExecute(session, tableExecuteHandle);
         }
     }
 
diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergMetadata.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergMetadata.java
@@ -1972,31 +1972,31 @@ private static void commitTransaction(Transaction transaction, String operation)
     }
 
     @Override
-    public void executeTableExecute(ConnectorSession session, ConnectorTableExecuteHandle tableExecuteHandle)
+    public Map<String, Long> executeTableExecute(ConnectorSession session, ConnectorTableExecuteHandle tableExecuteHandle)
     {
         IcebergTableExecuteHandle executeHandle = (IcebergTableExecuteHandle) tableExecuteHandle;
         switch (executeHandle.procedureId()) {
             case OPTIMIZE_MANIFESTS:
                 executeOptimizeManifests(session, executeHandle);
-                return;
+                return ImmutableMap.of();
             case DROP_EXTENDED_STATS:
                 executeDropExtendedStats(session, executeHandle);
-                return;
+                return ImmutableMap.of();
             case ROLLBACK_TO_SNAPSHOT:
                 executeRollbackToSnapshot(session, executeHandle);
-                return;
+                return ImmutableMap.of();
             case EXPIRE_SNAPSHOTS:
                 executeExpireSnapshots(session, executeHandle);
-                return;
+                return ImmutableMap.of();
             case REMOVE_ORPHAN_FILES:
                 executeRemoveOrphanFiles(session, executeHandle);
-                return;
+                return ImmutableMap.of();
             case ADD_FILES:
                 executeAddFiles(session, executeHandle);
-                return;
+                return ImmutableMap.of();
             case ADD_FILES_FROM_TABLE:
                 executeAddFilesFromTable(session, executeHandle);
-                return;
+                return ImmutableMap.of();
             default:
                 throw new IllegalArgumentException("Unknown procedure '" + executeHandle.procedureId() + "'");
         }
diff --git a/plugin/trino-lakehouse/src/main/java/io/trino/plugin/lakehouse/LakehouseMetadata.java b/plugin/trino-lakehouse/src/main/java/io/trino/plugin/lakehouse/LakehouseMetadata.java
@@ -198,9 +198,9 @@ public void finishTableExecute(ConnectorSession session, ConnectorTableExecuteHa
     }
 
     @Override
-    public void executeTableExecute(ConnectorSession session, ConnectorTableExecuteHandle tableExecuteHandle)
+    public Map<String, Long> executeTableExecute(ConnectorSession session, ConnectorTableExecuteHandle tableExecuteHandle)
     {
-        forHandle(tableExecuteHandle).executeTableExecute(session, tableExecuteHandle);
+        return forHandle(tableExecuteHandle).executeTableExecute(session, tableExecuteHandle);
     }
 
     @Override

Original file line number	Diff line number	Diff line change
`@@ -362,11 +362,11 @@ public void finishTableExecute(Session session, TableExecuteHandle tableExecuteH`
`362`	`362`	`}`
`363`	`363`
`364`	`364`	`@Override`
`365`		`- public void executeTableExecute(Session session, TableExecuteHandle tableExecuteHandle)`
	`365`	`+ public Map<String, Long> executeTableExecute(Session session, TableExecuteHandle tableExecuteHandle)`
`366`	`366`	`{`
`367`	`367`	`CatalogHandle catalogHandle = tableExecuteHandle.catalogHandle();`
`368`	`368`	`ConnectorMetadata metadata = getMetadata(session, catalogHandle);`
`369`		`- metadata.executeTableExecute(session.toConnectorSession(catalogHandle), tableExecuteHandle.connectorHandle());`
	`369`	`+ return metadata.executeTableExecute(session.toConnectorSession(catalogHandle), tableExecuteHandle.connectorHandle());`
`370`	`370`	`}`
`371`	`371`
`372`	`372`	`@Override`
Original file line number	Diff line number	Diff line change
`@@ -1323,6 +1323,9 @@ protected Scope visitTableExecute(TableExecute node, Optional<Scope> scope)`
`1323`	`1323`	`analysis.setUpdateType("ALTER TABLE EXECUTE");`
`1324`	`1324`	`analysis.setUpdateTarget(executeHandle.catalogHandle().getVersion(), tableName, Optional.of(table), Optional.empty());`
`1325`	`1325`
	`1326`	`+ if (!procedureMetadata.getExecutionMode().isReadsData()) {`
	`1327`	`+ return createAndAssignScope(node, scope, Field.newUnqualified("metric_name", VARCHAR), Field.newUnqualified("metric_value", BIGINT));`
	`1328`	`+ }`
`1326`	`1329`	`return createAndAssignScope(node, scope, Field.newUnqualified("rows", BIGINT));`
`1327`	`1330`	`}`
`1328`	`1331`
Original file line number	Diff line number	Diff line change
`@@ -3476,7 +3476,8 @@ public PhysicalOperation visitSimpleTableExecuteNode(SimpleTableExecuteNode node`
`3476`	`3476`	`node.getId(),`
`3477`	`3477`	`metadata,`
`3478`	`3478`	`session,`
`3479`		`- node.getExecuteHandle());`
	`3479`	`+ node.getExecuteHandle(),`
	`3480`	`+ getSymbolTypes(node.getOutputSymbols()));`
`3480`	`3481`
`3481`	`3482`	`return new PhysicalOperation(operatorFactory, makeLayout(node));`
`3482`	`3483`	`}`
Original file line number	Diff line number	Diff line change
`@@ -687,12 +687,11 @@ public PlanAndMappings visitSimpleTableExecuteNode(SimpleTableExecuteNode node,`
`687`	`687`	`{`
`688`	`688`	`Map<Symbol, Symbol> mapping = new HashMap<>(context.getCorrelationMapping());`
`689`	`689`	`SymbolMapper mapper = symbolMapper(mapping);`
`690`		`- Symbol newOutput = mapper.map(node.getOutput());`
`691`	`690`
`692`	`691`	`return new PlanAndMappings(`
`693`	`692`	`new SimpleTableExecuteNode(`
`694`	`693`	`node.getId(),`
`695`		`- newOutput,`
	`694`	`+ mapper.map(node.getOutputSymbols()),`
`696`	`695`	`node.getExecuteHandle()),`
`697`	`696`	`mapping);`
`698`	`697`	`}`
Original file line number	Diff line number	Diff line change
`@@ -181,11 +181,11 @@ public void finishTableExecute(ConnectorSession session, ConnectorTableExecuteHa`
`181`	`181`	`}`
`182`	`182`
`183`	`183`	`@Override`
`184`		`- public void executeTableExecute(ConnectorSession session, ConnectorTableExecuteHandle tableExecuteHandle)`
	`184`	`+ public Map<String, Long> executeTableExecute(ConnectorSession session, ConnectorTableExecuteHandle tableExecuteHandle)`
`185`	`185`	`{`
`186`	`186`	`Span span = startSpan("executeTableExecute", tableExecuteHandle);`
`187`	`187`	`try (var _ = scopedSpan(span)) {`
`188`		`- delegate.executeTableExecute(session, tableExecuteHandle);`
	`188`	`+ return delegate.executeTableExecute(session, tableExecuteHandle);`
`189`	`189`	`}`
`190`	`190`	`}`
`191`	`191`
Original file line number	Diff line number	Diff line change
`@@ -236,11 +236,11 @@ public void finishTableExecute(Session session, TableExecuteHandle handle, Colle`
`236`	`236`	`}`
`237`	`237`
`238`	`238`	`@Override`
`239`		`- public void executeTableExecute(Session session, TableExecuteHandle handle)`
	`239`	`+ public Map<String, Long> executeTableExecute(Session session, TableExecuteHandle handle)`
`240`	`240`	`{`
`241`	`241`	`Span span = startSpan("executeTableExecute", handle);`
`242`	`242`	`try (var _ = scopedSpan(span)) {`
`243`		`- delegate.executeTableExecute(session, handle);`
	`243`	`+ return delegate.executeTableExecute(session, handle);`
`244`	`244`	`}`
`245`	`245`	`}`
`246`	`246`