Cohort sampling

.gitignore

@@ -1,3 +1,4 @@
+WebAPIConfig/
 *application.properties
 .idea/
 .metadata/

pom.xml

@@ -205,6 +205,8 @@
     <logging.level.org.ohdsi>info</logging.level.org.ohdsi>
     <logging.level.org.springframework.orm>info</logging.level.org.springframework.orm>
     <logging.level.org.springframework.jdbc>info</logging.level.org.springframework.jdbc>
+    <logging.level.org.springframework.web>info</logging.level.org.springframework.web>
+    <logging.level.org.hibernate>info</logging.level.org.hibernate>
     <logging.level.org.apache.shiro>warn</logging.level.org.apache.shiro>
 
     <spring.batch.taskExecutor.corePoolSize>10</spring.batch.taskExecutor.corePoolSize>
@@ -949,6 +951,11 @@
         <version>2.0.1</version>
         <scope>test</scope>
     </dependency>
+    <dependency>
+      <groupId>log4j</groupId>
+      <artifactId>log4j</artifactId>
+      <version>1.2.17</version>
+    </dependency>
   </dependencies>
   <profiles>
     <profile>

src/main/java/org/ohdsi/webapi/cohortsample/CleanupCohortSamplesTasklet.java

@@ -0,0 +1,145 @@
+package org.ohdsi.webapi.cohortsample;
+
+import org.ohdsi.webapi.cohortdefinition.CleanupCohortTasklet;
+import org.ohdsi.webapi.job.JobExecutionResource;
+import org.ohdsi.webapi.job.JobTemplate;
+import org.ohdsi.webapi.source.Source;
+import org.ohdsi.webapi.source.SourceDaimon;
+import org.ohdsi.webapi.source.SourceRepository;
+import org.ohdsi.webapi.util.PreparedStatementRenderer;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.springframework.batch.core.Job;
+import org.springframework.batch.core.JobParameters;
+import org.springframework.batch.core.JobParametersBuilder;
+import org.springframework.batch.core.Step;
+import org.springframework.batch.core.StepContribution;
+import org.springframework.batch.core.configuration.annotation.JobBuilderFactory;
+import org.springframework.batch.core.configuration.annotation.StepBuilderFactory;
+import org.springframework.batch.core.job.builder.SimpleJobBuilder;
+import org.springframework.batch.core.scope.context.ChunkContext;
+import org.springframework.batch.core.step.tasklet.Tasklet;
+import org.springframework.batch.repeat.RepeatStatus;
+import org.springframework.transaction.support.TransactionTemplate;
+
+import java.util.List;
+import java.util.Map;
+
+import static org.ohdsi.webapi.Constants.Params.COHORT_DEFINITION_ID;
+import static org.ohdsi.webapi.Constants.Params.JOB_NAME;
+import static org.ohdsi.webapi.Constants.Params.SOURCE_ID;
+
+public class CleanupCohortSamplesTasklet implements Tasklet {
+	private static final Logger log = LoggerFactory.getLogger(CleanupCohortTasklet.class);
+
+	private final TransactionTemplate transactionTemplate;
+	private final SourceRepository sourceRepository;
+	private final CohortSamplingService samplingService;
+	private final CohortSampleRepository sampleRepository;
+
+	public CleanupCohortSamplesTasklet(
+			final TransactionTemplate transactionTemplate,
+			final SourceRepository sourceRepository,
+			CohortSamplingService samplingService,
+			CohortSampleRepository sampleRepository
+	) {
+		this.transactionTemplate = transactionTemplate;
+		this.sourceRepository = sourceRepository;
+		this.samplingService = samplingService;
+		this.sampleRepository = sampleRepository;
+	}
+
+	private Integer doTask(ChunkContext chunkContext) {
+		Map<String, Object> jobParams = chunkContext.getStepContext().getJobParameters();
+		int cohortDefinitionId = Integer.parseInt(jobParams.get(COHORT_DEFINITION_ID).toString());
+
+		if (jobParams.containsKey(SOURCE_ID)) {
+			int sourceId = Integer.parseInt(jobParams.get(SOURCE_ID).toString());
+			Source source = this.sourceRepository.findOne(sourceId);
+			if (source != null) {
+				return mapSource(source, cohortDefinitionId);
+			} else {
+				return 0;
+			}
+		} else {
+			return this.sourceRepository.findAll().stream()
+					.filter(source-> source.getDaimons()
+							.stream()
+							.anyMatch(daimon -> daimon.getDaimonType() == SourceDaimon.DaimonType.Results))
+					.mapToInt(source -> mapSource(source, cohortDefinitionId))
+					.sum();
+		}
+	}
+
+	private int mapSource(Source source, int cohortDefinitionId) {
+		try {
+			String resultSchema = source.getTableQualifier(SourceDaimon.DaimonType.Results);
+			return transactionTemplate.execute(transactionStatus -> {
+				List<CohortSample> samples = sampleRepository.findByCohortDefinitionIdAndSourceId(cohortDefinitionId, source.getId());
+				if (samples.isEmpty()) {
+					return 0;
+				}
+
+				sampleRepository.delete(samples);
+
+				int[] cohortSampleIds = samples.stream()
+						.mapToInt(CohortSample::getId)
+						.toArray();
+
+				PreparedStatementRenderer renderer = new PreparedStatementRenderer(
+						source,
+						"/resources/cohortsample/sql/deleteSampleElementsById.sql",
+						"results_schema",
+						resultSchema,
+						"cohortSampleId",
+						cohortSampleIds);
+
+				samplingService.getSourceJdbcTemplate(source)
+						.update(renderer.getSql(), renderer.getOrderedParams());
+				return cohortSampleIds.length;
+			});
+		} catch (Exception e) {
+			log.error("Error deleting samples for cohort: {}, cause: {}", cohortDefinitionId, e.getMessage());
+			return 0;
+		}
+	}
+
+	@Override
+	public RepeatStatus execute(final StepContribution contribution, final ChunkContext chunkContext) throws Exception {
+		this.transactionTemplate.execute(status -> doTask(chunkContext));
+
+		return RepeatStatus.FINISHED;
+	}
+
+	public JobExecutionResource launch(JobBuilderFactory jobBuilders, StepBuilderFactory stepBuilders, JobTemplate jobTemplate, int cohortDefinitionId) {
+		JobParametersBuilder builder = new JobParametersBuilder();
+		builder.addString(JOB_NAME, String.format("Cleanup cohort samples of cohort definition %d.", cohortDefinitionId));
+		builder.addString(COHORT_DEFINITION_ID, String.valueOf(cohortDefinitionId));
+
+		log.info("Beginning cohort cleanup for cohort definition id: {}", cohortDefinitionId);
+		return launch(jobBuilders, stepBuilders, jobTemplate, builder.toJobParameters());
+	}
+
+	public JobExecutionResource launch(JobBuilderFactory jobBuilders, StepBuilderFactory stepBuilders, JobTemplate jobTemplate, int cohortDefinitionId, int sourceId) {
+		JobParametersBuilder builder = new JobParametersBuilder();
+		builder.addString(JOB_NAME, String.format("Cleanup cohort samples of cohort definition %d.", cohortDefinitionId));
+		builder.addString(COHORT_DEFINITION_ID, String.valueOf(cohortDefinitionId));
+		builder.addString(SOURCE_ID, String.valueOf(sourceId));
+
+		log.info("Beginning cohort cleanup for cohort definition id {} and source ID {}", cohortDefinitionId, sourceId);
+		return launch(jobBuilders, stepBuilders, jobTemplate, builder.toJobParameters());
+	}
+
+	private JobExecutionResource launch(JobBuilderFactory jobBuilders, StepBuilderFactory stepBuilders, JobTemplate jobTemplate, JobParameters jobParameters) {
+		Step cleanupStep = stepBuilders.get("cohortSample.cleanupSamples")
+				.tasklet(this)
+				.build();
+
+		SimpleJobBuilder cleanupJobBuilder = jobBuilders.get("cleanupSamples")
+				.start(cleanupStep);
+
+		Job cleanupCohortJob = cleanupJobBuilder.build();
+
+		return jobTemplate.launch(cleanupCohortJob, jobParameters);
+	}
+}

