Introduction
When you’re using JDBC or if you are generating SQL statements by hand, you always know what statements are sent to the database server. Although there are situations when a native query is the most obvious solution to a given business use case, most statements are simple enough to be generated automatically That’s exactly what JPA and Hibernate do, and the application developer can focus on entity state transitions instead.
Nevertheless, the application developer must always assert that Hibernate generates the expected statements, as well as the number of statements being generated (to avoid N+1 query issues).
Proxying the underlying JDBC Driver or DataSource
In production, it’s very common to Proxy the underlying Driver Connection providing mechanism so that the application benefits from connection pooling, or for monitoring connection pool usage.
For this purpose, the underlying JDBC Driver or DataSource can be proxied using tools such as P6spy or datasource-proxy.
In fact, this is also a very convenient way of logging JDBC statements along with their bind parameters.
While for many application, it’s not an issue to add yet another dependency when you are developing an open source framework you strive to minimize the number of dependencies your project needs to depend on. Luckily, for Hibernate, we don’t even need to use an external dependency for intercepting JDBC statements, and this post is going to show you how easily you can tackle this requirement.
StatementInspector
For many use cases, the StatementInspector is the only thing you need to capture all SQL statements that are executed by Hibernate.
The StatementInspector must be provided during SessionFactory bootstrapping as follows:
public class SQLStatementInterceptor {
private final LinkedList<String> sqlQueries = new LinkedList<>();
public SQLStatementInterceptor(SessionFactoryBuilder sessionFactoryBuilder) {
sessionFactoryBuilder.applyStatementInspector(
(StatementInspector) sql -> {
sqlQueries.add( sql );
return sql;
} );
}
public LinkedList<String> getSqlQueries() {
return sqlQueries;
}
}With this utility we can easily verify the Oracle follow-on-locking mechanism which is caused by the FOR UPDATE clause restrictions imposed by the database engine:
sqlStatementInterceptor.getSqlQueries().clear();
List<Product> products = session.createQuery(
"select p from Product p", Product.class )
.setLockOptions( new LockOptions( LockMode.PESSIMISTIC_WRITE ) )
.setFirstResult( 40 )
.setMaxResults( 10 )
.getResultList();
assertEquals( 10, products.size() );
assertEquals( 11, sqlStatementInterceptor.getSqlQueries().size() );So far, so good. But as simple as the StatementInspector may be, it does not mix well with JDBC batching.
StatementInspector intercepts the prepare phase, whereas for batching we need to intercept the addBatch and executeBatch method calls.
Even without native support for such a feature, we can easily design a custom ConnectionProvider that can intercept all PreparedStatement method calls.
First, we start with the ConnectionProviderDelegate which is capable of substituting any other ConnectionProvider
that would otherwise be picked by Hibernate (e.g. DatasourceConnectionProviderImpl, DriverManagerConnectionProviderImpl, HikariCPConnectionProvider)
for the current configuration properties.
public class ConnectionProviderDelegate implements
ConnectionProvider,
Configurable,
ServiceRegistryAwareService {
private ServiceRegistryImplementor serviceRegistry;
private ConnectionProvider connectionProvider;
@Override
public void injectServices(ServiceRegistryImplementor serviceRegistry) {
this.serviceRegistry = serviceRegistry;
}
@Override
public void configure(Map configurationValues) {
@SuppressWarnings("unchecked")
Map<String, Object> settings = new HashMap<>( configurationValues );
settings.remove( AvailableSettings.CONNECTION_PROVIDER );
connectionProvider = ConnectionProviderInitiator.INSTANCE.initiateService(
settings,
serviceRegistry
);
if ( connectionProvider instanceof Configurable ) {
Configurable configurableConnectionProvider = (Configurable) connectionProvider;
configurableConnectionProvider.configure( settings );
}
}
@Override
public Connection getConnection() throws SQLException {
return connectionProvider.getConnection();
}
@Override
public void closeConnection(Connection conn) throws SQLException {
connectionProvider.closeConnection( conn );
}
@Override
public boolean supportsAggressiveRelease() {
return connectionProvider.supportsAggressiveRelease();
}
@Override
public boolean isUnwrappableAs(Class unwrapType) {
return connectionProvider.isUnwrappableAs( unwrapType );
}
@Override
public <T> T unwrap(Class<T> unwrapType) {
return connectionProvider.unwrap( unwrapType );
}
}With the ConnectionProviderDelegate in place, we can now implement the PreparedStatementSpyConnectionProvider which, using Mockito, it returns a Connection spy instead of an actual JDBC Driver Connection object:
public class PreparedStatementSpyConnectionProvider
extends ConnectionProviderDelegate {
private final Map<PreparedStatement, String> preparedStatementMap = new LinkedHashMap<>();
@Override
public Connection getConnection() throws SQLException {
Connection connection = super.getConnection();
return spy( connection );
}
private Connection spy(Connection connection) {
if ( new MockUtil().isMock( connection ) ) {
return connection;
}
Connection connectionSpy = Mockito.spy( connection );
try {
doAnswer( invocation -> {
PreparedStatement statement = (PreparedStatement) invocation.callRealMethod();
PreparedStatement statementSpy = Mockito.spy( statement );
String sql = (String) invocation.getArguments()[0];
preparedStatementMap.put( statementSpy, sql );
return statementSpy;
} ).when( connectionSpy ).prepareStatement( anyString() );
}
catch ( SQLException e ) {
throw new IllegalArgumentException( e );
}
return connectionSpy;
}
/**
* Clears the recorded PreparedStatements and reset the associated Mocks.
*/
public void clear() {
preparedStatementMap.keySet().forEach( Mockito::reset );
preparedStatementMap.clear();
}
/**
* Get one and only one PreparedStatement associated to the given SQL statement.
*
* @param sql SQL statement.
*
* @return matching PreparedStatement.
*
* @throws IllegalArgumentException If there is no matching PreparedStatement or multiple instances, an exception is being thrown.
*/
public PreparedStatement getPreparedStatement(String sql) {
List<PreparedStatement> preparedStatements = getPreparedStatements( sql );
if ( preparedStatements.isEmpty() ) {
throw new IllegalArgumentException(
"There is no PreparedStatement for this SQL statement " + sql );
}
else if ( preparedStatements.size() > 1 ) {
throw new IllegalArgumentException( "There are " + preparedStatements
.size() + " PreparedStatements for this SQL statement " + sql );
}
return preparedStatements.get( 0 );
}
/**
* Get the PreparedStatements that are associated to the following SQL statement.
*
* @param sql SQL statement.
*
* @return list of recorded PreparedStatements matching the SQL statement.
*/
public List<PreparedStatement> getPreparedStatements(String sql) {
return preparedStatementMap.entrySet()
.stream()
.filter( entry -> entry.getValue().equals( sql ) )
.map( Map.Entry::getKey )
.collect( Collectors.toList() );
}
/**
* Get the PreparedStatements that were executed since the last clear operation.
*
* @return list of recorded PreparedStatements.
*/
public List<PreparedStatement> getPreparedStatements() {
return new ArrayList<>( preparedStatementMap.keySet() );
}
}To use this custom provider, we just need to provide an instance via the hibernate.connection.provider_class configuration property:
private PreparedStatementSpyConnectionProvider connectionProvider =
new PreparedStatementSpyConnectionProvider();
@Override
protected void addSettings(Map settings) {
settings.put(
AvailableSettings.CONNECTION_PROVIDER,
connectionProvider
);
}Now, we can assert that the underlying PreparedStatement is batching statements according to our expectations:
Session session = sessionFactory().openSession();
session.setJdbcBatchSize( 3 );
session.beginTransaction();
try {
for ( long i = 0; i < 5; i++ ) {
Event event = new Event();
event.id = id++;
event.name = "Event " + i;
session.persist( event );
}
}
finally {
connectionProvider.clear();
session.getTransaction().commit();
session.close();
}
PreparedStatement preparedStatement = connectionProvider.getPreparedStatement(
"insert into Event (name, id) values (?, ?)" );
verify(preparedStatement, times( 5 )).addBatch();
verify(preparedStatement, times( 2 )).executeBatch();The PreparedStatement is not a mock but a real object spy, which can intercept method call while also propagating the call to the underlying actual JDBC Driver PreparedStatement object.
Although getting the PreparedStatement by its associated SQL String is useful for the aforementioned test case, we can also get all executed PreparedStatements like this:
List<PreparedStatement> preparedStatements = connectionProvider.getPreparedStatements();
assertEquals(1, preparedStatements.size());
preparedStatement = preparedStatements.get( 0 );
verify(preparedStatement, times( 5 )).addBatch();
verify(preparedStatement, times( 2 )).executeBatch();