DatabaseOperator { public <T> T transaction(final SQLTransaction<T> operations) throws SQLException { Connection conn = null; try { conn = this.queryRunner.getDataSource().getConnection(); conn.setAutoCommit(false); final DatabaseTransOperator transOperator = new DatabaseTransOperator(this.queryRunner, conn); final T res = operations.execute(transOperator); conn.commit(); return res; } catch (final SQLException ex) { logger.error("transaction failed", ex); if (this.dbMetrics != null) { this.dbMetrics.markDBFailTransaction(); } throw ex; } finally { DbUtils.closeQuietly(conn); } } @Inject  DatabaseOperator(final QueryRunner queryRunner); T query(final String baseQuery, final ResultSetHandler<T> resultHandler,
      final Object... params); T transaction(final SQLTransaction<T> operations); int update(final String updateClause, final Object... params); int[] batch(final String sqlCommand, final Object[]... params); AzkabanDataSource getDataSource();  }
@Test public void testTransaction() throws Exception { when(this.queryRunner.update(this.conn, "update blah set ? = ?", "1", 26)).thenReturn(1); when(this.queryRunner.query(this.conn, "select * from blah where ? = ?", this.handler, "id", 1)) .thenReturn(26); final SQLTransaction<Integer> transaction = transOperator -> { transOperator.update("update blah set ? = ?", "1", 26); return transOperator.query("select * from blah where ? = ?", this.handler, "id", 1); }; final int res = this.dbOperator.transaction(transaction); Assert.assertEquals(26, res); }