Draft for DB2-Trustedcontext implementation

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

@@ -0,0 +1,11 @@
+package io.ebean.datasource;
+
+import java.sql.Connection;
+
+/**
+ * @author Roland Praml, Foconis Analytics GmbH
+ */
+public interface DataSourceConnection extends Connection {
+
+  void clearPreparedStatementCache();
+}

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

@@ -1,6 +1,7 @@
 package io.ebean.datasource;
 
 import java.sql.Connection;
+import java.sql.SQLException;
 
 
 /**
@@ -17,11 +18,27 @@ public interface DataSourcePoolListener {
   /**
    * Called after a connection has been retrieved from the connection pool
    */
+  default void onAfterBorrowConnection(DataSourcePool pool, DataSourceConnection connection) {
+    onAfterBorrowConnection(connection);
+  }
+
+  /**
+   * @deprecated implement {@link #onAfterBorrowConnection(DataSourcePool, DataSourceConnection)}
+   */
+  @Deprecated
   default void onAfterBorrowConnection(Connection connection) {}
 
   /**
    * Called before a connection will be put back to the connection pool
    */
+  default void onBeforeReturnConnection(DataSourcePool pool, DataSourceConnection connection) {
+    onBeforeReturnConnection(connection);
+  }
+
+  /**
+   * @deprecated implement {@link #onBeforeReturnConnection(DataSourcePool, DataSourceConnection)}
+   */
+  @Deprecated
   default void onBeforeReturnConnection(Connection connection) {}
 
 

ebean-datasource/pom.xml

@@ -1,5 +1,6 @@
 <?xml version="1.0" encoding="UTF-8"?>
-<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
+<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+         xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
   <modelVersion>4.0.0</modelVersion>
   <parent>
     <groupId>io.ebean</groupId>
@@ -26,6 +27,13 @@
       <scope>test</scope>
     </dependency>
 
+    <dependency>
+      <groupId>org.junit.jupiter</groupId>
+      <artifactId>junit-jupiter-params</artifactId>
+      <version>5.10.2</version>
+      <scope>test</scope>
+    </dependency>
+
     <dependency>
       <groupId>io.ebean</groupId>
       <artifactId>ebean-test-containers</artifactId>

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

@@ -559,7 +559,7 @@ void removeClosedConnection(PooledConnection pooledConnection) {
    */
   private void returnTheConnection(PooledConnection pooledConnection, boolean forceClose) {
     if (poolListener != null && !forceClose) {
-      poolListener.onBeforeReturnConnection(pooledConnection);
+      poolListener.onBeforeReturnConnection(this, pooledConnection);
     }
     queue.returnPooledConnection(pooledConnection, forceClose);
   }
@@ -631,7 +631,7 @@ private PooledConnection getPooledConnection() throws SQLException {
       c.setStackTrace(Thread.currentThread().getStackTrace());
     }
     if (poolListener != null) {
-      poolListener.onAfterBorrowConnection(c);
+      poolListener.onAfterBorrowConnection(this, c);
     }
     return c;
   }

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

@@ -1,5 +1,7 @@
 package io.ebean.datasource.pool;
 
+import io.ebean.datasource.DataSourceConnection;
+
 import java.sql.*;
 import java.util.ArrayList;
 import java.util.Map;
@@ -17,7 +19,7 @@
  * It has caching of Statements and PreparedStatements. Remembers the last
  * statement that was executed. Keeps statistics on how long it is in use.
  */
-final class PooledConnection extends ConnectionDelegator {
+final class PooledConnection extends ConnectionDelegator implements DataSourceConnection {
 
   private static final String IDLE_CONNECTION_ACCESSED_ERROR = "Pooled Connection has been accessed whilst idle in the pool, via method: ";
 
@@ -974,4 +976,13 @@ private String stackTraceAsString(StackTraceElement[] stackTrace) {
     return filteredList.toString();
   }
 
+  @Override
+  public void clearPreparedStatementCache() {
+    lock.lock();
+    try {
+      pstmtCache.clear();
+    } finally {
+      lock.unlock();
+    }
+  }
 }

ebean-datasource/src/test/java/io/ebean/datasource/tcdriver/TrustedContextListener.java

@@ -0,0 +1,38 @@
+package io.ebean.datasource.tcdriver;
+
+import io.ebean.datasource.DataSourceConnection;
+import io.ebean.datasource.DataSourcePool;
+import io.ebean.datasource.DataSourcePoolListener;
+
+import java.sql.SQLException;
+
+/**
+ * Listener, that sets up TrustedConnection properly
+ *
+ * @author Roland Praml, Foconis Analytics GmbH
+ */
+public class TrustedContextListener implements DataSourcePoolListener {
+  private ThreadLocal<String> user = new ThreadLocal<>();
+  private ThreadLocal<String> pass = new ThreadLocal<>();
+  private ThreadLocal<String> schema = new ThreadLocal<>();
+
+  @Override
+  public void onAfterBorrowConnection(DataSourcePool pool, DataSourceConnection connection) {
+    try {
+      TrustedDb2Connection trustedDb2Connection = connection.unwrap(TrustedDb2Connection.class);
+      if (trustedDb2Connection.switchUser(user.get(), pass.get())) {
+        trustedDb2Connection.setSchema(schema.get());
+        connection.clearPreparedStatementCache();
+      }
+      //System.out.println("Switched to " + user.get() + ", Schema: " + schema.get());
+    } catch (SQLException e) {
+      throw new RuntimeException(e); // TODO: Allow throwing sqlException here
+    }
+  }
+
+  public void setContext(String user, String pass, String schema) {
+    this.user.set(user);
+    this.pass.set(pass);
+    this.schema.set(schema);
+  }
+}