[REST Catalog API] Rework catalog instance rotation (#338)

Author: george-zubrienko

build.sbt

@@ -43,6 +43,7 @@ lazy val root = (project in file("."))
     // Framework dependencies
     libraryDependencies += "dev.zio" %% "zio" % "2.1.24",
     libraryDependencies += "dev.zio" %% "zio-streams" % "2.1.24",
+      libraryDependencies += "dev.zio" %% "zio-cache" % "0.2.8",
     libraryDependencies += "com.microsoft.sqlserver" % "mssql-jdbc" % "12.8.2.jre11",
     libraryDependencies += "software.amazon.awssdk" % "s3" % "2.33.13",
     libraryDependencies += "com.lihaoyi" %% "upickle" % "4.4.3",

src/main/scala/models/settings/iceberg/DefaultIcebergStagingSettings.scala

@@ -1,17 +1,18 @@
 package com.sneaksanddata.arcane.framework
 package models.settings.iceberg
 
+import models.serialization.ZIODurationRW.*
 import services.iceberg.IcebergCatalogCredential
 import services.iceberg.base.S3CatalogFileIO
 
 import upickle.ReadWriter
-import upickle.implicits.key
 
 case class DefaultIcebergStagingSettings(
     catalogProperties: Map[String, String],
     override val namespace: String,
     override val catalogUri: String,
-    override val warehouse: String
+    override val warehouse: String,
+    override val maxCatalogInstanceLifetime: zio.Duration
 ) extends IcebergCatalogSettings derives ReadWriter:
   /** Important to note that currently we do not provide separation between Sink and Staging catalog auth and FileIO
     * implementations. This should be fixed in the future.

src/main/scala/models/settings/iceberg/IcebergCatalogSettings.scala

@@ -19,3 +19,7 @@ trait IcebergCatalogSettings:
   /** The catalog additional properties.
     */
   val additionalProperties: Map[String, String]
+
+  /** Maximum lifetime of a RESTSessionCatalog instance
+    */
+  val maxCatalogInstanceLifetime: zio.Duration

src/main/scala/models/settings/sink/DefaultIcebergSinkSettings.scala

@@ -1,18 +1,19 @@
 package com.sneaksanddata.arcane.framework
 package models.settings.sink
 
+import models.serialization.ZIODurationRW.*
 import models.settings.iceberg.IcebergCatalogSettings
 import services.iceberg.IcebergCatalogCredential
 import services.iceberg.base.S3CatalogFileIO
 
 import upickle.ReadWriter
-import upickle.implicits.key
 
 case class DefaultIcebergSinkSettings(
     catalogProperties: Map[String, String],
     override val namespace: String,
     override val catalogUri: String,
-    override val warehouse: String
+    override val warehouse: String,
+    override val maxCatalogInstanceLifetime: zio.Duration
 ) extends IcebergCatalogSettings derives ReadWriter:
   /** Important to note that currently we do not provide separation between Sink and Staging catalog auth and FileIO
     * implementations. This should be fixed in the future.

src/main/scala/services/iceberg/IcebergCatalogFactory.scala

@@ -10,26 +10,18 @@ import org.apache.iceberg.aws.s3.S3FileIO
 import org.apache.iceberg.catalog.SessionCatalog.SessionContext
 import org.apache.iceberg.rest.auth.OAuth2Properties
 import org.apache.iceberg.rest.{HTTPClient, RESTSessionCatalog}
-import zio.{Task, ZIO}
+import zio.cache.Cache
+import zio.{Cached, Schedule, Scope, Task, ZIO, ZLayer}
 
 import java.time.Instant
 import scala.collection.concurrent.TrieMap
 import scala.jdk.CollectionConverters.*
 import scala.language.implicitConversions
 
-final class IcebergCatalogFactory(icebergCatalogSettings: IcebergCatalogSettings) extends CatalogFactory:
-
-  private val catalogProperties: Map[String, String] =
-    Map(
-      CatalogProperties.WAREHOUSE_LOCATION -> icebergCatalogSettings.warehouse,
-      CatalogProperties.URI                -> icebergCatalogSettings.catalogUri,
-      "rest-metrics-reporting-enabled"     -> "false",
-      "view-endpoints-supported"           -> "false"
-    ) ++ icebergCatalogSettings.additionalProperties
-
-  private val maxCatalogLifetime =
-    zio.Duration.fromSeconds(10 * 60).toSeconds // limit catalog instance lifetime to 10 minutes
-  private val catalogs: TrieMap[String, (RESTSessionCatalog, Long)] = TrieMap()
+final class IcebergCatalogFactory(
+    icebergCatalogSettings: IcebergCatalogSettings,
+    cachedRef: Cached[Throwable, RESTSessionCatalog]
+) extends CatalogFactory:
 
   def getSessionContext: SessionContext =
     SessionContext(
@@ -41,38 +33,46 @@ final class IcebergCatalogFactory(icebergCatalogSettings: IcebergCatalogSettings
       Map().asJava
     );
 
-  def newCatalog: Task[RESTSessionCatalog] = ZIO.scoped {
-    for
+  def getCatalog: Task[RESTSessionCatalog] = cachedRef.get
+
+object IcebergCatalogFactory:
+  private def getCatalogProperties(icebergCatalogSettings: IcebergCatalogSettings): Map[String, String] =
+    Map(
+      CatalogProperties.WAREHOUSE_LOCATION -> icebergCatalogSettings.warehouse,
+      CatalogProperties.URI                -> icebergCatalogSettings.catalogUri,
+      "rest-metrics-reporting-enabled"     -> "false",
+      "view-endpoints-supported"           -> "false"
+    ) ++ icebergCatalogSettings.additionalProperties
+
+  private def newCatalog(icebergCatalogSettings: IcebergCatalogSettings): ZIO[Scope, Throwable, RESTSessionCatalog] =
+    ZIO.acquireRelease(for
+      properties <- ZIO.succeed(getCatalogProperties(icebergCatalogSettings))
       result <- ZIO.succeed(
         new RESTSessionCatalog(
           config => HTTPClient.builder(config).uri(config.get(CatalogProperties.URI)).build(),
           (_, _) =>
             val baseIO = new S3FileIO()
-            baseIO.initialize(catalogProperties.asJava)
+            baseIO.initialize(properties.asJava)
             baseIO
         )
       )
       name <- ZIO.succeed(java.util.UUID.randomUUID().toString)
       _    <- zlog("Creating new Iceberg RESTSessionCatalog instance with id %s", name)
       _ <- ZIO.attemptBlocking(
-        result.initialize(name, (catalogProperties ++ Map("header.X-Arcane-Runner-Identifier" -> name)).asJava)
+        result.initialize(name, (properties ++ Map("header.X-Arcane-Runner-Identifier" -> name)).asJava)
       )
-      _ <- ZIO.attempt(catalogs.addOne((name, (result, Instant.now.getEpochSecond))))
-    yield result
-  }
+    yield result)(catalog =>
+      zlog("Closing RESTSessionCatalog %s due to cache expiry or application shutdown", catalog.name()) *> ZIO
+        .attemptBlocking(catalog.close())
+        .orDieWith(e => new Throwable("Unable to close RESTSessionCatalog instance correctly", e))
+    )
 
-  def getCatalog: Task[RESTSessionCatalog] = for
-    catalogInfo <- ZIO.attempt(catalogs.find { case (_, (_, duration)) =>
-      Instant.now.getEpochSecond - duration < maxCatalogLifetime
-    })
-    selected <- catalogInfo match {
-      case Some((_, (catalog, _))) => ZIO.succeed(catalog)
-      case None =>
-        ZIO
-          .attempt(catalogs.foreach { case (_, (catalog, _)) =>
-            catalog.close()
-          })
-          .map(_ => catalogs.clear())
-          .flatMap(_ => newCatalog)
-    }
-  yield selected
+  /** Provide an instance of IcebergCatalogFactory which rotates catalog instance automatically
+    * @return
+    */
+  def live(icebergCatalogSettings: IcebergCatalogSettings): ZIO[Scope, Throwable, IcebergCatalogFactory] =
+    for cachedRef <- Cached.auto(
+        acquire = newCatalog(icebergCatalogSettings),
+        policy = Schedule.spaced(icebergCatalogSettings.maxCatalogInstanceLifetime)
+      )
+    yield new IcebergCatalogFactory(icebergCatalogSettings, cachedRef)

src/main/scala/services/iceberg/IcebergEntityManager.scala

@@ -10,14 +10,13 @@ import services.iceberg.SchemaConversions.toIcebergType
 import services.iceberg.base.{CatalogEntityManager, SinkEntityManager, StagingEntityManager}
 
 import org.apache.iceberg.catalog.{Namespace, TableIdentifier}
-import org.apache.iceberg.{PartitionSpec, SortOrder}
+import org.apache.iceberg.{PartitionSpec, SortOrder, Table}
 import zio.{Task, ZIO, ZLayer}
 
 import scala.jdk.CollectionConverters.*
 
-trait IcebergEntityManager(catalogSettings: IcebergCatalogSettings) extends CatalogEntityManager:
-  override val catalogFactory = new IcebergCatalogFactory(catalogSettings)
-
+trait IcebergEntityManager(catalogSettings: IcebergCatalogSettings, catalogFactory: IcebergCatalogFactory)
+    extends CatalogEntityManager:
   private def delete(tableId: TableIdentifier): Task[Boolean] = for
     _       <- zlog("Deleting table %s", tableId.name())
     catalog <- catalogFactory.getCatalog
@@ -36,6 +35,18 @@ trait IcebergEntityManager(catalogSettings: IcebergCatalogSettings) extends Cata
     result  <- delete(tableId)
   yield result
 
+  override def getTableRef(tableName: String): Task[Table] = for
+    tableId  <- ZIO.succeed(TableIdentifier.of(catalogSettings.namespace, tableName))
+    catalog  <- catalogFactory.getCatalog
+    tableRef <- ZIO.attemptBlocking(catalog.loadTable(catalogFactory.getSessionContext, tableId))
+  yield tableRef
+
+  override def tableExists(tableName: String): Task[Boolean] = for
+    tableId <- ZIO.succeed(TableIdentifier.of(catalogSettings.namespace, tableName))
+    catalog <- catalogFactory.getCatalog
+    result  <- ZIO.attemptBlocking(catalog.tableExists(catalogFactory.getSessionContext, tableId))
+  yield result
+
   /** Creates a new table in the Iceberg catalog, using the provided schema
     * @return
     */
@@ -92,21 +103,25 @@ trait IcebergEntityManager(catalogSettings: IcebergCatalogSettings) extends Cata
     }
   yield ()
 
