JpaJobRepository

package com.trendmicro.grid.acl.ds.jpa;

import com.trendmicro.grid.acl.ds.JobRepository;
import com.trendmicro.grid.acl.ds.datatypes.SharedJob;
import com.trendmicro.grid.acl.ds.datatypes.SharedJobDetails;
import com.trendmicro.grid.acl.ds.datatypes.SharedSource;
import com.trendmicro.grid.acl.ds.jpa.entities.JpaJob;
import com.trendmicro.grid.acl.ds.jpa.entities.JpaJobSources;
import com.trendmicro.grid.acl.ds.jpa.entities.JpaSource;
import com.trendmicro.grid.acl.l0.datatypes.*;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Repository;
import org.springframework.transaction.annotation.Transactional;

import javax.annotation.Resource;
import javax.persistence.EntityManager;
import javax.persistence.PersistenceContext;
import javax.persistence.TypedQuery;
import java.awt.*;
import java.util.*;
import java.util.List;

import static com.trendmicro.grid.acl.ds.DSUtil.cast;
import static com.trendmicro.grid.acl.ds.jpa.util.JpaUtils.*;
import static com.trendmicro.grid.acl.ds.jpa.util.JpaUtils.toListPage;
import static java.lang.System.currentTimeMillis;
import static net.sf.tinyjee.util.Assert.assertNotNull;

/**
 * Implements JobRepository using JPA.
 *
 * @author juergen_kellerer, 2010-06-01
 * @version 1.0
 */
@Repository
@Transactional(readOnly = true)
public class JpaJobRepository implements JobRepository {

	private static final Logger log = LoggerFactory.getLogger(JpaJobRepository.class);

	private static int pageSize = 25;

	public static int getPageSize() {
		return pageSize;
	}

	public static void setPageSize(int pageSize) {
		JpaJobRepository.pageSize = pageSize;
	}

	static final Callback<byte[], UUID> bytesToUUIDConverter = new Callback<byte[], UUID>() {
		public UUID call(byte[] element) {
			return fromBytes(element);
		}
	};

	@PersistenceContext(unitName = "CoreDB")
	EntityManager em;

	@Resource
	JpaSourceFetcher sourceFetcher;

	/**
	 * {@inheritDoc}
	 */
	@Override
	@Transactional(readOnly = false)
	public UUID createJob() {
		JpaJob job = new JpaJob(UUID.randomUUID(), null);
		em.persist(job);
		em.flush();
		return job.getJobId();
	}

	/**
	 * {@inheritDoc}
	 */
	@Override
	@Transactional(readOnly = false)
	public UUID createSubJob(UUID parentJobId) {
		JpaJob parent = getJob(parentJobId);
		JpaJob job = new JpaJob(UUID.randomUUID(), parent);
		em.persist(job);
		parent.getSubJobs().add(job);
		return job.getJobId();
	}

	/**
	 * {@inheritDoc}
	 */
	@Override
	@Transactional(readOnly = false)
	public void addJobSources(UUID jobId, List<Source> sources) {
		addJobSources(getJob(jobId), sources);
	}

	/**
	 * {@inheritDoc}
	 */
	@Override
	@Transactional(readOnly = false)
	public void addJobSources(Job job, List<Source> sources) {
		JpaJob jpaJob = toJpaJob(job);
		for (Source source : sources) {
			JpaSource jpaSource = sourceFetcher.fetch(source.getSourceInformation().getIdentifier());
			em.persist(new JpaJobSources(jpaJob, jpaSource));
		}
	}

	/**
	 * {@inheritDoc}
	 */
	@Override
	@Transactional(readOnly = false)
	public void updateJob(Job job) {
		assertNotNull("job", job);

		if (job instanceof JpaJob) {
			job.setLastUpdated(new Date());
			em.merge(job);
		} else {
			JpaJob j = getJob(job.getJobId());
			if (j == null)
				throw new IllegalArgumentException("Cannot update unexisting job " + job);

			j.setLastUpdated(new Date());
			j.setFinished(job.getFinished());
			j.setMetadata(job.getMetadata());
			j.setStagesPassed(job.getStagesPassed());
			j.setState(job.getState());
		}
	}

	/**
	 * {@inheritDoc}
	 */
	@Override
	@Transactional(readOnly = false)
	public void finalizeJob(UUID jobId, Job.State finalJobState) throws IllegalArgumentException {
		finalizeJob(getJob(jobId), finalJobState);
	}

