Don't use raw Metric type outside of Metrics class

trinodb · Jul 30, 2021 · 1f896ea · 1f896ea
1 parent f5fd024
commit 1f896ea
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 14 deletions.
diff --git a/lib/trino-plugin-toolkit/src/test/java/io/trino/plugin/base/metrics/TestMetrics.java b/lib/trino-plugin-toolkit/src/test/java/io/trino/plugin/base/metrics/TestMetrics.java
@@ -38,7 +38,7 @@ public void testMergeCount()
                 "b", new LongCount(3),
                 "c", new LongCount(4)));
         Metrics merged = merge(m1, m2);
-        Map<String, Metric> expectedMap = ImmutableMap.of(
+        Map<String, Metric<?>> expectedMap = ImmutableMap.of(
                 "a", new LongCount(1),
                 "b", new LongCount(5),
                 "c", new LongCount(4));

diff --git a/plugin/trino-memory/src/test/java/io/trino/plugin/memory/TestMemoryConnectorTest.java b/plugin/trino-memory/src/test/java/io/trino/plugin/memory/TestMemoryConnectorTest.java
@@ -23,7 +23,6 @@
 import io.trino.plugin.base.metrics.LongCount;
 import io.trino.spi.QueryId;
 import io.trino.spi.metrics.Count;
-import io.trino.spi.metrics.Metric;
 import io.trino.spi.metrics.Metrics;
 import io.trino.testing.BaseConnectorTest;
 import io.trino.testing.DistributedQueryRunner;
@@ -40,7 +39,6 @@
 import org.testng.annotations.Test;
 
 import java.util.List;
-import java.util.Map;
 
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static io.trino.SystemSessionProperties.ENABLE_LARGE_DYNAMIC_FILTERS;
@@ -152,22 +150,22 @@ public void testSelect()
     @Test
     public void testCustomMetricsScanFilter()
     {
-        Map<String, Metric> metrics = collectCustomMetrics("SELECT partkey FROM part WHERE partkey % 1000 > 0");
-        assertThat(metrics.get("rows")).isEqualTo(new LongCount(PART_COUNT));
-        assertThat(metrics.get("started")).isEqualTo(metrics.get("finished"));
-        assertThat(((Count) metrics.get("finished")).getTotal()).isGreaterThan(0);
+        Metrics metrics = collectCustomMetrics("SELECT partkey FROM part WHERE partkey % 1000 > 0");
+        assertThat(metrics.getMetrics().get("rows")).isEqualTo(new LongCount(PART_COUNT));
+        assertThat(metrics.getMetrics().get("started")).isEqualTo(metrics.getMetrics().get("finished"));
+        assertThat(((Count) metrics.getMetrics().get("finished")).getTotal()).isGreaterThan(0);
     }
 
     @Test
     public void testCustomMetricsScanOnly()
     {
-        Map<String, Metric> metrics = collectCustomMetrics("SELECT partkey FROM part");
-        assertThat(metrics.get("rows")).isEqualTo(new LongCount(PART_COUNT));
-        assertThat(metrics.get("started")).isEqualTo(metrics.get("finished"));
-        assertThat(((Count) metrics.get("finished")).getTotal()).isGreaterThan(0);
+        Metrics metrics = collectCustomMetrics("SELECT partkey FROM part");
+        assertThat(metrics.getMetrics().get("rows")).isEqualTo(new LongCount(PART_COUNT));
+        assertThat(metrics.getMetrics().get("started")).isEqualTo(metrics.getMetrics().get("finished"));
+        assertThat(((Count) metrics.getMetrics().get("finished")).getTotal()).isGreaterThan(0);
     }
 
-    private Map<String, Metric> collectCustomMetrics(String sql)
+    private Metrics collectCustomMetrics(String sql)
     {
         DistributedQueryRunner runner = (DistributedQueryRunner) getQueryRunner();
         ResultWithQueryId<MaterializedResult> result = runner.executeWithQueryId(getSession(), sql);
@@ -179,8 +177,7 @@ private Map<String, Metric> collectCustomMetrics(String sql)
                 .getOperatorSummaries()
                 .stream()
                 .map(OperatorStats::getMetrics)
-                .reduce(Metrics.EMPTY, Metrics::mergeWith)
-                .getMetrics();
+                .reduce(Metrics.EMPTY, Metrics::mergeWith);
     }
 
     @Test(timeOut = 30_000)