AuditDAO.java

/**
 * This Source Code Form is subject to the terms of the Mozilla Public License,
 * v. 2.0. If a copy of the MPL was not distributed with this file, You can
 * obtain one at http://mozilla.org/MPL/2.0/. OpenMRS is also distributed under
 * the terms of the Healthcare Disclaimer located at http://openmrs.org/license.
 *
 * Copyright (C) OpenMRS Inc. OpenMRS is a registered trademark and the OpenMRS
 * graphic logo is a trademark of OpenMRS Inc.
 */
package org.openmrs.module.eversauditing.api.dao;

import org.hibernate.Session;
import org.hibernate.SessionFactory;
import org.hibernate.envers.AuditReader;
import org.hibernate.envers.AuditReaderFactory;
import org.hibernate.envers.RevisionType;
import org.hibernate.envers.query.AuditQuery;
import org.openmrs.api.context.Context;
import org.openmrs.api.db.hibernate.envers.OpenmrsRevisionEntity;
import org.openmrs.module.eversauditing.AuditEntity;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Repository;
import org.springframework.transaction.annotation.Transactional;

import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collector;
import java.util.stream.Collectors;

@Repository
public class AuditDao {
	
	private final SessionFactory sessionFactory;
	
	private final Logger logger = LoggerFactory.getLogger(AuditDao.class);
	
	public AuditDao(SessionFactory sessionFactory) {
		this.sessionFactory = sessionFactory;
	}
	
	@SuppressWarnings("unchecked")
    public <T> List<AuditEntity<T>> getAllRevisions(Class<T> entityClass) {
		AuditReader auditReader = AuditReaderFactory.get(sessionFactory.getCurrentSession());;
        AuditQuery auditQuery = auditReader.createQuery().forRevisionsOfEntity(entityClass, false, true);
        return (List<AuditEntity<T>>) auditQuery.getResultList().stream()
                .map(result -> {
                    Object[] array = (Object[]) result;
                    T entity = entityClass.cast(array[0]);
                    OpenmrsRevisionEntity revisionEntity = (OpenmrsRevisionEntity) array[1];
                    RevisionType revisionType = (RevisionType) array[2];
					String changedBy = Context.getUserService().getUser(revisionEntity.getChangedBy()).toString();
                    return new AuditEntity<>(entity, revisionEntity, revisionType, changedBy);
                })
                .collect(Collectors.toList());
    }
	
	public <T> T getRevisionById(Class<T> entityClass, int entityId, int revisionId) {
		AuditReader auditReader = AuditReaderFactory.get(sessionFactory.getCurrentSession());
		T entity = auditReader.find(entityClass, entityId, revisionId);
		return entity;
	}
	
	@Transactional
	public void migrateAuditData() {
		Session session = sessionFactory.getCurrentSession();
		String currentSchema = (String) session.createNativeQuery("SELECT DATABASE()").getSingleResult();
		String getAuditClassesQuery = "SELECT table_name FROM information_schema.tables WHERE table_name LIKE '%_aud' and table_schema = :currentSchema";
		List<String> auditClassNames = session.createNativeQuery(getAuditClassesQuery)
		        .setParameter("currentSchema", currentSchema).getResultList();
		
		for (String auditEntity : auditClassNames) {
			String entity = auditEntity.split("_aud")[0];
			try {

				String tableExistsQuery = "SELECT COUNT(*) FROM information_schema.tables WHERE table_name = :entity AND table_schema = :currentSchema";
				Long tableExists = ((Number) session.createNativeQuery(tableExistsQuery)
						.setParameter("entity", entity)
						.setParameter("currentSchema", currentSchema)
						.getSingleResult()).longValue();
						
				if(tableExists == 0){
					logger.warn("Skipping migration: Table {} not found", entity);
					continue;		
				}

				String getCountQuery = "SELECT COUNT(*) FROM ";
				int count = Integer.parseInt(session.createNativeQuery(getCountQuery + entity).getSingleResult().toString());
				int auditCount = Integer.parseInt(session.createNativeQuery(getCountQuery + auditEntity).getSingleResult()
				        .toString());
				
				if (count != 0 && auditCount == 0) {
					
					String getAuditColumnsQuery = "SELECT column_name FROM information_schema.columns WHERE table_name = :auditEntity and table_schema = :currentSchema";
					
					List<String> auditColumns = session.createNativeQuery(getAuditColumnsQuery).setParameter("auditEntity", auditEntity)
					        .setParameter("currentSchema", currentSchema).getResultList();
					auditColumns.removeIf(c -> c.equalsIgnoreCase("REV") || c.equalsIgnoreCase("REVTYPE"));
					
					List<String> columns = new ArrayList<>(auditColumns);
					
					auditColumns.add("rev");
					auditColumns.add("revtype");

					String getRevIdQuery = "SELECT COALESCE(MAX(id), 0) FROM revision_entity";
					String updateRevisionEntity = "INSERT INTO revision_entity (id, timestamp) VALUES (:revId, :timestamp)";
					String migrateQuery = "INSERT INTO " + auditEntity + " (" + String.join(",", auditColumns) + ") " + "SELECT " + String.join(",", columns)
					        + ", :revId AS rev, 0 AS revtype FROM " + entity;
					
					Long currentTimestamp = -111111111111111L;
					Long revId = Long.parseLong(session.createNativeQuery(getRevIdQuery).getSingleResult().toString()) + 1;
					
					session.createNativeQuery(updateRevisionEntity).setParameter("revId", revId)
					        .setParameter("timestamp", currentTimestamp).executeUpdate();
					session.createNativeQuery(migrateQuery).setParameter("revId", revId).executeUpdate();
					logger.info("migration completed for: {}", entity);
					
				}
			}
			catch (Exception e) {
				logger.warn("Exception occurred while migrating: {} for {}, rollbacking all the transactions", e.getMessage(), entity);
			}
		}
	}
}