Simplify ConnectionFactory chain of responsibility

trinodb · Oct 17, 2023 · 1206ec6 · 1206ec6
1 parent 4f0f7fb
commit 1206ec6
Show file tree

Hide file tree

Showing 11 changed files with 23 additions and 62 deletions.
diff --git a/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/ForLazyConnectionFactory.java b/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/ForLazyConnectionFactory.java
diff --git a/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/JdbcDiagnosticModule.java b/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/JdbcDiagnosticModule.java
@@ -18,6 +18,7 @@
 import com.google.inject.Module;
 import com.google.inject.Provider;
 import com.google.inject.Provides;
+import com.google.inject.Scopes;
 import com.google.inject.Singleton;
 import io.airlift.log.Logger;
 import io.trino.plugin.base.CatalogName;
@@ -39,11 +40,12 @@ public void configure(Binder binder)
     {
         binder.install(new MBeanServerModule());
         binder.install(new MBeanModule());
+        binder.bind(StatisticsAwareConnectionFactory.class).in(Scopes.SINGLETON);
 
         Provider<CatalogName> catalogName = binder.getProvider(CatalogName.class);
         newExporter(binder).export(Key.get(JdbcClient.class, StatsCollecting.class))
                 .as(generator -> generator.generatedNameOf(JdbcClient.class, catalogName.get().toString()));
-        newExporter(binder).export(Key.get(ConnectionFactory.class, StatsCollecting.class))
+        newExporter(binder).export(StatisticsAwareConnectionFactory.class)
                 .as(generator -> generator.generatedNameOf(ConnectionFactory.class, catalogName.get().toString()));
         newExporter(binder).export(JdbcClient.class)
                 .as(generator -> generator.generatedNameOf(CachingJdbcClient.class, catalogName.get().toString()));
@@ -65,12 +67,4 @@ public JdbcClient createJdbcClientWithStats(@ForBaseJdbc JdbcClient client, Cata
             return client;
         }));
     }
-
-    @Provides
-    @Singleton
-    @StatsCollecting
-    public static ConnectionFactory createConnectionFactoryWithStats(@ForBaseJdbc ConnectionFactory connectionFactory)
-    {
-        return new StatisticsAwareConnectionFactory(connectionFactory);
-    }
 }
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
@@ -88,10 +88,6 @@ public void setup(Binder binder)
 
         newSetBinder(binder, ConnectorTableFunction.class);
 
