Resetting Schema and Catalog, when connection returend to pool

When a connection is taken from the pool and schema and/or catalog will
be changed, the schema/catalog is not reset. So the next one will get a
connection, that is not in the expected schema/catalog

Author: rPraml

ebean-datasource/src/main/java/io/ebean/datasource/pool/ConnectionDelegator.java

@@ -7,7 +7,6 @@
 abstract class ConnectionDelegator implements Connection {
 
   private final Connection delegate;
-  protected String currentSchema;
 
   ConnectionDelegator(Connection delegate) {
     this.delegate = delegate;
@@ -20,12 +19,6 @@ Connection delegate() {
     return delegate;
   }
 
-  @Override
-  public final void setSchema(String schema) throws SQLException {
-    delegate.setSchema(schema);
-    currentSchema = schema;
-  }
-
   @Override
   public final String getSchema() throws SQLException {
     return delegate.getSchema();

ebean-datasource/src/main/java/io/ebean/datasource/pool/PooledConnection.java

@@ -82,6 +82,13 @@ final class PooledConnection extends ConnectionDelegator {
    */
   private boolean failoverToReadOnly;
   private boolean resetAutoCommit;
+  private boolean resetSchema;
+  private boolean resetCatalog;
+  private String currentSchema;
+  private String currentCatalog;
+  private final String originalSchema;
+  private final String originalCatalog;
+
   private long startUseTime;
   private long lastUseTime;
   /**
@@ -106,7 +113,7 @@ final class PooledConnection extends ConnectionDelegator {
    * close() will return the connection back to the pool , while
    * closeDestroy() will close() the underlining connection properly.
    */
-  PooledConnection(ConnectionPool pool, int uniqueId, Connection connection) {
+  PooledConnection(ConnectionPool pool, int uniqueId, Connection connection) throws SQLException {
     super(connection);
     this.pool = pool;
     this.connection = connection;
@@ -115,6 +122,8 @@ final class PooledConnection extends ConnectionDelegator {
     this.maxStackTrace = pool.maxStackTraceSize();
     this.creationTime = System.currentTimeMillis();
     this.lastUseTime = creationTime;
+    this.currentSchema = this.originalSchema = connection.getSchema();
+    this.currentCatalog = this.originalCatalog = connection.getCatalog();
     pool.inc();
   }
 
@@ -130,6 +139,8 @@ final class PooledConnection extends ConnectionDelegator {
     this.maxStackTrace = 0;
     this.creationTime = System.currentTimeMillis();
     this.lastUseTime = creationTime;
+    this.currentSchema = this.originalSchema = "DEFAULT";
+    this.currentCatalog = this.originalCatalog = "DEFAULT";
   }
 
   /**
@@ -272,7 +283,7 @@ void returnPreparedStatement(ExtendedPreparedStatement pstmt) {
    */
   @Override
   public PreparedStatement prepareStatement(String sql, int returnKeysFlag) throws SQLException {
-    String key = sql + ':' + currentSchema + ':' + returnKeysFlag;
+    String key = sql + ':' + currentSchema + ':' + currentCatalog + ':' + returnKeysFlag;
     return prepareStatement(sql, true, returnKeysFlag, key);
   }
 
@@ -281,7 +292,7 @@ public PreparedStatement prepareStatement(String sql, int returnKeysFlag) throws
    */
   @Override
   public PreparedStatement prepareStatement(String sql) throws SQLException {
-    String key = sql + ':' + currentSchema;
+    String key = sql + ':' + currentSchema + ':' + currentCatalog;
     return prepareStatement(sql, false, 0, key);
   }
 
@@ -411,6 +422,16 @@ public void close() throws SQLException {
         resetIsolationReadOnlyRequired = false;
       }
 
+      if (resetSchema) {
+        connection.setSchema(originalSchema);
+        resetSchema = false;
+      }
+
+      if (resetCatalog) {
+        connection.setCatalog(originalCatalog);
+        resetCatalog = false;
+      }
+
       // the connection is assumed GOOD so put it back in the pool
       lastUseTime = System.currentTimeMillis();
       connection.clearWarnings();
@@ -511,6 +532,7 @@ public void setReadOnly(boolean readOnly) throws SQLException {
     connection.setReadOnly(readOnly);
   }
 
+
   /**
    * Also note the Isolation level needs to be reset when put back into the pool.
    */
@@ -673,11 +695,23 @@ public void setAutoCommit(boolean autoCommit) throws SQLException {
     }
   }
 
+  @Override
+  public void setSchema(String schema) throws SQLException {
+    if (status == STATUS_IDLE) {
+      throw new SQLException(IDLE_CONNECTION_ACCESSED_ERROR + "setSchema()");
+    }
+    currentSchema = schema;
+    resetSchema = true;
+    connection.setSchema(schema);
+  }
+
   @Override
   public void setCatalog(String catalog) throws SQLException {
     if (status == STATUS_IDLE) {
       throw new SQLException(IDLE_CONNECTION_ACCESSED_ERROR + "setCatalog()");
     }
+    currentCatalog = catalog;
+    resetCatalog = true;
     connection.setCatalog(catalog);
   }
 

ebean-datasource/src/test/java/io/ebean/datasource/pool/SchemaTest.java

@@ -0,0 +1,92 @@
+package io.ebean.datasource.pool;
+
+import io.ebean.datasource.DataSourceConfig;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import java.sql.Connection;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.sql.Statement;
+
+import static org.assertj.core.api.AssertionsForClassTypes.assertThat;
+
+class SchemaTest {
+
+  private final ConnectionPool pool;
+
+  SchemaTest() {
+    pool = createPool();
+  }
+
+  private ConnectionPool createPool() {
+    DataSourceConfig config = new DataSourceConfig();
+    config.setUrl("jdbc:h2:mem:tests");
+    config.setUsername("sa");
+    config.setPassword("");
+    config.setMinConnections(2);
+    config.setMaxConnections(4);
+
+    return new ConnectionPool("test", config);
+  }
+
+  @BeforeEach
+  void initTables() throws SQLException {
+    try (Connection conn = pool.getConnection()) {
+      Statement statement = conn.createStatement();
+      statement.execute("CREATE TABLE test (name VARCHAR(255))");
+      statement.execute("INSERT INTO test (name) VALUES ('default schema')");
+      statement.execute("CREATE SCHEMA SCHEMA1");
+      statement.execute("CREATE SCHEMA SCHEMA2");
+      conn.commit();
+    }
+    try (Connection conn = pool.getConnection()) {
+      String orig = conn.getSchema();
+      conn.setSchema("SCHEMA1");
+      Statement statement = conn.createStatement();
+      statement.execute("CREATE TABLE test (name VARCHAR(255))");
+      statement.execute("INSERT INTO test (name) VALUES ('schema1')");
+      conn.setSchema("SCHEMA2");
+      statement.execute("CREATE TABLE test (name VARCHAR(255))");
+      statement.execute("INSERT INTO test (name) VALUES ('schema2')");
+      conn.setSchema(orig);
+      conn.commit();
+    }
+  }
+
+  @AfterEach
+  void after() {
+    pool.shutdown();
+  }
+
+  @Test
+  void getConnectionWithSchemaSwitch() throws SQLException {
+    try (Connection conn = pool.getConnection()) {
+      Statement statement = conn.createStatement();
+      ResultSet rs = statement.executeQuery("SELECT name FROM test");
+      assertThat(rs.next()).isTrue();
+      assertThat(rs.getString("name")).isEqualTo("default schema");
+    }
+    try (Connection conn = pool.getConnection()) {
+      conn.setSchema("SCHEMA1");
+      Statement statement = conn.createStatement();
+      ResultSet rs = statement.executeQuery("SELECT name FROM test");
+      assertThat(rs.next()).isTrue();
+      assertThat(rs.getString("name")).isEqualTo("schema1");
+    }
+    try (Connection conn = pool.getConnection()) {
+      conn.setSchema("SCHEMA2");
+      Statement statement = conn.createStatement();
+      ResultSet rs = statement.executeQuery("SELECT name FROM test");
+      assertThat(rs.next()).isTrue();
+      assertThat(rs.getString("name")).isEqualTo("schema2");
+    }
+    try (Connection conn = pool.getConnection()) {
+      Statement statement = conn.createStatement();
+      ResultSet rs = statement.executeQuery("SELECT name FROM test");
+      assertThat(rs.next()).isTrue();
+      assertThat(rs.getString("name")).isEqualTo("default schema");
+    }
+  }
+}