exception simulator manager

aws · Jul 7, 2023 · fd70e5c · fd70e5c
1 parent e23a5e3
commit fd70e5c
Show file tree

Hide file tree

Showing 7 changed files with 177 additions and 28 deletions.
diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/dev/DeveloperConnectionPlugin.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/dev/DeveloperConnectionPlugin.java
@@ -24,10 +24,8 @@
 import java.util.Set;
 import java.util.logging.Logger;
 import org.checkerframework.checker.nullness.qual.NonNull;
-import software.amazon.jdbc.HostListProviderService;
 import software.amazon.jdbc.HostSpec;
 import software.amazon.jdbc.JdbcCallable;
-import software.amazon.jdbc.PluginService;
 import software.amazon.jdbc.plugin.AbstractConnectionPlugin;
 import software.amazon.jdbc.util.StringUtils;
 import software.amazon.jdbc.util.WrapperUtils;
@@ -44,17 +42,9 @@ public class DeveloperConnectionPlugin extends AbstractConnectionPlugin implemen
 
   private String nextMethodName;
   private Throwable nextException;
-  private ExceptionSimulatorCallback exceptionSimulatorCallback;
-  private final PluginService pluginService;
-  private final Properties props;
+  private ExceptionSimulatorExecuteJdbcMethodCallback exceptionSimulatorExecuteJdbcMethodCallback;
 
-  public DeveloperConnectionPlugin(
-      final PluginService pluginService,
-      final Properties props) {
-
-    this.pluginService = pluginService;
-    this.props = props;
-  }
+  public DeveloperConnectionPlugin() { }
 
   @Override
   public Set<String> getSubscribedMethods() {
@@ -77,8 +67,8 @@ public void raiseExceptionOnNextCall(final @NonNull String methodName, final Thr
   }
 
   @Override
-  public void setCallback(final ExceptionSimulatorCallback exceptionSimulatorCallback) {
-    this.exceptionSimulatorCallback = exceptionSimulatorCallback;
+  public void setCallback(final ExceptionSimulatorExecuteJdbcMethodCallback exceptionSimulatorExecuteJdbcMethodCallback) {
+    this.exceptionSimulatorExecuteJdbcMethodCallback = exceptionSimulatorExecuteJdbcMethodCallback;
   }
 
   @Override
@@ -109,34 +99,108 @@ protected <T, E extends Exception> void raiseExceptionIfNeeded(
 
     if (this.nextException != null) {
       if (ALL_METHODS.equals(this.nextMethodName) || methodName.equals(this.nextMethodName)) {
-        this.raiseException(exceptionClass, this.nextException);
+        this.raiseException(exceptionClass, this.nextException, methodName);
       }
-    } else if (this.exceptionSimulatorCallback != null) {
-      Throwable userException = this.exceptionSimulatorCallback.getExceptionToRaise(
-          resultClass,
+    } else if (this.exceptionSimulatorExecuteJdbcMethodCallback != null) {
+      this.raiseException(
           exceptionClass,
-          methodName,
-          jdbcMethodArgs);
-
-      if (userException != null) {
-        this.raiseException(exceptionClass, this.nextException);
-      }
+          this.exceptionSimulatorExecuteJdbcMethodCallback.getExceptionToRaise(
+              resultClass,
+              exceptionClass,
+              methodName,
+              jdbcMethodArgs),
+          methodName);
     }
   }
 
   protected <E extends Exception> void raiseException(
       final Class<E> exceptionClass,
-      final Throwable throwable)
+      final Throwable throwable,
+      final String methodName)
       throws E {
 
+    if (throwable == null) {
+      return;
+    }
+
     if (throwable instanceof RuntimeException) {
       this.nextException = null;
       this.nextMethodName = null;
+      LOGGER.finest(() -> String.format("Raise an exception %s while executing %s.",
+          throwable.getClass().getName(), methodName));
       throw (RuntimeException) throwable;
     } else {
       E resulException = WrapperUtils.wrapExceptionIfNeeded(exceptionClass, throwable);
       this.nextException = null;
       this.nextMethodName = null;
+      LOGGER.finest(() -> String.format("Raise an exception %s while executing %s.",
+          resulException.getClass().getName(), methodName));
+      throw resulException;
+    }
+  }
+
+  @Override
+  public Connection connect(
+      final String driverProtocol,
+      final HostSpec hostSpec,
+      final Properties props,
+      final boolean isInitialConnection,
+      final JdbcCallable<Connection, SQLException> connectFunc)
+      throws SQLException {
+
+    this.raiseExceptionOnConnectIfNeeded(driverProtocol, hostSpec, props, isInitialConnection);
+    return super.connect(driverProtocol, hostSpec, props, isInitialConnection, connectFunc);
+  }
+
+  @Override
+  public Connection forceConnect(
+      final String driverProtocol,
+      final HostSpec hostSpec,
+      final Properties props,
+      final boolean isInitialConnection,
+      final JdbcCallable<Connection, SQLException> forceConnectFunc)
+      throws SQLException {
+
+    this.raiseExceptionOnConnectIfNeeded(driverProtocol, hostSpec, props, isInitialConnection);
+    return super.connect(driverProtocol, hostSpec, props, isInitialConnection, forceConnectFunc);
+  }
+
+  protected void raiseExceptionOnConnectIfNeeded(
+      final String driverProtocol,
+      final HostSpec hostSpec,
+      final Properties props,
+      final boolean isInitialConnection)
+      throws SQLException {
+
+    if (ExceptionSimulatorManager.nextException != null) {
+      this.raiseExceptionOnConnect(ExceptionSimulatorManager.nextException);
+    } else if (ExceptionSimulatorManager.connectCallback != null) {
+      this.raiseExceptionOnConnect(
+          ExceptionSimulatorManager.connectCallback.getExceptionToRaise(
+            driverProtocol,
+            hostSpec,
+            props,
+            isInitialConnection));
+    }
+  }
+
+  protected void raiseExceptionOnConnect(final Throwable throwable)
+      throws SQLException {
+
+    if (throwable == null) {
+      return;
+    }
+
+    if (throwable instanceof RuntimeException) {
+      ExceptionSimulatorManager.nextException = null;
+      LOGGER.finest(() -> String.format("Raise an exception %s while opening a new connection.",
+          ExceptionSimulatorManager.nextException.getClass().getName()));
+      throw (RuntimeException) throwable;
+    } else {
+      SQLException resulException = WrapperUtils.wrapExceptionIfNeeded(SQLException.class, throwable);
+      ExceptionSimulatorManager.nextException = null;
+      LOGGER.finest(() -> String.format("Raise an exception %s while opening a new connection.",
+          resulException.getClass().getName()));
       throw resulException;
     }
   }

diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/dev/DeveloperConnectionPluginFactory.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/dev/DeveloperConnectionPluginFactory.java
@@ -25,6 +25,6 @@ public class DeveloperConnectionPluginFactory implements ConnectionPluginFactory
 
   @Override
   public ConnectionPlugin getInstance(PluginService pluginService, Properties props) {
-    return new DeveloperConnectionPlugin(pluginService, props);
+    return new DeveloperConnectionPlugin();
   }
 }
diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/dev/ExceptionSimulator.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/dev/ExceptionSimulator.java
@@ -22,5 +22,5 @@ public interface ExceptionSimulator {
 
   void raiseExceptionOnNextCall(final String methodName, final Throwable throwable);
 
-  void setCallback(final ExceptionSimulatorCallback exceptionSimulatorCallback);
+  void setCallback(final ExceptionSimulatorExecuteJdbcMethodCallback exceptionSimulatorExecuteJdbcMethodCallback);
 }
diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/dev/ExceptionSimulatorConnectCallback.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/dev/ExceptionSimulatorConnectCallback.java
@@ -0,0 +1,13 @@
+package software.amazon.jdbc.plugin.dev;
+
+import java.sql.SQLException;
+import java.util.Properties;
+import software.amazon.jdbc.HostSpec;
+
+public interface ExceptionSimulatorConnectCallback {
+  SQLException getExceptionToRaise(
+      final String driverProtocol,
+      final HostSpec hostSpec,
+      final Properties props,
+      final boolean isInitialConnection);
+}
diff --git a/...lugin/dev/ExceptionSimulatorCallback.java → ...onSimulatorExecuteJdbcMethodCallback.java b/...lugin/dev/ExceptionSimulatorCallback.java → ...onSimulatorExecuteJdbcMethodCallback.java
@@ -16,7 +16,7 @@
 
 package software.amazon.jdbc.plugin.dev;
 
