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 deleted file mode 100644 index 21ee99114bd5..000000000000 --- a/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/ForLazyConnectionFactory.java +++ /dev/null @@ -1,31 +0,0 @@ -/* - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package io.trino.plugin.jdbc; - -import com.google.inject.BindingAnnotation; - -import java.lang.annotation.Retention; -import java.lang.annotation.Target; - -import static java.lang.annotation.ElementType.FIELD; -import static java.lang.annotation.ElementType.METHOD; -import static java.lang.annotation.ElementType.PARAMETER; -import static java.lang.annotation.RetentionPolicy.RUNTIME; - -@Retention(RUNTIME) -@Target({FIELD, PARAMETER, METHOD}) -@BindingAnnotation -public @interface ForLazyConnectionFactory -{ -} 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 index 557095b85053..93bed3852ec8 100644 --- 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 = 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 index fbb3d98760aa..ba18738c0d96 100644 --- 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 @@ -50,6 +50,7 @@ public void setup(Binder binder) install(new JdbcDiagnosticModule()); install(new IdentifierMappingModule()); install(new RemoteQueryModifierModule()); + install(new RetryingConnectionFactoryModule()); newOptionalBinder(binder, ConnectorAccessControl.class); newOptionalBinder(binder, QueryBuilder.class).setDefault().to(DefaultQueryBuilder.class).in(Scopes.SINGLETON); @@ -88,10 +89,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 index f056e07c7b95..284cc1ef8d01 100644 --- 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 @@ -32,7 +32,7 @@ public final class LazyConnectionFactory private final ConnectionFactory delegate; @Inject - public LazyConnectionFactory(@ForLazyConnectionFactory ConnectionFactory delegate) + public LazyConnectionFactory(RetryingConnectionFactory 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 index 6063feaa0ecf..0993cdb26e59 100644 --- 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 @@ -14,15 +14,17 @@ package io.trino.plugin.jdbc; import com.google.common.base.Throwables; +import com.google.inject.Inject; 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; import java.sql.Connection; import java.sql.SQLException; -import java.sql.SQLRecoverableException; +import java.sql.SQLTransientException; import static java.time.temporal.ChronoUnit.MILLIS; import static java.time.temporal.ChronoUnit.SECONDS; @@ -31,19 +33,22 @@ public class RetryingConnectionFactory implements ConnectionFactory { - private static final RetryPolicy RETRY_POLICY = RetryPolicy.builder() - .withMaxDuration(java.time.Duration.of(30, SECONDS)) - .withMaxAttempts(5) - .withBackoff(50, 5_000, MILLIS, 4) - .handleIf(RetryingConnectionFactory::isSqlRecoverableException) - .abortOn(TrinoException.class) - .build(); + private final RetryPolicy retryPolicy; private final ConnectionFactory delegate; - public RetryingConnectionFactory(ConnectionFactory delegate) + @Inject + public RetryingConnectionFactory(StatisticsAwareConnectionFactory delegate, RetryStrategy retryStrategy) { + requireNonNull(retryStrategy); this.delegate = requireNonNull(delegate, "delegate is null"); + this.retryPolicy = RetryPolicy.builder() + .withMaxDuration(java.time.Duration.of(30, SECONDS)) + .withMaxAttempts(5) + .withBackoff(50, 5_000, MILLIS, 4) + .handleIf(retryStrategy::isExceptionRecoverable) + .abortOn(TrinoException.class) + .build(); } @Override @@ -51,7 +56,7 @@ public Connection openConnection(ConnectorSession session) throws SQLException { try { - return Failsafe.with(RETRY_POLICY) + return Failsafe.with(retryPolicy) .get(() -> delegate.openConnection(session)); } catch (FailsafeException ex) { @@ -69,9 +74,19 @@ public void close() delegate.close(); } - private static boolean isSqlRecoverableException(Throwable exception) + public interface RetryStrategy { - return Throwables.getCausalChain(exception).stream() - .anyMatch(SQLRecoverableException.class::isInstance); + boolean isExceptionRecoverable(Throwable exception); + } + + public static class DefaultRetryStrategy + implements RetryStrategy + { + @Override + public boolean isExceptionRecoverable(Throwable exception) + { + return Throwables.getCausalChain(exception).stream() + .anyMatch(SQLTransientException.class::isInstance); + } } } diff --git a/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/RetryingConnectionFactoryModule.java b/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/RetryingConnectionFactoryModule.java new file mode 100644 index 000000000000..a0815d38e84c --- /dev/null +++ b/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/RetryingConnectionFactoryModule.java @@ -0,0 +1,35 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.jdbc; + +import com.google.inject.AbstractModule; +import com.google.inject.Scopes; +import io.trino.plugin.jdbc.RetryingConnectionFactory.DefaultRetryStrategy; +import io.trino.plugin.jdbc.RetryingConnectionFactory.RetryStrategy; + +import static com.google.inject.multibindings.OptionalBinder.newOptionalBinder; + +public class RetryingConnectionFactoryModule + extends AbstractModule +{ + @Override + public void configure() + { + bind(RetryingConnectionFactory.class).in(Scopes.SINGLETON); + newOptionalBinder(binder(), RetryStrategy.class) + .setDefault() + .to(DefaultRetryStrategy.class) + .in(Scopes.SINGLETON); + } +} diff --git a/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/jmx/StatisticsAwareConnectionFactory.java b/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/jmx/StatisticsAwareConnectionFactory.java index 6c9153997a06..9f59238ba790 100644 --- a/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/jmx/StatisticsAwareConnectionFactory.java +++ b/plugin/trino-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 index 1f6ddb9a9d0a..cd0fb3dd1d6e 100644 --- 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 @@ -13,6 +13,8 @@ */ package io.trino.plugin.jdbc; +import com.google.inject.Guice; +import com.google.inject.Injector; import io.trino.plugin.jdbc.credential.EmptyCredentialProvider; import org.h2.Driver; import org.junit.jupiter.api.Test; @@ -30,11 +32,15 @@ public class TestLazyConnectionFactory public void testNoConnectionIsCreated() throws Exception { - ConnectionFactory failingConnectionFactory = session -> { - throw new AssertionError("Expected no connection creation"); - }; + Injector injector = Guice.createInjector(binder -> { + binder.bind(ConnectionFactory.class).annotatedWith(ForBaseJdbc.class).toInstance( + session -> { + throw new AssertionError("Expected no connection creation"); + }); + binder.install(new RetryingConnectionFactoryModule()); + }); - try (LazyConnectionFactory lazyConnectionFactory = new LazyConnectionFactory(failingConnectionFactory); + try (LazyConnectionFactory lazyConnectionFactory = injector.getInstance(LazyConnectionFactory.class); Connection ignored = lazyConnectionFactory.openConnection(SESSION)) { // no-op } @@ -47,8 +53,13 @@ public void testConnectionCannotBeReusedAfterClose() BaseJdbcConfig config = new BaseJdbcConfig() .setConnectionUrl(format("jdbc:h2:mem:test%s;DB_CLOSE_DELAY=-1", System.nanoTime() + ThreadLocalRandom.current().nextLong())); - try (DriverConnectionFactory h2ConnectionFactory = new DriverConnectionFactory(new Driver(), config, new EmptyCredentialProvider()); - LazyConnectionFactory lazyConnectionFactory = new LazyConnectionFactory(h2ConnectionFactory)) { + Injector injector = Guice.createInjector(binder -> { + binder.bind(ConnectionFactory.class).annotatedWith(ForBaseJdbc.class).toInstance( + new DriverConnectionFactory(new Driver(), config, new EmptyCredentialProvider())); + binder.install(new RetryingConnectionFactoryModule()); + }); + + try (LazyConnectionFactory lazyConnectionFactory = injector.getInstance(LazyConnectionFactory.class)) { Connection connection = lazyConnectionFactory.openConnection(SESSION); connection.close(); assertThatThrownBy(() -> connection.createStatement()) 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 index 985136fc422e..d85c1c5d1ef8 100644 --- 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 @@ -13,6 +13,13 @@ */ package io.trino.plugin.jdbc; +import com.google.common.base.Throwables; +import com.google.inject.Guice; +import com.google.inject.Inject; +import com.google.inject.Injector; +import com.google.inject.Key; +import com.google.inject.Scopes; +import io.trino.plugin.jdbc.RetryingConnectionFactory.RetryStrategy; import io.trino.spi.StandardErrorCode; import io.trino.spi.TrinoException; import io.trino.spi.connector.ConnectorSession; @@ -21,17 +28,20 @@ import java.sql.Connection; import java.sql.SQLException; import java.sql.SQLRecoverableException; +import java.sql.SQLTransientException; import java.util.ArrayDeque; import java.util.Deque; import java.util.stream.Stream; import static com.google.common.reflect.Reflection.newProxy; +import static com.google.inject.multibindings.OptionalBinder.newOptionalBinder; import static io.trino.plugin.jdbc.TestRetryingConnectionFactory.MockConnectorFactory.Action.RETURN; import static io.trino.plugin.jdbc.TestRetryingConnectionFactory.MockConnectorFactory.Action.THROW_NPE; import static io.trino.plugin.jdbc.TestRetryingConnectionFactory.MockConnectorFactory.Action.THROW_SQL_EXCEPTION; import static io.trino.plugin.jdbc.TestRetryingConnectionFactory.MockConnectorFactory.Action.THROW_SQL_RECOVERABLE_EXCEPTION; +import static io.trino.plugin.jdbc.TestRetryingConnectionFactory.MockConnectorFactory.Action.THROW_SQL_TRANSIENT_EXCEPTION; import static io.trino.plugin.jdbc.TestRetryingConnectionFactory.MockConnectorFactory.Action.THROW_TRINO_EXCEPTION; -import static io.trino.plugin.jdbc.TestRetryingConnectionFactory.MockConnectorFactory.Action.THROW_WRAPPED_SQL_RECOVERABLE_EXCEPTION; +import static io.trino.plugin.jdbc.TestRetryingConnectionFactory.MockConnectorFactory.Action.THROW_WRAPPED_SQL_TRANSIENT_EXCEPTION; import static io.trino.spi.block.TestingSession.SESSION; import static io.trino.spi.testing.InterfaceTestUtils.assertAllMethodsOverridden; import static java.util.Objects.requireNonNull; @@ -50,42 +60,55 @@ public void testEverythingImplemented() public void testSimplyReturnConnection() throws Exception { - MockConnectorFactory mock = new MockConnectorFactory(RETURN); - ConnectionFactory factory = new RetryingConnectionFactory(mock); - assertThat(factory.openConnection(SESSION)).isNotNull(); + Injector injector = createInjector(RETURN); + ConnectionFactory factory = injector.getInstance(RetryingConnectionFactory.class); + MockConnectorFactory mock = injector.getInstance(MockConnectorFactory.class); + + Connection connection = factory.openConnection(SESSION); + + assertThat(connection).isNotNull(); assertThat(mock.getCallCount()).isEqualTo(1); } @Test public void testRetryAndStopOnTrinoException() { - MockConnectorFactory mock = new MockConnectorFactory(THROW_SQL_RECOVERABLE_EXCEPTION, THROW_TRINO_EXCEPTION); - ConnectionFactory factory = new RetryingConnectionFactory(mock); + Injector injector = createInjector(THROW_SQL_TRANSIENT_EXCEPTION, THROW_TRINO_EXCEPTION); + ConnectionFactory factory = injector.getInstance(RetryingConnectionFactory.class); + MockConnectorFactory mock = injector.getInstance(MockConnectorFactory.class); + assertThatThrownBy(() -> factory.openConnection(SESSION)) .isInstanceOf(TrinoException.class) .hasMessage("Testing Trino exception"); + assertThat(mock.getCallCount()).isEqualTo(2); } @Test public void testRetryAndStopOnSqlException() { - MockConnectorFactory mock = new MockConnectorFactory(THROW_SQL_RECOVERABLE_EXCEPTION, THROW_SQL_EXCEPTION); - ConnectionFactory factory = new RetryingConnectionFactory(mock); + Injector injector = createInjector(THROW_SQL_TRANSIENT_EXCEPTION, THROW_SQL_EXCEPTION); + ConnectionFactory factory = injector.getInstance(RetryingConnectionFactory.class); + MockConnectorFactory mock = injector.getInstance(MockConnectorFactory.class); + assertThatThrownBy(() -> factory.openConnection(SESSION)) .isInstanceOf(SQLException.class) .hasMessage("Testing sql exception"); + assertThat(mock.getCallCount()).isEqualTo(2); } @Test public void testNullPointerException() { - MockConnectorFactory mock = new MockConnectorFactory(THROW_NPE); - ConnectionFactory factory = new RetryingConnectionFactory(mock); + Injector injector = createInjector(THROW_NPE); + ConnectionFactory factory = injector.getInstance(RetryingConnectionFactory.class); + MockConnectorFactory mock = injector.getInstance(MockConnectorFactory.class); + assertThatThrownBy(() -> factory.openConnection(SESSION)) .isInstanceOf(NullPointerException.class) .hasMessage("Testing NPE"); + assertThat(mock.getCallCount()).isEqualTo(1); } @@ -93,9 +116,13 @@ public void testNullPointerException() public void testRetryAndReturn() throws Exception { - MockConnectorFactory mock = new MockConnectorFactory(THROW_SQL_RECOVERABLE_EXCEPTION, RETURN); - ConnectionFactory factory = new RetryingConnectionFactory(mock); - assertThat(factory.openConnection(SESSION)).isNotNull(); + Injector injector = createInjector(THROW_SQL_TRANSIENT_EXCEPTION, RETURN); + ConnectionFactory factory = injector.getInstance(RetryingConnectionFactory.class); + MockConnectorFactory mock = injector.getInstance(MockConnectorFactory.class); + + Connection connection = factory.openConnection(SESSION); + + assertThat(connection).isNotNull(); assertThat(mock.getCallCount()).isEqualTo(2); } @@ -103,18 +130,69 @@ public void testRetryAndReturn() public void testRetryOnWrappedAndReturn() throws Exception { - MockConnectorFactory mock = new MockConnectorFactory(THROW_WRAPPED_SQL_RECOVERABLE_EXCEPTION, RETURN); - ConnectionFactory factory = new RetryingConnectionFactory(mock); - assertThat(factory.openConnection(SESSION)).isNotNull(); + Injector injector = createInjector(THROW_WRAPPED_SQL_TRANSIENT_EXCEPTION, RETURN); + ConnectionFactory factory = injector.getInstance(RetryingConnectionFactory.class); + MockConnectorFactory mock = injector.getInstance(MockConnectorFactory.class); + + Connection connection = factory.openConnection(SESSION); + + assertThat(connection).isNotNull(); + assertThat(mock.getCallCount()).isEqualTo(2); + } + + @Test + public void testOverridingRetryStrategyWorks() + throws Exception + { + Injector injector = createInjectorWithOverridenStrategy(THROW_SQL_RECOVERABLE_EXCEPTION, RETURN); + ConnectionFactory factory = injector.getInstance(RetryingConnectionFactory.class); + MockConnectorFactory mock = injector.getInstance(MockConnectorFactory.class); + + Connection connection = factory.openConnection(SESSION); + + assertThat(connection).isNotNull(); assertThat(mock.getCallCount()).isEqualTo(2); } + 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(ForBaseJdbc.class).to(Key.get(MockConnectorFactory.class)); + binder.install(new RetryingConnectionFactoryModule()); + }); + } + + private static Injector createInjectorWithOverridenStrategy(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(ForBaseJdbc.class).to(Key.get(MockConnectorFactory.class)); + binder.install(new RetryingConnectionFactoryModule()); + newOptionalBinder(binder, RetryStrategy.class).setBinding().to(OverrideRetryStrategy.class).in(Scopes.SINGLETON); + }); + } + + private static class OverrideRetryStrategy + implements RetryStrategy + { + @Override + public boolean isExceptionRecoverable(Throwable exception) + { + return Throwables.getCausalChain(exception).stream() + .anyMatch(SQLRecoverableException.class::isInstance); + } + } + public static class MockConnectorFactory implements ConnectionFactory { private final Deque actions = new ArrayDeque<>(); private int callCount; + @Inject public MockConnectorFactory(Action... actions) { Stream.of(actions) @@ -145,6 +223,10 @@ public Connection openConnection(ConnectorSession session) throw new SQLRecoverableException("Testing sql recoverable exception"); case THROW_WRAPPED_SQL_RECOVERABLE_EXCEPTION: throw new RuntimeException(new SQLRecoverableException("Testing sql recoverable exception")); + case THROW_SQL_TRANSIENT_EXCEPTION: + throw new SQLTransientException("Testing sql transient exception"); + case THROW_WRAPPED_SQL_TRANSIENT_EXCEPTION: + throw new RuntimeException(new SQLTransientException("Testing sql transient exception")); } throw new IllegalStateException("Unsupported action:" + action); } @@ -155,6 +237,8 @@ public enum Action THROW_SQL_EXCEPTION, THROW_SQL_RECOVERABLE_EXCEPTION, THROW_WRAPPED_SQL_RECOVERABLE_EXCEPTION, + THROW_SQL_TRANSIENT_EXCEPTION, + THROW_WRAPPED_SQL_TRANSIENT_EXCEPTION, THROW_NPE, RETURN, } 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 index af36cb507e88..d3c2f545861c 100644 --- 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 @@ -13,6 +13,7 @@ */ package io.trino.plugin.oracle; +import com.google.common.base.Throwables; import com.google.inject.Binder; import com.google.inject.Key; import com.google.inject.Module; @@ -26,7 +27,7 @@ import io.trino.plugin.jdbc.ForBaseJdbc; import io.trino.plugin.jdbc.JdbcClient; import io.trino.plugin.jdbc.MaxDomainCompactionThreshold; -import io.trino.plugin.jdbc.RetryingConnectionFactory; +import io.trino.plugin.jdbc.RetryingConnectionFactory.RetryStrategy; import io.trino.plugin.jdbc.credential.CredentialProvider; import io.trino.plugin.jdbc.ptf.Query; import io.trino.spi.function.table.ConnectorTableFunction; @@ -34,6 +35,7 @@ import oracle.jdbc.OracleDriver; import java.sql.SQLException; +import java.sql.SQLRecoverableException; import java.util.Properties; import static com.google.inject.multibindings.Multibinder.newSetBinder; @@ -53,6 +55,7 @@ public void configure(Binder binder) configBinder(binder).bindConfig(OracleConfig.class); newOptionalBinder(binder, Key.get(int.class, MaxDomainCompactionThreshold.class)).setBinding().toInstance(ORACLE_MAX_LIST_EXPRESSIONS); newSetBinder(binder, ConnectorTableFunction.class).addBinding().toProvider(Query.class).in(Scopes.SINGLETON); + newOptionalBinder(binder, RetryStrategy.class).setBinding().to(OracleRetryStrategy.class).in(Scopes.SINGLETON); } @Provides @@ -76,11 +79,22 @@ public static ConnectionFactory connectionFactory(BaseJdbcConfig config, Credent openTelemetry); } - return new RetryingConnectionFactory(new DriverConnectionFactory( + return new DriverConnectionFactory( new OracleDriver(), config.getConnectionUrl(), connectionProperties, credentialProvider, - openTelemetry)); + openTelemetry); + } + + private static class OracleRetryStrategy + implements RetryStrategy + { + @Override + public boolean isExceptionRecoverable(Throwable exception) + { + return Throwables.getCausalChain(exception).stream() + .anyMatch(SQLRecoverableException.class::isInstance); + } } } 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 index f73ae90215f3..6bf8e1ed04d9 100644 --- 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 @@ -22,7 +22,9 @@ import io.trino.plugin.jdbc.ConnectionFactory; import io.trino.plugin.jdbc.DriverConnectionFactory; import io.trino.plugin.jdbc.RetryingConnectionFactory; +import io.trino.plugin.jdbc.RetryingConnectionFactory.DefaultRetryStrategy; 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,11 +127,11 @@ 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)); - return new RetryingConnectionFactory(connectionFactory); + StaticCredentialProvider.of(username, password))); + return new RetryingConnectionFactory(connectionFactory, new DefaultRetryStrategy()); } @Override 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 index 852d76723aa4..c18bdc26d041 100644 --- 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; @@ -48,6 +47,7 @@ import io.trino.plugin.jdbc.LazyConnectionFactory; import io.trino.plugin.jdbc.MaxDomainCompactionThreshold; import io.trino.plugin.jdbc.QueryBuilder; +import io.trino.plugin.jdbc.RetryingConnectionFactoryModule; import io.trino.plugin.jdbc.ReusableConnectionFactoryModule; import io.trino.plugin.jdbc.StatsCollecting; import io.trino.plugin.jdbc.TypeHandlingJdbcConfig; @@ -97,6 +97,7 @@ public PhoenixClientModule(String catalogName) protected void setup(Binder binder) { install(new RemoteQueryModifierModule()); + install(new RetryingConnectionFactoryModule()); binder.bind(ConnectorSplitManager.class).annotatedWith(ForJdbcDynamicFiltering.class).to(PhoenixSplitManager.class).in(Scopes.SINGLETON); binder.bind(ConnectorSplitManager.class).annotatedWith(ForClassLoaderSafe.class).to(JdbcDynamicFilteringSplitManager.class).in(Scopes.SINGLETON); binder.bind(ConnectorSplitManager.class).to(ClassLoaderSafeConnectorSplitManager.class).in(Scopes.SINGLETON); @@ -130,10 +131,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,