feat: implement leftover variant query parts #498

lapis2/src/main/antlr/org/genspectrum/lapis/model/variantqueryparser/VariantQuery.g4

@@ -45,7 +45,8 @@ nOfMatchExactly: 'EXACTLY-' | 'exactly-';
 nOfNumberOfMatchers: NUMBER+;
 nOfExprs: expr (',' expr)*;
 
-nucleotideInsertionQuery: insertionKeyword position ':' (possibleAmbiguousNucleotideSymbol | '?')+;
+nucleotideInsertionQuery: insertionKeyword position ':' nucleotideInsertionSymbol+;
+nucleotideInsertionSymbol: possibleAmbiguousNucleotideSymbol | '?';
 insertionKeyword: 'ins_' | 'INS_';
 
 aaMutationQuery: gene ':' aaSymbol? position possiblyAmbiguousAaSymbol?;
@@ -55,7 +56,8 @@ possiblyAmbiguousAaSymbol: aaSymbol | ambiguousAaSymbol;
 gene: covidGene;
 covidGene : E | M | N | S | ORF;
 
-aaInsertionQuery: insertionKeyword gene ':' position ':' (possiblyAmbiguousAaSymbol | '?')+;
+aaInsertionQuery: insertionKeyword gene ':' position ':' aaInsertionSymbol+;
+aaInsertionSymbol: possiblyAmbiguousAaSymbol | '?';
 
 nextcladePangolineageQuery: nextcladePangoLineagePrefix pangolineageQuery;
 nextcladePangoLineagePrefix: 'nextcladePangoLineage:' | 'NEXTCLADEPANGOLINEAGE:';

lapis2/src/main/kotlin/org/genspectrum/lapis/model/CovidSpecifics.kt

@@ -0,0 +1,12 @@
+package org.genspectrum.lapis.model
+
+/**
+ * Those are special values used for the COVID instance that supports advanced "variant queries".
+ */
+
+const val PANGO_LINEAGE_COLUMN = "pangoLineage"
+const val NEXTCLADE_PANGO_LINEAGE_COLUMN = "nextcladePangoLineage"
+const val NEXTSTRAIN_CLADE_COLUMN = "nextstrainClade"
+const val GISAID_CLADE_COLUMN = "gisaidClade"
+
+const val NEXTSTRAIN_CLADE_RECOMBINANT = "RECOMBINANT"

lapis2/src/main/kotlin/org/genspectrum/lapis/model/VariantQueryCustomListener.kt

@@ -1,10 +1,10 @@
 package org.genspectrum.lapis.model
 
 import VariantQueryBaseListener
+import VariantQueryParser
 import VariantQueryParser.AaInsertionQueryContext
 import VariantQueryParser.AaMutationQueryContext
 import VariantQueryParser.AndContext
-import VariantQueryParser.GisaidCladeLineageQueryContext
 import VariantQueryParser.MaybeContext
 import VariantQueryParser.NOfQueryContext
 import VariantQueryParser.NextcladePangolineageQueryContext
@@ -14,18 +14,24 @@ import VariantQueryParser.NucleotideInsertionQueryContext
 import VariantQueryParser.NucleotideMutationQueryContext
 import VariantQueryParser.OrContext
 import VariantQueryParser.PangolineageQueryContext
+import org.antlr.v4.runtime.RuleContext
 import org.antlr.v4.runtime.tree.ParseTreeListener
+import org.genspectrum.lapis.request.LAPIS_INSERTION_AMBIGUITY_SYMBOL
+import org.genspectrum.lapis.request.SILO_INSERTION_AMBIGUITY_SYMBOL
+import org.genspectrum.lapis.silo.AminoAcidInsertionContains
 import org.genspectrum.lapis.silo.AminoAcidSymbolEquals
 import org.genspectrum.lapis.silo.And
 import org.genspectrum.lapis.silo.HasAminoAcidMutation
 import org.genspectrum.lapis.silo.HasNucleotideMutation
 import org.genspectrum.lapis.silo.Maybe
 import org.genspectrum.lapis.silo.NOf
 import org.genspectrum.lapis.silo.Not
