Merge 651f3bd into d625347

src/main/scala/org/broadinstitute/dsde/firecloud/model/ModelSchema.scala

@@ -15,7 +15,7 @@ import org.broadinstitute.dsde.firecloud.FireCloudConfig
  */
 
 object ModelSchema {
-
+  
   def getTypeSchema(entityType: String): Try[EntityMetadata] = {
     EntityTypes.types.get(entityType) match {
       case Some(schema) => Success(schema)
@@ -59,3 +59,42 @@ case class EntityMetadata(
 )
 
 case class EntityModel(schema : Map[String, EntityMetadata]) //entity name -> stuff about it
+
+object FlexibleModelSchema {
+  def memberTypeFromEntityType(entityType: String): Try[String] = {
+    val memberType = EntityTypes.types.get(entityType) match {
+      case Some(schema) => schema.memberType
+      case None => if (entityType.endsWith("_set"))
+        Some(entityType.replace("_set", "")) else
+        None
+    }
+    Try(memberType.get)
+  }
+
+  def pluralizeEntityType(entityType: String): String =  {
+    EntityTypes.types.get(entityType) match {
+      case Some(schema) => schema.plural
+      case None => entityType + "s"
+    }
+  }
+
+//  def getCollectionMemberType(collectionEntityType: String): Option[String] = {
+//    EntityTypes.types.get(collectionEntityType) match {
+//      case Some(schema) => schema.memberType
+//      case None => collectionEntityType.endsWith("_set")
+//    }
+//    getTypeSchema(collectionEntityType).map(_.memberType)
+//  }
+
+  def isCollectionType(entityType: String): Boolean = {
+    EntityTypes.types.get(entityType) match {
+      case Some(schema) => schema.memberType.isDefined
+      case None => entityType.endsWith("_set")
+    }
+  }
+
+  def isUsingFirecloudSchema(entityType: String): Boolean = {
+    EntityTypes.types.get(entityType).isDefined
+  }
+
+}

src/main/scala/org/broadinstitute/dsde/firecloud/service/ExportEntitiesByTypeActor.scala

@@ -89,7 +89,7 @@ class ExportEntitiesByTypeActor(rawlsDAO: RawlsDAO,
   def streamEntities(): Future[Unit] = {
     entityTypeMetadata flatMap { metadata =>
       val entityQueries = getEntityQueries(metadata, entityType)
-      if (TSVFormatter.isCollectionType(entityType)) {
+      if (FlexibleModelSchema.isCollectionType(entityType)) {
         streamCollectionType(entityQueries, metadata)
       } else {
         val headers = TSVFormatter.makeEntityHeaders(entityType, metadata.attributeNames, attributeNames)
@@ -123,6 +123,7 @@ class ExportEntitiesByTypeActor(rawlsDAO: RawlsDAO,
 
   // Stream exceptions have to be handled by directly closing out the RequestContext responder stream
   private def handleStreamException(t: Throwable): Unit = {
+    logger.info("handling exception", t)
     val message = t match {
       case f: FireCloudExceptionWithErrorReport => s"FireCloudException: Error generating entity download: ${f.errorReport.message}"
       case _ => s"FireCloudException: Error generating entity download: ${t.getMessage}"

src/main/scala/org/broadinstitute/dsde/firecloud/utils/TSVFormatter.scala

@@ -3,6 +3,7 @@ package org.broadinstitute.dsde.firecloud.utils
 import org.broadinstitute.dsde.rawls.model._
 import org.broadinstitute.dsde.firecloud.model._
 import org.broadinstitute.dsde.firecloud.service.TsvTypes
+import scala.util.{Failure, Success}
 
 import spray.json.JsValue
 
@@ -84,7 +85,7 @@ object TSVFormatter {
     * @return IndexedSeq of header Strings
     */
   def makeMembershipHeaders(entityType: String): IndexedSeq[String] = {
-    IndexedSeq[String](s"${TsvTypes.MEMBERSHIP}:${entityType}_id", memberTypeFromEntityType(entityType))
+    IndexedSeq[String](s"${TsvTypes.MEMBERSHIP}:${entityType}_id", FlexibleModelSchema.memberTypeFromEntityType(entityType).getOrElse(""))
   }
 
   /**
@@ -95,13 +96,12 @@ object TSVFormatter {
     * @return Ordered list of rows
     */
   def makeMembershipRows(entityType: String, entities: Seq[Entity]): Seq[IndexedSeq[String]] = {
-    val memberPlural = pluralizeMemberType(memberTypeFromEntityType(entityType))
     entities.filter { _.entityType == entityType }.flatMap {
       entity =>
         entity.attributes.filter {
           // To make the membership file, we need the array of elements that correspond to the set type.
           // All other top-level properties are not necessary and are only used for the data load file.
-          case (attributeName, _) => attributeName.equals(AttributeName.withDefaultNS(memberPlural))
+          case (attributeName, _) => attributeName.equals(AttributeName.withDefaultNS(FlexibleModelSchema.pluralizeEntityType(FlexibleModelSchema.memberTypeFromEntityType(entityType).get)))
         }.flatMap {
           case (_, AttributeEntityReference(`entityType`, entityName)) => Seq(IndexedSeq[String](entity.name, entityName))
           case (_, AttributeEntityReferenceList(refs)) => refs.map( ref => IndexedSeq[String](entity.name, ref.entityName) )
@@ -119,7 +119,7 @@ object TSVFormatter {
     * @return Entity name as first column header, followed by matching entity attribute labels
     */
   def makeEntityHeaders(entityType: String, allHeaders: Seq[String], requestedHeaders: Option[IndexedSeq[String]]): IndexedSeq[String] = {
-    val memberPlural = pluralizeMemberType(memberTypeFromEntityType(entityType))
+    val memberPlural = FlexibleModelSchema.pluralizeEntityType(FlexibleModelSchema.memberTypeFromEntityType(entityType).getOrElse(""))
 
     val requestedHeadersSansId = requestedHeaders.
       // remove empty strings
@@ -129,16 +129,16 @@ object TSVFormatter {
       // entity id always needs to be first and is handled differently so remove it from requestedHeaders
       map(_.filterNot(_.equalsIgnoreCase(entityType + "_id"))).
       // filter out member attribute if a set type
-      map { h => if (isCollectionType(entityType)) h.filterNot(_.equals(memberPlural)) else h }
+      map { h => if (FlexibleModelSchema.isCollectionType(entityType)) h.filterNot(_.equals(memberPlural)) else h }
 
-    val filteredAllHeaders = if (isCollectionType(entityType)) {
+    val filteredAllHeaders = if (FlexibleModelSchema.isCollectionType(entityType)) {
       allHeaders.filterNot(_.equals(memberPlural))
     } else {
       allHeaders
     }
 
     val entityHeader: String = requestedHeadersSansId match {
-      case Some(headers) if !ModelSchema.getRequiredAttributes(entityType).get.keySet.forall(headers.contains) => s"${TsvTypes.UPDATE}:${entityType}_id"
+      case Some(headers) if FlexibleModelSchema.isUsingFirecloudSchema(entityType) && !ModelSchema.getRequiredAttributes(entityType).get.keySet.forall(headers.contains) => s"${TsvTypes.UPDATE}:${entityType}_id"
       case _ => s"${TsvTypes.ENTITY}:${entityType}_id"
     }
     (entityHeader +: requestedHeadersSansId.getOrElse(filteredAllHeaders)).toIndexedSeq
@@ -153,11 +153,13 @@ object TSVFormatter {
     * @return Ordered list of rows, each row entry value ordered by its corresponding header position
     */
   def makeEntityRows(entityType: String, entities: Seq[Entity], headers: IndexedSeq[String]): IndexedSeq[IndexedSeq[String]] = {
-    val memberPlural = pluralizeMemberType(memberTypeFromEntityType(entityType))
     // if we have a set entity, we need to filter out the attribute array of the members so that we only
     // have top-level attributes to construct columns from.
-    val filteredEntities = if (isCollectionType(entityType)) {
-      filterAttributeFromEntities(entities, memberPlural)
+    val filteredEntities = if (FlexibleModelSchema.isCollectionType(entityType)) {
+      FlexibleModelSchema.memberTypeFromEntityType(entityType) match {
+        case Success(memberType) => filterAttributeFromEntities(entities, FlexibleModelSchema.pluralizeEntityType(memberType))
+        case Failure(regret) => ??? // should not happen?
+      }
     } else {
       entities
     }
@@ -168,10 +170,5 @@ object TSVFormatter {
       .toIndexedSeq
   }
 
-  def memberTypeFromEntityType(entityType: String): String = ModelSchema.getCollectionMemberType(entityType).get.getOrElse(entityType.replace("_set", ""))
-
-  def pluralizeMemberType(memberType: String): String = ModelSchema.getPlural(memberType).getOrElse(memberType + "s")
-
-  def isCollectionType(entityType: String): Boolean = ModelSchema.isCollectionType(entityType).getOrElse(false)
 
 }

src/test/scala/org/broadinstitute/dsde/firecloud/dataaccess/MockRawlsDAO.scala

@@ -88,8 +88,32 @@ object MockRawlsDAO {
       idName = "sample_set_id",
       attributeNames = largeSampleSetAttributes.map(_._1.name).toSeq))
 
+
+  val validBigQueryEntities = List(
+    Entity("shakespeare", "bigQuery", Map(AttributeName.withDefaultNS("query_str") -> AttributeString("SELECT * FROM [bigquery-public-data:samples.shakespeare] LIMIT 1000"))),
+    Entity("king", "bigQuery", Map(AttributeName.withDefaultNS("query_str") -> AttributeString("SELECT * FROM [bigquery-public-data:samples.king] LIMIT 1000")))
+  )
+
+  val validBigQuerySetEntities = List(
+    Entity("settest", "bigQuery_set", Map(AttributeName.withDefaultNS("bigQuerys") -> AttributeEntityReferenceList(Seq(
+      AttributeEntityReference("bigQuery", "shakespeare"),
+      AttributeEntityReference("bigQuery", "king")))))
+  )
+
+  val nonModelBigQueryMetadata = Map(
+    "bigQuery" -> EntityTypeMetadata(
+      count = 2,
+      idName = "bigQuery_id",
+      attributeNames = Seq("query_str")))
+
+  val nonModelBigQuerySetMetadata = Map(
+    "bigQuery_set" -> EntityTypeMetadata(
+      count = 1,
+      idName = "bigQuery_set_id",
+      attributeNames = Seq("bigQuerys")))
 }
 
+
 /**
   * Created by davidan on 9/28/16.
   *
@@ -304,6 +328,20 @@ class MockRawlsDAO extends RawlsDAO with MockGroupSupport {
         results = sampleRange
       )
       Future.successful(queryResponse)
+    } else if (workspaceName == "nonModel") {
+      val queryResponse: EntityQueryResponse = EntityQueryResponse(
+        parameters = query,
+        resultMetadata = EntityQueryResultMetadata(unfilteredCount = 2, filteredCount = 2, filteredPageCount = 1),
+        results = validBigQueryEntities
+      )
+      Future.successful(queryResponse)
+    } else if (workspaceName == "nonModelSet") {
+      val queryResponse: EntityQueryResponse = EntityQueryResponse(
+        parameters = query,
+        resultMetadata = EntityQueryResultMetadata(unfilteredCount = 1, filteredCount = 1, filteredPageCount = 1),
+        results = validBigQuerySetEntities
+      )
+      Future.successful(queryResponse)
     } else {
       val queryResponse: EntityQueryResponse = EntityQueryResponse(
         parameters = query,
@@ -323,6 +361,10 @@ class MockRawlsDAO extends RawlsDAO with MockGroupSupport {
       Future.successful(largeSampleSetMetadata)
     } else if (workspaceName == "large" || workspaceName == "page3exception") {
       Future.successful(largeSampleMetadata)
+    } else if (workspaceName == "nonModel") {
+      Future.successful(nonModelBigQueryMetadata)
+    } else if (workspaceName == "nonModelSet") {
+      Future.successful(nonModelBigQuerySetMetadata)
     } else {
       Future.successful(validEntitiesMetadata)
     }

src/test/scala/org/broadinstitute/dsde/firecloud/mock/MockTSV.scala

@@ -241,6 +241,12 @@ object MockTSVStrings {
     List("abc", "123", "").tabbed
   ).newlineSeparated
 
+  val nonModelBigQuery = List(
+    List("bigQuery_id", "query_str").tabbed,
+    List("shakespeare", "SELECT * FROM [bigquery-public-data:samples.shakespeare] LIMIT 1000").tabbed,
+    List("king", "SELECT * FROM [bigquery-public-data:samples.king] LIMIT 1000").tabbed
+  ).newlineSeparated
+
 }
 
 object MockTSVLoadFiles {

src/test/scala/org/broadinstitute/dsde/firecloud/service/EntityServiceSpec.scala

@@ -56,6 +56,9 @@ class EntityServiceSpec extends BaseServiceSpec with EntityService {
   )
   val validSampleEntities = List(Entity("sample_01", "sample", sampleAtts))
 
+  val nonModelBigQueryEntities = List(Entity("shakespeare", "bigQuery", Map(AttributeName.withDefaultNS("query_str") -> AttributeString(""))))
+//  val nonModelBigQuerySetEntities = List(Entity("testset", "bigQuery_set", Map(AttributeName.withDefaultNS())))
+
   override def beforeAll(): Unit = {
     workspaceServer = startClientAndServer(MockUtils.workspaceServerPort)
     // Valid Entities by sample type case

src/test/scala/org/broadinstitute/dsde/firecloud/service/ExportEntitiesByTypeServiceSpec.scala

@@ -32,6 +32,8 @@ class ExportEntitiesByTypeServiceSpec extends BaseServiceSpec with ExportEntitie
   val invalidFireCloudEntitiesParticipantSetTSVPath = "/api/workspaces/broad-dsde-dev/invalid/entities/participant_set/tsv"
   val exceptionFireCloudEntitiesSampleTSVPath = "/api/workspaces/broad-dsde-dev/exception/entities/sample/tsv"
   val page3ExceptionFireCloudEntitiesSampleTSVPath = "/api/workspaces/broad-dsde-dev/page3exception/entities/sample/tsv"
+  val nonModelEntitiesBigQueryTSVPath = "/api/workspaces/broad-dsde-dev/nonModel/entities/bigQuery/tsv"
+  val nonModelEntitiesBigQuerySetTSVPath = "/api/workspaces/broad-dsde-dev/nonModelSet/entities/bigQuery_set/tsv"
 
   // Pick the first few headers from the list of available sample headers:
   val filterProps: Seq[String] = MockRawlsDAO.largeSampleHeaders.take(5).map(_.name)
@@ -77,6 +79,33 @@ class ExportEntitiesByTypeServiceSpec extends BaseServiceSpec with ExportEntitie
       }
     }
 
+    "when calling GET on exporting a non-FC model entity type with all attributes" - {
+      "OK response is returned and attributes are included" in {
+        Get(nonModelEntitiesBigQueryTSVPath) ~> dummyUserIdHeaders("1234") ~> sealRoute(exportEntitiesRoutes) ~> check {
+          handled should be(true)
+          status should be(OK)
+          entity shouldNot be(empty) // Entity is the first line of content as output by StreamingActor
+          chunks shouldNot be(empty) // Chunks has all of the rest of the content, as output by StreamingActor
+          headers.contains(HttpHeaders.Connection("Keep-Alive")) should be(true)
+          headers.contains(HttpHeaders.`Content-Disposition`.apply("attachment", Map("filename" -> "bigQuery.txt"))) should be(true)
+          validateLineCount(chunks, 2)
+          entity.asString.contains("query_str") should be(true)
+        }
+      }
+    }
+
+    "when calling GET on exporting a non-FC model entity set type with all attributes" - {
+      "OK response is returned" in {
+        Get(nonModelEntitiesBigQuerySetTSVPath) ~> dummyUserIdHeaders("1234") ~> sealRoute(exportEntitiesRoutes) ~> check {
+          handled should be(true)
+          status should be(OK)
+          entity shouldNot be(empty) // Entity is the first line of content as output by StreamingActor
+          headers.contains(HttpHeaders.Connection("Keep-Alive")) should be(true)
+          headers.contains(HttpHeaders.`Content-Disposition`.apply("attachment", Map("filename" -> "bigQuery_set.zip"))) should be(true)
+        }
+      }
+    }
+
     "when calling GET on exporting LARGE (20K) sample TSV" - {
       "OK response is returned" in {
         Get(largeFireCloudEntitiesSampleTSVPath) ~> dummyUserIdHeaders("1234") ~> sealRoute(exportEntitiesRoutes) ~> check {