Ignore exceptions when computing stats after query finished

core/trino-main/src/main/java/io/trino/metadata/MetadataManager.java

@@ -24,18 +24,16 @@
 import com.google.inject.Inject;
 import io.airlift.log.Logger;
 import io.airlift.slice.Slice;
-import io.trino.FeaturesConfig;
 import io.trino.Session;
 import io.trino.connector.system.GlobalSystemConnector;
+import io.trino.execution.QueryManager;
 import io.trino.metadata.LanguageFunctionManager.RunAsIdentityLoader;
 import io.trino.security.AccessControl;
-import io.trino.security.AllowAllAccessControl;
 import io.trino.security.InjectedConnectorAccessControl;
 import io.trino.spi.ErrorCode;
 import io.trino.spi.QueryId;
 import io.trino.spi.RefreshType;
 import io.trino.spi.TrinoException;
-import io.trino.spi.block.BlockEncodingSerde;
 import io.trino.spi.connector.AggregateFunction;
 import io.trino.spi.connector.AggregationApplicationResult;
 import io.trino.spi.connector.Assignment;
@@ -115,12 +113,9 @@
 import io.trino.spi.type.Type;
 import io.trino.spi.type.TypeManager;
 import io.trino.spi.type.TypeNotFoundException;
-import io.trino.spi.type.TypeOperators;
 import io.trino.sql.analyzer.TypeSignatureProvider;
-import io.trino.sql.parser.SqlParser;
 import io.trino.sql.planner.PartitioningHandle;
 import io.trino.transaction.TransactionManager;
-import io.trino.type.BlockTypeOperators;
 import io.trino.type.TypeCoercion;
 
 import java.util.ArrayList;
@@ -132,6 +127,7 @@
 import java.util.Locale;
 import java.util.Map;
 import java.util.Map.Entry;
+import java.util.NoSuchElementException;
 import java.util.Optional;
 import java.util.OptionalInt;
 import java.util.OptionalLong;
@@ -152,7 +148,6 @@
 import static com.google.common.util.concurrent.MoreExecutors.directExecutor;
 import static io.airlift.concurrent.MoreFutures.toListenableFuture;
 import static io.trino.SystemSessionProperties.getRetryPolicy;
-import static io.trino.client.NodeVersion.UNKNOWN;
 import static io.trino.metadata.CatalogMetadata.SecurityManagement.CONNECTOR;
 import static io.trino.metadata.CatalogMetadata.SecurityManagement.SYSTEM;
 import static io.trino.metadata.GlobalFunctionCatalog.BUILTIN_SCHEMA;
@@ -175,8 +170,6 @@
 import static io.trino.spi.StandardErrorCode.TABLE_REDIRECTION_ERROR;
 import static io.trino.spi.StandardErrorCode.UNSUPPORTED_TABLE_TYPE;
 import static io.trino.spi.connector.MaterializedViewFreshness.Freshness.STALE;
-import static io.trino.transaction.InMemoryTransactionManager.createTestTransactionManager;
-import static io.trino.type.InternalTypeManager.TESTING_TYPE_MANAGER;
 import static java.lang.String.format;
 import static java.util.Collections.singletonList;
 import static java.util.Locale.ENGLISH;
@@ -199,6 +192,7 @@ public final class MetadataManager
     private final TableFunctionRegistry tableFunctionRegistry;
     private final TypeManager typeManager;
     private final TypeCoercion typeCoercion;
+    private final QueryManager queryManager;
 
     private final ConcurrentMap<QueryId, QueryCatalogs> catalogsByQueryId = new ConcurrentHashMap<>();
 
