Align closingBinder API with Guice

Accept `Key`.
HotSushi · Apr 10, 2024 · 1c8eadf · 1c8eadf
1 parent 6afe0d8
commit 1c8eadf
Show file tree

Hide file tree

Showing 7 changed files with 22 additions and 21 deletions.
diff --git a/core/trino-main/src/main/java/io/trino/server/CoordinatorModule.java b/core/trino-main/src/main/java/io/trino/server/CoordinatorModule.java
@@ -16,6 +16,7 @@
 import com.google.common.collect.ImmutableList;
 import com.google.inject.Binder;
 import com.google.inject.Inject;
+import com.google.inject.Key;
 import com.google.inject.Provides;
 import com.google.inject.Scopes;
 import com.google.inject.Singleton;
@@ -365,10 +366,10 @@ List<OutputStatsEstimatorFactory> getCompositeOutputDataSizeEstimatorDelegateFac
 
         // cleanup
         closingBinder(binder)
-                .registerExecutor(ExecutorService.class, ForStatementResource.class)
-                .registerExecutor(ScheduledExecutorService.class, ForStatementResource.class)
-                .registerExecutor(ExecutorService.class, ForQueryExecution.class)
-                .registerExecutor(ScheduledExecutorService.class, ForScheduler.class);
+                .registerExecutor(Key.get(ExecutorService.class, ForStatementResource.class))
+                .registerExecutor(Key.get(ScheduledExecutorService.class, ForStatementResource.class))
+                .registerExecutor(Key.get(ExecutorService.class, ForQueryExecution.class))
+                .registerExecutor(Key.get(ScheduledExecutorService.class, ForScheduler.class));
     }
 
     // working around circular dependency Metadata <-> PlannerContext

diff --git a/core/trino-main/src/main/java/io/trino/server/ServerMainModule.java b/core/trino-main/src/main/java/io/trino/server/ServerMainModule.java
@@ -15,6 +15,7 @@
 
 import com.google.inject.Binder;
 import com.google.inject.Inject;
+import com.google.inject.Key;
 import com.google.inject.Provides;
 import com.google.inject.Scopes;
 import com.google.inject.Singleton;
@@ -496,9 +497,9 @@ protected void setup(Binder binder)
 
         // cleanup
         closingBinder(binder)
-                .registerExecutor(ScheduledExecutorService.class, ForExchange.class)
-                .registerExecutor(ExecutorService.class, ForAsyncHttp.class)
-                .registerExecutor(ScheduledExecutorService.class, ForAsyncHttp.class);
+                .registerExecutor(Key.get(ScheduledExecutorService.class, ForExchange.class))
+                .registerExecutor(Key.get(ExecutorService.class, ForAsyncHttp.class))
+                .registerExecutor(Key.get(ScheduledExecutorService.class, ForAsyncHttp.class));
     }
 
     private static class RegisterFunctionBundles

diff --git a/core/trino-main/src/main/java/io/trino/transaction/InMemoryTransactionManagerModule.java b/core/trino-main/src/main/java/io/trino/transaction/InMemoryTransactionManagerModule.java
@@ -14,6 +14,7 @@
 package io.trino.transaction;
 
 import com.google.inject.Binder;
+import com.google.inject.Key;
 import com.google.inject.Module;
 import com.google.inject.Provides;
 import com.google.inject.Singleton;
@@ -37,8 +38,8 @@ public void configure(Binder binder)
     {
         configBinder(binder).bindConfig(TransactionManagerConfig.class);
         closingBinder(binder)
-                .registerExecutor(ExecutorService.class, ForTransactionManager.class)
-                .registerExecutor(ScheduledExecutorService.class, ForTransactionManager.class);
+                .registerExecutor(Key.get(ExecutorService.class, ForTransactionManager.class))
+                .registerExecutor(Key.get(ScheduledExecutorService.class, ForTransactionManager.class));
     }
 
     @Provides

diff --git a/lib/trino-plugin-toolkit/src/main/java/io/trino/plugin/base/ClosingBinder.java b/lib/trino-plugin-toolkit/src/main/java/io/trino/plugin/base/ClosingBinder.java
@@ -24,7 +24,6 @@
 
 import java.io.Closeable;
 import java.io.IOException;
-import java.lang.annotation.Annotation;
 import java.lang.annotation.Retention;
 import java.lang.annotation.Target;
 import java.util.Set;
@@ -35,6 +34,7 @@
 import static java.lang.annotation.ElementType.METHOD;
 import static java.lang.annotation.ElementType.PARAMETER;
 import static java.lang.annotation.RetentionPolicy.RUNTIME;
+import static java.util.Objects.requireNonNull;
 
 public class ClosingBinder
 {
@@ -55,25 +55,23 @@ private ClosingBinder(Binder binder)
 
     public ClosingBinder registerExecutor(Class<? extends ExecutorService> type)
     {
-        executors.addBinding().to(type);
-        return this;
+        return registerExecutor(Key.get(type));
     }
 
-    public ClosingBinder registerExecutor(Class<? extends ExecutorService> type, Class<? extends Annotation> annotation)
+    public ClosingBinder registerExecutor(Key<? extends ExecutorService> key)
     {
-        executors.addBinding().to(Key.get(type, annotation));
+        executors.addBinding().to(requireNonNull(key, "key is null"));
         return this;
     }
 
     public ClosingBinder registerCloseable(Class<? extends Closeable> type)
     {
-        closeables.addBinding().to(type);
-        return this;
+        return registerCloseable(Key.get(type));
     }
 
-    public ClosingBinder registerCloseable(Class<? extends Closeable> type, Class<? extends Annotation> annotation)
+    public ClosingBinder registerCloseable(Key<? extends Closeable> key)
     {
-        closeables.addBinding().to(Key.get(type, annotation));
+        closeables.addBinding().to(key);
         return this;
     }
 

diff --git a/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/JdbcModule.java b/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/JdbcModule.java
@@ -106,7 +106,7 @@ public void setup(Binder binder)
         newSetBinder(binder, JdbcQueryEventListener.class);
 
         closingBinder(binder)
-                .registerExecutor(ExecutorService.class, ForRecordCursor.class);
+                .registerExecutor(Key.get(ExecutorService.class, ForRecordCursor.class));
     }
 
     public static Multibinder<SessionPropertiesProvider> sessionPropertiesProviderBinder(Binder binder)

diff --git a/...trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/ThriftMetastoreModule.java b/...trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/ThriftMetastoreModule.java
@@ -93,7 +93,7 @@ protected void setup(Binder binder)
         binder.bind(Key.get(boolean.class, AllowHiveTableRename.class)).toInstance(true);
 
         closingBinder(binder)
-                .registerExecutor(ExecutorService.class, ThriftHiveWriteStatisticsExecutor.class);
+                .registerExecutor(Key.get(ExecutorService.class, ThriftHiveWriteStatisticsExecutor.class));
     }
 
     private static class ThriftHiveMetastoreStatisticExecutorProvider

diff --git a/plugin/trino-phoenix5/src/main/java/io/trino/plugin/phoenix5/PhoenixClientModule.java b/plugin/trino-phoenix5/src/main/java/io/trino/plugin/phoenix5/PhoenixClientModule.java
@@ -151,7 +151,7 @@ protected void setup(Binder binder)
         install(new DecimalModule());
 
         closingBinder(binder)
-                .registerExecutor(ExecutorService.class, ForRecordCursor.class);
+                .registerExecutor(Key.get(ExecutorService.class, ForRecordCursor.class));
     }
 
     private void checkConfiguration(String connectionUrl)