From 9974464a09a01e89f5ca04da169938924fe24235 Mon Sep 17 00:00:00 2001
From: Janet Gainer-Dewar <jdewar@broadinstitute.org>
Date: Tue, 13 Feb 2024 17:01:28 -0500
Subject: [PATCH] Use custom role when writing large objects

---
 .../slick/MetadataSlickDatabase.scala         |  9 ++-----
 .../database/slick/SlickDatabase.scala        | 25 ++++++++++++++++++-
 2 files changed, 26 insertions(+), 8 deletions(-)

diff --git a/database/sql/src/main/scala/cromwell/database/slick/MetadataSlickDatabase.scala b/database/sql/src/main/scala/cromwell/database/slick/MetadataSlickDatabase.scala
index c0723af6795..fbab736ae26 100644
--- a/database/sql/src/main/scala/cromwell/database/slick/MetadataSlickDatabase.scala
+++ b/database/sql/src/main/scala/cromwell/database/slick/MetadataSlickDatabase.scala
@@ -14,7 +14,6 @@ import cromwell.database.sql.tables.{
   MetadataEntry,
   WorkflowMetadataSummaryEntry
 }
-import net.ceedubs.ficus.Ficus._
 import slick.basic.DatabasePublisher
 import slick.jdbc.{ResultSetConcurrency, ResultSetType}
 
@@ -76,8 +75,6 @@ class MetadataSlickDatabase(originalDatabaseConfig: Config)
   import dataAccess.driver.api._
   import MetadataSlickDatabase._
 
-  lazy val pgLargeObjectWriteRole: Option[String] = originalDatabaseConfig.as[Option[String]]("pgLargeObjectWriteRole")
-
   override def existsMetadataEntries()(implicit ec: ExecutionContext): Future[Boolean] = {
     val action = dataAccess.metadataEntriesExists.result
     runTransaction(action)
@@ -106,8 +103,6 @@ class MetadataSlickDatabase(originalDatabaseConfig: Config)
       labelMetadataKey
     )
 
-    val roleSet = pgLargeObjectWriteRole.map(role => sqlu"""SET ROLE TO "#$role"""")
-
     // These entries also require a write to the summary queue.
     def writeSummarizable(): Future[Unit] = if (partitioned.summarizableMetadata.isEmpty) Future.successful(())
     else {
@@ -116,7 +111,7 @@ class MetadataSlickDatabase(originalDatabaseConfig: Config)
         val insertMetadata = dataAccess.metadataEntryIdsAutoInc ++= batch
         insertMetadata.flatMap(ids => writeSummaryQueueEntries(ids))
       }
-      runTransaction(DBIO.sequence(roleSet ++ insertActions)).void
+      runTransaction(DBIO.sequence(insertActions)).void
     }
 
     // Non-summarizable metadata that only needs to go to the metadata table can be written much more efficiently
@@ -124,7 +119,7 @@ class MetadataSlickDatabase(originalDatabaseConfig: Config)
     def writeNonSummarizable(): Future[Unit] = if (partitioned.nonSummarizableMetadata.isEmpty) Future.successful(())
     else {
       val action = DBIO.sequence(
-        roleSet ++ partitioned.nonSummarizableMetadata.grouped(insertBatchSize).map(dataAccess.metadataEntries ++= _)
+        partitioned.nonSummarizableMetadata.grouped(insertBatchSize).map(dataAccess.metadataEntries ++= _)
       )
       runLobAction(action).void
     }
diff --git a/database/sql/src/main/scala/cromwell/database/slick/SlickDatabase.scala b/database/sql/src/main/scala/cromwell/database/slick/SlickDatabase.scala
index 55c408f944f..bc9b085b1f1 100644
--- a/database/sql/src/main/scala/cromwell/database/slick/SlickDatabase.scala
+++ b/database/sql/src/main/scala/cromwell/database/slick/SlickDatabase.scala
@@ -72,6 +72,19 @@ abstract class SlickDatabase(override val originalDatabaseConfig: Config) extend
   // NOTE: if you want to refactor database is inner-class type: this.dataAccess.driver.backend.DatabaseFactory
   val database = slickConfig.db
 
+  /*
+  In some cases we want to write Postgres Large Objects (corresponding to Clob/Blob in Slick)
+  with a role other than the database user we are authenticated as. This is important
+  when we want multiple login users to be able to access the records. We can SET ROLE to
+  a role granted to all of them, and they'll all be able to access the Large Objects.
+  This SO thread also has a good explanation:
+  https://dba.stackexchange.com/questions/147607/postgres-large-objects-multiple-users
+   */
+  private lazy val pgLargeObjectWriteRole: Option[String] =
+    originalDatabaseConfig.as[Option[String]]("pgLargeObjectWriteRole")
+  private lazy val roleSetCmd =
+    pgLargeObjectWriteRole.map(role => sqlu"""SET LOCAL ROLE TO "#$role"""")
+
   override lazy val connectionDescription: String = databaseConfig.getString(urlKey)
 
   SlickDatabase.log.info(s"Running with database $urlKey = $connectionDescription")
@@ -167,7 +180,17 @@ abstract class SlickDatabase(override val originalDatabaseConfig: Config) extend
                                         isolationLevel: TransactionIsolation = TransactionIsolation.RepeatableRead,
                                         timeout: Duration = Duration.Inf
   ): Future[R] =
-    runActionInternal(action.transactionally.withTransactionIsolation(isolationLevel), timeout = timeout)
+    runActionInternal(withLobRole(action).transactionally.withTransactionIsolation(isolationLevel), timeout = timeout)
+
+  /*
+  If we're using Postgres and have been configured to do so, set the desired role on the transaction.
+  See comments on `roleSetCmd` above for more information.
+   */
+  private def withLobRole[R](action: DBIO[R]): DBIO[R] =
+    (dataAccess.driver, roleSetCmd) match {
+      case (PostgresProfile, Some(roleSet)) => roleSet.andThen(action)
+      case _ => action
+    }
 
   /* Note that this is only appropriate for actions that do not involve Blob
    * or Clob fields in Postgres, since large object support requires running