Skip to content

Commit

Permalink
Use custom role when writing large objects
Browse files Browse the repository at this point in the history
  • Loading branch information
jgainerdewar committed Feb 13, 2024
1 parent 6994608 commit 9974464
Show file tree
Hide file tree
Showing 2 changed files with 26 additions and 8 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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}

Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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 {
Expand All @@ -116,15 +111,15 @@ 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
// than summarizable metadata.
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
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down Expand Up @@ -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
Expand Down

0 comments on commit 9974464

Please sign in to comment.