Merge pull request #105 from ebean-orm/feature/change-schema-reset

Use pool schema/catalog for original values, lazy initialise originalSchema if required

Author: rbygrave

ebean-datasource-api/src/main/java/io/ebean/datasource/DataSourceBuilder.java

@@ -249,6 +249,11 @@ default DataSourceBuilder schema(String schema) {
     return setSchema(schema);
   }
 
+  /**
+   * Set the default database catalog to use.
+   */
+  DataSourceBuilder catalog(String catalog);
+
   /**
    * @deprecated - migrate to {@link #driver(String)}.
    */
@@ -812,6 +817,11 @@ interface Settings extends DataSourceBuilder {
      */
     String getSchema();
 
+    /**
+     * Return the database catalog.
+     */
+    String catalog();
+
     /**
      * Return the driver instance to use.
      */

ebean-datasource-api/src/main/java/io/ebean/datasource/DataSourceConfig.java

@@ -36,6 +36,7 @@ public class DataSourceConfig implements DataSourceBuilder.Settings {
   private String password;
   private String password2;
   private String schema;
+  private String catalog;
   private Driver driver;
   private Class<? extends Driver> driverClass;
   private String driverClassName;
@@ -103,6 +104,7 @@ public DataSourceConfig copy() {
     copy.password = password;
     copy.password2 = password2;
     copy.schema = schema;
+    copy.catalog = catalog;
     copy.platform = platform;
     copy.ownerUsername = ownerUsername;
     copy.ownerPassword = ownerPassword;
@@ -171,6 +173,9 @@ public DataSourceConfig setDefaults(DataSourceBuilder builder) {
     if (schema == null) {
       schema = other.getSchema();
     }
+    if (catalog == null) {
+      catalog = other.catalog();
+    }
     if (minConnections == 2 && other.getMinConnections() < 2) {
       minConnections = other.getMinConnections();
     }
@@ -307,6 +312,17 @@ public DataSourceConfig setSchema(String schema) {
     return this;
   }
 
+  @Override
+  public String catalog() {
+    return catalog;
+  }
+
+  @Override
+  public DataSourceConfig catalog(String catalog) {
+    this.catalog = catalog;
+    return this;
+  }
+
   @Override
   public String getDriver() {
     return driverClassName;
@@ -740,6 +756,7 @@ private void loadSettings(ConfigPropertiesHelper properties) {
     password = properties.get("password", password);
     password2 = properties.get("password2", password2);
     schema = properties.get("schema", schema);
+    catalog = properties.get("catalog", catalog);
     platform = properties.get("platform", platform);
     ownerUsername = properties.get("ownerUsername", ownerUsername);
     ownerPassword = properties.get("ownerPassword", ownerPassword);

ebean-datasource-api/src/test/java/io/ebean/datasource/DataSourceConfigTest.java

@@ -77,6 +77,7 @@ public void copy() {
     source.setUrl("url");
     source.setReadOnlyUrl("readOnlyUrl");
     source.setSchema("sch");
+    source.catalog("cat");
 
     Map<String,String> customSource = new LinkedHashMap<>();
     customSource.put("a","a");
@@ -90,6 +91,7 @@ public void copy() {
     assertEquals("url", copy.getUrl());
     assertEquals("readOnlyUrl", copy.getReadOnlyUrl());
     assertEquals("sch", copy.getSchema());
+    assertEquals("cat", copy.catalog());
     assertEquals(42, copy.getMinConnections());
     assertEquals(45, copy.getMaxConnections());
 
@@ -113,6 +115,7 @@ public void defaults() {
     assertThat(readOnly.getUsername()).isEqualTo(config.getUsername());
     assertThat(readOnly.getPassword()).isEqualTo(config.getPassword());
     assertThat(readOnly.getSchema()).isEqualTo(config.getSchema());
+    assertThat(readOnly.catalog()).isEqualTo(config.catalog());
     assertThat(readOnly.getMinConnections()).isEqualTo(config.getMinConnections());
     assertThat(readOnly.getCustomProperties()).containsKeys("useSSL");
   }
@@ -258,6 +261,7 @@ private static void assertConfigValues(DataSourceBuilder.Settings config) {
     assertThat(config.getUsername()).isEqualTo("myusername");
     assertThat(config.getPassword()).isEqualTo("mypassword");
     assertThat(config.getSchema()).isEqualTo("myschema");
+    assertThat(config.catalog()).isEqualTo("mycat");
     assertThat(config.getApplicationName()).isEqualTo("myApp");
     Properties clientInfo = config.getClientInfo();
     assertThat(clientInfo.getProperty("ClientUser")).isEqualTo("ciu");

ebean-datasource-api/src/test/resources/example.properties

@@ -1,6 +1,7 @@
 datasource.foo.username=myusername
 datasource.foo.password=mypassword
 datasource.foo.schema=myschema
+datasource.foo.catalog=mycat
 datasource.foo.url=myUrl
 datasource.foo.readOnlyUrl=myReadOnlyUrl
 datasource.foo.applicationName=myApp

ebean-datasource-api/src/test/resources/example2.properties

@@ -1,6 +1,7 @@
 bar.username=myusername
 bar.password=mypassword
 bar.schema=myschema
+bar.catalog=mycat
 bar.url=myUrl
 bar.readOnlyUrl=myReadOnlyUrl
 bar.applicationName=myApp

ebean-datasource-api/src/test/resources/example3.properties

@@ -1,6 +1,7 @@
 username=myusername
 password=mypassword
 schema=myschema
+catalog=mycat
 url=myUrl
 readOnlyUrl=myReadOnlyUrl
 applicationName=myApp

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

@@ -42,6 +42,7 @@ final class ConnectionPool implements DataSourcePool {
   private final List<String> initSql;
   private final String user;
   private final String schema;
+  private final String catalog;
   private final String heartbeatSql;
   private final int heartbeatFreqSecs;
   private final int heartbeatTimeoutSeconds;
@@ -119,6 +120,7 @@ final class ConnectionPool implements DataSourcePool {
     this.clientInfo = params.getClientInfo();
     this.queue = new PooledConnectionQueue(this);
     this.schema = params.getSchema();
+    this.catalog = params.catalog();
     this.user = params.getUsername();
     this.shutdownOnJvmExit = params.isShutdownOnJvmExit();
     this.source = DriverDataSource.of(name, params);
@@ -251,6 +253,14 @@ public int size() {
     return size.get();
   }
 
+  String schema() {
+    return schema;
+  }
+
+  String catalog() {
+    return catalog;
+  }
+
   /**
    * Increment the current pool size.
    */
@@ -388,6 +398,9 @@ private Connection initConnection(Connection conn) throws SQLException {
     if (readOnly) {
       conn.setReadOnly(true);
     }
+    if (catalog != null) {
+      conn.setCatalog(catalog);
+    }
     if (schema != null) {
       conn.setSchema(schema);
     }

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

@@ -57,6 +57,32 @@ final class PooledConnection extends ConnectionDelegator {
 
   private static final int RO_MYSQL_1290 = 1290;
 
+  /**
+   * Constant for schema/catalog, when we are in SCHEMA_CATALOG_UNKNOWN state
+   * This is used for correct cache key computation. (We cannot use 'null'
+   * here, as we might get a collision in edge cases)
+   */
+  private static final String UNKNOWN = "@unknown";
+
+  /**
+   * The schema/catalog is unknown, that means, we have not yet touched the
+   * value on the underlying connection.
+   * We do not have to restore it.
+   */
+  private static final int SCHEMA_CATALOG_UNKNOWN = 0;
+
+  /**
+   * The schema/catalog is changed. The original value has to be restored on
+   * close()
+   */
+  private static final int SCHEMA_CATALOG_CHANGED = 1;
+
+  /**
+   * We know the original value of the underlying connection, but there is no
+   * demand to restore it.
+   */
+  private static final int SCHEMA_CATALOG_KNOWN = 2;
+
   private final String name;
   private final ConnectionPool pool;
   private final Connection connection;
@@ -82,12 +108,16 @@ 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 int schemaState = SCHEMA_CATALOG_UNKNOWN;
+  private int catalogState = SCHEMA_CATALOG_UNKNOWN;
+
+  // this is used for cache computation
+  private String cacheKeySchema = UNKNOWN;
+  private String cacheKeyCatalog = UNKNOWN;
+
+  // original values are lazily initialized and restored on close()
+  private String originalSchema;
+  private String originalCatalog;
 
   private long startUseTime;
   private long lastUseTime;
@@ -118,12 +148,20 @@ final class PooledConnection extends ConnectionDelegator {
     this.pool = pool;
     this.connection = connection;
     this.name = pool.name() + uniqueId;
+    this.originalSchema = pool.schema();
+    this.originalCatalog = pool.catalog();
+    if (originalSchema != null) {
+      this.schemaState = SCHEMA_CATALOG_KNOWN;
+      this.cacheKeySchema = originalSchema;
+    }
+    if (originalCatalog != null) {
+      this.catalogState = SCHEMA_CATALOG_KNOWN;
+      this.cacheKeyCatalog = originalCatalog;
+    }
     this.pstmtCache = new PstmtCache(pool.pstmtCacheSize());
     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();
   }
 
@@ -139,8 +177,6 @@ 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";
   }
 
   /**
@@ -283,7 +319,7 @@ void returnPreparedStatement(ExtendedPreparedStatement pstmt) {
    */
   @Override
   public PreparedStatement prepareStatement(String sql, int returnKeysFlag) throws SQLException {
-    String key = sql + ':' + currentSchema + ':' + currentCatalog + ':' + returnKeysFlag;
+    String key = sql + ':' + cacheKeySchema + ':' + cacheKeyCatalog + ':' + returnKeysFlag;
     return prepareStatement(sql, true, returnKeysFlag, key);
   }
 
@@ -292,7 +328,7 @@ public PreparedStatement prepareStatement(String sql, int returnKeysFlag) throws
    */
   @Override
   public PreparedStatement prepareStatement(String sql) throws SQLException {
-    String key = sql + ':' + currentSchema + ':' + currentCatalog;
+    String key = sql + ':' + cacheKeySchema + ':' + cacheKeyCatalog;
     return prepareStatement(sql, false, 0, key);
   }
 
@@ -422,14 +458,17 @@ public void close() throws SQLException {
         resetIsolationReadOnlyRequired = false;
       }
 
-      if (resetSchema) {
-        connection.setSchema(originalSchema);
-        resetSchema = false;
+      if (catalogState == SCHEMA_CATALOG_CHANGED) {
+        connection.setCatalog(originalCatalog);
+        cacheKeyCatalog = originalCatalog;
+        catalogState = SCHEMA_CATALOG_KNOWN;
       }
 
-      if (resetCatalog) {
-        connection.setCatalog(originalCatalog);
-        resetCatalog = false;
+      if (schemaState == SCHEMA_CATALOG_CHANGED) {
+        connection.setSchema(originalSchema);
+        // we can use original value for cache computation from now on
+        cacheKeySchema = originalSchema;
+        schemaState = SCHEMA_CATALOG_KNOWN;
       }
 
       // the connection is assumed GOOD so put it back in the pool
@@ -700,8 +739,13 @@ public void setSchema(String schema) throws SQLException {
     if (status == STATUS_IDLE) {
       throw new SQLException(IDLE_CONNECTION_ACCESSED_ERROR + "setSchema()");
     }
-    currentSchema = schema;
-    resetSchema = true;
+    if (schemaState == SCHEMA_CATALOG_UNKNOWN) {
+      // lazily initialise the originalSchema
+      originalSchema = getSchema();
+      // state would be KNOWN here
+    }
+    schemaState = SCHEMA_CATALOG_CHANGED;
+    cacheKeySchema = schema;
     connection.setSchema(schema);
   }
 
@@ -710,8 +754,11 @@ public void setCatalog(String catalog) throws SQLException {
     if (status == STATUS_IDLE) {
       throw new SQLException(IDLE_CONNECTION_ACCESSED_ERROR + "setCatalog()");
     }
-    currentCatalog = catalog;
-    resetCatalog = true;
+    if (schemaState == SCHEMA_CATALOG_UNKNOWN) {
+      originalCatalog = getCatalog();
+    }
+    catalogState = SCHEMA_CATALOG_CHANGED;
+    cacheKeyCatalog = catalog;
     connection.setCatalog(catalog);
   }
 

ebean-datasource/src/test/java/io/ebean/datasource/test/PostgresInitTest.java

@@ -81,7 +81,7 @@ void test_with_clientInfo() throws SQLException {
   void test_with_applicationNameAndSchema() throws SQLException {
     DataSourceConfig ds = new DataSourceConfig();
     ds.setUrl("jdbc:postgresql://127.0.0.1:9999/app");
-    ds.setSchema("fred");
+    ds.setSchema("public");
     ds.setUsername("db_owner");
     ds.setPassword("test");
     ds.setApplicationName("my-application-name");
@@ -112,7 +112,7 @@ void test_with_applicationNameAndSchema() throws SQLException {
   void test_password2() throws SQLException {
     DataSourceConfig ds = new DataSourceConfig();
     ds.setUrl("jdbc:postgresql://127.0.0.1:9999/app");
-    ds.setSchema("fred");
+    ds.setSchema("public");
     ds.setUsername("db_owner");
     ds.setPassword("test");
     ds.setPassword2("newRolledPassword");