-public interface ExceptionSimulatorCallback {
+public interface ExceptionSimulatorExecuteJdbcMethodCallback {
   <T, E extends Exception> Throwable getExceptionToRaise(
       final Class<T> resultClass,
       final Class<E> exceptionClass,

diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/dev/ExceptionSimulatorManager.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/dev/ExceptionSimulatorManager.java
@@ -0,0 +1,15 @@
+package software.amazon.jdbc.plugin.dev;
+
+public class ExceptionSimulatorManager {
+
+  static Throwable nextException;
+  static ExceptionSimulatorConnectCallback connectCallback;
+
+  public static void raiseExceptionOnNextConnect(final Throwable throwable) {
+    nextException = throwable;
+  }
+
+  public static void setCallback(final ExceptionSimulatorConnectCallback exceptionSimulatorConnectCallback) {
+    connectCallback = exceptionSimulatorConnectCallback;
+  }
+}
diff --git a/wrapper/src/test/java/software/amazon/jdbc/plugin/dev/DeveloperConnectionPluginTest.java b/wrapper/src/test/java/software/amazon/jdbc/plugin/dev/DeveloperConnectionPluginTest.java
@@ -23,6 +23,7 @@
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.Mockito.when;
 
 import java.sql.Connection;
@@ -43,6 +44,7 @@ public class DeveloperConnectionPluginTest {
 
   @Mock ConnectionProvider mockConnectionProvider;
   @Mock Connection mockConnection;
+  @Mock ExceptionSimulatorConnectCallback mockConnectCallback;
 
   private AutoCloseable closeable;
 
@@ -56,6 +58,7 @@ void init() throws SQLException {
     closeable = MockitoAnnotations.openMocks(this);
 
     when(mockConnectionProvider.connect(any(), any(), any(), any())).thenReturn(mockConnection);
+    when(mockConnectCallback.getExceptionToRaise(any(), any(), any(), anyBoolean())).thenReturn(null);
   }
 
   @Test
@@ -201,4 +204,58 @@ public void test_RaiseUnexpectedExceptionClass() throws SQLException {
       assertDoesNotThrow(() -> wrapper.createStatement());
     }
   }
+
+  @Test
+  public void test_RaiseExceptionOnConnect() {
+
+    final Properties props = new Properties();
+    props.put(PropertyDefinition.PLUGINS.name, "dev");
+    props.put(DialectManager.DIALECT.name, DialectCodes.PG);
+
+    final SQLException exception = new SQLException("test");
+    ExceptionSimulatorManager.raiseExceptionOnNextConnect(exception);
+
+    Throwable thrownException = assertThrows(
+        SQLException.class,
+        () -> new ConnectionWrapper(props, "any-protocol://any-host/", mockConnectionProvider));
+    assertSame(exception, thrownException);
+
+    assertDoesNotThrow(
+        () -> new ConnectionWrapper(props, "any-protocol://any-host/", mockConnectionProvider));
+  }
+
+  @Test
+  public void test_NoExceptionOnConnectWithCallback() {
+
+    final Properties props = new Properties();
+    props.put(PropertyDefinition.PLUGINS.name, "dev");
+    props.put(DialectManager.DIALECT.name, DialectCodes.PG);
+
+    ExceptionSimulatorManager.setCallback(mockConnectCallback);
+
+    assertDoesNotThrow(
+        () -> new ConnectionWrapper(props, "any-protocol://any-host/", mockConnectionProvider));
+  }
+
+  @Test
+  public void test_RaiseExceptionOnConnectWithCallback() {
+
+    final Properties props = new Properties();
+    props.put(PropertyDefinition.PLUGINS.name, "dev");
+    props.put(DialectManager.DIALECT.name, DialectCodes.PG);
+
+    final SQLException exception = new SQLException("test");
+    when(mockConnectCallback.getExceptionToRaise(any(), any(), any(), anyBoolean()))
+        .thenReturn(exception)
+        .thenReturn(null);
+    ExceptionSimulatorManager.setCallback(mockConnectCallback);
+
+    Throwable thrownException = assertThrows(
+        SQLException.class,
+        () -> new ConnectionWrapper(props, "any-protocol://any-host/", mockConnectionProvider));
+    assertSame(exception, thrownException);
+
+    assertDoesNotThrow(
+        () -> new ConnectionWrapper(props, "any-protocol://any-host/", mockConnectionProvider));
+  }
 }