+import org.genspectrum.lapis.silo.NucleotideInsertionContains
 import org.genspectrum.lapis.silo.NucleotideSymbolEquals
 import org.genspectrum.lapis.silo.Or
 import org.genspectrum.lapis.silo.PangoLineageEquals
 import org.genspectrum.lapis.silo.SiloFilterExpression
+import org.genspectrum.lapis.silo.StringEquals
 
 class VariantQueryCustomListener : VariantQueryBaseListener(), ParseTreeListener {
     private val expressionStack = ArrayDeque<SiloFilterExpression>()
@@ -48,15 +54,8 @@ class VariantQueryCustomListener : VariantQueryBaseListener(), ParseTreeListener
         expressionStack.addLast(expression)
     }
 
-    override fun enterPangolineageQuery(ctx: PangolineageQueryContext?) {
-        if (ctx == null) {
-            return
-        }
-        val pangolineage = ctx.pangolineage().text
-        val includeSublineages = ctx.pangolineageIncludeSublineages() != null
-
-        val expr = PangoLineageEquals("pango_lineage", pangolineage, includeSublineages)
-        expressionStack.addLast(expr)
+    override fun enterPangolineageQuery(ctx: PangolineageQueryContext) {
+        addPangoLineage(ctx, PANGO_LINEAGE_COLUMN)
     }
 
     override fun exitAnd(ctx: AndContext?) {
@@ -95,8 +94,14 @@ class VariantQueryCustomListener : VariantQueryBaseListener(), ParseTreeListener
         expressionStack.addLast(NOf(n, matchExactly, children.reversed()))
     }
 
-    override fun enterNucleotideInsertionQuery(ctx: NucleotideInsertionQueryContext?) {
-        throw SiloNotImplementedError("Nucleotide insertions are not supported yet.", NotImplementedError())
+    override fun enterNucleotideInsertionQuery(ctx: NucleotideInsertionQueryContext) {
+        val value = ctx.nucleotideInsertionSymbol().joinToString("", transform = ::mapInsertionSymbol)
+        expressionStack.addLast(
+            NucleotideInsertionContains(
+                ctx.position().text.toInt(),
+                value,
+            ),
+        )
     }
 
     override fun enterAaMutationQuery(ctx: AaMutationQueryContext?) {
@@ -113,21 +118,49 @@ class VariantQueryCustomListener : VariantQueryBaseListener(), ParseTreeListener
         expressionStack.addLast(expression)
     }
 
-    override fun enterAaInsertionQuery(ctx: AaInsertionQueryContext?) {
-        throw SiloNotImplementedError("Amino acid insertions are not supported yet.", NotImplementedError())
+    override fun enterAaInsertionQuery(ctx: AaInsertionQueryContext) {
+        val value = ctx.aaInsertionSymbol().joinToString("", transform = ::mapInsertionSymbol)
+        expressionStack.addLast(
+            AminoAcidInsertionContains(
+                ctx.position().text.toInt(),
+                value,
+                ctx.gene().text,
+            ),
+        )
+    }
+
+    override fun enterNextcladePangolineageQuery(ctx: NextcladePangolineageQueryContext) {
+        addPangoLineage(ctx.pangolineageQuery(), NEXTCLADE_PANGO_LINEAGE_COLUMN)
     }
 
-    override fun enterNextcladePangolineageQuery(ctx: NextcladePangolineageQueryContext?) {
-        throw SiloNotImplementedError("Nextclade pango lineages are not supported yet.", NotImplementedError())
+    override fun enterNextstrainCladeQuery(ctx: NextstrainCladeQueryContext) {
+        val value = when (ctx.text) {
+            NEXTSTRAIN_CLADE_RECOMBINANT -> ctx.text.lowercase()
+            else -> ctx.text
+        }
+        expressionStack.addLast(StringEquals(NEXTSTRAIN_CLADE_COLUMN, value))
     }
 
-    override fun enterNextstrainCladeQuery(ctx: NextstrainCladeQueryContext?) {
-        throw SiloNotImplementedError("Nextstrain clade lineages are not supported yet.", NotImplementedError())
+    override fun enterGisaidCladeNomenclature(ctx: VariantQueryParser.GisaidCladeNomenclatureContext) {
+        expressionStack.addLast(StringEquals(GISAID_CLADE_COLUMN, ctx.text))
     }
 
-    override fun enterGisaidCladeLineageQuery(ctx: GisaidCladeLineageQueryContext?) {
-        throw SiloNotImplementedError("Gisaid clade lineages are not supported yet.", NotImplementedError())
+    private fun addPangoLineage(
+        ctx: PangolineageQueryContext,
+        pangoLineageColumnName: String,
+    ) {
+        val pangolineage = ctx.pangolineage().text
+        val includeSublineages = ctx.pangolineageIncludeSublineages() != null
+
+        val expr = PangoLineageEquals(pangoLineageColumnName, pangolineage, includeSublineages)
+        expressionStack.addLast(expr)
     }
 }
 
+fun mapInsertionSymbol(ctx: RuleContext): String =
+    when (ctx.text) {
+        LAPIS_INSERTION_AMBIGUITY_SYMBOL -> SILO_INSERTION_AMBIGUITY_SYMBOL
+        else -> ctx.text
+    }
+
 class SiloNotImplementedError(message: String?, cause: Throwable?) : Exception(message, cause)

lapis2/src/main/kotlin/org/genspectrum/lapis/request/AminoAcidInsertion.kt

@@ -26,7 +26,10 @@ data class AminoAcidInsertion(val position: Int, val gene: String, val insertion
                     "Invalid amino acid insertion: $aminoAcidInsertion: Did not find gene",
                 )
 
-            val insertions = matchGroups["insertions"]?.value?.replace("?", ".*")
+            val insertions = matchGroups["insertions"]?.value?.replace(
+                LAPIS_INSERTION_AMBIGUITY_SYMBOL,
+                SILO_INSERTION_AMBIGUITY_SYMBOL,
+            )
                 ?: throw BadRequestException(
                     "Invalid amino acid insertion: $aminoAcidInsertion: Did not find insertions",
                 )

lapis2/src/main/kotlin/org/genspectrum/lapis/request/NucleotideInsertion.kt

@@ -8,6 +8,9 @@ import org.springframework.boot.jackson.JsonComponent
 import org.springframework.core.convert.converter.Converter
 import org.springframework.stereotype.Component
 
+const val LAPIS_INSERTION_AMBIGUITY_SYMBOL = "?"
+const val SILO_INSERTION_AMBIGUITY_SYMBOL = ".*"
+
 data class NucleotideInsertion(val position: Int, val insertions: String, val segment: String?) {
     companion object {
         fun fromString(nucleotideInsertion: String): NucleotideInsertion {
@@ -21,7 +24,10 @@ data class NucleotideInsertion(val position: Int, val insertions: String, val se
                     "Invalid nucleotide insertion: $nucleotideInsertion: Did not find position",
                 )
 
-            val insertions = matchGroups["insertions"]?.value?.replace("?", ".*")
+            val insertions = matchGroups["insertions"]?.value?.replace(
+                LAPIS_INSERTION_AMBIGUITY_SYMBOL,
+                SILO_INSERTION_AMBIGUITY_SYMBOL,
+            )
                 ?: throw BadRequestException(
                     "Invalid nucleotide insertion: $nucleotideInsertion: Did not find insertions",
                 )

lapis2/src/test/kotlin/org/genspectrum/lapis/model/VariantQueryFacadeTest.kt

@@ -1,20 +1,22 @@
 package org.genspectrum.lapis.model
 
+import org.genspectrum.lapis.silo.AminoAcidInsertionContains
 import org.genspectrum.lapis.silo.AminoAcidSymbolEquals
 import org.genspectrum.lapis.silo.And
 import org.genspectrum.lapis.silo.HasAminoAcidMutation
 import org.genspectrum.lapis.silo.HasNucleotideMutation
 import org.genspectrum.lapis.silo.Maybe
 import org.genspectrum.lapis.silo.NOf
 import org.genspectrum.lapis.silo.Not
+import org.genspectrum.lapis.silo.NucleotideInsertionContains
 import org.genspectrum.lapis.silo.NucleotideSymbolEquals
 import org.genspectrum.lapis.silo.Or
 import org.genspectrum.lapis.silo.PangoLineageEquals
+import org.genspectrum.lapis.silo.StringEquals
 import org.hamcrest.MatcherAssert.assertThat
 import org.hamcrest.Matchers.equalTo
 import org.junit.jupiter.api.BeforeEach
 import org.junit.jupiter.api.Test
-import org.junit.jupiter.api.assertThrows
 
 class VariantQueryFacadeTest {
     private lateinit var underTest: VariantQueryFacade
@@ -74,7 +76,7 @@ class VariantQueryFacadeTest {
                             ),
                         ),
                     ),
-                    PangoLineageEquals("pango_lineage", "A.1.2.3", true),
+                    PangoLineageEquals(PANGO_LINEAGE_COLUMN, "A.1.2.3", true),
                 ),
             )
 
@@ -196,7 +198,7 @@ class VariantQueryFacadeTest {
 
         val result = underTest.map(variantQuery)
 
-        val expectedResult = PangoLineageEquals("pango_lineage", "A.1.2.3", false)
+        val expectedResult = PangoLineageEquals(PANGO_LINEAGE_COLUMN, "A.1.2.3", false)
         assertThat(result, equalTo(expectedResult))
     }
 
@@ -207,7 +209,18 @@ class VariantQueryFacadeTest {
 
         val result = underTest.map(variantQuery)
 
-        val expectedResult = PangoLineageEquals("pango_lineage", "A.1.2.3", true)
+        val expectedResult = PangoLineageEquals(PANGO_LINEAGE_COLUMN, "A.1.2.3", true)
+        assertThat(result, equalTo(expectedResult))
+    }
+
+    @Test
+    @Suppress("ktlint:standard:max-line-length")
+    fun `given a variantQuery with a 'NextcladePangolineage' expression then map should return the corresponding SiloQuery`() {
+        val variantQuery = "nextcladePangoLineage:A.1.2.3*"
+
+        val result = underTest.map(variantQuery)
+
+        val expectedResult = PangoLineageEquals(NEXTCLADE_PANGO_LINEAGE_COLUMN, "A.1.2.3", true)
         assertThat(result, equalTo(expectedResult))
     }
 
@@ -248,19 +261,25 @@ class VariantQueryFacadeTest {
     }
 
     @Test
-    fun `given a variantQuery with a 'Insertion' expression then map should throw an error`() {
+    fun `given a variantQuery with a 'Insertion' expression then returns SILO query`() {
         val variantQuery = "ins_1234:GAG"
 
-        val exception = assertThrows<SiloNotImplementedError> { underTest.map(variantQuery) }
+        val result = underTest.map(variantQuery)
+
+        assertThat(result, equalTo(NucleotideInsertionContains(1234, "GAG")))
+    }
 
-        assertThat(
-            exception.message,
-            equalTo("Nucleotide insertions are not supported yet."),
-        )
+    @Test
+    fun `given a variantQuery with a 'Insertion' with wildcard expression then returns SILO query`() {
+        val variantQuery = "ins_1234:G?A?G"
+
+        val result = underTest.map(variantQuery)
+
+        assertThat(result, equalTo(NucleotideInsertionContains(1234, "G.*A.*G")))
     }
 
     @Test
-    fun `given amino acidAA mutation expression then should map to AminoAcidSymbolEquals`() {
+    fun `given amino acid mutation expression then should map to AminoAcidSymbolEquals`() {
         val variantQuery = "S:N501Y"
 
         val result = underTest.map(variantQuery)
@@ -296,50 +315,56 @@ class VariantQueryFacadeTest {
     }
 
     @Test
-    fun `given a valid variantQuery with a 'AA insertion' expression then map should throw an error`() {
+    fun `given a valid variantQuery with a 'AA insertion' expression then returns SILO query`() {
         val variantQuery = "ins_S:501:EPE"
 
-        val exception = assertThrows<SiloNotImplementedError> { underTest.map(variantQuery) }
+        val result = underTest.map(variantQuery)
 
-        assertThat(
-            exception.message,
-            equalTo("Amino acid insertions are not supported yet."),
-        )
+        assertThat(result, equalTo(AminoAcidInsertionContains(501, "EPE", "S")))
     }
 
     @Test
-    fun `given a valid variantQuery with a 'nextclade pango lineage' expression then map should throw an error`() {
-        val variantQuery = "nextcladePangoLineage:BA.5*"
+    fun `given a valid variantQuery with a 'AA insertion' with wildcard then returns SILO query`() {
+        val variantQuery = "ins_S:501:E?E?"
 
-        val exception = assertThrows<SiloNotImplementedError> { underTest.map(variantQuery) }
+        val result = underTest.map(variantQuery)
 
-        assertThat(
-            exception.message,
-            equalTo("Nextclade pango lineages are not supported yet."),
-        )
+        assertThat(result, equalTo(AminoAcidInsertionContains(501, "E.*E.*", "S")))
     }
 
     @Test
-    fun `given a valid variantQuery with a 'Nextstrain clade lineage' expression then map should throw an error`() {
+    fun `given a valid variantQuery with a 'NextstrainCladeLineage' expression then returns SILO query`() {
         val variantQuery = "nextstrainClade:22B"
 
-        val exception = assertThrows<SiloNotImplementedError> { underTest.map(variantQuery) }
+        val result = underTest.map(variantQuery)
 
-        assertThat(
-            exception.message,
-            equalTo("Nextstrain clade lineages are not supported yet."),
-        )
+        assertThat(result, equalTo(StringEquals(NEXTSTRAIN_CLADE_COLUMN, "22B")))
+    }
+
+    @Test
+    fun `given a valid variantQuery with a 'NextstrainCladeLineage' recombinant expression then returns SILO query`() {
+        val variantQuery = "nextstrainClade:RECOMBINANT"
+
+        val result = underTest.map(variantQuery)
+
+        assertThat(result, equalTo(StringEquals(NEXTSTRAIN_CLADE_COLUMN, "recombinant")))
     }
 
     @Test
-    fun `given a valid variantQuery with a 'Gisaid clade lineage' expression then map should throw an error`() {
+    fun `given a valid variantQuery with a single letter 'GisaidCladeLineage' expression then returns SILO query`() {
+        val variantQuery = "gisaid:X"
+
+        val result = underTest.map(variantQuery)
+
+        assertThat(result, equalTo(StringEquals(GISAID_CLADE_COLUMN, "X")))
+    }
+
+    @Test
+    fun `given a valid variantQuery with a 'GisaidCladeLineage' expression then returns SILO query`() {
         val variantQuery = "gisaid:AB"
 
-        val exception = assertThrows<SiloNotImplementedError> { underTest.map(variantQuery) }
+        val result = underTest.map(variantQuery)
 
-        assertThat(
-            exception.message,
-            equalTo("Gisaid clade lineages are not supported yet."),
-        )
+        assertThat(result, equalTo(StringEquals(GISAID_CLADE_COLUMN, "AB")))
     }
 }