Include additional keys in workflow query response (#2971)

* Add additionalKeys param to query workflows

Add an additionalKeys parameter that allows users
to specify additional keys that should be included
in the query result json for each workflow.

* Include additionalKeys in query response

h/t to @kshakir and @ruchim for restructuring
the code to handle futures.

* Query for labels in additionalKeys

* Test changes to query endpoint

Add tests to check that the key
specified in the additionalKeys
parameter for the query endpoint is
returned in the workflow query metadata.

* Fix function ordering

* Validate additionalKeys value

Throw a validation error if an
unexpected key is included in the
additionalKeys parameter for the
query endpoint. Currently only
"labels" and "parentWorkflowId" are
supported.

* Rename additionalKeys to additionalQueryResultFields

* Code refactoring/clean-up

* Fix renaming additionalKeys parameter

Author: samanehsan

engine/src/main/resources/swagger/cromwell.yaml

@@ -356,6 +356,15 @@ paths:
             returns workflows with all of the specified label keys. Specify the label key
             and label value pair as separated with
             "label-key:label-value"
+        - name: additionalQueryResultFields
+          required: false
+          in: query
+          type: array
+          items:
+            type: string
+          collectionFormat: multi
+          description: >
+            Includes the specified keys in the metadata for the returned workflows.
       tags:
         - Workflows
       responses:

engine/src/main/scala/cromwell/webservice/WorkflowJsonSupport.scala

@@ -50,6 +50,6 @@ object WorkflowJsonSupport extends DefaultJsonProtocol {
     }
   }
 
-  implicit val workflowQueryResult = jsonFormat5(WorkflowQueryResult)
+  implicit val workflowQueryResult = jsonFormat7(WorkflowQueryResult)
   implicit val workflowQueryResponse = jsonFormat1(WorkflowQueryResponse)
 }

engine/src/test/scala/cromwell/webservice/CromwellApiServiceSpec.scala

@@ -20,6 +20,7 @@ import cromwell.services.healthmonitor.HealthMonitorServiceActor.{GetCurrentStat
 import cromwell.services.metadata.MetadataService._
 import cromwell.services.metadata._
 import cromwell.util.SampleWdl.HelloWorld
+import mouse.boolean._
 import org.scalatest.{AsyncFlatSpec, Matchers}
 import spray.json.DefaultJsonProtocol._
 import spray.json._
@@ -488,6 +489,33 @@ class CromwellApiServiceSpec extends AsyncFlatSpec with ScalatestRouteTest with
         }
     }
 
+    it should "return labels if specified in additionalQueryResultFields param" in {
+      Get(s"/workflows/$version/query?additionalQueryResultFields=labels&id=${CromwellApiServiceSpec.ExistingWorkflowId}") ~>
+        akkaHttpService.workflowRoutes ~>
+        check {
+          status should be(StatusCodes.OK)
+          contentType should be(ContentTypes.`application/json`)
+          val results = responseAs[JsObject].fields("results").convertTo[Seq[JsObject]]
+          val fields = results.head.fields
+          fields("id") should be(JsString(CromwellApiServiceSpec.ExistingWorkflowId.toString))
+          fields(WorkflowMetadataKeys.Labels).asJsObject.fields("key1") should be(JsString("label1"))
+          fields(WorkflowMetadataKeys.Labels).asJsObject.fields("key2") should be(JsString("label2"))
+        }
+    }
+
+    it should "return parentWorkflowId if specified in additionalQueryResultFields param" in {
+      Get(s"/workflows/$version/query?additionalQueryResultFields=parentWorkflowId&id=${CromwellApiServiceSpec.ExistingWorkflowId}") ~>
+        akkaHttpService.workflowRoutes ~>
+        check {
+          status should be(StatusCodes.OK)
+          contentType should be(ContentTypes.`application/json`)
+          val results = responseAs[JsObject].fields("results").convertTo[Seq[JsObject]]
+          val fields = results.head.fields
+          fields("id") should be(JsString(CromwellApiServiceSpec.ExistingWorkflowId.toString))
+          fields(WorkflowMetadataKeys.ParentWorkflowId) should be(JsString("pid"))
+        }
+    }
+
     behavior of "REST API /query POST endpoint"
     it should "return good results for a good query map body" in {
       Post(s"/workflows/$version/query", HttpEntity(ContentTypes.`application/json`, """[{"status":"Succeeded"}]""")) ~>
@@ -502,6 +530,33 @@ class CromwellApiServiceSpec extends AsyncFlatSpec with ScalatestRouteTest with
         }
     }
 
