From 4e46c0dc656a2da259ee296cf68012dc198d265a Mon Sep 17 00:00:00 2001 From: David An Date: Thu, 5 Dec 2024 15:22:58 -0500 Subject: [PATCH 1/3] connect batchUpsert tracing contexts --- .../dataaccess/slick/AttributeComponent.scala | 134 +++++++++--------- .../dataaccess/slick/EntityComponent.scala | 14 +- .../entities/local/LocalEntityProvider.scala | 10 +- 3 files changed, 83 insertions(+), 75 deletions(-) diff --git a/core/src/main/scala/org/broadinstitute/dsde/rawls/dataaccess/slick/AttributeComponent.scala b/core/src/main/scala/org/broadinstitute/dsde/rawls/dataaccess/slick/AttributeComponent.scala index 9250151c55..220c36f34c 100644 --- a/core/src/main/scala/org/broadinstitute/dsde/rawls/dataaccess/slick/AttributeComponent.scala +++ b/core/src/main/scala/org/broadinstitute/dsde/rawls/dataaccess/slick/AttributeComponent.scala @@ -695,81 +695,81 @@ trait AttributeComponent { */ def rewriteAttrsAction(attributesToSave: Traversable[RECORD], existingAttributes: Traversable[RECORD], - insertFunction: Seq[RECORD] => () => WriteAction[Int] + insertFunction: Seq[RECORD] => () => WriteAction[Int], + parentContext: RawlsTracingContext = RawlsTracingContext() ): ReadWriteAction[Set[OWNER_ID]] = - traceDBIOWithParent("AttributeComponent.rewriteAttrsAction", RawlsTracingContext(Option.empty)) { - tracingContext => - val toSaveAttrMap = toPrimaryKeyMap(attributesToSave) - val existingAttrMap = toPrimaryKeyMap(existingAttributes) - - // note that currently-existing attributes will have a populated id e.g. "1234", but to-save will have an id of "0" - // therefore, we use this ComparableRecord class which omits the id when checking equality between existing and to-save. - object ComparableRecord { - def fromRecord(rec: RECORD): ComparableRecord = - new ComparableRecord( - rec.ownerId, - rec.namespace, - rec.name, - rec.valueString, - rec.valueNumber, - rec.valueBoolean, - rec.valueJson, - rec.valueEntityRef, - rec.listIndex, - rec.listLength, - rec.deleted, - rec.deletedDate - ) - } + traceDBIOWithParent("AttributeComponent.rewriteAttrsAction", parentContext) { tracingContext => + val toSaveAttrMap = toPrimaryKeyMap(attributesToSave) + val existingAttrMap = toPrimaryKeyMap(existingAttributes) + + // note that currently-existing attributes will have a populated id e.g. "1234", but to-save will have an id of "0" + // therefore, we use this ComparableRecord class which omits the id when checking equality between existing and to-save. + object ComparableRecord { + def fromRecord(rec: RECORD): ComparableRecord = + new ComparableRecord( + rec.ownerId, + rec.namespace, + rec.name, + rec.valueString, + rec.valueNumber, + rec.valueBoolean, + rec.valueJson, + rec.valueEntityRef, + rec.listIndex, + rec.listLength, + rec.deleted, + rec.deletedDate + ) + } - case class ComparableRecord( - ownerId: OWNER_ID, - namespace: String, - name: String, - valueString: Option[String], - valueNumber: Option[Double], - valueBoolean: Option[Boolean], - valueJson: Option[String], - valueEntityRef: Option[Long], - listIndex: Option[Int], - listLength: Option[Int], - deleted: Boolean, - deletedDate: Option[Timestamp] - ) + case class ComparableRecord( + ownerId: OWNER_ID, + namespace: String, + name: String, + valueString: Option[String], + valueNumber: Option[Double], + valueBoolean: Option[Boolean], + valueJson: Option[String], + valueEntityRef: Option[Long], + listIndex: Option[Int], + listLength: Option[Int], + deleted: Boolean, + deletedDate: Option[Timestamp] + ) - // create a set of ComparableRecords representing the existing attributes - val existingAttributesSet: Set[ComparableRecord] = - existingAttributes.toSet.map(r => ComparableRecord.fromRecord(r)) + // create a set of ComparableRecords representing the existing attributes + val existingAttributesSet: Set[ComparableRecord] = + existingAttributes.toSet.map(r => ComparableRecord.fromRecord(r)) - val existingKeys = existingAttrMap.keySet + val existingKeys = existingAttrMap.keySet - // insert attributes which are in save but not exists - val attributesToInsert = toSaveAttrMap.filterKeys(!existingKeys.contains(_)) + // insert attributes which are in save but not exists + val attributesToInsert = toSaveAttrMap.filterKeys(!existingKeys.contains(_)) - // delete attributes which are in exists but not save - val attributesToDelete = existingAttrMap.filterKeys(!toSaveAttrMap.keySet.contains(_)) + // delete attributes which are in exists but not save + val attributesToDelete = existingAttrMap.filterKeys(!toSaveAttrMap.keySet.contains(_)) - val attributesToUpdate = toSaveAttrMap.filter { case (k, v) => - existingKeys.contains(k) && // if the attribute doesn't already exist, don't attempt to update it - !existingAttributesSet.contains( - ComparableRecord.fromRecord(v) - ) // if the attribute exists and is unchanged, don't update it - } + val attributesToUpdate = toSaveAttrMap.filter { case (k, v) => + existingKeys.contains(k) && // if the attribute doesn't already exist, don't attempt to update it + !existingAttributesSet.contains( + ComparableRecord.fromRecord(v) + ) // if the attribute exists and is unchanged, don't update it + } - // collect the parent objects (e.g. entity, workspace) that have writes, so we know which object rows to re-calculate - val ownersWithWrites: Set[OWNER_ID] = (attributesToInsert.values.map(_.ownerId) ++ - attributesToUpdate.values.map(_.ownerId) ++ - attributesToDelete.values.map(_.ownerId)).toSet - - // perform the inserts/updates/deletes - patchAttributesAction( - attributesToInsert.values, - attributesToUpdate.values, - attributesToDelete.values.map(_.id), - insertFunction, - tracingContext - ) - .map(_ => ownersWithWrites) + // collect the parent objects (e.g. entity, workspace) that have writes, so we know which object rows to re-calculate + val ownersWithWrites: Set[OWNER_ID] = (attributesToInsert.values.map(_.ownerId) ++ + attributesToUpdate.values.map(_.ownerId) ++ + attributesToDelete.values.map(_.ownerId)).toSet + + // perform the inserts/updates/deletes + patchAttributesAction( + attributesToInsert.values, + attributesToUpdate.values, + attributesToDelete.values.map(_.id), + insertFunction, + tracingContext + ) + .map(_ => ownersWithWrites) } // noinspection SqlDialectInspection diff --git a/core/src/main/scala/org/broadinstitute/dsde/rawls/dataaccess/slick/EntityComponent.scala b/core/src/main/scala/org/broadinstitute/dsde/rawls/dataaccess/slick/EntityComponent.scala index 3650186422..6d571c0e61 100644 --- a/core/src/main/scala/org/broadinstitute/dsde/rawls/dataaccess/slick/EntityComponent.scala +++ b/core/src/main/scala/org/broadinstitute/dsde/rawls/dataaccess/slick/EntityComponent.scala @@ -1026,7 +1026,10 @@ trait EntityComponent { def save(workspaceContext: Workspace, entity: Entity): ReadWriteAction[Entity] = save(workspaceContext, Seq(entity)).map(_.head) - def save(workspaceContext: Workspace, entities: Traversable[Entity]): ReadWriteAction[Traversable[Entity]] = { + def save(workspaceContext: Workspace, + entities: Traversable[Entity], + parentContext: RawlsTracingContext = RawlsTracingContext() + ): ReadWriteAction[Traversable[Entity]] = { entities.foreach(validateEntity) for { @@ -1044,7 +1047,8 @@ trait EntityComponent { workspaceContext.workspaceIdAsUUID, entities, savingEntityRecs.map(_.id), - referencedAndSavingEntityRecs.map(e => e.toReference -> e.id).toMap + referencedAndSavingEntityRecs.map(e => e.toReference -> e.id).toMap, + parentContext ) // find the pre-existing records that we updated actuallyUpdatedPreExistingEntityRecs = preExistingEntityRecs.filter(e => @@ -1112,7 +1116,8 @@ trait EntityComponent { private def rewriteAttributes(workspaceId: UUID, entitiesToSave: Traversable[Entity], entityIds: Seq[Long], - entityIdsByRef: Map[AttributeEntityReference, Long] + entityIdsByRef: Map[AttributeEntityReference, Long], + parentContext: RawlsTracingContext = RawlsTracingContext() ) = { val attributesToSave = for { entity <- entitiesToSave @@ -1127,7 +1132,8 @@ trait EntityComponent { entityAttributeShardQuery(workspaceId).findByOwnerQuery(entityIds).result flatMap { existingAttributes => entityAttributeShardQuery(workspaceId).rewriteAttrsAction(attributesToSave, existingAttributes, - entityAttributeTempQuery.insertScratchAttributes + entityAttributeTempQuery.insertScratchAttributes, + parentContext ) } } diff --git a/core/src/main/scala/org/broadinstitute/dsde/rawls/entities/local/LocalEntityProvider.scala b/core/src/main/scala/org/broadinstitute/dsde/rawls/entities/local/LocalEntityProvider.scala index 8c95b29038..6417e93c45 100644 --- a/core/src/main/scala/org/broadinstitute/dsde/rawls/entities/local/LocalEntityProvider.scala +++ b/core/src/main/scala/org/broadinstitute/dsde/rawls/entities/local/LocalEntityProvider.scala @@ -497,13 +497,15 @@ class LocalEntityProvider(requestArguments: EntityRequestArguments, } else { val t = updateTrials.collect { case (entityUpdate, Success(entity)) => entity } - dataAccess.entityQuery - .save(workspaceContext, t) - .withStatementParameters(statementInit = _.setQueryTimeout(queryTimeoutSeconds)) + traceDBIOWithParent("saveAction", localContext) { subContext => + dataAccess.entityQuery + .save(workspaceContext, t, subContext.toTracingContext) + .withStatementParameters(statementInit = _.setQueryTimeout(queryTimeoutSeconds)) + } } } - traceDBIOWithParent("saveAction", localContext)(_ => saveAction) + saveAction } recover { case icve: java.sql.SQLIntegrityConstraintViolationException => val userMessage = From e4bb6b7b1b49986f18c3f7c9fe978e13b3d251a0 Mon Sep 17 00:00:00 2001 From: David An Date: Thu, 5 Dec 2024 15:40:13 -0500 Subject: [PATCH 2/3] fix ordering of traces in updateAction --- .../dsde/rawls/dataaccess/slick/AttributeComponent.scala | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/core/src/main/scala/org/broadinstitute/dsde/rawls/dataaccess/slick/AttributeComponent.scala b/core/src/main/scala/org/broadinstitute/dsde/rawls/dataaccess/slick/AttributeComponent.scala index 220c36f34c..b4f067cc4e 100644 --- a/core/src/main/scala/org/broadinstitute/dsde/rawls/dataaccess/slick/AttributeComponent.scala +++ b/core/src/main/scala/org/broadinstitute/dsde/rawls/dataaccess/slick/AttributeComponent.scala @@ -804,8 +804,10 @@ trait AttributeComponent { def updateAction(insertIntoScratchFunction: () => WriteAction[Int], tracingContext: RawlsTracingContext) = traceDBIOWithParent("updateAction", tracingContext) { span => - traceDBIOWithParent("insertIntoScratchFunction", span)(_ => insertIntoScratchFunction()) andThen - traceDBIOWithParent("updateInMasterAction", span)(_ => updateInMasterAction()) + for { + _ <- traceDBIOWithParent("insertIntoScratchFunction", span)(_ => insertIntoScratchFunction()) + numUpdates <- traceDBIOWithParent("updateInMasterAction", span)(_ => updateInMasterAction()) + } yield numUpdates } } From 28b516dfa85af8b20b4be4bc0014a8fa98eab6d8 Mon Sep 17 00:00:00 2001 From: David An Date: Thu, 5 Dec 2024 15:45:56 -0500 Subject: [PATCH 3/3] add more detail --- .../dataaccess/slick/EntityComponent.scala | 28 ++++++++++++------- 1 file changed, 18 insertions(+), 10 deletions(-) diff --git a/core/src/main/scala/org/broadinstitute/dsde/rawls/dataaccess/slick/EntityComponent.scala b/core/src/main/scala/org/broadinstitute/dsde/rawls/dataaccess/slick/EntityComponent.scala index 6d571c0e61..94407ce768 100644 --- a/core/src/main/scala/org/broadinstitute/dsde/rawls/dataaccess/slick/EntityComponent.scala +++ b/core/src/main/scala/org/broadinstitute/dsde/rawls/dataaccess/slick/EntityComponent.scala @@ -1033,15 +1033,21 @@ trait EntityComponent { entities.foreach(validateEntity) for { - _ <- workspaceQuery.updateLastModified(workspaceContext.workspaceIdAsUUID) - preExistingEntityRecs <- getEntityRecords(workspaceContext.workspaceIdAsUUID, entities.map(_.toReference).toSet) - savingEntityRecs <- entityQueryWithInlineAttributes - .insertNewEntities(workspaceContext, entities, preExistingEntityRecs.map(_.toReference)) - .map(_ ++ preExistingEntityRecs) - referencedAndSavingEntityRecs <- lookupNotYetLoadedReferences(workspaceContext, - entities, - savingEntityRecs.map(_.toReference) - ).map(_ ++ savingEntityRecs) + _ <- traceDBIOWithParent("updateLastModified", parentContext)(_ => + workspaceQuery.updateLastModified(workspaceContext.workspaceIdAsUUID) + ) + preExistingEntityRecs <- traceDBIOWithParent("getEntityRecords", parentContext)(_ => + getEntityRecords(workspaceContext.workspaceIdAsUUID, entities.map(_.toReference).toSet) + ) + savingEntityRecs <- traceDBIOWithParent("insertNewEntities", parentContext)(_ => + entityQueryWithInlineAttributes + .insertNewEntities(workspaceContext, entities, preExistingEntityRecs.map(_.toReference)) + .map(_ ++ preExistingEntityRecs) + ) + referencedAndSavingEntityRecs <- traceDBIOWithParent("lookupNotYetLoadedReferences", parentContext)(_ => + lookupNotYetLoadedReferences(workspaceContext, entities, savingEntityRecs.map(_.toReference)) + .map(_ ++ savingEntityRecs) + ) actuallyUpdatedEntityIds <- rewriteAttributes( workspaceContext.workspaceIdAsUUID, @@ -1068,7 +1074,9 @@ trait EntityComponent { // 2) were repeated in the input payload, causing one insert and subsequent update(s) recsToUpdate = (actuallyUpdatedPreExistingEntityRecs ++ insertedRepeats).distinct - _ <- entityQueryWithInlineAttributes.optimisticLockUpdate(recsToUpdate, entities) + _ <- traceDBIOWithParent("optimisticLockUpdate", parentContext)(_ => + entityQueryWithInlineAttributes.optimisticLockUpdate(recsToUpdate, entities) + ) } yield entities }