src/main/java/org/ohdsi/webapi/cohortsample/CohortSample.java

@@ -0,0 +1,139 @@
+package org.ohdsi.webapi.cohortsample;
+
+import org.hibernate.annotations.GenericGenerator;
+import org.hibernate.annotations.Parameter;
+import org.ohdsi.webapi.model.CommonEntity;
+
+import javax.persistence.Column;
+import javax.persistence.Entity;
+import javax.persistence.GeneratedValue;
+import javax.persistence.Id;
+import javax.persistence.Table;
+import javax.persistence.Transient;
+import java.util.List;
+
+/**
+ * Cohort sample details.
+ */
+@Entity(name = "CohortSample")
+@Table(name = "cohort_sample")
+public class CohortSample extends CommonEntity<Integer> {
+	@Id
+	@GenericGenerator(
+			name = "cohort_sample_generator",
+			strategy = "org.hibernate.id.enhanced.SequenceStyleGenerator",
+			parameters = {
+					@Parameter(name = "sequence_name", value = "cohort_sample_sequence"),
+					@Parameter(name = "increment_size", value = "1")
+			}
+	)
+	@GeneratedValue(generator = "cohort_sample_generator")
+	private Integer id;
+
+	@Column
+	private String name;
+
+	@Column(name = "cohort_definition_id")
+	private int cohortDefinitionId;
+
+	@Column(name = "source_id")
+	private int sourceId;
+
+	@Column(name = "age_mode")
+	private String ageMode;
+
+	@Column(name = "age_min")
+	private Integer ageMin;
+
+	@Column(name = "age_max")
+	private Integer ageMax;
+
+	@Column(name = "gender_concept_ids")
+	private String genderConceptIds;
+
+	@Column
+	private int size;
+
+	@Transient
+	private List<SampleElement> elements;
+
+	public Integer getId() {
+		return this.id;
+	}
+
+	public void setId(Integer id) {
+		this.id = id;
+	}
+
+	public int getCohortDefinitionId() {
+		return cohortDefinitionId;
+	}
+
+	public void setCohortDefinitionId(int cohortDefinitionId) {
+		this.cohortDefinitionId = cohortDefinitionId;
+	}
+
+	public List<SampleElement> getElements() {
+		return elements;
+	}
+
+	public void setElements(List<SampleElement> elements) {
+		this.elements = elements;
+	}
+
+	public int getSize() {
+		return size;
+	}
+
+	public void setSize(int size) {
+		this.size = size;
+	}
+
+	public Integer getAgeMin() {
+		return ageMin;
+	}
+
+	public void setAgeMin(Integer ageMin) {
+		this.ageMin = ageMin;
+	}
+
+	public Integer getAgeMax() {
+		return ageMax;
+	}
+
+	public void setAgeMax(Integer ageMax) {
+		this.ageMax = ageMax;
+	}
+
+	public String getGenderConceptIds() {
+		return genderConceptIds;
+	}
+
+	public void setGenderConceptIds(String genderConceptIds) {
+		this.genderConceptIds = genderConceptIds;
+	}
+
+	public int getSourceId() {
+		return sourceId;
+	}
+
+	public void setSourceId(int sourceId) {
+		this.sourceId = sourceId;
+	}
+
+	public void setName(String name) {
+		this.name = name;
+	}
+
+	public String getName() {
+		return name;
+	}
+
+	public String getAgeMode() {
+		return ageMode;
+	}
+
+	public void setAgeMode(String ageMode) {
+		this.ageMode = ageMode;
+	}
+}

