got WSM endpoint

broadinstitute · Oct 19, 2023 · aa1bc52 · aa1bc52
1 parent 06531d7
commit aa1bc52
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 3 deletions.
diff --git a/filesystems/blob/src/main/scala/cromwell/filesystems/blob/BlobPathBuilder.scala b/filesystems/blob/src/main/scala/cromwell/filesystems/blob/BlobPathBuilder.scala
@@ -194,9 +194,8 @@ case class BlobPath private[blob](pathString: String, endpoint: EndpointURL, con
   }
 
   def containerWSMResourceId: Try[UUID] = {
-
     val wsmGenerator: Option[WSMBlobSasTokenGenerator] = fsm.blobTokenGenerator match {
-      case wsmGenerator: WSMBlobSasTokenGenerator    => Option(wsmGenerator)
+      case wsmGenerator: WSMBlobSasTokenGenerator => Option(wsmGenerator)
       case _: Any => None
     }
     val workspaceId: Try[UUID] = parseTerraWorkspaceIdFromPath
@@ -205,6 +204,20 @@ case class BlobPath private[blob](pathString: String, endpoint: EndpointURL, con
     Try(wsmGenerator.get.getContainerResourceId(workspaceId.get, container, wsmAuth.get)).flatten
   }
 
+  def wsmEndpoint: Try[String] = {
+    val wsmGenerator: Option[WSMBlobSasTokenGenerator] = fsm.blobTokenGenerator match {
+      case wsmGenerator: WSMBlobSasTokenGenerator => Option(wsmGenerator)
+      case _: Any => None
+    }
+    val maybeEndpoint = wsmGenerator.map{generator =>
+      generator.wsmClientProvider.getBaseWorkspaceManagerUrl
+    }
+    maybeEndpoint match {
+      case endpoint: Some[String] => Try(endpoint.value)
+      case _ => Failure(new NoSuchElementException("Could not determine WSM API endpoint."))
+    }
+  }
+
   override def getSymlinkSafePath(options: LinkOption*): Path  = toAbsolutePath
 
 }
diff --git a/...ems/blob/src/main/scala/cromwell/filesystems/blob/WorkspaceManagerApiClientProvider.scala b/...ems/blob/src/main/scala/cromwell/filesystems/blob/WorkspaceManagerApiClientProvider.scala
@@ -20,6 +20,7 @@ import scala.util.Try
 trait WorkspaceManagerApiClientProvider {
   def getControlledAzureResourceApi(token: String): WsmControlledAzureResourceApi
   def getResourceApi(token: String): WsmResourceApi
+  def getBaseWorkspaceManagerUrl: String
 }
 
 class HttpWorkspaceManagerClientProvider(baseWorkspaceManagerUrl: WorkspaceManagerURL) extends WorkspaceManagerApiClientProvider {
@@ -40,6 +41,10 @@ class HttpWorkspaceManagerClientProvider(baseWorkspaceManagerUrl: WorkspaceManag
     apiClient.setAccessToken(token)
     WsmControlledAzureResourceApi(new ControlledAzureResourceApi(apiClient))
   }
+
+  def getBaseWorkspaceManagerUrl: String = {
+    baseWorkspaceManagerUrl.value
+  }
 }
 
 case class WsmResourceApi(resourcesApi : ResourceApi) {

diff --git a/...kends/tes/src/main/scala/cromwell/backend/impl/tes/TesAsyncBackendJobExecutionActor.scala b/...kends/tes/src/main/scala/cromwell/backend/impl/tes/TesAsyncBackendJobExecutionActor.scala
@@ -29,6 +29,7 @@ import wom.values.WomFile
 
 import java.io.FileNotFoundException
 import java.nio.file.FileAlreadyExistsException
+import java.util.UUID
 import scala.concurrent.Future
 import scala.util.{Failure, Success, Try}
 sealed trait TesRunStatus {
@@ -67,6 +68,12 @@ object TesAsyncBackendJobExecutionActor {
        |CONTAINER_RESOURCE_ID="${sasParams.containerResourceId}"
        |""".stripMargin
   }
+
+  /* Under certain situations (and only on Terra), we want the VM running a TES task to have the ability to acquire a
+   * fresh SaS token for itself. In order to be able to do this, we must provide the task execution script with a
+   * WSM endpoint, WorkspaceID, and container resource ID. The task VM will use the user assigned managed identity that
+   * it is running as in order to authenticate.
+   */
   def getLocalizedSasTokenParams(taskInputs: List[Input], pathGetter: String => Try[Path]): Option[LocalizedSasTokenParams] = {
     val shouldLocalizeSas = true //TODO: Make this a Workflow Option or come from the WDL
     if (!shouldLocalizeSas || taskInputs.isEmpty) return None
@@ -89,9 +96,12 @@ object TesAsyncBackendJobExecutionActor {
       case _: Any => None
     }
 
+    val uuid: Try[UUID] = blobPath.get.containerWSMResourceId
+    if(!uuid.isSuccess) return None
     val container = templateBlobFile.container
     val maybeWorkspaceId = blobPath
-    val wsmEndpoint = "1234"
+    val tryWsm = blobPath.get.wsmEndpoint
+    val wsmEndpoint = tryWsm.getOrElse("invalid")
     maybeWorkspaceId.map(workspaceId => LocalizedSasTokenParams(wsmEndpoint, container.value, blobPath.get.containerWSMResourceId.get.toString))
   }
 }