diff --git a/src/main/java/com/google/cloud/genomics/utils/Contig.java b/src/main/java/com/google/cloud/genomics/utils/Contig.java index aa20f62..4891e90 100644 --- a/src/main/java/com/google/cloud/genomics/utils/Contig.java +++ b/src/main/java/com/google/cloud/genomics/utils/Contig.java @@ -40,7 +40,7 @@ public class Contig implements Serializable { private static final long serialVersionUID = -1730387112193404207L; - public static final long DEFAULT_NUMBER_OF_BASES_PER_SHARD = 100000; + public static final long DEFAULT_NUMBER_OF_BASES_PER_SHARD = 1000000; public enum SexChromosomeFilter { INCLUDE_XY, EXCLUDE_XY } diff --git a/src/main/java/com/google/cloud/genomics/utils/Paginator.java b/src/main/java/com/google/cloud/genomics/utils/Paginator.java index 20b6d07..11ea559 100644 --- a/src/main/java/com/google/cloud/genomics/utils/Paginator.java +++ b/src/main/java/com/google/cloud/genomics/utils/Paginator.java @@ -13,6 +13,12 @@ */ package com.google.cloud.genomics.utils; +import java.io.IOException; +import java.io.Serializable; +import java.util.Collections; +import java.util.Iterator; +import java.util.Map; + import com.google.api.services.genomics.Genomics; import com.google.api.services.genomics.GenomicsRequest; import com.google.api.services.genomics.model.Annotation; @@ -51,17 +57,16 @@ import com.google.api.services.genomics.model.Variant; import com.google.api.services.genomics.model.VariantSet; import com.google.common.base.Function; +import com.google.common.base.Joiner; import com.google.common.base.Optional; +import com.google.common.base.Preconditions; import com.google.common.base.Predicate; import com.google.common.collect.AbstractSequentialIterator; import com.google.common.collect.FluentIterable; +import com.google.common.collect.ImmutableMap; import com.google.common.collect.Iterables; import com.google.common.collect.Lists; - -import java.io.IOException; -import java.io.Serializable; -import java.util.Collections; -import java.util.Iterator; +import com.google.common.collect.Maps; /** * An abstraction that understands the {@code pageToken} / {@code nextPageToken} protocol for paging @@ -319,6 +324,13 @@ public static class Reads extends Paginator< SearchReadsResponse, Read> { + // TODO: When this is supported server-side, these additional fields are no longer required + private static ImmutableMap REQUIRED_STRICT_SHARD_FIELDS = + ImmutableMap.builder() + .putAll(REQUIRED_FIELDS) + .put("alignment", ".*\\p{Punct}alignment\\p{Punct}.*") + .put("position", ".*\\p{Punct}position\\p{Punct}.*") + .build(); private final ShardBoundary shardBoundary; private Predicate shardPredicate = null; @@ -369,6 +381,14 @@ public boolean apply(Read read) { .or(request)); } + @Override + public GenomicsRequestInitializer> setFieldsInitializer(final String fields) { + if(shardBoundary == ShardBoundary.STRICT) { + return new GenomicsSearchFieldRequestInitializer(fields, REQUIRED_STRICT_SHARD_FIELDS); + } + return new GenomicsSearchFieldRequestInitializer(fields, REQUIRED_FIELDS); + } + @Override Genomics.Reads getApi(Genomics genomics) { return genomics.reads(); } @@ -396,6 +416,13 @@ public static class Annotations extends Paginator< SearchAnnotationsResponse, Annotation> { + // TODO: When this is supported server-side, these additional fields are no longer required + private static ImmutableMap REQUIRED_STRICT_SHARD_FIELDS = + ImmutableMap.builder() + .putAll(REQUIRED_FIELDS) + .put("position", ".*\\p{Punct}position\\p{Punct}.*") + .put("start", ".*\\p{Punct}start\\p{Punct}.*") + .build(); private final ShardBoundary shardBoundary; private Predicate shardPredicate = null; @@ -445,6 +472,14 @@ public boolean apply(Annotation anno) { .or(request)); } + @Override + public GenomicsRequestInitializer> setFieldsInitializer(final String fields) { + if(shardBoundary == ShardBoundary.STRICT) { + return new GenomicsSearchFieldRequestInitializer(fields, REQUIRED_STRICT_SHARD_FIELDS); + } + return new GenomicsSearchFieldRequestInitializer(fields, REQUIRED_FIELDS); + } + @Override Genomics.Annotations getApi(Genomics genomics) { return genomics.annotations(); } @@ -784,6 +819,12 @@ public static class Variants extends Paginator< SearchVariantsResponse, Variant> { + // TODO: When this is supported server-side, these additional fields are no longer required + private static ImmutableMap REQUIRED_STRICT_SHARD_FIELDS = + ImmutableMap.builder() + .putAll(REQUIRED_FIELDS) + .put("start", ".*\\p{Punct}start\\p{Punct}.*") + .build(); private final ShardBoundary shardBoundary; private Predicate shardPredicate = null; @@ -834,6 +875,14 @@ public boolean apply(Variant variant) { .or(request)); } + @Override + public GenomicsRequestInitializer> setFieldsInitializer(final String fields) { + if(shardBoundary == ShardBoundary.STRICT) { + return new GenomicsSearchFieldRequestInitializer(fields, REQUIRED_STRICT_SHARD_FIELDS); + } + return new GenomicsSearchFieldRequestInitializer(fields, REQUIRED_FIELDS); + } + @Override Genomics.Variants getApi(Genomics genomics) { return genomics.variants(); } @@ -905,17 +954,46 @@ private Variantsets(Genomics genomics) { @Override public void initialize(GenomicsRequest search) {} }; - public static GenomicsRequestInitializer> setFieldsInitializer( - final String fields) { - return new GenomicsRequestInitializer>() { - @Override public void initialize(GenomicsRequest search) { - if (null != fields) { - search.setFields(fields); - } - } - }; + protected static final ImmutableMap REQUIRED_FIELDS = ImmutableMap + .builder() + .put("nextPageToken", "(.+\\p{Punct})?nextPageToken(\\p{Punct}.+)?"). + build(); + + public GenomicsRequestInitializer> setFieldsInitializer(final String fields) { + return new GenomicsSearchFieldRequestInitializer(fields, REQUIRED_FIELDS); } + public static class GenomicsSearchFieldRequestInitializer implements + GenomicsRequestInitializer> { + private final String fields; + private final ImmutableMap requiredFields; + + public GenomicsSearchFieldRequestInitializer(String fields, + ImmutableMap requiredFields) { + this.fields = fields; + this.requiredFields = requiredFields; + } + + @Override + public void initialize(GenomicsRequest search) { + if (null != fields) { + Map missingFields = + Maps.filterValues(requiredFields, new Predicate() { + @Override + public boolean apply(String input) { + return !fields.matches(input); + } + }); + Preconditions + .checkArgument( + missingFields.isEmpty(), + "Required field(s) missing: '%s' Add this field to the list of fields to be returned in the partial response.", + Joiner.on(",").join(missingFields.keySet())); + search.setFields(fields); + } + } + } + private final Genomics genomics; public Paginator(Genomics genomics) { diff --git a/src/test/java/com/google/cloud/genomics/utils/PaginatorTest.java b/src/test/java/com/google/cloud/genomics/utils/PaginatorTest.java index 51b799e..bc6a2fd 100644 --- a/src/test/java/com/google/cloud/genomics/utils/PaginatorTest.java +++ b/src/test/java/com/google/cloud/genomics/utils/PaginatorTest.java @@ -32,7 +32,10 @@ import org.hamcrest.CoreMatchers; import org.junit.Before; +import org.junit.Ignore; +import org.junit.Rule; import org.junit.Test; +import org.junit.rules.ExpectedException; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; import org.mockito.Mock; @@ -55,6 +58,9 @@ public class PaginatorTest { @Mock Genomics.Reads reads; @Mock Genomics.Reads.Search readsSearch; + @Rule + public ExpectedException thrown = ExpectedException.none(); + @Before public void initMocks() { MockitoAnnotations.initMocks(this); @@ -120,16 +126,26 @@ public void testFields() throws Exception { Paginator.ReadGroupSets paginator = Paginator.ReadGroupSets.create(genomics); List ids = Lists.newArrayList(); for (ReadGroupSet set : paginator.search( - new SearchReadGroupSetsRequest().setName("HG"), "readGroupSets(id,name)")) { + new SearchReadGroupSetsRequest().setName("HG"), "nextPageToken,readGroupSets(id,name)")) { ids.add(set.getId()); } assertEquals(Lists.newArrayList("r1"), ids); // Make sure the fields parameter actually gets passed along - Mockito.verify(readGroupSetSearch, Mockito.atLeastOnce()).setFields("readGroupSets(id,name)"); + Mockito.verify(readGroupSetSearch, Mockito.atLeastOnce()).setFields("nextPageToken,readGroupSets(id,name)"); } + @Test + public void testFieldsMissingNextPageToken() throws Exception { + Mockito.when(readGroupSets.search(new SearchReadGroupSetsRequest().setName("HG"))) + .thenReturn(readGroupSetSearch); + + Paginator.ReadGroupSets paginator = Paginator.ReadGroupSets.create(genomics); + thrown.expect(IllegalArgumentException.class); + paginator.search(new SearchReadGroupSetsRequest().setName("HG"), "readGroupSets(id,name)").iterator().next(); + } + @Test public void testVariantPagination() throws Exception { @@ -157,6 +173,17 @@ public void testVariantPagination() throws Exception { assertThat(filteredVariants, CoreMatchers.hasItems(atStartWithinExtent, atStartOverlapExtent, beyondStartWithinExtent, beyondOverlapExtent)); + // Ensure searches with fields verify the preconditions for strict shards. + final String nullFields = null; + assertNotNull(filteredPaginator.search(request, nullFields).iterator().next()); + assertNotNull(filteredPaginator.search(request, "nextPageToken,variants(start,id,calls(genotype,callSetName))").iterator().next()); + assertNotNull(filteredPaginator.search(request, "id,nextPageToken,variants(start,id,calls(genotype,callSetName))").iterator().next()); + assertNotNull(filteredPaginator.search(request, "variants(start,id,calls(genotype,callSetName)),nextPageToken").iterator().next()); + try { + filteredPaginator.search(request, "nextPageToken,variants(id,calls(genotype,callSetName))").iterator().next(); + fail("should have thrown an IllegalArgumentxception"); + } catch (IllegalArgumentException e) {} + Paginator.Variants overlappingPaginator = Paginator.Variants.create(genomics, ShardBoundary.OVERLAPS); List overlappingVariants = Lists.newArrayList(); for (Variant variant : overlappingPaginator.search(request)) { @@ -164,8 +191,15 @@ public void testVariantPagination() throws Exception { } assertEquals(6, overlappingVariants.size()); assertThat(overlappingVariants, CoreMatchers.hasItems(input)); + + // Ensure searches with fields verify the preconditions for overlapping shards. + assertNotNull(overlappingPaginator.search(request, nullFields).iterator().next()); + assertNotNull(overlappingPaginator.search(request, "nextPageToken,variants(start,id,calls(genotype,callSetName))").iterator().next()); + assertNotNull(overlappingPaginator.search(request, "nextPageToken,variants(id,calls(genotype,callSetName))").iterator().next()); + assertNotNull(overlappingPaginator.search(request, "id,nextPageToken,variants(start,id,calls(genotype,callSetName))").iterator().next()); + assertNotNull(overlappingPaginator.search(request, "variants(id,calls(genotype,callSetName)),nextPageToken").iterator().next()); } - + @Test public void testVariantPaginationEmptyShard() throws Exception { @@ -178,8 +212,8 @@ public void testVariantPaginationEmptyShard() throws Exception { Paginator.Variants filteredPaginator = Paginator.Variants.create(genomics, ShardBoundary.STRICT); assertNotNull(filteredPaginator.search(request)); } - - Read readHelper(int start, int end) { + + static Read readHelper(int start, int end) { Position position = new Position().setPosition((long) start); LinearAlignment alignment = new LinearAlignment().setPosition(position); return new Read().setAlignment(alignment).setFragmentLength(end-start); @@ -221,4 +255,34 @@ public void testReadPagination() throws Exception { assertThat(overlappingReads, CoreMatchers.hasItems(input)); } + @Test + public void testReadPaginationStrictShardPrecondition() throws Exception { + SearchReadsRequest request = new SearchReadsRequest().setStart(1000L).setEnd(2000L); + Mockito.when(reads.search(request)).thenReturn(readsSearch); + + Paginator.Reads filteredPaginator = Paginator.Reads.create(genomics, ShardBoundary.STRICT); + thrown.expect(IllegalArgumentException.class); + filteredPaginator.search(request, "nextPageToken,reads(id,alignment(cigar))").iterator().next(); + } + + @Test + public void testStrictReadPaginationNextPageTokenPrecondition() throws Exception { + SearchReadsRequest request = new SearchReadsRequest().setStart(1000L).setEnd(2000L); + Mockito.when(reads.search(request)).thenReturn(readsSearch); + + Paginator.Reads filteredPaginator = Paginator.Reads.create(genomics, ShardBoundary.STRICT); + thrown.expect(IllegalArgumentException.class); + filteredPaginator.search(request, "reads(id,alignment(cigar,position))").iterator().next(); + } + + @Test + public void testOverlappingReadPaginationNextPageTokenPrecondition() throws Exception { + SearchReadsRequest request = new SearchReadsRequest().setStart(1000L).setEnd(2000L); + Mockito.when(reads.search(request)).thenReturn(readsSearch); + + Paginator.Reads overlappingPaginator = Paginator.Reads.create(genomics, ShardBoundary.OVERLAPS); + thrown.expect(IllegalArgumentException.class); + overlappingPaginator.search(request, "reads(id,alignment(cigar,position))").iterator().next(); + } + }