diff --git a/ebean-datasource/src/main/java/io/ebean/datasource/pool/ConnectionDelegator.java b/ebean-datasource/src/main/java/io/ebean/datasource/pool/ConnectionDelegator.java index 0634929..6fd3a67 100644 --- a/ebean-datasource/src/main/java/io/ebean/datasource/pool/ConnectionDelegator.java +++ b/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(); diff --git a/ebean-datasource/src/main/java/io/ebean/datasource/pool/PooledConnection.java b/ebean-datasource/src/main/java/io/ebean/datasource/pool/PooledConnection.java index fcb0d58..fba3005 100644 --- a/ebean-datasource/src/main/java/io/ebean/datasource/pool/PooledConnection.java +++ b/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); } diff --git a/ebean-datasource/src/test/java/io/ebean/datasource/pool/SchemaTest.java b/ebean-datasource/src/test/java/io/ebean/datasource/pool/SchemaTest.java new file mode 100644 index 0000000..8116f1e --- /dev/null +++ b/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"); + } + } +}