Add balances to ACS snapshots

apps/common/src/main/resources/db/migration/canton-network/postgres/stable/V044__acs_snapshot_balance.sql

@@ -0,0 +1,2 @@
+alter table acs_snapshot add column unlocked_amulet_balance numeric;
+alter table acs_snapshot add column locked_amulet_balance numeric;

apps/common/src/main/scala/org/lfdecentralizedtrust/splice/store/db/AcsJdbcTypes.scala

@@ -19,7 +19,12 @@ import io.circe.parser.parse as circeParse
 import org.lfdecentralizedtrust.splice.store.db.AcsQueries.AcsStoreId
 import org.lfdecentralizedtrust.splice.store.db.TxLogQueries.TxLogStoreId
 import org.lfdecentralizedtrust.splice.util.Contract.Companion
-import org.lfdecentralizedtrust.splice.util.{Contract, LegacyOffset, QualifiedName}
+import org.lfdecentralizedtrust.splice.util.{
+  Contract,
+  LegacyOffset,
+  PackageQualifiedName,
+  QualifiedName,
+}
 import slick.ast.FieldSymbol
 import slick.jdbc.{GetResult, JdbcType, PositionedParameters, PositionedResult, SetParameter}
 import com.digitalasset.canton.resource.DbParameterUtils
@@ -203,6 +208,9 @@ trait AcsJdbcTypes {
       }
   }
 
+  protected implicit lazy val packageQualifiedNameSetParameter: SetParameter[PackageQualifiedName] =
+    SetParameter.SetString.contramap(_.toString)
+
   protected implicit lazy val qualifiedNameSetParameter: SetParameter[QualifiedName] =
     (v1: QualifiedName, v2: PositionedParameters) =>
       implicitly[SetParameter[String2066]].apply(lengthLimited(v1.toString()), v2)

apps/common/src/test/scala/org/lfdecentralizedtrust/splice/store/StoreTest.scala

