From 613d9ff47e04a46ccbf2935b8c4bd88bd386ee22 Mon Sep 17 00:00:00 2001
From: Yuya Ebihara <ebyhry@gmail.com>
Date: Mon, 21 Aug 2023 08:40:25 +0900
Subject: [PATCH] Add support for OpenTelemetry in MongoDB connector

---
 plugin/trino-mongodb/pom.xml                                 | 5 +++++
 .../main/java/io/trino/plugin/mongodb/MongoClientModule.java | 5 ++++-
 .../java/io/trino/plugin/mongodb/MongoConnectorFactory.java  | 4 +++-
 3 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/plugin/trino-mongodb/pom.xml b/plugin/trino-mongodb/pom.xml
index a56131844f61..08211c4915ed 100644
--- a/plugin/trino-mongodb/pom.xml
+++ b/plugin/trino-mongodb/pom.xml
@@ -48,6 +48,11 @@
             <artifactId>log</artifactId>
         </dependency>
 
+        <dependency>
+            <groupId>io.opentelemetry.instrumentation</groupId>
+            <artifactId>opentelemetry-mongo-3.1</artifactId>
+        </dependency>
+
         <dependency>
             <groupId>io.trino</groupId>
             <artifactId>trino-cache</artifactId>
diff --git a/plugin/trino-mongodb/src/main/java/io/trino/plugin/mongodb/MongoClientModule.java b/plugin/trino-mongodb/src/main/java/io/trino/plugin/mongodb/MongoClientModule.java
index 60441e7ecca7..81cfcf57294d 100644
--- a/plugin/trino-mongodb/src/main/java/io/trino/plugin/mongodb/MongoClientModule.java
+++ b/plugin/trino-mongodb/src/main/java/io/trino/plugin/mongodb/MongoClientModule.java
@@ -23,6 +23,8 @@
 import com.mongodb.client.MongoClient;
 import com.mongodb.client.MongoClients;
 import io.airlift.configuration.AbstractConfigurationAwareModule;
+import io.opentelemetry.api.OpenTelemetry;
+import io.opentelemetry.instrumentation.mongo.v3_1.MongoTelemetry;
 import io.trino.plugin.base.session.SessionPropertiesProvider;
 import io.trino.plugin.mongodb.ptf.Query;
 import io.trino.spi.function.table.ConnectorTableFunction;
@@ -60,10 +62,11 @@ public void setup(Binder binder)
 
     @Singleton
     @Provides
-    public static MongoSession createMongoSession(TypeManager typeManager, MongoClientConfig config, Set<MongoClientSettingConfigurator> configurators)
+    public static MongoSession createMongoSession(TypeManager typeManager, MongoClientConfig config, Set<MongoClientSettingConfigurator> configurators, OpenTelemetry openTelemetry)
     {
         MongoClientSettings.Builder options = MongoClientSettings.builder();
         configurators.forEach(configurator -> configurator.configure(options));
+        options.addCommandListener(MongoTelemetry.builder(openTelemetry).build().newCommandListener());
         MongoClient client = MongoClients.create(options.build());
 
         return new MongoSession(
diff --git a/plugin/trino-mongodb/src/main/java/io/trino/plugin/mongodb/MongoConnectorFactory.java b/plugin/trino-mongodb/src/main/java/io/trino/plugin/mongodb/MongoConnectorFactory.java
index 6960ead92ff3..4d61c6340bca 100644
--- a/plugin/trino-mongodb/src/main/java/io/trino/plugin/mongodb/MongoConnectorFactory.java
+++ b/plugin/trino-mongodb/src/main/java/io/trino/plugin/mongodb/MongoConnectorFactory.java
@@ -16,6 +16,7 @@
 import com.google.inject.Injector;
 import io.airlift.bootstrap.Bootstrap;
 import io.airlift.json.JsonModule;
+import io.opentelemetry.api.OpenTelemetry;
 import io.trino.spi.connector.Connector;
 import io.trino.spi.connector.ConnectorContext;
 import io.trino.spi.connector.ConnectorFactory;
@@ -54,7 +55,8 @@ public Connector create(String catalogName, Map<String, String> config, Connecto
         Bootstrap app = new Bootstrap(
                 new JsonModule(),
                 new MongoClientModule(),
-                binder -> binder.bind(TypeManager.class).toInstance(context.getTypeManager()));
+                binder -> binder.bind(TypeManager.class).toInstance(context.getTypeManager()),
+                binder -> binder.bind(OpenTelemetry.class).toInstance(context.getOpenTelemetry()));
 
         Injector injector = app.doNotInitializeLogging()
                 .setRequiredConfigurationProperties(config)