feat(lapis2): stream data from SILO

Refs: #744
GenSpectrum · Apr 22, 2024 · d8ade36 · d8ade36
1 parent f6a6184
commit d8ade36
Show file tree

Hide file tree

Showing 34 changed files with 563 additions and 423 deletions.
diff --git a/lapis2/src/main/kotlin/org/genspectrum/lapis/controller/CompressionFilter.kt b/lapis2/src/main/kotlin/org/genspectrum/lapis/controller/CompressionFilter.kt
@@ -25,6 +25,7 @@ import org.springframework.http.converter.json.MappingJackson2HttpMessageConvert
 import org.springframework.stereotype.Component
 import org.springframework.web.context.annotation.RequestScope
 import org.springframework.web.filter.OncePerRequestFilter
+import java.io.IOException
 import java.io.OutputStream
 import java.nio.charset.Charset
 import java.util.Enumeration
@@ -154,8 +155,12 @@ class CompressionFilter(val objectMapper: ObjectMapper, val requestCompression:
             maybeCompressingResponse,
         )
 
-        maybeCompressingResponse.outputStream.flush()
-        maybeCompressingResponse.outputStream.close()
+        try {
+            maybeCompressingResponse.outputStream.flush()
+            maybeCompressingResponse.outputStream.close()
+        } catch (e: IOException) {
+            log.debug { "Failed to flush and close the compressing output stream: ${e.message}" }
+        }
     }
 
     private fun getValidatedCompressionProperty(reReadableRequest: CachedBodyHttpServletRequest): Compression? {

diff --git a/lapis2/src/main/kotlin/org/genspectrum/lapis/controller/ControllerDescriptions.kt b/lapis2/src/main/kotlin/org/genspectrum/lapis/controller/ControllerDescriptions.kt
@@ -1,6 +1,6 @@
 package org.genspectrum.lapis.controller
 
-import org.genspectrum.lapis.controller.LapisMediaType.TEXT_CSV_WITHOUT_HEADERS
+import org.genspectrum.lapis.controller.LapisMediaType.TEXT_CSV_WITHOUT_HEADERS_VALUE
 
 const val DETAILS_ENDPOINT_DESCRIPTION = """Returns the specified metadata fields of sequences matching the filter."""
 const val AGGREGATED_ENDPOINT_DESCRIPTION =
@@ -53,7 +53,7 @@ const val OFFSET_DESCRIPTION =
     This is useful for pagination in combination with \"limit\"."""
 const val FORMAT_DESCRIPTION = """The data format of the response.
 Alternatively, the data format can be specified by setting the \"Accept\"-header.
-You can include the parameter to return the CSV/TSV without headers: "$TEXT_CSV_WITHOUT_HEADERS". 
+You can include the parameter to return the CSV/TSV without headers: "$TEXT_CSV_WITHOUT_HEADERS_VALUE". 
 When both are specified, the request parameter takes precedence over the header."""
 
 private const val MAYBE_DESCRIPTION = """

diff --git a/lapis2/src/main/kotlin/org/genspectrum/lapis/controller/CsvWriter.kt b/lapis2/src/main/kotlin/org/genspectrum/lapis/controller/CsvWriter.kt
@@ -4,42 +4,42 @@ import com.fasterxml.jackson.annotation.JsonIgnore
 import org.apache.commons.csv.CSVFormat
 import org.apache.commons.csv.CSVPrinter
 import org.springframework.stereotype.Component
-import java.io.StringWriter
+import java.util.stream.Stream
 
 interface CsvRecord {
     @JsonIgnore
     fun getValuesList(): List<String?>
 
     @JsonIgnore
-    fun getHeader(): Array<String>
+    fun getHeader(): Iterable<String>
 }
 
 @Component
 class CsvWriter {
     fun write(
-        headers: Array<String>?,
-        data: List<CsvRecord>,
+        appendable: Appendable,
+        includeHeaders: Boolean,
+        data: Stream<out CsvRecord>,
         delimiter: Delimiter,
-    ): String {
-        val stringWriter = StringWriter()
+    ) {
+        var shouldWriteHeaders = includeHeaders
+
         CSVPrinter(
-            stringWriter,
+            appendable,
             CSVFormat.DEFAULT.builder()
                 .setRecordSeparator("\n")
                 .setDelimiter(delimiter.value)
                 .setNullString("")
-                .also {
-                    when {
-                        headers != null -> it.setHeader(*headers)
-                    }
-                }
                 .build(),
         ).use {
             for (datum in data) {
+                if (shouldWriteHeaders) {
+                    it.printRecord(datum.getHeader())
+                    shouldWriteHeaders = false
+                }
                 it.printRecord(datum.getValuesList())
             }
         }
-        return stringWriter.toString().trimEnd('\n')
     }
 }
 

diff --git a/lapis2/src/main/kotlin/org/genspectrum/lapis/controller/DataFormatParameterFilter.kt b/lapis2/src/main/kotlin/org/genspectrum/lapis/controller/DataFormatParameterFilter.kt
@@ -7,10 +7,8 @@ import jakarta.servlet.http.HttpServletResponse
 import mu.KotlinLogging
 import org.genspectrum.lapis.util.CachedBodyHttpServletRequest
 import org.genspectrum.lapis.util.HeaderModifyingRequestWrapper
-import org.genspectrum.lapis.util.ResponseWithContentType
 import org.springframework.core.annotation.Order
 import org.springframework.http.HttpHeaders.ACCEPT
-import org.springframework.http.InvalidMediaTypeException
 import org.springframework.http.MediaType
 import org.springframework.stereotype.Component
 import org.springframework.web.filter.OncePerRequestFilter
@@ -44,37 +42,16 @@ class DataFormatParameterFilter(val objectMapper: ObjectMapper) : OncePerRequest
 
         filterChain.doFilter(
             requestWithModifiedAcceptHeader,
-            when (isCsvOrTsvWithoutHeaders(requestWithModifiedAcceptHeader)) {
-                true -> ResponseWithContentType(response, MediaType.TEXT_PLAIN_VALUE)
-                false -> response
-            },
+            response,
         )
     }
 
     private fun findAcceptHeaderOverwriteValue(reReadableRequest: CachedBodyHttpServletRequest) =
         when (reReadableRequest.getStringField(FORMAT_PROPERTY)?.uppercase()) {
-            DataFormat.CSV -> LapisMediaType.TEXT_CSV
-            DataFormat.CSV_WITHOUT_HEADERS -> LapisMediaType.TEXT_CSV_WITHOUT_HEADERS
-            DataFormat.TSV -> LapisMediaType.TEXT_TSV
+            DataFormat.CSV -> LapisMediaType.TEXT_CSV_VALUE
+            DataFormat.CSV_WITHOUT_HEADERS -> LapisMediaType.TEXT_CSV_WITHOUT_HEADERS_VALUE
+            DataFormat.TSV -> LapisMediaType.TEXT_TSV_VALUE
             DataFormat.JSON -> MediaType.APPLICATION_JSON_VALUE
             else -> null
         }
-
-    private fun isCsvOrTsvWithoutHeaders(requestWithModifiedAcceptHeader: HeaderModifyingRequestWrapper): Boolean {
-        val acceptHeader = requestWithModifiedAcceptHeader.getHeader(ACCEPT) ?: return false
-
-        val acceptMediaType = try {
-            MediaType.parseMediaType(acceptHeader)
-        } catch (e: InvalidMediaTypeException) {
-            log.info { "failed to parse accept header: " + e.message }
-            return false
-        }
-
-        if (acceptMediaType.parameters[HEADERS_ACCEPT_HEADER_PARAMETER] == "false") {
-            log.info { "Setting response content type to plain text due to 'headers=false'" }
-            return true
-        }
-
-        return false
-    }
 }
diff --git a/lapis2/src/main/kotlin/org/genspectrum/lapis/controller/DownloadAsFileFilter.kt b/lapis2/src/main/kotlin/org/genspectrum/lapis/controller/DownloadAsFileFilter.kt
@@ -4,8 +4,8 @@ import com.fasterxml.jackson.databind.ObjectMapper
 import jakarta.servlet.FilterChain
 import jakarta.servlet.http.HttpServletRequest
 import jakarta.servlet.http.HttpServletResponse
-import org.genspectrum.lapis.controller.LapisMediaType.TEXT_CSV
-import org.genspectrum.lapis.controller.LapisMediaType.TEXT_TSV
+import org.genspectrum.lapis.controller.LapisMediaType.TEXT_CSV_VALUE
+import org.genspectrum.lapis.controller.LapisMediaType.TEXT_TSV_VALUE
 import org.genspectrum.lapis.util.CachedBodyHttpServletRequest
 import org.springframework.core.annotation.Order
 import org.springframework.http.HttpHeaders.ACCEPT
@@ -45,8 +45,8 @@ class DownloadAsFileFilter(
         }
 
         val fileEnding = when (request.getHeader(ACCEPT)) {
-            TEXT_CSV -> "csv"
-            TEXT_TSV -> "tsv"
+            TEXT_CSV_VALUE -> "csv"
+            TEXT_TSV_VALUE -> "tsv"
             else -> when (matchingRoute?.servesFasta) {
                 true -> "fasta"
                 else -> "json"

diff --git a/lapis2/src/main/kotlin/org/genspectrum/lapis/controller/Headers.kt b/lapis2/src/main/kotlin/org/genspectrum/lapis/controller/Headers.kt
@@ -1,14 +1,17 @@
 package org.genspectrum.lapis.controller
 
+import org.springframework.http.MediaType
+
 object LapisHeaders {
     const val REQUEST_ID = "X-Request-ID"
     const val LAPIS_DATA_VERSION = "Lapis-Data-Version"
 }
 
 object LapisMediaType {
-    const val TEXT_X_FASTA = "text/x-fasta"
-    const val TEXT_CSV = "text/csv"
-    const val TEXT_CSV_WITHOUT_HEADERS = "text/csv;$HEADERS_ACCEPT_HEADER_PARAMETER=false"
-    const val TEXT_TSV = "text/tab-separated-values"
-    const val APPLICATION_YAML = "application/yaml"
+    const val TEXT_X_FASTA_VALUE = "text/x-fasta"
+    val TEXT_X_FASTA = MediaType.parseMediaType(TEXT_X_FASTA_VALUE)
+    const val TEXT_CSV_VALUE = "text/csv"
+    const val TEXT_CSV_WITHOUT_HEADERS_VALUE = "text/csv;$HEADERS_ACCEPT_HEADER_PARAMETER=false"
+    const val TEXT_TSV_VALUE = "text/tab-separated-values"
+    const val APPLICATION_YAML_VALUE = "application/yaml"
 }
diff --git a/lapis2/src/main/kotlin/org/genspectrum/lapis/controller/InfoController.kt b/lapis2/src/main/kotlin/org/genspectrum/lapis/controller/InfoController.kt
@@ -3,7 +3,7 @@ package org.genspectrum.lapis.controller
 import io.swagger.v3.oas.annotations.Operation
 import org.genspectrum.lapis.config.DatabaseConfig
 import org.genspectrum.lapis.config.ReferenceGenome
-import org.genspectrum.lapis.controller.LapisMediaType.APPLICATION_YAML
+import org.genspectrum.lapis.controller.LapisMediaType.APPLICATION_YAML_VALUE
 import org.genspectrum.lapis.model.SiloQueryModel
 import org.genspectrum.lapis.request.LapisInfo
 import org.springframework.http.MediaType
@@ -29,7 +29,7 @@ class InfoController(
         return LapisInfo(siloInfo.dataVersion)
     }
 
-    @GetMapping(DATABASE_CONFIG_ROUTE, produces = [MediaType.APPLICATION_JSON_VALUE, APPLICATION_YAML])
+    @GetMapping(DATABASE_CONFIG_ROUTE, produces = [MediaType.APPLICATION_JSON_VALUE, APPLICATION_YAML_VALUE])
     @Operation(description = DATABASE_CONFIG_ENDPOINT_DESCRIPTION)
     fun getDatabaseConfigAsJson(): DatabaseConfig = databaseConfig