Test S3 request count in Hive connector

trinodb · Jul 24, 2023 · 7a5f728 · 7a5f728
1 parent c3a9278
commit 7a5f728
Show file tree

Hide file tree

Showing 6 changed files with 274 additions and 10 deletions.
diff --git a/plugin/trino-hive/pom.xml b/plugin/trino-hive/pom.xml
@@ -308,6 +308,12 @@
             <scope>runtime</scope>
         </dependency>
 
+        <dependency>
+            <groupId>io.opentelemetry</groupId>
+            <artifactId>opentelemetry-context</artifactId>
+            <scope>runtime</scope>
+        </dependency>
+
         <dependency>
             <groupId>io.trino</groupId>
             <artifactId>trino-hadoop-toolkit</artifactId>
@@ -348,6 +354,24 @@
             </exclusions>
         </dependency>
 
+        <dependency>
+            <groupId>io.opentelemetry</groupId>
+            <artifactId>opentelemetry-sdk</artifactId>
+            <scope>test</scope>
+        </dependency>
+
+        <dependency>
+            <groupId>io.opentelemetry</groupId>
+            <artifactId>opentelemetry-sdk-testing</artifactId>
+            <scope>test</scope>
+        </dependency>
+
+        <dependency>
+            <groupId>io.opentelemetry</groupId>
+            <artifactId>opentelemetry-sdk-trace</artifactId>
+            <scope>test</scope>
+        </dependency>
+
         <!-- for benchmark -->
         <dependency>
             <groupId>io.trino</groupId>
@@ -368,6 +392,13 @@
             <scope>test</scope>
         </dependency>
 
+        <dependency>
+            <groupId>io.trino</groupId>
+            <artifactId>trino-filesystem</artifactId>
+            <type>test-jar</type>
+            <scope>test</scope>
+        </dependency>
+
         <dependency>
             <groupId>io.trino</groupId>
             <artifactId>trino-main</artifactId>

diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/InternalHiveConnectorFactory.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/InternalHiveConnectorFactory.java
@@ -87,7 +87,7 @@ private InternalHiveConnectorFactory() {}
 
     public static Connector createConnector(String catalogName, Map<String, String> config, ConnectorContext context, Module module)
     {
-        return createConnector(catalogName, config, context, module, Optional.empty(), Optional.empty(), Optional.empty());
+        return createConnector(catalogName, config, context, module, Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty());
     }
 
     public static Connector createConnector(
@@ -97,6 +97,7 @@ public static Connector createConnector(
             Module module,
             Optional<HiveMetastore> metastore,
             Optional<TrinoFileSystemFactory> fileSystemFactory,
+            Optional<OpenTelemetry> openTelemetry,
             Optional<DirectoryLister> directoryLister)
     {
         requireNonNull(config, "config is null");
@@ -127,7 +128,7 @@ public static Connector createConnector(
                     new HiveProcedureModule(),
                     new MBeanServerModule(),
                     binder -> {
-                        binder.bind(OpenTelemetry.class).toInstance(context.getOpenTelemetry());
+                        binder.bind(OpenTelemetry.class).toInstance(openTelemetry.orElse(context.getOpenTelemetry()));
                         binder.bind(Tracer.class).toInstance(context.getTracer());
                         binder.bind(NodeVersion.class).toInstance(new NodeVersion(context.getNodeManager().getCurrentNode().getVersion()));
                         binder.bind(NodeManager.class).toInstance(context.getNodeManager());

diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/HiveQueryRunner.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/HiveQueryRunner.java
@@ -19,6 +19,7 @@
 import com.google.inject.Module;
 import io.airlift.log.Logger;
 import io.airlift.log.Logging;
+import io.opentelemetry.api.OpenTelemetry;
 import io.trino.Session;
 import io.trino.metadata.QualifiedObjectName;
 import io.trino.plugin.hive.fs.DirectoryLister;
@@ -105,6 +106,7 @@ public static class Builder<SELF extends Builder<?>>
             File baseDir = queryRunner.getCoordinator().getBaseDataDir().resolve("hive_data").toFile();
             return createTestingFileHiveMetastore(baseDir);
         };
+        private Optional<OpenTelemetry> openTelemetry = Optional.empty();
         private Module module = EMPTY_MODULE;
         private Optional<DirectoryLister> directoryLister = Optional.empty();
         private boolean tpcdsCatalogEnabled;
@@ -173,6 +175,13 @@ public SELF setMetastore(Function<DistributedQueryRunner, HiveMetastore> metasto
             return self();
         }
 
+        @CanIgnoreReturnValue
+        public SELF setOpenTelemetry(OpenTelemetry openTelemetry)
+        {
+            this.openTelemetry = Optional.of(openTelemetry);
+            return self();
+        }
+
         @CanIgnoreReturnValue
         public SELF setModule(Module module)
         {
@@ -244,7 +253,7 @@ public DistributedQueryRunner build()
                 }
 
                 HiveMetastore metastore = this.metastore.apply(queryRunner);
-                queryRunner.installPlugin(new TestingHivePlugin(Optional.of(metastore), module, directoryLister));
+                queryRunner.installPlugin(new TestingHivePlugin(Optional.of(metastore), openTelemetry, module, directoryLister));
 
                 Map<String, String> hiveProperties = new HashMap<>();
                 if (!skipTimezoneSetup) {

diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestingHiveConnectorFactory.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestingHiveConnectorFactory.java
@@ -14,6 +14,7 @@
 package io.trino.plugin.hive;
 
 import com.google.inject.Module;
+import io.opentelemetry.api.OpenTelemetry;
 import io.trino.plugin.hive.fs.DirectoryLister;
 import io.trino.plugin.hive.metastore.HiveMetastore;
 import io.trino.spi.connector.Connector;
@@ -31,17 +32,23 @@ public class TestingHiveConnectorFactory
         implements ConnectorFactory
 {
     private final Optional<HiveMetastore> metastore;
+    private final Optional<OpenTelemetry> openTelemetry;
     private final Module module;
     private final Optional<DirectoryLister> directoryLister;
 
     public TestingHiveConnectorFactory(HiveMetastore metastore)
     {
-        this(Optional.of(metastore), EMPTY_MODULE, Optional.empty());
+        this(Optional.of(metastore), Optional.empty(), EMPTY_MODULE, Optional.empty());
     }
 
-    public TestingHiveConnectorFactory(Optional<HiveMetastore> metastore, Module module, Optional<DirectoryLister> directoryLister)
+    public TestingHiveConnectorFactory(
+            Optional<HiveMetastore> metastore,
+            Optional<OpenTelemetry> openTelemetry,
+            Module module,
+            Optional<DirectoryLister> directoryLister)
     {
         this.metastore = requireNonNull(metastore, "metastore is null");
+        this.openTelemetry = requireNonNull(openTelemetry, "openTelemetry is null");
         this.module = requireNonNull(module, "module is null");
         this.directoryLister = requireNonNull(directoryLister, "directoryLister is null");
     }
@@ -55,6 +62,6 @@ public String getName()
     @Override
     public Connector create(String catalogName, Map<String, String> config, ConnectorContext context)
     {
-        return createConnector(catalogName, config, context, module, metastore, Optional.empty(), directoryLister);
+        return createConnector(catalogName, config, context, module, metastore, Optional.empty(), openTelemetry, directoryLister);
     }
 }
diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestingHivePlugin.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestingHivePlugin.java
@@ -15,6 +15,7 @@
 
 import com.google.common.collect.ImmutableList;
 import com.google.inject.Module;
+import io.opentelemetry.api.OpenTelemetry;
 import io.trino.plugin.hive.fs.DirectoryLister;
 import io.trino.plugin.hive.metastore.HiveMetastore;
 import io.trino.spi.Plugin;
@@ -29,29 +30,31 @@ public class TestingHivePlugin
         implements Plugin
 {
     private final Optional<HiveMetastore> metastore;
+    private final Optional<OpenTelemetry> openTelemetry;
     private final Module module;
     private final Optional<DirectoryLister> directoryLister;
 
     public TestingHivePlugin()
     {
-        this(Optional.empty(), EMPTY_MODULE, Optional.empty());
+        this(Optional.empty(), Optional.empty(), EMPTY_MODULE, Optional.empty());
     }
 
     public TestingHivePlugin(HiveMetastore metastore)
     {
-        this(Optional.of(metastore), EMPTY_MODULE, Optional.empty());
+        this(Optional.of(metastore), Optional.empty(), EMPTY_MODULE, Optional.empty());
     }
 
-    public TestingHivePlugin(Optional<HiveMetastore> metastore, Module module, Optional<DirectoryLister> directoryLister)
+    public TestingHivePlugin(Optional<HiveMetastore> metastore, Optional<OpenTelemetry> openTelemetry, Module module, Optional<DirectoryLister> directoryLister)
     {
         this.metastore = requireNonNull(metastore, "metastore is null");
+        this.openTelemetry = requireNonNull(openTelemetry, "openTelemetry is null");
         this.module = requireNonNull(module, "module is null");
         this.directoryLister = requireNonNull(directoryLister, "directoryLister is null");
     }
 
     @Override
     public Iterable<ConnectorFactory> getConnectorFactories()
     {
-        return ImmutableList.of(new TestingHiveConnectorFactory(metastore, module, directoryLister));
+        return ImmutableList.of(new TestingHiveConnectorFactory(metastore, openTelemetry, module, directoryLister));
     }
 }