	/**
	 * {@inheritDoc}
	 */
	@Override
	@Transactional(readOnly = false)
	public void finalizeJob(Job job, Job.State finalJobState) throws IllegalArgumentException {
		JpaJob jpaJob = toJpaJob(job);

		if (jpaJob == null)
			return;

		if (jpaJob.getState() == Job.State.PREPARED) {
			if (finalJobState != null) {
				throw new IllegalArgumentException(
						"finalJobState is not 'null' of a job that is still in prepared state.");
			}

			em.remove(jpaJob);
		} else {
			if (finalJobState == null || !Job.FINAL_STATES.contains(finalJobState))
				throw new IllegalArgumentException("finalJobState not in " + Job.FINAL_STATES);

			jpaJob.setState(finalJobState);
			jpaJob.setLastUpdated(new Date());
			jpaJob.setFinished(jpaJob.getLastUpdated());
		}
	}

	private JpaJob toJpaJob(Job job) {
		JpaJob jpaJob = null;
		if (job != null) {
			if (job instanceof JpaJob && em.contains(job))
				jpaJob = (JpaJob) job;
			else
				jpaJob = getJob(job.getJobId());
		}
		return jpaJob;
	}

	/**
	 * {@inheritDoc}
	 */
	@Override
	public UUIDListPage getJobsByStateInRange(Job.State jobState,
											  Date updatedFromDate, Date updatedToDate, int pageNumber) {
		final boolean rangeQuery = updatedFromDate != null || updatedToDate != null;
		final TypedQuery<byte[]> q = em.createNamedQuery(rangeQuery ?
				"Job.SelectJobIdsByStateInRange" : "Job.SelectJobIdsByState", byte[].class);

		if (rangeQuery) {
			q.setParameter("fromDate", updatedFromDate == null ? new Date(0) : updatedFromDate);
			q.setParameter("toDate", updatedToDate == null ? new Date(currentTimeMillis() + 1000) : updatedToDate);
		}

		applyPage(q, pageNumber, pageSize).setParameter("state", jobState);
		return toListPage(q, bytesToUUIDConverter, new UUIDListPage());
	}

	/**
	 * {@inheritDoc}
	 */
	public Job.State getJobState(UUID jobId) {
		List<Job.State> result = em.createNamedQuery("Job.SelectJobStateById", Job.State.class).
				setParameter("publicGUID", toBytes(jobId)).getResultList();

		if (result.isEmpty()) {
			if (log.isTraceEnabled())
				log.trace("Queried state for non-existing job '{}'", jobId);
			return null;
		}

		return result.get(0);
	}

	JpaJob getJob(UUID jobId) {
		return jobId == null ? null : getJob(em.createNamedQuery("Job.SelectById", JpaJob.class), jobId);
	}

	<E extends Job> E getJob(TypedQuery<E> query, UUID jobId) {
		List<E> result = query.setParameter("publicGUID", toBytes(jobId)).getResultList();

		if (result.isEmpty()) {
			if (log.isTraceEnabled())
				log.trace("Queried non-existing job '{}'", jobId);
			return null;
		}

		return result.get(0);
	}

	/**
	 * {@inheritDoc}
	 */
	public Collection<SharedJob> getJobs(Collection<UUID> jobIds) {
		final Collection<SharedJob> result = new ArrayList<SharedJob>(jobIds.size());
		final TypedQuery<SharedJob> query = em.createNamedQuery("Job.SelectById", SharedJob.class);

		for (UUID id : jobIds)
			result.add(getJob(query, id));

		return result;
	}

	/**
	 * {@inheritDoc}
	 */
	public SharedJobDetails getJobDetails(UUID jobId) {
		JpaJob job = getJob(jobId);
		if (job == null)
			return null;

		List<byte[]> guids = em.createNamedQuery("JobSources.SelectAssignedSourceIds", byte[].class).
				setParameter("primaryKey", job.getPrimaryKey()).getResultList();
		List<SharedSource> assignedSources = new ArrayList<SharedSource>(guids.size());
		for (byte[] guid : guids)
			assignedSources.add(sourceFetcher.fetch(new SourceIdentifier(guid)));

		return new SharedJobDetails(job, cast(job.getSubJobs(), SharedJob.class), assignedSources);
	}
}