@@ -210,13 +204,15 @@ public MetadataManager(
             GlobalFunctionCatalog globalFunctionCatalog,
             LanguageFunctionManager languageFunctionManager,
             TableFunctionRegistry tableFunctionRegistry,
-            TypeManager typeManager)
+            TypeManager typeManager,
+            QueryManager queryManager)
     {
         this.accessControl = requireNonNull(accessControl, "accessControl is null");
         this.typeManager = requireNonNull(typeManager, "typeManager is null");
         functions = requireNonNull(globalFunctionCatalog, "globalFunctionCatalog is null");
         functionResolver = new BuiltinFunctionResolver(this, typeManager, globalFunctionCatalog);
         this.typeCoercion = new TypeCoercion(typeManager::getType);
+        this.queryManager = requireNonNull(queryManager, "queryManager is null");
 
         this.systemSecurityMetadata = requireNonNull(systemSecurityMetadata, "systemSecurityMetadata is null");
         this.transactionManager = requireNonNull(transactionManager, "transactionManager is null");
@@ -482,11 +478,33 @@ public TableMetadata getTableMetadata(Session session, TableHandle tableHandle)
     @Override
     public TableStatistics getTableStatistics(Session session, TableHandle tableHandle)
     {
-        CatalogHandle catalogHandle = tableHandle.catalogHandle();
-        ConnectorMetadata metadata = getMetadata(session, catalogHandle);
-        TableStatistics tableStatistics = metadata.getTableStatistics(session.toConnectorSession(catalogHandle), tableHandle.connectorHandle());
-        verifyNotNull(tableStatistics, "%s returned null tableStatistics for %s", metadata, tableHandle);
-        return tableStatistics;
+        try {
+            CatalogHandle catalogHandle = tableHandle.catalogHandle();
+            ConnectorMetadata metadata = getMetadata(session, catalogHandle);
+            TableStatistics tableStatistics = metadata.getTableStatistics(session.toConnectorSession(catalogHandle), tableHandle.connectorHandle());
+            verifyNotNull(tableStatistics, "%s returned null tableStatistics for %s", metadata, tableHandle);
+            return tableStatistics;
+        }
+        catch (RuntimeException e) {
+            if (isQueryDone(session)) {
+                // getting statistics for finished query may result in many different execeptions being thrown.
+                // As we do not care about the result anyway mask it by returning empty statistics.
+                return TableStatistics.empty();
+            }
+            throw e;
+        }
+    }
+
+    private boolean isQueryDone(Session session)
+    {
+        boolean done;
+        try {
+            done = queryManager.getQueryState(session.getQueryId()).isDone();
+        }
+        catch (NoSuchElementException ex) {
+            done = true;
+        }
+        return done;
     }
 
     @Override
@@ -2864,82 +2882,4 @@ private static boolean cannotExist(QualifiedObjectName name)
     {
         return name.catalogName().isEmpty() || name.schemaName().isEmpty() || name.objectName().isEmpty();
     }
-
-    public static MetadataManager createTestMetadataManager()
-    {
-        return testMetadataManagerBuilder().build();
-    }
-
-    public static TestMetadataManagerBuilder testMetadataManagerBuilder()
-    {
-        return new TestMetadataManagerBuilder();
-    }
-
-    public static class TestMetadataManagerBuilder
-    {
-        private TransactionManager transactionManager;
-        private TypeManager typeManager = TESTING_TYPE_MANAGER;
-        private GlobalFunctionCatalog globalFunctionCatalog;
-        private LanguageFunctionManager languageFunctionManager;
-
-        private TestMetadataManagerBuilder() {}
-
-        public TestMetadataManagerBuilder withTransactionManager(TransactionManager transactionManager)
-        {
-            this.transactionManager = transactionManager;
-            return this;
-        }
-
-        public TestMetadataManagerBuilder withTypeManager(TypeManager typeManager)
-        {
-            this.typeManager = requireNonNull(typeManager, "typeManager is null");
-            return this;
-        }
-
-        public TestMetadataManagerBuilder withGlobalFunctionCatalog(GlobalFunctionCatalog globalFunctionCatalog)
-        {
-            this.globalFunctionCatalog = globalFunctionCatalog;
-            return this;
-        }
-
-        public TestMetadataManagerBuilder withLanguageFunctionManager(LanguageFunctionManager languageFunctionManager)
-        {
-            this.languageFunctionManager = languageFunctionManager;
-            return this;
-        }
-
-        public MetadataManager build()
-        {
-            TransactionManager transactionManager = this.transactionManager;
-            if (transactionManager == null) {
-                transactionManager = createTestTransactionManager();
-            }
-
-            GlobalFunctionCatalog globalFunctionCatalog = this.globalFunctionCatalog;
-            if (globalFunctionCatalog == null) {
-                globalFunctionCatalog = new GlobalFunctionCatalog(
-                        () -> { throw new UnsupportedOperationException(); },
-                        () -> { throw new UnsupportedOperationException(); },
-                        () -> { throw new UnsupportedOperationException(); });
-                TypeOperators typeOperators = new TypeOperators();
-                globalFunctionCatalog.addFunctions(SystemFunctionBundle.create(new FeaturesConfig(), typeOperators, new BlockTypeOperators(typeOperators), UNKNOWN));
-            }
-
-            if (languageFunctionManager == null) {
-                BlockEncodingSerde blockEncodingSerde = new InternalBlockEncodingSerde(new BlockEncodingManager(), typeManager);
-                languageFunctionManager = new LanguageFunctionManager(new SqlParser(), typeManager, user -> ImmutableSet.of(), blockEncodingSerde);
-            }
-
-            TableFunctionRegistry tableFunctionRegistry = new TableFunctionRegistry(_ -> new CatalogTableFunctions(ImmutableList.of()));
-
-            return new MetadataManager(
-                    new AllowAllAccessControl(),
-                    new DisabledSystemSecurityMetadata(),
-                    transactionManager,
-                    globalFunctionCatalog,
-                    languageFunctionManager,
-                    tableFunctionRegistry,
-                    typeManager);
-        }
-    }
 }

core/trino-main/src/main/java/io/trino/testing/NotImplementedQueryManager.java