@@ -320,13 +320,14 @@ abstract class StoreTest extends AsyncWordSpec with BaseTest {
       createdAtRound: Long,
       ratePerRound: BigDecimal,
       version: DarResource = DarResources.amulet_current,
+      dso: PartyId = dsoParty,
   ) = {
     val templateId = new Identifier(
       version.packageId,
       amuletCodegen.LockedAmulet.TEMPLATE_ID.getModuleName,
       amuletCodegen.LockedAmulet.TEMPLATE_ID.getEntityName,
     )
-    val amuletTemplate = amulet(owner, amount, createdAtRound, ratePerRound).payload
+    val amuletTemplate = amulet(owner, amount, createdAtRound, ratePerRound, version, dso).payload
     val template = new amuletCodegen.LockedAmulet(
       amuletTemplate,
       new expiryCodegen.TimeLock(

apps/scan/src/main/scala/org/lfdecentralizedtrust/splice/scan/ScanApp.scala

@@ -188,6 +188,7 @@ class ScanApp(
       acsSnapshotStore = AcsSnapshotStore(
         storage,
         store.updateHistory,
+        dsoParty,
         migrationInfo.currentMigrationId,
         loggerFactory,
       )

apps/scan/src/main/scala/org/lfdecentralizedtrust/splice/scan/store/AcsSnapshotStore.scala

@@ -9,6 +9,8 @@ import org.lfdecentralizedtrust.splice.codegen.java.splice.amulet.{Amulet, Locke
 import org.lfdecentralizedtrust.splice.scan.store.AcsSnapshotStore.{
   AcsSnapshot,
   QueryAcsSnapshotResult,
+  amuletQualifiedName,
+  lockedAmuletQualifiedName,
 }
 import org.lfdecentralizedtrust.splice.store.UpdateHistory.SelectFromCreateEvents
 import org.lfdecentralizedtrust.splice.store.{HardLimit, Limit, LimitHelpers, UpdateHistory}
@@ -33,6 +35,7 @@ import scala.concurrent.{ExecutionContext, Future}
 class AcsSnapshotStore(
     storage: DbStorage,
     val updateHistory: UpdateHistory,
+    dsoParty: PartyId,
     val currentMigrationId: Long,
     override protected val loggerFactory: NamedLoggerFactory,
 )(implicit ec: ExecutionContext, closeContext: CloseContext)
@@ -53,7 +56,7 @@ class AcsSnapshotStore(
   )(implicit tc: TraceContext): Future[Option[AcsSnapshot]] = {
     storage
       .querySingle(
-        sql"""select snapshot_record_time, migration_id, history_id, first_row_id, last_row_id
+        sql"""select snapshot_record_time, migration_id, history_id, first_row_id, last_row_id, unlocked_amulet_balance, locked_amulet_balance
             from acs_snapshot
             where snapshot_record_time <= $before
               and migration_id = $migrationId
@@ -79,7 +82,7 @@ class AcsSnapshotStore(
     }.flatMap { _ =>
       val from = lastSnapshot.map(_.snapshotRecordTime).getOrElse(CantonTimestamp.MinValue)
       val previousSnapshotDataFilter = lastSnapshot match {
-        case Some(AcsSnapshot(_, _, _, firstRowId, lastRowId)) =>
+        case Some(AcsSnapshot(_, _, _, firstRowId, lastRowId, _, _)) =>
           sql"where snapshot.row_id >= $firstRowId and snapshot.row_id <= $lastRowId"
         case None =>
           sql"where false"
@@ -92,12 +95,11 @@ class AcsSnapshotStore(
             and #$tableAlias.record_time < $until
            """
       val statement = (sql"""
-        with inserted_rows as (
             with previous_snapshot_data as (select contract_id
                                             from acs_snapshot_data snapshot
                                                      join update_history_creates creates on snapshot.create_id = creates.row_id
                                             """ ++ previousSnapshotDataFilter ++
-        sql"""),
+        sql"""      ),
                 new_creates as (select contract_id
                                 from update_history_creates creates
                                 """ ++ recordTimeFilter("creates") ++ sql"""
@@ -124,28 +126,37 @@ class AcsSnapshotStore(
                                                           history_id,
                                                           migration_id,
                                                           created_at,
-                                                          creates.contract_id
+                                                          creates.contract_id,
+                                                          create_arguments
                                                    from contracts_to_insert contracts
                                                             join update_history_creates creates
-                                                            on contracts.contract_id = creates.contract_id)
-
-                insert into acs_snapshot_data (create_id, template_id, stakeholder)
-                select row_id,
-                       concat(package_name, ':', template_id_module_name, ':', template_id_entity_name),
-                       stakeholder
-                from creates_to_insert
-                         cross join unnest(array_cat(signatories, observers)) as stakeholders(stakeholder)
-                where history_id = $historyId
-                  and migration_id = $migrationId
-                -- consistent ordering across SVs
-                order by created_at, contract_id
-                returning row_id
-        )
+                                                            on contracts.contract_id = creates.contract_id),
+                inserted_rows as (insert into acs_snapshot_data (create_id, template_id, stakeholder)
+                                  select row_id,
+                                         concat(package_name, ':', template_id_module_name, ':', template_id_entity_name),
+                                         stakeholder
+                                  from creates_to_insert
+                                           cross join unnest(array_cat(signatories, observers)) as stakeholders(stakeholder)
+                                  where history_id = $historyId
+                                    and migration_id = $migrationId
+                                  -- consistent ordering across SVs
+                                  order by created_at, contract_id
+                                  returning row_id, create_id, template_id, stakeholder
+                )
         insert
-        into acs_snapshot (snapshot_record_time, migration_id, history_id, first_row_id, last_row_id)
-        select $until, $migrationId, $historyId, min(row_id), max(row_id)
+        into acs_snapshot (snapshot_record_time, migration_id, history_id, first_row_id, last_row_id, unlocked_amulet_balance, locked_amulet_balance)
+        select
+          $until,
+          $migrationId,
+          $historyId,
+          min(inserted_rows.row_id),
+          max(inserted_rows.row_id),
+          -- the stakeholder filter ensures that we don't double-count amulet amounts
+          sum(case when inserted_rows.template_id = $amuletQualifiedName and stakeholder=$dsoParty then (create_arguments->'record'->'fields'->2->'value'->'record'->'fields'->0->'value'->>'numeric')::numeric else 0 end),
+          sum(case when inserted_rows.template_id = $lockedAmuletQualifiedName and stakeholder=$dsoParty then (create_arguments->'record'->'fields'->0->'value'->'record'->'fields'->2->'value'->'record'->'fields'->0->'value'->>'numeric')::numeric else 0 end)
         from inserted_rows
-        having min(row_id) is not null;
+        join creates_to_insert on inserted_rows.create_id = creates_to_insert.row_id
+        having min(inserted_rows.row_id) is not null;
              """).toActionBuilder.asUpdate
       storage.update(statement, "insertNewSnapshot")
     }.andThen { _ =>
@@ -176,7 +187,7 @@ class AcsSnapshotStore(
     for {
       snapshot <- storage
         .querySingle(
-          sql"""select snapshot_record_time, migration_id, history_id, first_row_id, last_row_id
+          sql"""select snapshot_record_time, migration_id, history_id, first_row_id, last_row_id, unlocked_amulet_balance, locked_amulet_balance
             from acs_snapshot
             where snapshot_record_time = $snapshot
               and migration_id = $migrationId
@@ -350,6 +361,8 @@ object AcsSnapshotStore {
       historyId: Long,
       firstRowId: Long,
       lastRowId: Long,
+      unlockedAmuletBalance: Option[BigDecimal],
+      lockedAmuletBalance: Option[BigDecimal],
   ) extends PrettyPrinting {
     import org.lfdecentralizedtrust.splice.util.PrettyInstances.*
     override def pretty: Pretty[this.type] = prettyOfClass(
@@ -358,6 +371,8 @@ object AcsSnapshotStore {
       param("historyId", _.historyId),
       param("firstRowId", _.firstRowId),
       param("lastRowId", _.lastRowId),
+      param("unlockedAmuletBalance", _.unlockedAmuletBalance),
+      param("lockedAmuletBalance", _.lockedAmuletBalance),
     )
   }
 
@@ -369,6 +384,8 @@ object AcsSnapshotStore {
         historyId = r.<<[Long],
         firstRowId = r.<<[Long],
         lastRowId = r.<<[Long],
+        unlockedAmuletBalance = r.<<[Option[BigDecimal]],
+        lockedAmuletBalance = r.<<[Option[BigDecimal]],
       )
     )
   }
@@ -380,10 +397,11 @@ object AcsSnapshotStore {
       afterToken: Option[Long],
   )
 
-  private val holdingsTemplates =
-    Vector(Amulet.TEMPLATE_ID_WITH_PACKAGE_ID, LockedAmulet.TEMPLATE_ID_WITH_PACKAGE_ID).map(
-      PackageQualifiedName.getFromResources
-    )
+  private val amuletQualifiedName =
+    PackageQualifiedName.getFromResources(Amulet.TEMPLATE_ID_WITH_PACKAGE_ID)
+  private val lockedAmuletQualifiedName =
+    PackageQualifiedName.getFromResources(LockedAmulet.TEMPLATE_ID_WITH_PACKAGE_ID)
+  private val holdingsTemplates = Vector(amuletQualifiedName, lockedAmuletQualifiedName)
 
   private def decodeHoldingContract(createdEvent: CreatedEvent): Either[
     Contract[LockedAmulet.ContractId, LockedAmulet],
@@ -437,11 +455,13 @@ object AcsSnapshotStore {
   def apply(
       storage: Storage,
       updateHistory: UpdateHistory,
+      dsoParty: PartyId,
       migrationId: Long,
       loggerFactory: NamedLoggerFactory,
   )(implicit ec: ExecutionContext, closeContext: CloseContext): AcsSnapshotStore =
     storage match {
-      case db: DbStorage => new AcsSnapshotStore(db, updateHistory, migrationId, loggerFactory)
+      case db: DbStorage =>
+        new AcsSnapshotStore(db, updateHistory, dsoParty, migrationId, loggerFactory)
       case storageType => throw new RuntimeException(s"Unsupported storage type $storageType")
     }
 

apps/scan/src/test/scala/org/lfdecentralizedtrust/splice/scan/automation/AcsSnapshotTriggerTest.scala

@@ -628,7 +628,7 @@ class AcsSnapshotTriggerTest
         time: CantonTimestamp,
         migrationId: Long = currentMigrationId,
     ): AcsSnapshot = {
-      val lastSnapshot = AcsSnapshot(time, migrationId, historyId, 0, 100)
+      val lastSnapshot = AcsSnapshot(time, migrationId, historyId, 0, 100, None, None)
       when(
         store.lookupSnapshotBefore(eqTo(migrationId), eqTo(CantonTimestamp.MaxValue))(
           any[TraceContext]