package me.jasoncampos.inject.persist.hibernate;

import javax.inject.Inject;
import javax.inject.Provider;
import javax.inject.Singleton;

import org.hibernate.Session;

import com.google.common.base.Preconditions;
import com.google.inject.persist.UnitOfWork;

/**
 * Hibernate {@code UnitOfWork} implementation which manages a ThreadLocal {@code Session} instance. Invoking
 * {@link #begin()} will open a new {@code Session} if one is not already open. Subsequent invocations of {@link #get()}
 * will return the same {@code Session} instance until {@link #end()} is invoked. <br />
 * <br />
 *
 * This class may also be used as a {@code Provider<Session>} without the {@code UnitOfWork} context. Note that any
 * session not created during a {@link #begin()} invocation *must* be manually closed. Invoking {@link #get()} and then
 * {@link #end()} will *not* close the session created from {@link #get()}.
 *
 * @author Jason Campos <jcmapos8782@gmail.com>
 */
@Singleton
public class HibernateUnitOfWork implements Provider<Session>, UnitOfWork {

	private final ThreadLocal<Session> sessions = new ThreadLocal<>();
	private final HibernatePersistService sessionFactory;

	@Inject
	public HibernateUnitOfWork(final HibernatePersistService sessionFactory) {
		this.sessionFactory = sessionFactory;
	}

	@Override
	public Session get() {
		// If a unit of work has begun, return that session.
		// Otherwise, assume the session will be manually managed.
		if (isWorking()) {
			return sessions.get();
		}

		return sessionFactory.get().openSession();
	}

	public boolean isWorking() {
		return sessions.get() != null;
	}

	@Override
	public void begin() {
		if (!isWorking()) {
			sessions.set(get());
		}
	}

	@Override
	public void end() {
		Preconditions.checkState(isWorking(), "UnitOfWork.end() invoked with no corresponding UnitOfWork.begin()");
		final Session session = sessions.get();
		try {
			if (session.isOpen() && session.isJoinedToTransaction()) {
				session.flush();
			}
		} finally {
			if (session.isOpen()) {
				session.close();
			}
			sessions.remove();
		}
	}
}

@@ -165,24 +165,83 @@
 te.Session;%0D
+%0Aimport org.slf4j.Logger;%0D%0Aimport org.slf4j.LoggerFactory;%0D
 %0A%0D%0Aimport co
@@ -1129,16 +1129,107 @@
 fWork %7B%0D
+%0A%09private static final Logger logger = LoggerFactory.getLogger(HibernateUnitOfWork.class);%0D
 %0A%0D%0A%09priv
@@ -1691,16 +1691,140 @@
 ();%0D%0A%09%09%7D
+ else %7B%0D%0A%09%09%09logger.warn(%22Opening hibernate Session with no current UnitOfWork. This session must be manually closed.%22);%0D%0A%09%09%7D
 %0D%0A%0D%0A%09%09re
@@ -2008,26 +2008,124 @@
 %0A%09%09%09
-sessions.set(get()
+final Session session = get();%0D%0A%09%09%09sessions.set(session);%0D%0A%09%09%09logger.debug(%22UnitOfWork started. session=%7B%7D%22, session
 );%0D%0A
@@ -2494,16 +2494,16 @@
 %0D%0A%09%09%09%7D%0D%0A
-
 %09%09%09sessi
@@ -2513,24 +2513,86 @@
 .remove();%0D%0A
+%09%09%09logger.debug(%22UnitOfWork complete. session=%7B%7D%22, session);%0D%0A
 %09%09%7D%0D%0A%09%7D%0D%0A%7D%0D%0A

Add logging around begin() and end() unit of work