-class IcebergSinkEntityManager(catalogSettings: IcebergCatalogSettings)
-    extends IcebergEntityManager(catalogSettings)
+class IcebergSinkEntityManager(catalogSettings: IcebergCatalogSettings, catalogFactory: IcebergCatalogFactory)
+    extends IcebergEntityManager(catalogSettings, catalogFactory)
     with SinkEntityManager
 
-class IcebergStagingEntityManager(catalogSettings: IcebergCatalogSettings)
-    extends IcebergEntityManager(catalogSettings)
+class IcebergStagingEntityManager(catalogSettings: IcebergCatalogSettings, catalogFactory: IcebergCatalogFactory)
+    extends IcebergEntityManager(catalogSettings, catalogFactory)
     with StagingEntityManager
 
 object IcebergEntityManager:
-  val sinkLayer = ZLayer {
-    for context <- ZIO.service[PluginStreamContext]
-    yield IcebergSinkEntityManager(context.sink.icebergCatalog)
+  val sinkLayer: ZLayer[PluginStreamContext, Throwable, IcebergSinkEntityManager] = ZLayer.scoped {
+    for
+      context <- ZIO.service[PluginStreamContext]
+      factory <- IcebergCatalogFactory.live(context.sink.icebergCatalog)
+    yield IcebergSinkEntityManager(context.sink.icebergCatalog, factory)
   }
 
-  val stagingLayer = ZLayer {
-    for context <- ZIO.service[PluginStreamContext]
-    yield IcebergStagingEntityManager(context.staging.icebergCatalog)
+  val stagingLayer: ZLayer[PluginStreamContext, Throwable, IcebergStagingEntityManager] = ZLayer.scoped {
+    for
+      context <- ZIO.service[PluginStreamContext]
+      factory <- IcebergCatalogFactory.live(context.staging.icebergCatalog)
+    yield IcebergStagingEntityManager(context.staging.icebergCatalog, factory)
   }

src/main/scala/services/iceberg/IcebergS3CatalogWriter.scala

@@ -28,7 +28,7 @@ given Conversion[ArcaneSchema, Schema] with
   def apply(schema: ArcaneSchema): Schema = SchemaConversions.toIcebergSchema(schema)
 
 // https://www.tabular.io/blog/java-api-part-3/
-class IcebergS3CatalogWriter(entityManager: CatalogEntityManager, stagingSettings: StagingSettings)
+class IcebergS3CatalogWriter(entityManager: StagingEntityManager, stagingSettings: StagingSettings)
     extends CatalogWriter[RESTCatalog, Table, Schema]:
 
   private val maxRowsPerFile = stagingSettings.table.maxRowsPerFile.getOrElse(10000)
@@ -101,25 +101,13 @@ class IcebergS3CatalogWriter(entityManager: CatalogEntityManager, stagingSetting
       logAnnotations: Seq[(LogAnnotation[String], String)]
   ): Task[Table] =
     for
-      _       <- entityManager.createTable(CreateTableRequest(name, schema, false))
-      _       <- zlog("Created a staging table %s, waiting for commit", logAnnotations, name)
-      catalog <- entityManager.catalogFactory.getCatalog
+      _ <- entityManager.createTable(CreateTableRequest(name, schema, false))
+      _ <- zlog("Created a staging table %s, waiting for commit", logAnnotations, name)
       _ <- ZIO
         .sleep(zio.Duration.fromSeconds(1))
-        .repeatUntil(_ =>
-          catalog
-            .tableExists(
-              entityManager.catalogFactory.getSessionContext,
-              TableIdentifier.of(stagingSettings.icebergCatalog.namespace, name)
-            )
-        )
-      _ <- zlog("Staging table %s created, appending data", logAnnotations, name)
-      table <- ZIO.attemptBlocking(
-        catalog.loadTable(
-          entityManager.catalogFactory.getSessionContext,
-          TableIdentifier.of(stagingSettings.icebergCatalog.namespace, name)
-        )
-      )
+        .repeatUntilZIO(_ => entityManager.tableExists(name).orDie)
+      _            <- zlog("Staging table %s created, appending data", logAnnotations, name)
+      table        <- entityManager.getTableRef(name)
       updatedTable <- appendData(data, schema, false, table, logAnnotations)
       _            <- zlog("Staging table %s ready for merge", logAnnotations, name)
     yield updatedTable
@@ -130,9 +118,7 @@ class IcebergS3CatalogWriter(entityManager: CatalogEntityManager, stagingSetting
       schema: Schema,
       logAnnotations: Seq[(LogAnnotation[String], String)]
   ): Task[Table] = for
-    tableId      <- ZIO.succeed(TableIdentifier.of(stagingSettings.icebergCatalog.namespace, name))
-    catalog      <- entityManager.catalogFactory.getCatalog
-    table        <- ZIO.attemptBlocking(catalog.loadTable(entityManager.catalogFactory.getSessionContext, tableId))
+    table        <- entityManager.getTableRef(name)
     updatedTable <- appendData(data, schema, false, table, logAnnotations)
   yield updatedTable
 
@@ -142,10 +128,10 @@ object IcebergS3CatalogWriter:
     * @return
     *   The initialized IcebergS3CatalogWriter instance
     */
-  def apply(entityManager: CatalogEntityManager, stagingSettings: StagingSettings): IcebergS3CatalogWriter =
+  def apply(entityManager: StagingEntityManager, stagingSettings: StagingSettings): IcebergS3CatalogWriter =
     new IcebergS3CatalogWriter(entityManager, stagingSettings)
 
-  def layer = ZLayer {
+  def layer: ZLayer[StagingEntityManager & PluginStreamContext, Nothing, IcebergS3CatalogWriter] = ZLayer.scoped {
     for
       stagingEntityManager <- ZIO.service[StagingEntityManager]
       context              <- ZIO.service[PluginStreamContext]

src/main/scala/services/iceberg/IcebergTablePropertyManager.scala

@@ -12,8 +12,8 @@ import zio.{Task, ZIO, ZLayer}
 import scala.jdk.CollectionConverters.*
 import scala.language.implicitConversions
 
-trait IcebergTablePropertyManager(catalogSettings: IcebergCatalogSettings) extends TablePropertyManager:
-  override val catalogFactory = new IcebergCatalogFactory(catalogSettings)
+trait IcebergTablePropertyManager(catalogSettings: IcebergCatalogSettings, catalogFactory: IcebergCatalogFactory)
+    extends TablePropertyManager:
 
   private def loadTable(tableName: String): Task[Table] = for
     tableId <- ZIO.succeed(
@@ -72,24 +72,28 @@ trait IcebergTablePropertyManager(catalogSettings: IcebergCatalogSettings) exten
   override def getTableSchema(tableName: String): Task[Schema] = for table <- loadTable(tableName)
   yield table.schema()
 
-class IcebergSinkTablePropertyManager(catalogSettings: IcebergCatalogSettings)
-    extends IcebergTablePropertyManager(catalogSettings)
+class IcebergSinkTablePropertyManager(catalogSettings: IcebergCatalogSettings, catalogFactory: IcebergCatalogFactory)
+    extends IcebergTablePropertyManager(catalogSettings, catalogFactory)
     with SinkPropertyManager
 
-class IcebergStagingTablePropertyManager(catalogSettings: IcebergCatalogSettings)
-    extends IcebergTablePropertyManager(catalogSettings)
+class IcebergStagingTablePropertyManager(catalogSettings: IcebergCatalogSettings, catalogFactory: IcebergCatalogFactory)
+    extends IcebergTablePropertyManager(catalogSettings, catalogFactory)
     with StagingPropertyManager
 
 object IcebergTablePropertyManager:
 
-  val sinkLayer =
-    ZLayer {
-      for context <- ZIO.service[PluginStreamContext]
-      yield IcebergSinkTablePropertyManager(context.sink.icebergCatalog)
+  val sinkLayer: ZLayer[PluginStreamContext, Throwable, IcebergSinkTablePropertyManager] =
+    ZLayer.scoped {
+      for
+        context <- ZIO.service[PluginStreamContext]
+        factory <- IcebergCatalogFactory.live(context.sink.icebergCatalog)
+      yield IcebergSinkTablePropertyManager(context.sink.icebergCatalog, factory)
     }
 
-  val stagingLayer =
-    ZLayer {
-      for context <- ZIO.service[PluginStreamContext]
-      yield IcebergStagingTablePropertyManager(context.staging.icebergCatalog)
+  val stagingLayer: ZLayer[PluginStreamContext, Throwable, IcebergStagingTablePropertyManager] =
+    ZLayer.scoped {
+      for
+        context <- ZIO.service[PluginStreamContext]
+        factory <- IcebergCatalogFactory.live(context.staging.icebergCatalog)
+      yield IcebergStagingTablePropertyManager(context.staging.icebergCatalog, factory)
     }

src/main/scala/services/iceberg/base/CatalogEntityManager.scala

@@ -8,10 +8,6 @@ import org.apache.iceberg.Table
 import zio.Task
 
 trait CatalogEntityManager:
-  /** Catalog connection factory
-    */
-  val catalogFactory: CatalogFactory
-
   /** Deletes the specified table from the catalog
     *
     * @param tableName
@@ -38,6 +34,16 @@ trait CatalogEntityManager:
     */
   def migrateSchema(oldSchema: ArcaneSchema, newSchema: ArcaneSchema, tableName: String): Task[Unit]
 
+  /** Retrieve catalog factory used by this entity manager
+    * @return
+    */
+  def getTableRef(tableName: String): Task[Table]
+
+  /** Check if a specified table exists in the catalog
+    * @return
+    */
+  def tableExists(tableName: String): Task[Boolean]
+
 /** Entity manager for sink catalog
   */
 trait SinkEntityManager extends CatalogEntityManager

src/main/scala/services/iceberg/base/CatalogFactory.scala

@@ -3,7 +3,7 @@ package services.iceberg.base
 
 import org.apache.iceberg.catalog.SessionCatalog.SessionContext
 import org.apache.iceberg.rest.RESTSessionCatalog
-import zio.Task
+import zio.{Scope, Task, ZIO}
 
 /** Object responsible for creating catalog clients. Takes care of recycling expired instances and provided the current
   * active one to the caller
@@ -14,11 +14,6 @@ trait CatalogFactory:
     */
   def getSessionContext: SessionContext
 
-  /** Create a new RESTSessionCatalog isntance
-    * @return
-    */
-  def newCatalog: Task[RESTSessionCatalog]
-
   /** Retrieve current active RESTSessionCatalog instance
     * @return
     */