Shutdown MongoDB connector properly

Author: wendigo

plugin/trino-mongodb/src/main/java/io/trino/plugin/mongodb/MongoClientModule.java

@@ -37,6 +37,7 @@
 import static com.google.inject.multibindings.OptionalBinder.newOptionalBinder;
 import static io.airlift.configuration.ConditionalModule.conditionalModule;
 import static io.airlift.configuration.ConfigBinder.configBinder;
+import static io.trino.plugin.base.ClosingBinder.closingBinder;
 import static java.util.concurrent.TimeUnit.MILLISECONDS;
 
 public class MongoClientModule
@@ -67,6 +68,7 @@ public void setup(Binder binder)
                 internalBinder -> internalBinder.bind(MongoServerDetailsProvider.class).toInstance(ImmutableList::of),
                 internalBinder -> internalBinder.bind(MongoServerDetailsProvider.class).to(SessionBasedMongoServerDetailsProvider.class).in(Scopes.SINGLETON)));
 
+        closingBinder(binder).registerCloseable(MongoSession.class);
         newSetBinder(binder, ConnectorTableFunction.class).addBinding().toProvider(Query.class).in(Scopes.SINGLETON);
     }
 

plugin/trino-mongodb/src/main/java/io/trino/plugin/mongodb/MongoConnector.java

@@ -15,6 +15,7 @@
 
 import com.google.common.collect.ImmutableSet;
 import com.google.inject.Inject;
+import io.airlift.bootstrap.LifeCycleManager;
 import io.trino.plugin.base.session.SessionPropertiesProvider;
 import io.trino.spi.connector.Connector;
 import io.trino.spi.connector.ConnectorMetadata;
@@ -36,7 +37,7 @@
 public class MongoConnector
         implements Connector
 {
-    private final MongoSession mongoSession;
+    private final LifeCycleManager lifeCycleManager;
     private final MongoTransactionManager transactionManager;
     private final MongoSplitManager splitManager;
     private final MongoPageSourceProvider pageSourceProvider;
@@ -46,15 +47,15 @@ public class MongoConnector
 
     @Inject
     public MongoConnector(
-            MongoSession mongoSession,
+            LifeCycleManager lifeCycleManager,
             MongoTransactionManager transactionManager,
             MongoSplitManager splitManager,
             MongoPageSourceProvider pageSourceProvider,
             MongoPageSinkProvider pageSinkProvider,
             Set<ConnectorTableFunction> connectorTableFunctions,
             Set<SessionPropertiesProvider> sessionPropertiesProviders)
     {
-        this.mongoSession = mongoSession;
+        this.lifeCycleManager = requireNonNull(lifeCycleManager, "lifeCycleManager is null");
         this.transactionManager = requireNonNull(transactionManager, "transactionManager is null");
         this.splitManager = requireNonNull(splitManager, "splitManager is null");
         this.pageSourceProvider = requireNonNull(pageSourceProvider, "pageSourceProvider is null");
@@ -122,6 +123,6 @@ public List<PropertyMetadata<?>> getSessionProperties()
     @Override
     public void shutdown()
     {
-        mongoSession.shutdown();
+        lifeCycleManager.stop();
     }
 }

plugin/trino-mongodb/src/main/java/io/trino/plugin/mongodb/MongoSession.java

@@ -69,6 +69,7 @@
 import org.bson.types.Decimal128;
 import org.bson.types.ObjectId;
 
+import java.io.Closeable;
 import java.math.BigDecimal;
 import java.time.Instant;
 import java.time.LocalDate;
@@ -130,6 +131,7 @@
 import static java.util.stream.Collectors.toSet;
 
 public class MongoSession
+        implements Closeable
 {
     private static final Logger log = Logger.get(MongoSession.class);
     private static final Set<String> SYSTEM_DATABASES = Set.of("admin", "local", "config");
@@ -195,7 +197,8 @@ public MongoSession(TypeManager typeManager, MongoClient client, MongoClientConf
                 .build();
     }
 
-    public void shutdown()
+    @Override
+    public void close()
     {
         client.close();
     }