+    it should "return labels if specified in additionalQueryResultFields param" in {
+      Post(s"/workflows/$version/query", HttpEntity(ContentTypes.`application/json`, """[{"additionalQueryResultFields":"labels"}]""")) ~>
+        akkaHttpService.workflowRoutes ~>
+        check {
+          assertResult(StatusCodes.OK) {
+            status
+          }
+          val results = responseAs[JsObject].fields("results").convertTo[Seq[JsObject]]
+          val fields = results.head.fields
+          fields(WorkflowMetadataKeys.Labels).asJsObject.fields("key1") should be(JsString("label1"))
+          fields(WorkflowMetadataKeys.Labels).asJsObject.fields("key2") should be(JsString("label2"))
+        }
+    }
+
+    it should "return parentWorkflowId if specified in additionalQueryResultFields param" in {
+      Post(s"/workflows/$version/query", HttpEntity(ContentTypes.`application/json`, """[{"additionalQueryResultFields":"parentWorkflowId"}]""")) ~>
+        akkaHttpService.workflowRoutes ~>
+        check {
+          assertResult(StatusCodes.OK) {
+            status
+          }
+          assertResult(true) {
+            entityAs[String].contains("\"parentWorkflowId\":\"pid\"")
+          }
+        }
+    }
+
     behavior of "REST API /labels GET endpoint"
     it should "return labels for a workflow ID" in {
       Get(s"/workflows/$version/${CromwellApiServiceSpec.ExistingWorkflowId}/labels") ~>
@@ -588,9 +643,16 @@ object CromwellApiServiceSpec {
   class MockServiceRegistryActor extends Actor {
     import MockServiceRegistryActor._
     override def receive = {
-      case WorkflowQuery(_) =>
+      case WorkflowQuery(parameters) =>
+        val labels: Option[Map[String, String]] = {
+          parameters.contains(("additionalQueryResultFields", "labels")).option(
+            Map("key1" -> "label1", "key2" -> "label2"))
+        }
+        val parentWorkflowId: Option[String] =  {
+          parameters.contains(("additionalQueryResultFields", "parentWorkflowId")).option("pid")
+        }
         val response = WorkflowQuerySuccess(WorkflowQueryResponse(List(WorkflowQueryResult(ExistingWorkflowId.toString,
-          None, Some(WorkflowSucceeded.toString), None, None))), None)
+          None, Some(WorkflowSucceeded.toString), None, None, labels, parentWorkflowId))), None)
         sender ! response
       case ValidateWorkflowId(id) =>
         if (RecognizedWorkflowIds.contains(id)) sender ! MetadataService.RecognizedWorkflowId

services/src/main/scala/cromwell/services/metadata/MetadataService.scala