-        binder.bind(ConnectionFactory.class)
-                .annotatedWith(ForLazyConnectionFactory.class)
-                .to(Key.get(ConnectionFactory.class, StatsCollecting.class))
-                .in(Scopes.SINGLETON);
         install(conditionalModule(
                 QueryConfig.class,
                 QueryConfig::isReuseConnection,

diff --git a/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/LazyConnectionFactory.java b/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/LazyConnectionFactory.java
@@ -16,6 +16,7 @@
 import com.google.errorprone.annotations.ThreadSafe;
 import com.google.errorprone.annotations.concurrent.GuardedBy;
 import com.google.inject.Inject;
+import io.trino.plugin.jdbc.jmx.StatisticsAwareConnectionFactory;
 import io.trino.spi.connector.ConnectorSession;
 import jakarta.annotation.Nullable;
 
@@ -29,10 +30,10 @@
 public final class LazyConnectionFactory
         implements ConnectionFactory
 {
-    private final ConnectionFactory delegate;
+    private final StatisticsAwareConnectionFactory delegate;
 
     @Inject
-    public LazyConnectionFactory(@ForLazyConnectionFactory ConnectionFactory delegate)
+    public LazyConnectionFactory(StatisticsAwareConnectionFactory delegate)
     {
         this.delegate = requireNonNull(delegate, "delegate is null");
     }

diff --git a/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/RetryingConnectionFactory.java b/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/RetryingConnectionFactory.java
@@ -18,6 +18,7 @@
 import dev.failsafe.Failsafe;
 import dev.failsafe.FailsafeException;
 import dev.failsafe.RetryPolicy;
+import io.trino.plugin.jdbc.jmx.StatisticsAwareConnectionFactory;
 import io.trino.spi.TrinoException;
 import io.trino.spi.connector.ConnectorSession;
 
@@ -40,10 +41,10 @@ public class RetryingConnectionFactory
             .abortOn(TrinoException.class)
             .build();
 
-    private final ConnectionFactory delegate;
+    private final StatisticsAwareConnectionFactory delegate;
 
     @Inject
-    public RetryingConnectionFactory(@StatsCollecting ConnectionFactory delegate)
+    public RetryingConnectionFactory(StatisticsAwareConnectionFactory delegate)
     {
         this.delegate = requireNonNull(delegate, "delegate is null");
     }

diff --git a/...no-base-jdbc/src/main/java/io/trino/plugin/jdbc/jmx/StatisticsAwareConnectionFactory.java b/...no-base-jdbc/src/main/java/io/trino/plugin/jdbc/jmx/StatisticsAwareConnectionFactory.java
@@ -13,7 +13,9 @@
  */
 package io.trino.plugin.jdbc.jmx;
 
+import com.google.inject.Inject;
 import io.trino.plugin.jdbc.ConnectionFactory;
+import io.trino.plugin.jdbc.ForBaseJdbc;
 import io.trino.spi.connector.ConnectorSession;
 import org.weakref.jmx.Managed;
 import org.weakref.jmx.Nested;
@@ -30,7 +32,8 @@ public class StatisticsAwareConnectionFactory
     private final JdbcApiStats closeConnection = new JdbcApiStats();
     private final ConnectionFactory delegate;
 
-    public StatisticsAwareConnectionFactory(ConnectionFactory delegate)
+    @Inject
+    public StatisticsAwareConnectionFactory(@ForBaseJdbc ConnectionFactory delegate)
     {
         this.delegate = requireNonNull(delegate, "delegate is null");
     }

diff --git a/plugin/trino-base-jdbc/src/test/java/io/trino/plugin/jdbc/TestLazyConnectionFactory.java b/plugin/trino-base-jdbc/src/test/java/io/trino/plugin/jdbc/TestLazyConnectionFactory.java
@@ -32,7 +32,7 @@ public class TestLazyConnectionFactory
     public void testNoConnectionIsCreated()
             throws Exception
     {
-        Injector injector = Guice.createInjector(binder -> binder.bind(ConnectionFactory.class).annotatedWith(ForLazyConnectionFactory.class).toInstance(
+        Injector injector = Guice.createInjector(binder -> binder.bind(ConnectionFactory.class).annotatedWith(ForBaseJdbc.class).toInstance(
                 session -> {
                     throw new AssertionError("Expected no connection creation");
                 }));
@@ -50,8 +50,8 @@ public void testConnectionCannotBeReusedAfterClose()
         BaseJdbcConfig config = new BaseJdbcConfig()
                 .setConnectionUrl(format("jdbc:h2:mem:test%s;DB_CLOSE_DELAY=-1", System.nanoTime() + ThreadLocalRandom.current().nextLong()));
 
-        Injector injector = Guice.createInjector(binder -> binder.bind(ConnectionFactory.class).annotatedWith(ForLazyConnectionFactory.class).toInstance(
-            new DriverConnectionFactory(new Driver(), config, new EmptyCredentialProvider())));
+        Injector injector = Guice.createInjector(binder -> binder.bind(ConnectionFactory.class).annotatedWith(ForBaseJdbc.class).toInstance(
+                new DriverConnectionFactory(new Driver(), config, new EmptyCredentialProvider())));
 
         try (LazyConnectionFactory lazyConnectionFactory = injector.getInstance(LazyConnectionFactory.class)) {
             Connection connection = lazyConnectionFactory.openConnection(SESSION);

diff --git a/plugin/trino-base-jdbc/src/test/java/io/trino/plugin/jdbc/TestRetryingConnectionFactory.java b/plugin/trino-base-jdbc/src/test/java/io/trino/plugin/jdbc/TestRetryingConnectionFactory.java
@@ -140,7 +140,7 @@ private static Injector createInjector(MockConnectorFactory.Action... actions)
         return Guice.createInjector(binder -> {
             binder.bind(MockConnectorFactory.Action[].class).toInstance(actions);
             binder.bind(MockConnectorFactory.class).in(Scopes.SINGLETON);
-            binder.bind(ConnectionFactory.class).annotatedWith(StatsCollecting.class).to(Key.get(MockConnectorFactory.class));
+            binder.bind(ConnectionFactory.class).annotatedWith(ForBaseJdbc.class).to(Key.get(MockConnectorFactory.class));
         });
     }
 

diff --git a/plugin/trino-oracle/src/main/java/io/trino/plugin/oracle/OracleClientModule.java b/plugin/trino-oracle/src/main/java/io/trino/plugin/oracle/OracleClientModule.java
@@ -28,6 +28,7 @@
 import io.trino.plugin.jdbc.MaxDomainCompactionThreshold;
 import io.trino.plugin.jdbc.RetryingConnectionFactory;
 import io.trino.plugin.jdbc.credential.CredentialProvider;
+import io.trino.plugin.jdbc.jmx.StatisticsAwareConnectionFactory;
 import io.trino.plugin.jdbc.ptf.Query;
 import io.trino.spi.function.table.ConnectorTableFunction;
 import oracle.jdbc.OracleConnection;
@@ -76,11 +77,11 @@ public static ConnectionFactory connectionFactory(BaseJdbcConfig config, Credent
                     openTelemetry);
         }
 
-        return new RetryingConnectionFactory(new DriverConnectionFactory(
+        return new RetryingConnectionFactory(new StatisticsAwareConnectionFactory(new DriverConnectionFactory(
                 new OracleDriver(),
                 config.getConnectionUrl(),
                 connectionProperties,
                 credentialProvider,
-                openTelemetry));
+                openTelemetry)));
     }
 }
diff --git a/plugin/trino-oracle/src/test/java/io/trino/plugin/oracle/TestingOracleServer.java b/plugin/trino-oracle/src/test/java/io/trino/plugin/oracle/TestingOracleServer.java
@@ -23,6 +23,7 @@
 import io.trino.plugin.jdbc.DriverConnectionFactory;
 import io.trino.plugin.jdbc.RetryingConnectionFactory;
 import io.trino.plugin.jdbc.credential.StaticCredentialProvider;
+import io.trino.plugin.jdbc.jmx.StatisticsAwareConnectionFactory;
 import io.trino.testing.ResourcePresence;
 import oracle.jdbc.OracleDriver;
 import org.testcontainers.containers.OracleContainer;
@@ -125,10 +126,10 @@ public void execute(String sql, String user, String password)
 
     private ConnectionFactory getConnectionFactory(String connectionUrl, String username, String password)
     {
-        DriverConnectionFactory connectionFactory = new DriverConnectionFactory(
+        StatisticsAwareConnectionFactory connectionFactory = new StatisticsAwareConnectionFactory(new DriverConnectionFactory(
                 new OracleDriver(),
                 new BaseJdbcConfig().setConnectionUrl(connectionUrl),
-                StaticCredentialProvider.of(username, password));
+                StaticCredentialProvider.of(username, password)));
         return new RetryingConnectionFactory(connectionFactory);
     }
 

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
@@ -34,7 +34,6 @@
 import io.trino.plugin.jdbc.DynamicFilteringStats;
 import io.trino.plugin.jdbc.ForBaseJdbc;
 import io.trino.plugin.jdbc.ForJdbcDynamicFiltering;
-import io.trino.plugin.jdbc.ForLazyConnectionFactory;
 import io.trino.plugin.jdbc.ForRecordCursor;
 import io.trino.plugin.jdbc.JdbcClient;
 import io.trino.plugin.jdbc.JdbcDiagnosticModule;
@@ -130,10 +129,6 @@ protected void setup(Binder binder)
         binder.bind(ConnectorMetadata.class).annotatedWith(ForClassLoaderSafe.class).to(PhoenixMetadata.class).in(Scopes.SINGLETON);
         binder.bind(ConnectorMetadata.class).to(ClassLoaderSafeConnectorMetadata.class).in(Scopes.SINGLETON);
 
-        binder.bind(ConnectionFactory.class)
-                .annotatedWith(ForLazyConnectionFactory.class)
-                .to(Key.get(ConnectionFactory.class, StatsCollecting.class))
-                .in(Scopes.SINGLETON);
         install(conditionalModule(
                 PhoenixConfig.class,
                 PhoenixConfig::isReuseConnection,