Use tracing for testing ConnectorMetadata invocations

This refactors `TestInformationSchemaConnector.testMetadataCalls` to
leverage `TracingConnectorMetadata` instead of manually performed counts
within `CountingMockConnector. Among other things this exposes
`listViews` calls which were previously not counted.

The differences in expected counts in the test come from the fact how
the counting was done previously. For example
`MockConnectorMetadata.listTables` calls `listSchemaNames` before
`listTables`, so this single metadata call was being recorded as both
`listSchemaNames` and `listTables` (zero or more times).

client/trino-jdbc/src/test/java/io/trino/jdbc/TestTrinoDatabaseMetaData.java

@@ -150,6 +150,7 @@ public void tearDownServer()
     {
         server.close();
         server = null;
+        countingMockConnector.close();
         countingMockConnector = null;
     }
 

core/trino-main/src/test/java/io/trino/connector/MockConnector.java

@@ -131,6 +131,7 @@ public class MockConnector
     private static final String UPDATE_ROW_ID = "update_row_id";
     private static final String MERGE_ROW_ID = "merge_row_id";
 
+    private final Function<ConnectorMetadata, ConnectorMetadata> metadataWrapper;
     private final Function<ConnectorSession, List<String>> listSchemaNames;
     private final BiFunction<ConnectorSession, String, List<String>> listTables;
     private final Optional<BiFunction<ConnectorSession, SchemaTablePrefix, Iterator<TableColumnsMetadata>>> streamTableColumns;
@@ -176,6 +177,7 @@ public class MockConnector
     private final BiFunction<ConnectorSession, ConnectorTableExecuteHandle, Optional<ConnectorTableLayout>> getLayoutForTableExecute;
 
     MockConnector(
+            Function<ConnectorMetadata, ConnectorMetadata> metadataWrapper,
             List<PropertyMetadata<?>> sessionProperties,
             Function<ConnectorSession, List<String>> listSchemaNames,
             BiFunction<ConnectorSession, String, List<String>> listTables,
@@ -220,6 +222,7 @@ public class MockConnector
             OptionalInt maxWriterTasks,
             BiFunction<ConnectorSession, ConnectorTableExecuteHandle, Optional<ConnectorTableLayout>> getLayoutForTableExecute)
     {
+        this.metadataWrapper = requireNonNull(metadataWrapper, "metadataWrapper is null");
         this.sessionProperties = ImmutableList.copyOf(requireNonNull(sessionProperties, "sessionProperties is null"));
         this.listSchemaNames = requireNonNull(listSchemaNames, "listSchemaNames is null");
         this.listTables = requireNonNull(listTables, "listTables is null");
@@ -280,7 +283,7 @@ public ConnectorTransactionHandle beginTransaction(IsolationLevel isolationLevel
     @Override
     public ConnectorMetadata getMetadata(ConnectorSession session, ConnectorTransactionHandle transaction)
     {
-        return new MockConnectorMetadata();
+        return metadataWrapper.apply(new MockConnectorMetadata());
     }
 
     @Override

core/trino-main/src/test/java/io/trino/connector/MockConnectorFactory.java

@@ -26,6 +26,7 @@
 import io.trino.spi.connector.ConnectorContext;
 import io.trino.spi.connector.ConnectorFactory;
 import io.trino.spi.connector.ConnectorMaterializedViewDefinition;
+import io.trino.spi.connector.ConnectorMetadata;
 import io.trino.spi.connector.ConnectorNodePartitioningProvider;
 import io.trino.spi.connector.ConnectorSession;
 import io.trino.spi.connector.ConnectorSplitSource;
@@ -81,12 +82,14 @@
 import static io.trino.spi.statistics.TableStatistics.empty;
 import static io.trino.spi.type.VarcharType.createUnboundedVarcharType;
 import static java.util.Objects.requireNonNull;
+import static java.util.function.Function.identity;
 
 public class MockConnectorFactory
         implements ConnectorFactory
 {
     private final String name;
     private final List<PropertyMetadata<?>> sessionProperty;
+    private final Function<ConnectorMetadata, ConnectorMetadata> metadataWrapper;
     private final Function<ConnectorSession, List<String>> listSchemaNames;
     private final BiFunction<ConnectorSession, String, List<String>> listTables;
     private final Optional<BiFunction<ConnectorSession, SchemaTablePrefix, Iterator<TableColumnsMetadata>>> streamTableColumns;
@@ -135,6 +138,7 @@ public class MockConnectorFactory
     private MockConnectorFactory(
             String name,
             List<PropertyMetadata<?>> sessionProperty,
+            Function<ConnectorMetadata, ConnectorMetadata> metadataWrapper,
             Function<ConnectorSession, List<String>> listSchemaNames,
             BiFunction<ConnectorSession, String, List<String>> listTables,
             Optional<BiFunction<ConnectorSession, SchemaTablePrefix, Iterator<TableColumnsMetadata>>> streamTableColumns,
@@ -180,6 +184,7 @@ private MockConnectorFactory(
     {
         this.name = requireNonNull(name, "name is null");
         this.sessionProperty = ImmutableList.copyOf(requireNonNull(sessionProperty, "sessionProperty is null"));
+        this.metadataWrapper = requireNonNull(metadataWrapper, "metadataWrapper is null");
         this.listSchemaNames = requireNonNull(listSchemaNames, "listSchemaNames is null");
         this.listTables = requireNonNull(listTables, "listTables is null");
         this.streamTableColumns = requireNonNull(streamTableColumns, "streamTableColumns is null");
@@ -234,6 +239,7 @@ public String getName()
     public Connector create(String catalogName, Map<String, String> config, ConnectorContext context)
     {
         return new MockConnector(
+                metadataWrapper,
                 sessionProperty,
                 listSchemaNames,
                 listTables,
@@ -367,6 +373,7 @@ public static final class Builder
     {
         private String name = "mock";
         private final List<PropertyMetadata<?>> sessionProperties = new ArrayList<>();
+        private Function<ConnectorMetadata, ConnectorMetadata> metadataWrapper = identity();
         private Function<ConnectorSession, List<String>> listSchemaNames = defaultListSchemaNames();
         private BiFunction<ConnectorSession, String, List<String>> listTables = defaultListTables();
         private Optional<BiFunction<ConnectorSession, SchemaTablePrefix, Iterator<TableColumnsMetadata>>> streamTableColumns = Optional.empty();
@@ -438,6 +445,12 @@ public Builder withSessionProperties(Iterable<PropertyMetadata<?>> sessionProper
             return this;
         }
 
+        public Builder withMetadataWrapper(Function<ConnectorMetadata, ConnectorMetadata> metadataWrapper)
+        {
+            this.metadataWrapper = requireNonNull(metadataWrapper, "metadataWrapper is null");
+            return this;
+        }
+
         public Builder withListSchemaNames(Function<ConnectorSession, List<String>> listSchemaNames)
         {
             this.listSchemaNames = requireNonNull(listSchemaNames, "listSchemaNames is null");
@@ -733,6 +746,7 @@ public MockConnectorFactory build()
             return new MockConnectorFactory(
                     name,
                     sessionProperties,
+                    metadataWrapper,
                     listSchemaNames,
                     listTables,
                     streamTableColumns,

testing/trino-testing/pom.xml

@@ -86,6 +86,16 @@
             <artifactId>opentelemetry-api</artifactId>
         </dependency>
 
+        <dependency>
+            <groupId>io.opentelemetry</groupId>
+            <artifactId>opentelemetry-sdk-testing</artifactId>
+        </dependency>
+
+        <dependency>
+            <groupId>io.opentelemetry</groupId>
+            <artifactId>opentelemetry-sdk-trace</artifactId>
+        </dependency>
+
         <dependency>
             <groupId>io.trino</groupId>
             <artifactId>trino-client</artifactId>

testing/trino-testing/src/main/java/io/trino/testing/CountingMockConnector.java

@@ -14,12 +14,18 @@
 package io.trino.testing;
 
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.Multiset;
+import io.opentelemetry.sdk.testing.exporter.InMemorySpanExporter;
+import io.opentelemetry.sdk.trace.SdkTracerProvider;
+import io.opentelemetry.sdk.trace.export.SimpleSpanProcessor;
 import io.trino.connector.MockConnectorFactory;
 import io.trino.spi.Plugin;
 import io.trino.spi.connector.ConnectorFactory;
 import io.trino.spi.connector.SchemaTableName;
 import io.trino.spi.security.RoleGrant;
 import io.trino.spi.security.TrinoPrincipal;
+import io.trino.tracing.TracingConnectorMetadata;
+import io.trino.util.AutoCloseableCloser;
 
 import java.util.Objects;
 import java.util.Optional;
@@ -30,12 +36,16 @@
 import java.util.stream.Stream;
 
 import static com.google.common.base.MoreObjects.toStringHelper;
+import static com.google.common.collect.ImmutableMultiset.toImmutableMultiset;
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static io.trino.connector.MockConnectorFactory.Builder.defaultGetColumns;
 import static io.trino.connector.MockConnectorFactory.Builder.defaultGetTableHandle;
 import static io.trino.spi.security.PrincipalType.USER;
+import static java.util.Map.entry;
+import static java.util.stream.Collectors.joining;
 
 public class CountingMockConnector
+        implements AutoCloseable
 {
     private final Object lock = new Object();
 
@@ -51,12 +61,32 @@ public class CountingMockConnector
             .mapToObj(i -> new RoleGrant(new TrinoPrincipal(USER, "user" + (i == 0 ? "" : i)), "role" + i / 2, false))
             .collect(toImmutableSet());
 
+    private final AutoCloseableCloser closer = AutoCloseableCloser.create();
+
     private final AtomicLong listSchemasCallsCounter = new AtomicLong();
     private final AtomicLong listTablesCallsCounter = new AtomicLong();
     private final AtomicLong getTableHandleCallsCounter = new AtomicLong();
     private final AtomicLong getColumnsCallsCounter = new AtomicLong();
     private final ListRoleGrantsCounter listRoleGrantCounter = new ListRoleGrantsCounter();
 
+    private final InMemorySpanExporter spanExporter;
+    private final SdkTracerProvider tracerProvider;
+
+    public CountingMockConnector()
+    {
+        spanExporter = closer.register(InMemorySpanExporter.create());
+        tracerProvider = closer.register(SdkTracerProvider.builder()
+                .addSpanProcessor(SimpleSpanProcessor.create(spanExporter))
+                .build());
+    }
+
+    @Override
+    public void close()
+            throws Exception
+    {
+        closer.close();
+    }
+
     public Plugin getPlugin()
     {
         return new Plugin()
@@ -78,6 +108,10 @@ public Stream<SchemaTableName> getAllTables()
                         .map(tableName -> new SchemaTableName("test_schema2", tableName)));
     }
 
+    /**
+     * @deprecated Use {@link #runTracing}.
+     */
+    @Deprecated
     public MetadataCallsCount runCounting(Runnable runnable)
     {
         synchronized (lock) {
@@ -101,9 +135,34 @@ public MetadataCallsCount runCounting(Runnable runnable)
         }
     }
 
+    public Multiset<String> runTracing(Runnable runnable)
+    {
+        synchronized (lock) {
+            spanExporter.reset();
+
+            runnable.run();
+
+            return spanExporter.getFinishedSpanItems().stream()
+                    .map(span -> {
+                        String attributes = span.getAttributes().asMap().entrySet().stream()
+                                .map(entry -> entry(entry.getKey().getKey(), entry.getValue()))
+                                .filter(entry -> !entry.getKey().equals("trino.catalog"))
+                                .map(entry -> "%s=%s".formatted(entry.getKey().replaceFirst("^trino\\.", ""), entry.getValue()))
+                                .sorted()
+                                .collect(joining(", "));
+                        if (attributes.isEmpty()) {
+                            return span.getName();
+                        }
+                        return "%s(%s)".formatted(span.getName(), attributes);
+                    })
+                    .collect(toImmutableMultiset());
+        }
+    }
+
     private ConnectorFactory getConnectorFactory()
     {
         MockConnectorFactory mockConnectorFactory = MockConnectorFactory.builder()
+                .withMetadataWrapper(connectorMetadata -> new TracingConnectorMetadata(tracerProvider.get("test"), "mock", connectorMetadata))
                 .withListSchemaNames(connectorSession -> {
                     listSchemasCallsCounter.incrementAndGet();
                     return ImmutableList.of("test_schema1", "test_schema2");
@@ -150,6 +209,7 @@ private ConnectorFactory getConnectorFactory()
         return mockConnectorFactory;
     }
 
+    @Deprecated
     public static final class MetadataCallsCount
     {
         private final long listSchemasCount;