HHH-6594 - Fix and test

hibernate-core/src/main/java/org/hibernate/engine/jdbc/batch/internal/AbstractBatchImpl.java

@@ -42,6 +42,7 @@ import org.jboss.logging.Logger;
  * Convenience base class for implementors of the Batch interface.
  *
  * @author Steve Ebersole
+ * @author Lukasz Antoniak (lukasz dot antoniak at gmail dot com)
  */
 public abstract class AbstractBatchImpl implements Batch {
 
@@ -134,6 +135,7 @@ public abstract class AbstractBatchImpl implements Batch {
 	}
 
 	private PreparedStatement buildBatchStatement(String sql, boolean callable) {
+		sql = jdbcCoordinator.getTransactionCoordinator().getTransactionContext().onPrepareStatement( sql );
 		try {
 			if ( callable ) {
 				return jdbcCoordinator.getLogicalConnection().getShareableConnectionProxy().prepareCall( sql );

hibernate-core/src/main/java/org/hibernate/engine/jdbc/internal/StatementPreparerImpl.java

@@ -37,7 +37,8 @@ import java.sql.ResultSet;
 import java.sql.SQLException;
 
 /**
-* @author Steve Ebersole
+ * @author Steve Ebersole
+ * @author Lukasz Antoniak (lukasz dot antoniak at gmail dot com)
 */
 class StatementPreparerImpl implements StatementPreparer {
 	private long transactionTimeOut = -1;
@@ -169,7 +170,7 @@ class StatementPreparerImpl implements StatementPreparer {
 		protected final String sql;
 
 		protected StatementPreparationTemplate(String sql) {
-			this.sql = sql;
+			this.sql = jdbcCoordinator.getTransactionCoordinator().getTransactionContext().onPrepareStatement( sql );
 		}
 
 		public PreparedStatement prepareStatement() {

hibernate-core/src/main/java/org/hibernate/engine/transaction/spi/TransactionContext.java

@@ -112,5 +112,7 @@ public interface TransactionContext extends Serializable {
 
 	public void afterTransactionCompletion(TransactionImplementor hibernateTransaction, boolean successful);
 
+	public String onPrepareStatement(String sql); 
+
 	public JdbcConnectionAccess getJdbcConnectionAccess();
 }

hibernate-core/src/main/java/org/hibernate/internal/SessionImpl.java

@@ -45,6 +45,7 @@ import java.util.Set;
 
 import org.jboss.logging.Logger;
 
+import org.hibernate.AssertionFailure;
 import org.hibernate.CacheMode;
 import org.hibernate.ConnectionReleaseMode;
 import org.hibernate.Criteria;
@@ -537,6 +538,16 @@ public final class SessionImpl
 		}
 	}
 
+	@Override
+	public String onPrepareStatement(String sql) {
+		errorIfClosed();
+		sql = interceptor.onPrepareStatement( sql );
+		if ( sql == null || sql.length() == 0 ) {
+			throw new AssertionFailure( "Interceptor.onPrepareStatement() returned null or empty string." );
+		}
+		return sql;
+	}
+
 	/**
 	 * clear all the internal collections, just
 	 * to help the garbage collector, does not

hibernate-core/src/main/java/org/hibernate/internal/StatelessSessionImpl.java

@@ -361,6 +361,11 @@ public class StatelessSessionImpl extends AbstractSessionImpl implements Statele
 		// nothing to do here
 	}
 
+	@Override
+	public String onPrepareStatement(String sql) {
+		return sql;
+	}
+
 	public String bestGuessEntityName(Object object) {
 		if (object instanceof HibernateProxy) {
 			object = ( (HibernateProxy) object ).getHibernateLazyInitializer().getImplementation();

hibernate-core/src/test/java/org/hibernate/test/common/TransactionContextImpl.java

@@ -114,4 +114,9 @@ public class TransactionContextImpl implements TransactionContext {
 	@Override
 	public void afterTransactionCompletion(TransactionImplementor hibernateTransaction, boolean successful) {
 	}
+
+	@Override
+	public String onPrepareStatement(String sql) {
+		return sql;
+	}
 }

hibernate-core/src/test/java/org/hibernate/test/interceptor/InterceptorTest.java

@@ -23,9 +23,13 @@
  */
 package org.hibernate.test.interceptor;
 import java.io.Serializable;
+import java.util.LinkedList;
 import java.util.List;
+import java.util.Queue;
 
+import org.hibernate.AssertionFailure;
 import org.hibernate.EmptyInterceptor;
+import org.hibernate.Interceptor;
 import org.hibernate.Session;
 import org.hibernate.Transaction;
 import org.hibernate.type.Type;
@@ -39,9 +43,11 @@ import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertSame;
+import static org.junit.Assert.assertTrue;
 
 /**
  * @author Gavin King
+ * @author Lukasz Antoniak (lukasz dot antoniak at gmail dot com)
  */
 public class InterceptorTest extends BaseCoreFunctionalTestCase {
 	@Override
@@ -233,5 +239,78 @@ public class InterceptorTest extends BaseCoreFunctionalTestCase {
 		t.commit();
 		s.close();
 	}
+
+	@Test
+	@TestForIssue( jiraKey = "HHH-6594" )
+	public void testPrepareStatementIntercept() {
+		final Queue<String> expectedSQLs = new LinkedList<String>();
+		// Transaction 1
+		expectedSQLs.add( "insert" );
+		// Transaction 2
+		expectedSQLs.add( "select" );
+		expectedSQLs.add( "select" );
+		// Transaction 3
+		expectedSQLs.add( "select" );
+		expectedSQLs.add( "select" );
+		expectedSQLs.add( "update" );
+		// Transaction 4
+		expectedSQLs.add( "select" );
+		expectedSQLs.add( "delete" );
+
+		final Interceptor interceptor = new EmptyInterceptor() {
+			@Override
+			public String onPrepareStatement(String sql) {
+				assertNotNull( sql );
+				assertTrue( sql.toLowerCase().startsWith( expectedSQLs.poll().toLowerCase() ) );
+				return sql;
+			}
+		};
+
+		Session s = openSession(interceptor);
+		Transaction t = s.beginTransaction();
+		User u = new User( "Lukasz", "Antoniak" );
+		s.persist( u );
+		t.commit();
+		s.close();
+
+		s = openSession(interceptor);
+		t = s.beginTransaction();
+		s.get( User.class, "Lukasz" );
+		s.createQuery( "from User u" ).list();
+		t.commit();
+		s.close();
+
+		u.setPassword( "Kinga" );
+		s = openSession(interceptor);
+		t = s.beginTransaction();
+		s.merge( u );
+		t.commit();
+		s.close();
+
+		s = openSession(interceptor);
+		t = s.beginTransaction();
+		s.delete( u );
+		t.commit();
+		s.close();
+
+		assertTrue( expectedSQLs.isEmpty() );
+	}
+
+	@Test(expected = AssertionFailure.class)
+	public void testPrepareStatementFaultIntercept() {
+		final Interceptor interceptor = new EmptyInterceptor() {
+			@Override
+			public String onPrepareStatement(String sql) {
+				return null;
+			}
+		};
+
+		Session s = openSession(interceptor);
+		Transaction t = s.beginTransaction();
+		User u = new User( "Kinga", "Mroz" );
+		s.persist( u );
+		t.commit();
+		s.close();
+	}
 }