@@ -17,7 +17,7 @@ object MetadataService {
 
   final val MetadataServiceName = "MetadataService"
 
-  final case class WorkflowQueryResult(id: String, name: Option[String], status: Option[String], start: Option[OffsetDateTime], end: Option[OffsetDateTime])
+  final case class WorkflowQueryResult(id: String, name: Option[String], status: Option[String], start: Option[OffsetDateTime], end: Option[OffsetDateTime], labels: Option[Map[String, String]], parentWorkflowId: Option[String])
 
   final case class WorkflowQueryResponse(results: Seq[WorkflowQueryResult])
 

services/src/main/scala/cromwell/services/metadata/WorkflowQueryKey.scala

@@ -6,15 +6,16 @@ import cats.data
 import cats.syntax.traverse._
 import cats.syntax.validated._
 import cromwell.core.labels.Label
-import cromwell.core.{WorkflowId, WorkflowState}
+import cromwell.core.{WorkflowId, WorkflowMetadataKeys, WorkflowState}
 import common.validation.ErrorOr._
 import cats.data.Validated._
 import cats.instances.list._
+import mouse.boolean._
 
 import scala.util.{Success, Try}
 
 object WorkflowQueryKey {
-  val ValidKeys = Set(StartDate, EndDate, Name, Id, Status, LabelKeyValue, Page, PageSize) map { _.name }
+  val ValidKeys = Set(StartDate, EndDate, Name, Id, Status, LabelKeyValue, Page, PageSize, AdditionalQueryResultFields) map { _.name }
 
   case object StartDate extends DateTimeWorkflowQueryKey {
     override val name = "Start"
@@ -89,6 +90,19 @@ object WorkflowQueryKey {
       sequenceListOfValidatedNels("Unrecognized status values", nels)
     }
   }
+
+  case object AdditionalQueryResultFields extends SeqWorkflowQueryKey[String] {
+    override val name = "Additionalqueryresultfields"
+
+    override def validate(grouped: Map[String, Seq[(String, String)]]): ErrorOr[List[String]] = {
+      val values = valuesFromMap(grouped).toList
+      val allowedValues = Seq(WorkflowMetadataKeys.Labels, WorkflowMetadataKeys.ParentWorkflowId)
+      val nels: List[ErrorOr[String]] = values map { v => {
+        allowedValues.contains(v).fold(v.validNel[String], v.invalidNel[String])
+      }}
+      sequenceListOfValidatedNels("Unrecognized values", nels)
+    }
+  }
 }
 
 sealed trait WorkflowQueryKey[T] {

services/src/main/scala/cromwell/services/metadata/WorkflowQueryParameters.scala

@@ -17,7 +17,8 @@ case class WorkflowQueryParameters private(statuses: Set[String],
                                            startDate: Option[OffsetDateTime],
                                            endDate: Option[OffsetDateTime],
                                            page: Option[Int],
-                                           pageSize: Option[Int])
+                                           pageSize: Option[Int],
+                                           additionalQueryResultFields: Set[String])
 
 object WorkflowQueryParameters {
 
@@ -68,6 +69,7 @@ object WorkflowQueryParameters {
     val labelsValidation: ErrorOr[Set[Label]] = WorkflowQueryKey.LabelKeyValue.validate(valuesByCanonicalCapitalization).map(_.toSet)
     val pageValidation = Page.validate(valuesByCanonicalCapitalization)
     val pageSizeValidation = PageSize.validate(valuesByCanonicalCapitalization)
+    val additionalQueryResultFieldsValidation: ErrorOr[Set[String]] = AdditionalQueryResultFields.validate(valuesByCanonicalCapitalization).map(_.toSet)
 
     // Only validate start before end if both of the individual date parsing validations have already succeeded.
     val startBeforeEndValidation: ErrorOr[Unit] = (startDateValidation, endDateValidation) match {
@@ -84,7 +86,8 @@ object WorkflowQueryParameters {
       startDateValidation,
       endDateValidation,
       pageValidation,
-      pageSizeValidation) mapN { (_, _, statuses, names, ids, labels, startDate, endDate, page, pageSize) => WorkflowQueryParameters(statuses, names, ids, labels, startDate, endDate, page, pageSize) }
+      pageSizeValidation,
+      additionalQueryResultFieldsValidation) mapN { (_, _, statuses, names, ids, labels, startDate, endDate, page, pageSize, additionalQueryResultFields) => WorkflowQueryParameters(statuses, names, ids, labels, startDate, endDate, page, pageSize, additionalQueryResultFields) }
   }
 
   def apply(rawParameters: Seq[(String, String)]): WorkflowQueryParameters = {

services/src/main/scala/cromwell/services/metadata/impl/MetadataDatabaseAccess.scala

@@ -2,14 +2,20 @@ package cromwell.services.metadata.impl
 
 import cats.Semigroup
 import cats.data.NonEmptyList
+import cats.instances.future._
+import cats.instances.list._
+import cats.instances.option._
+import cats.syntax.functor._
 import cats.syntax.semigroup._
+import cats.syntax.traverse._
 import cromwell.core.{WorkflowId, WorkflowMetadataKeys, WorkflowState}
 import cromwell.database.sql.SqlConverters._
 import cromwell.database.sql.joins.{CallOrWorkflowQuery, CallQuery, WorkflowQuery}
 import cromwell.database.sql.tables.{MetadataEntry, WorkflowMetadataSummaryEntry}
 import cromwell.services.MetadataServicesStore
 import cromwell.services.metadata.MetadataService.{QueryMetadata, WorkflowQueryResponse}
 import cromwell.services.metadata._
+import mouse.boolean._
 
 import scala.concurrent.{ExecutionContext, Future}
 
@@ -35,7 +41,7 @@ object MetadataDatabaseAccess {
     }
   }
 
-  def baseSummary(workflowUuid: String) = WorkflowMetadataSummaryEntry(workflowUuid, None, None, None, None, None)
+  def baseSummary(workflowUuid: String) = WorkflowMetadataSummaryEntry(workflowUuid, None, None, None, None)
 
   // If visibility is made `private`, there's a bogus warning about this being unused.
   implicit class MetadatumEnhancer(val metadatum: MetadataEntry) extends AnyVal {
@@ -178,14 +184,60 @@ trait MetadataDatabaseAccess {
       queryParameters.statuses, queryParameters.names, queryParameters.ids.map(_.toString), queryParameters.labels.map(label => (label.key, label.value)),
       queryParameters.startDate.map(_.toSystemTimestamp), queryParameters.endDate.map(_.toSystemTimestamp))
 
-    workflowSummaryCount flatMap { count =>
-      workflowSummaries map { workflows =>
-        (WorkflowQueryResponse(workflows.toSeq map { workflow =>
-          MetadataService.WorkflowQueryResult(id = workflow.workflowExecutionUuid, name = workflow.workflowName, status = workflow.workflowStatus, start = workflow.startTimestamp map { _.toSystemOffsetDateTime }, end = workflow.endTimestamp map { _.toSystemOffsetDateTime })
-        }),
-          //only return metadata if page is defined
-          queryParameters.page map { _ => QueryMetadata(queryParameters.page, queryParameters.pageSize, Option(count)) })
+    def queryMetadata(count: Int): Option[QueryMetadata] = {
+      queryParameters.page.as(QueryMetadata(queryParameters.page, queryParameters.pageSize, Option(count)))
+    }
+
+    def summariesToQueryResults(workflows: Traversable[WorkflowMetadataSummaryEntry]): Future[List[MetadataService.WorkflowQueryResult]] = {
+      workflows.toList.traverse(summaryToQueryResult)
+    }
+
+    def summaryToQueryResult(workflow: WorkflowMetadataSummaryEntry): Future[MetadataService.WorkflowQueryResult] = {
+
+      def queryResult(labels: Map[String, String], parentId: Option[String]): MetadataService.WorkflowQueryResult = {
+        MetadataService.WorkflowQueryResult(
+          id = workflow.workflowExecutionUuid,
+          name = workflow.workflowName,
+          status = workflow.workflowStatus,
+          start = workflow.startTimestamp map { _.toSystemOffsetDateTime },
+          end = workflow.endTimestamp map { _.toSystemOffsetDateTime },
+          labels = labels.nonEmpty.option(labels),
+          parentWorkflowId = parentId
+        )
       }
+
+      def metadataEntriesToValue(entries: Seq[MetadataEntry]): Option[String] = {
+        entries.headOption.flatMap(_.metadataValue.toRawStringOption)
+      }
+
+      def keyToMetadataValue(key: String): Future[Option[String]] = {
+        val metadataEntries: Future[Seq[MetadataEntry]] = metadataDatabaseInterface.queryMetadataEntries(workflow.workflowExecutionUuid, key)
+        metadataEntries map metadataEntriesToValue
+      }
+
+      def getWorkflowLabels: Future[Map[String, String]] = {
+        queryParameters.additionalQueryResultFields.contains(WorkflowMetadataKeys.Labels).fold(
+          metadataDatabaseInterface.getWorkflowLabels(workflow.workflowExecutionUuid), Future.successful(Map.empty))
+      }
+
+      val getParentWorkflowId: Future[Option[String]] = {
+        queryParameters.additionalQueryResultFields.contains(WorkflowMetadataKeys.ParentWorkflowId).fold(
+          keyToMetadataValue(WorkflowMetadataKeys.ParentWorkflowId), Future.successful(None))
+      }
+
+      for {
+        labels <- getWorkflowLabels
+        parentWorkflowId <- getParentWorkflowId
+      } yield queryResult(labels, parentWorkflowId)
+
     }
+
+    for {
+      count <- workflowSummaryCount
+      workflows <- workflowSummaries
+      queryResults <- summariesToQueryResults(workflows)
+    } yield (WorkflowQueryResponse(queryResults.toSeq), queryMetadata(count))
+
   }
 }
+