@@ -0,0 +1,150 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.trino.testing;
+
+import com.google.common.util.concurrent.ListenableFuture;
+import io.trino.Session;
+import io.trino.execution.QueryExecution;
+import io.trino.execution.QueryInfo;
+import io.trino.execution.QueryManager;
+import io.trino.execution.QueryState;
+import io.trino.execution.StageId;
+import io.trino.execution.StateMachine;
+import io.trino.execution.TaskId;
+import io.trino.server.BasicQueryInfo;
+import io.trino.server.ResultQueryInfo;
+import io.trino.server.protocol.Slug;
+import io.trino.spi.QueryId;
+
+import java.util.List;
+import java.util.NoSuchElementException;
+import java.util.function.Consumer;
+
+public class NotImplementedQueryManager
+        implements QueryManager
+{
+    @Override
+    public List<BasicQueryInfo> getQueries()
+    {
+        throw new RuntimeException("not implemented");
+    }
+
+    @Override
+    public void setOutputInfoListener(QueryId queryId, Consumer<QueryExecution.QueryOutputInfo> listener)
+            throws NoSuchElementException
+    {
+        throw new RuntimeException("not implemented");
+    }
+
+    @Override
+    public void outputTaskFailed(TaskId taskId, Throwable failure)
+    {
+        throw new RuntimeException("not implemented");
+    }
+
+    @Override
+    public void resultsConsumed(QueryId queryId)
+    {
+        throw new RuntimeException("not implemented");
+    }
+
+    @Override
+    public void addStateChangeListener(QueryId queryId, StateMachine.StateChangeListener<QueryState> listener)
+            throws NoSuchElementException
+    {
+        throw new RuntimeException("not implemented");
+    }
+
+    @Override
+    public ListenableFuture<QueryState> getStateChange(QueryId queryId, QueryState currentState)
+    {
+        throw new RuntimeException("not implemented");
+    }
+
+    @Override
+    public BasicQueryInfo getQueryInfo(QueryId queryId)
+            throws NoSuchElementException
+    {
+        throw new RuntimeException("not implemented");
+    }
+
+    @Override
+    public QueryInfo getFullQueryInfo(QueryId queryId)
+            throws NoSuchElementException
+    {
+        throw new RuntimeException("not implemented");
+    }
+
+    @Override
+    public ResultQueryInfo getResultQueryInfo(QueryId queryId)
+            throws NoSuchElementException
+    {
+        throw new RuntimeException("not implemented");
+    }
+
+    @Override
+    public Session getQuerySession(QueryId queryId)
+    {
+        throw new RuntimeException("not implemented");
+    }
+
+    @Override
+    public Slug getQuerySlug(QueryId queryId)
+    {
+        throw new RuntimeException("not implemented");
+    }
+
+    @Override
+    public QueryState getQueryState(QueryId queryId)
+            throws NoSuchElementException
+    {
+        throw new RuntimeException("not implemented");
+    }
+
+    @Override
+    public boolean hasQuery(QueryId queryId)
+    {
+        throw new RuntimeException("not implemented");
+    }
+
+    @Override
+    public void recordHeartbeat(QueryId queryId)
+    {
+        throw new RuntimeException("not implemented");
+    }
+
+    @Override
+    public void createQuery(QueryExecution execution)
+    {
+        throw new RuntimeException("not implemented");
+    }
+
+    @Override
+    public void failQuery(QueryId queryId, Throwable cause)
+    {
+        throw new RuntimeException("not implemented");
+    }
+
+    @Override
+    public void cancelQuery(QueryId queryId)
+    {
+        throw new RuntimeException("not implemented");
+    }
+
+    @Override
+    public void cancelStage(StageId stageId)
+    {
+        throw new RuntimeException("not implemented");
+    }
+}

core/trino-main/src/main/java/io/trino/testing/PlanTester.java

@@ -366,7 +366,8 @@ private PlanTester(Session defaultSession, int nodeCountForStats)
                 globalFunctionCatalog,
                 languageFunctionManager,
                 tableFunctionRegistry,
-                typeManager);
+                typeManager,
+                new NotImplementedQueryManager());
         typeRegistry.addType(new JsonPath2016Type(new TypeDeserializer(typeManager), blockEncodingSerde));
         this.joinCompiler = new JoinCompiler(typeOperators);
         this.hashStrategyCompiler = new FlatHashStrategyCompiler(typeOperators);

core/trino-main/src/test/java/io/trino/cost/TestFilterStatsCalculator.java

@@ -17,8 +17,8 @@
 import io.airlift.slice.Slices;
 import io.trino.Session;
 import io.trino.metadata.Metadata;
-import io.trino.metadata.MetadataManager;
 import io.trino.metadata.ResolvedFunction;
+import io.trino.metadata.TestMetadataManager;
 import io.trino.metadata.TestingFunctionResolution;
 import io.trino.plugin.base.util.JsonTypeUtil;
 import io.trino.security.AllowAllAccessControl;
@@ -855,7 +855,7 @@ private PlanNodeStatsAssertion assertExpression(Expression expression, Session s
     private PlanNodeStatsAssertion assertExpression(Expression expression, Session session, PlanNodeStatsEstimate inputStatistics)
     {
         TransactionManager transactionManager = new TestingTransactionManager();
-        Metadata metadata = MetadataManager.testMetadataManagerBuilder().withTransactionManager(transactionManager).build();
+        Metadata metadata = TestMetadataManager.builder().withTransactionManager(transactionManager).build();
         return transaction(transactionManager, metadata, new AllowAllAccessControl())
                 .singleStatement()
                 .execute(session, transactionSession -> {