src/main/java/org/ohdsi/webapi/cohortsample/CohortSampleRepository.java

@@ -0,0 +1,26 @@
+package org.ohdsi.webapi.cohortsample;
+
+import org.springframework.data.jpa.repository.Query;
+import org.springframework.data.repository.CrudRepository;
+import org.springframework.data.repository.query.Param;
+import org.springframework.stereotype.Component;
+
+import java.util.List;
+
+/**
+ * Repository of samples. This does not fetch any sample elements.
+ */
+@Component
+public interface CohortSampleRepository extends CrudRepository<CohortSample, Integer> {
+	@Query("SELECT c FROM CohortSample c LEFT JOIN FETCH c.createdBy WHERE c.id = :id")
+	CohortSample findById(@Param("id") int cohortSampleId);
+
+	@Query("SELECT c FROM CohortSample c LEFT JOIN FETCH c.createdBy WHERE c.cohortDefinitionId = :cohortDefinitionId AND c.sourceId = :sourceId")
+	List<CohortSample> findByCohortDefinitionIdAndSourceId(@Param("cohortDefinitionId") int cohortDefinitionId, @Param("sourceId") int sourceId);
+
+	@Query("SELECT count(c.id) FROM CohortSample c WHERE c.cohortDefinitionId = :cohortDefinitionId")
+	int countSamples(@Param("cohortDefinitionId") int cohortDefinitionId);
+
+	@Query("SELECT count(c.id) FROM CohortSample c WHERE c.cohortDefinitionId = :cohortDefinitionId AND c.sourceId = :sourceId")
+	int countSamples(@Param("cohortDefinitionId") int cohortDefinitionId, @Param("sourceId") int sourceId);
+}