mongo transaction support, typed wrappers for MongoClient, MongoDatabase, ClientSession

Author: Roman Janusz

commons-macros/src/main/scala/com/avsystem/commons/macros/misc/MiscMacros.scala

@@ -75,6 +75,13 @@ final class MiscMacros(ctx: blackbox.Context) extends AbstractMacroCommons(ctx)
     q"new ${c.prefix}.ValName(${owner.asTerm.getter.name.decodedName.toString})"
   }
 
+  def optionalizeFirstArg(expr: Tree): Tree = expr match {
+    case Apply(MaybeTypeApply(Select(prefix, name: TermName), targs), head :: tail) =>
+      q"if($head ne null) $expr else ${c.untypecheck(prefix)}.$name[..$targs](..$tail)"
+    case _ =>
+      c.abort(expr.pos, "function application expected")
+  }
+
   def compilationError(error: Tree): Tree = error match {
     case StringLiteral(errorStr) =>
       abortAt(errorStr, error.pos)

commons-mongo/jvm/src/main/scala/com/avsystem/commons/mongo/typed/TypedClientSession.scala

@@ -0,0 +1,82 @@
+package com.avsystem.commons
+package mongo.typed
+
+import com.mongodb.reactivestreams.client.ClientSession
+import com.mongodb.session.ServerSession
+import com.mongodb.{ClientSessionOptions, ServerAddress, TransactionOptions}
+import monix.eval.Task
+import org.bson.{BsonDocument, BsonTimestamp}
+
+import java.io.Closeable
+
+class TypedClientSession(val nativeSession: ClientSession)
+  extends Closeable with TypedMongoUtils {
+
+  def hasActiveTransaction: Boolean =
+    nativeSession.hasActiveTransaction
+
+  def transactionOptions: TransactionOptions =
+    nativeSession.getTransactionOptions
+
+  def startTransaction(
+    transactionOptions: TransactionOptions = TransactionOptions.builder().build()
+  ): Unit =
+    nativeSession.startTransaction(transactionOptions)
+
+  def commitTransaction: Task[Unit] =
+    single(nativeSession.commitTransaction()).void
+
+  def abortTransaction: Task[Unit] =
+    single(nativeSession.abortTransaction()).void
+
+  def pinnedServerAddress: Option[ServerAddress] =
+    Option(nativeSession.getPinnedServerAddress)
+
+  def transactionContext(): Option[AnyRef] =
+    Option(nativeSession.getTransactionContext)
+
+  def setTransactionContext(address: ServerAddress, transactionContext: Any): Unit =
+    nativeSession.setTransactionContext(address, transactionContext)
+
+  def clearTransactionContext(): Unit =
+    nativeSession.clearTransactionContext()
+
+  def recoveryToken(): Option[BsonDocument] =
+    Option(nativeSession.getRecoveryToken)
+
+  def setRecoveryToken(recoverToken: BsonDocument): Unit =
+    nativeSession.setRecoveryToken(recoverToken)
+
+  def options: ClientSessionOptions =
+    nativeSession.getOptions
+
+  def casuallyConsistent: Boolean =
+    nativeSession.isCausallyConsistent
+
+  def originator: AnyRef =
+    nativeSession.getOriginator
+
+  def serverSession: ServerSession =
+    nativeSession.getServerSession
+
+  def operationTime: BsonTimestamp =
+    nativeSession.getOperationTime
+
+  def advanceOperationTime(operationTime: BsonTimestamp): Unit =
+    nativeSession.advanceOperationTime(operationTime)
+
+  def advanceClusterTime(clusterTime: BsonDocument): Unit =
+    nativeSession.advanceClusterTime(clusterTime)
+
+  def snapshotTimestamp: BsonTimestamp =
+    nativeSession.getSnapshotTimestamp
+
+  def setSnapshotTimestamp(snapshotTimestamp: BsonTimestamp): Unit =
+    nativeSession.setSnapshotTimestamp(snapshotTimestamp)
+
+  def clusterTime: BsonDocument =
+    nativeSession.getClusterTime
+
+  def close(): Unit =
+    nativeSession.close()
+}

commons-mongo/jvm/src/main/scala/com/avsystem/commons/mongo/typed/TypedMongoClient.scala

@@ -0,0 +1,77 @@
+package com.avsystem.commons
+package mongo.typed
+
+import com.avsystem.commons.mongo.BsonValueInput
+import com.avsystem.commons.serialization.GenCodec
+import com.mongodb.connection.ClusterDescription
+import com.mongodb.reactivestreams.client.{MongoClient, MongoClients}
+import com.mongodb.{ClientSessionOptions, ConnectionString, MongoClientSettings, MongoDriverInformation}
+import monix.eval.Task
+import monix.reactive.Observable
+import org.bson.Document
+
+import java.io.Closeable
+
+object TypedMongoClient {
+  def apply(): TypedMongoClient =
+    new TypedMongoClient(MongoClients.create())
+
+  def apply(connectionString: String): TypedMongoClient =
+    new TypedMongoClient(MongoClients.create(connectionString))
+
+  def apply(connectionString: ConnectionString): TypedMongoClient =
+    new TypedMongoClient(MongoClients.create(connectionString))
+
+  def apply(settings: MongoClientSettings): TypedMongoClient =
+    new TypedMongoClient(MongoClients.create(settings))
+
+  def apply(
+    connectionString: ConnectionString,
+    driverInformation: MongoDriverInformation,
+  ): TypedMongoClient =
+    new TypedMongoClient(MongoClients.create(connectionString, driverInformation))
+
+  def apply(
+    settings: MongoClientSettings,
+    driverInformation: MongoDriverInformation,
+  ): TypedMongoClient =
+    new TypedMongoClient(MongoClients.create(settings, driverInformation))
+}
+
+/**
+  * A better-typed wrapper over [[MongoClient]]. Uses Monix [[Task]] and [[Observable]] instead of
+  * [[org.reactivestreams.Publisher]]. Returns similar better-typed wrappers for database and client session objects.
+  */
+class TypedMongoClient(
+  val nativeClient: MongoClient,
+  val clientSession: OptArg[TypedClientSession] = OptArg.Empty
+) extends TypedMongoUtils with Closeable {
+  private val sessionOrNull = clientSession.toOpt.map(_.nativeSession).orNull
+
+  def withSession(session: TypedClientSession): TypedMongoClient =
+    new TypedMongoClient(nativeClient, session)
+
+  def getDatabase(name: String): TypedMongoDatabase =
+    new TypedMongoDatabase(nativeClient.getDatabase(name))
+
+  def listDatabaseNames: Observable[String] =
+    multi(optionalizeFirstArg(nativeClient.listDatabaseNames(sessionOrNull)))
+
+  def listDatabases: Observable[Document] =
+    multi(optionalizeFirstArg(nativeClient.listDatabases(sessionOrNull)))
+
+  def listDatabases[T: GenCodec]: Observable[T] =
+    listDatabases.map(doc => BsonValueInput.read[T](doc.toBsonDocument))
+
+  //TODO: `watch` methods
+
+  def startSession(
+    options: ClientSessionOptions = ClientSessionOptions.builder().build()
+  ): Task[TypedClientSession] =
+    single(nativeClient.startSession(options)).map(new TypedClientSession(_))
+
+  def clusterDescription(): ClusterDescription =
+    nativeClient.getClusterDescription
+
+  def close(): Unit = nativeClient.close()
+}