JpaPackageRepository

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

import com.trendmicro.grid.acl.ds.PackageProvider;
import com.trendmicro.grid.acl.ds.RemoteCacheable;
import com.trendmicro.grid.acl.ds.TagQueryExpression;
import com.trendmicro.grid.acl.ds.datatypes.SharedNamedFileIdentifier;
import com.trendmicro.grid.acl.ds.datatypes.SharedNamedFileIdentifierListPage;
import com.trendmicro.grid.acl.ds.datatypes.SharedPackageDetails;
import com.trendmicro.grid.acl.ds.datatypes.SharedPackageInformation;
import com.trendmicro.grid.acl.ds.jpa.entities.JpaFileInformation;
import com.trendmicro.grid.acl.ds.jpa.entities.JpaNamedFileIdentifier;
import com.trendmicro.grid.acl.ds.jpa.entities.JpaPackageDetails;
import com.trendmicro.grid.acl.ds.jpa.entities.JpaPackageInformation;
import com.trendmicro.grid.acl.ds.jpa.tagquery.TagQueryProviderSelector;
import com.trendmicro.grid.acl.ds.jpa.util.FileQueryConfigurator;
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.Query;
import javax.persistence.TypedQuery;
import java.util.*;

import static com.trendmicro.grid.acl.ds.jpa.util.JpaUtils.*;

/**
 * Implements FileProvider using JPA.
 *
 * @author juergen_kellerer, 2010-06-11
 * @version 1.0
 */
@Repository
@Transactional(readOnly = true)
public class JpaPackageRepository implements PackageProvider {

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

	private static int pageSize = 1000;

	public static int getPageSize() {
		return pageSize;
	}

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

	/**
	 * Assembles an SharedNamedFileIdentifier instance.
	 * <p/>
	 * Note: The elements returned by this implementation must not be of an
	 * instance of JpaNamedFileIdentifier as the equals and hashCode
	 * implementations differ and are not compatible.
	 */
	static final Callback<Object[], NamedFileIdentifier> namedFileIdentifierAssembler = new Callback<Object[], NamedFileIdentifier>() {
		@Override
		public NamedFileIdentifier call(Object[] columns) {
			final FileIdentifier identifier = (FileIdentifier) columns[1];
			if (columns[0] instanceof String) {
				final String fileName = (String) columns[0];
				return new SharedNamedFileIdentifier(identifier, fileName);
			} else
				return ((JpaNamedFileIdentifier) columns[0]).export(identifier);
		}
	};

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

	@Resource
	TagQueryProviderSelector providerSelector;

	Integer rootPackageKey;

	@Override
	public NameListPage getPackageNamesInFamily(String packageFamilyName, int pageNumber) {
		TypedQuery<String> query = em.createNamedQuery("Packages.GetPackageNamesInFamily", String.class).
				setParameter("familyName", packageFamilyName);
		return toNameListPage(query, pageNumber, pageSize);
	}

	@Override
	@RemoteCacheable
	public Collection<Boolean> isPackagesTaggedWithById(Collection<FileIdentifier> files, String[] tags) {
		// This is a un-optimized version of this method that should not be called.
		// The reason that this implementation exists is only to handle the case that caching is disabled.
		final Collection<Boolean> results = new ArrayList<Boolean>(files.size());
		for (SharedPackageInformation information : getPackageInformationListById(files))
			results.add(isTaggedWith(information, tags));

		return results;
	}

	@Override
	@RemoteCacheable
	public Collection<Boolean> isPackagesTaggedWithByName(Collection<String> packageNames, String[] tags) {
		// This is a un-optimized version of this method that should not be called.
		// The reason that this implementation exists is only to handle the case that caching is disabled.
		final Collection<Boolean> results = new ArrayList<Boolean>(packageNames.size());
		for (SharedPackageInformation information : getPackageInformationListByName(packageNames))
			results.add(isTaggedWith(information, tags));

		return results;
	}

	@Override
	public NameListPage getMatchingPackageNames(TagQueryExpression expression, Range range, int pageNumber) {
		if (range != null && !(range instanceof DaysRange)) {
			log.warn("TMACL-00830:Ignoring CompositeRange in JPA implementation. Such requests should get handled inside the Cache module.");
			return null;
		}

		final TagQueryProvider provider = providerSelector.getSelectedProvider(em);
		final Query query = provider.getMatchingPackageNamesQuery(em, expression, (DaysRange) range);

		return toNameListPage(query, pageNumber, pageSize);
	}

	@Override
	public SharedNamedFileIdentifierListPage getFilesContainedInPackageById(FileIdentifier packageFile, int pageNumber) {
		final FileQueryConfigurator queryConfigurator = new FileQueryConfigurator<Object>(em, null, "Packages.SelectFilesContainedInPackage");
		Query query = applyPage(queryConfigurator.getUntypedQuery(packageFile), pageNumber, pageSize);
		return toListPage(query, namedFileIdentifierAssembler, new SharedNamedFileIdentifierListPage());
	}

	@Override
	public SharedNamedFileIdentifierListPage getFilesContainedInPackageByName(String packageName, int pageNumber) {
		Query query = em.createNamedQuery("Packages.SelectFilesContainedInPackageByName").setParameter("name", packageName);
		applyPage(query, pageNumber, pageSize);
		return toListPage(query, namedFileIdentifierAssembler, new SharedNamedFileIdentifierListPage());
	}

	/**
	 * Returns a map of contained files for the given package details.
	 * <p/>
	 * Note: It is adviceable to use the values of the map only to satisfy foreign key dependencies
	 * as all operations cause lazy loading of file or package detail resources.
	 *
	 * @param packageDetails the package details to lookup all contained files for.
	 * @return a map of contained files for the given package details.
	 */
	public Map<NamedFileIdentifier, JpaNamedFileIdentifier> getFilesContainedInPackage(JpaPackageDetails packageDetails) {

		final List queryResults = em.createNamedQuery("Packages.SelectFilesContainedInPackageByPrimaryKey").
				setParameter("packageDetails", packageDetails).getResultList();
		final Map<NamedFileIdentifier, JpaNamedFileIdentifier> results = new HashMap<NamedFileIdentifier, JpaNamedFileIdentifier>(queryResults.size());

		for (Object queryResult : queryResults) {
			Object[] cols = (Object[]) queryResult;
			results.put(namedFileIdentifierAssembler.call(cols), (JpaNamedFileIdentifier) cols[0]);
		}

		return results;
	}

	@Override
	public NameListPage getReferencingPackageNamesById(FileIdentifier file, int pageNumber) {
		final FileQueryConfigurator<String> queryConfigurator =
				new FileQueryConfigurator<String>(em, String.class, "Packages.SelectReferencingPackageNames");
		return toNameListPage(queryConfigurator.getQuery(file), pageNumber, pageSize);
	}

	@Override
	public NameListPage getReferencingPackageNames(String packageName, int pageNumber) {
		final TypedQuery<String> query = em.createNamedQuery(
				"Packages.SelectReferencingPackageNamesByName", String.class).setParameter("name", packageName);
		return toNameListPage(query, pageNumber, pageSize);
	}

	@Override
	public NameListPage getPackagesContainedInPackageById(FileIdentifier packageFile, int pageNumber) {
		final FileQueryConfigurator<String> queryConfigurator =
				new FileQueryConfigurator<String>(em, String.class, "Packages.SelectPackagesContainedInPackage");
		return toNameListPage(queryConfigurator.getQuery(packageFile), pageNumber, pageSize);
	}

	@Override
	public NameListPage getPackagesContainedInPackageByName(String packageName, int pageNumber) {
		final TypedQuery<String> query = em.createNamedQuery(
				"Packages.SelectPackagesContainedInPackageByName", String.class).setParameter("name", packageName);
		return toNameListPage(query, pageNumber, pageSize);
	}

	@Override
	@RemoteCacheable
	public Collection<SharedPackageInformation> getPackageInformationListById(Collection<FileIdentifier> files) {
		final FileQueryConfigurator queryConfigurator = new FileQueryConfigurator<Object>(em, null, "Packages.SelectPackageInformation");

		final Collection<SharedPackageInformation> results = new ArrayList<SharedPackageInformation>(files.size());
		for (FileIdentifier file : files)
			doGetPackageInformationList(queryConfigurator.getUntypedQuery(file), results);

		return results;
	}

	@Override
	@RemoteCacheable
	public Collection<SharedPackageInformation> getPackageInformationListByName(Collection<String> packageNames) {
		final Query query = em.createNamedQuery("Packages.SelectPackageInformationByName");

		final Collection<SharedPackageInformation> results = new ArrayList<SharedPackageInformation>(packageNames.size());
		for (String name : packageNames)
			doGetPackageInformationList(query.setParameter("name", name), results);

		return results;
	}

	@SuppressWarnings("unchecked")
	private static void doGetPackageInformationList(final Query query, final Collection<SharedPackageInformation> results) {
		List<Object[]> rows = query.getResultList();
		if (rows.isEmpty())
			results.add(null);
		else {
			Object[] row = rows.get(0);
			((JpaPackageInformation) row[0]).setPackageFileInformation((JpaFileInformation) row[1]);
			results.add((SharedPackageInformation) row[0]);
		}
	}

	@Override
	@SuppressWarnings("unchecked")
	@RemoteCacheable(comment = "Note: Don't use this method for JPA purposes, " +
			"the returned entities will not be usable when returned from the remote cache.")
	public Collection<SharedPackageDetails> getPackageDetailsListById(Collection<FileIdentifier> files) {
		return (Collection) getJpaPackageDetailsListById(files);
	}

	/**
	 * Returns the package details on the given package file ids.
	 *
	 * @param files The SHA1 hashes of the package files.
	 * @return the package details on the given package file ids.
	 */
	public Collection<JpaPackageDetails> getJpaPackageDetailsListById(Collection<FileIdentifier> files) {
		final FileQueryConfigurator<Integer> queryConfigurator =
				new FileQueryConfigurator<Integer>(em, Integer.class, "Packages.SelectPackageDetailsPrimaryKey");

		final Collection<Integer> results = new ArrayList<Integer>(files.size());
		for (FileIdentifier file : files)
			appendSingleElementToList(results, queryConfigurator.getQuery(file));

		return toDetails(results);
	}

	@Override
	@SuppressWarnings("unchecked")
	@RemoteCacheable(comment = "Note: Don't use this method for JPA purposes, " +
			"the returned entities will not be usable when returned from the remote cache.")
	public Collection<SharedPackageDetails> getPackageDetailsListByName(Collection<String> packageNames) {
		return (Collection) getJpaPackageDetailsListByName(packageNames);
	}

	/**
	 * Returns the package details on the given package names.
	 *
	 * @param packageNames The names the packages to lookup.
	 * @return the package details on the given package names.
	 */
	public Collection<JpaPackageDetails> getJpaPackageDetailsListByName(Collection<String> packageNames) {
		final TypedQuery<Integer> query = em.createNamedQuery("Packages.SelectPackageDetailsPrimaryKeyByName", Integer.class);

		final Collection<Integer> results = new ArrayList<Integer>(packageNames.size());
		for (String name : packageNames)
			appendSingleElementToList(results, query.setParameter("name", name));

		return toDetails(results);
	}

	private Collection<JpaPackageDetails> toDetails(Collection<Integer> primaryKeys) {
		final Collection<JpaPackageDetails> results = new ArrayList<JpaPackageDetails>(primaryKeys.size());
		for (Integer pk : primaryKeys)
			results.add(pk == null ? null : em.find(JpaPackageDetails.class, pk));
		return results;
	}

	@Override
	@RemoteCacheable(comment = "Note: Don't use this method for JPA purposes, " +
			"the returned entities will not be usable when returned from the remote cache.")
	public Collection<FileIdentifier> getPackageFileIdentifiersByName(Collection<String> packageNames) {
		TypedQuery<FileIdentifier> query = em.createNamedQuery("Packages.GetPackageFileIdentifierByName", FileIdentifier.class);

		final Collection<FileIdentifier> results = new ArrayList<FileIdentifier>(packageNames.size());
		for (String name : packageNames)
			appendSingleElementToList(results, query.setParameter("name", name));

		return results;
	}

	/**
	 * Returns the primary key of the specified package or 'null' if it does not exist.
	 *
	 * @param packageName the name of the package to lookup.
	 * @return the primary key of the package or 'null' if it does not exist.
	 */
	public Integer getPackagePrimaryKeyByName(String packageName) {
		final boolean isRoot = JpaPackageInformation.ROOT_PACKAGE_NAME.equalsIgnoreCase(packageName);
		if (isRoot && rootPackageKey != null)
			return rootPackageKey;

		final List<Integer> r = em.createNamedQuery("Packages.SelectPackageDetailsPrimaryKeyByName", Integer.class).
				setParameter("name", packageName).getResultList();
		final Integer result = r.isEmpty() ? null : r.get(0);

		if (isRoot)
			rootPackageKey = result;

		return result;
	}
}