From 4e2deb8d37a0c095a7c0834cecbe0a0a0e5dcced Mon Sep 17 00:00:00 2001 From: Daniel Thom Date: Thu, 9 Jan 2025 08:30:47 -0700 Subject: [PATCH] Fix bulk update of supplemental attributes --- src/supplemental_attribute_manager.jl | 27 ++++++++- src/system_data.jl | 10 ---- test/test_supplemental_attributes.jl | 83 +++++++++++++++++++++++++++ 3 files changed, 109 insertions(+), 11 deletions(-) diff --git a/src/supplemental_attribute_manager.jl b/src/supplemental_attribute_manager.jl index 58fb64ce8..50808f3fb 100644 --- a/src/supplemental_attribute_manager.jl +++ b/src/supplemental_attribute_manager.jl @@ -1,7 +1,7 @@ const SupplementalAttributesByType = Dict{DataType, Dict{Base.UUID, <:SupplementalAttribute}} -struct SupplementalAttributeManager <: InfrastructureSystemsContainer +mutable struct SupplementalAttributeManager <: InfrastructureSystemsContainer data::SupplementalAttributesByType associations::SupplementalAttributeAssociations end @@ -15,6 +15,31 @@ end get_member_string(::SupplementalAttributeManager) = "supplemental attributes" +""" +Begin an update of supplemental attributes. Use this function when adding +or removing many supplemental attributes in order to improve performance. + +If an error occurs during the update, changes will be reverted. +""" +function begin_supplemental_attributes_update( + func::Function, + mgr::SupplementalAttributeManager, +) + orig_data = SupplementalAttributesByType() + for (key, val) in mgr.data + orig_data[key] = copy(val) + end + + try + SQLite.transaction(mgr.associations.db) do + func() + end + catch + mgr.data = orig_data + rethrow() + end +end + function add_supplemental_attribute!( mgr::SupplementalAttributeManager, component::InfrastructureSystemsComponent, diff --git a/src/system_data.jl b/src/system_data.jl index 2d4b99789..a9219beca 100644 --- a/src/system_data.jl +++ b/src/system_data.jl @@ -1152,16 +1152,6 @@ function add_supplemental_attribute!(data::SystemData, component, attribute; kwa return end -""" -Begin a transaction to add supplemental attributes. Use this function when adding -many supplemental attributes in order to improve performance. -""" -function begin_supplemental_attributes_transaction(func::Function, data::SystemData) - SQLite.transaction(data.supplemental_attribute_manager.associations.db) do - func() - end -end - function get_supplemental_attributes( filter_func::Function, ::Type{T}, diff --git a/test/test_supplemental_attributes.jl b/test/test_supplemental_attributes.jl index 090f70926..2d10546c4 100644 --- a/test/test_supplemental_attributes.jl +++ b/test/test_supplemental_attributes.jl @@ -13,6 +13,89 @@ ) end +@testset "Test bulk addition of supplemental attributes" begin + mgr = IS.SupplementalAttributeManager() + attr1 = IS.GeographicInfo(; geo_json = Dict("x" => 1.0)) + attr2 = IS.GeographicInfo(; geo_json = Dict("x" => 2.0)) + component = IS.TestComponent("component1", 1) + IS.begin_supplemental_attributes_update(mgr) do + IS.add_supplemental_attribute!(mgr, component, attr1) + IS.add_supplemental_attribute!(mgr, component, attr2) + end + @test length(mgr.data) == 1 + @test length(mgr.data[IS.GeographicInfo]) == 2 + @test IS.get_num_attributes(mgr.associations) == 2 +end + +@testset "Test bulk addition of supplemental attributes with error" begin + mgr = IS.SupplementalAttributeManager() + attr1 = IS.TestSupplemental(; value = 1.0) + attr2 = IS.GeographicInfo(; geo_json = Dict("x" => 2.0)) + component = IS.TestComponent("component1", 1) + @test_throws( + ArgumentError, + IS.begin_supplemental_attributes_update(mgr) do + IS.add_supplemental_attribute!(mgr, component, attr1) + IS.add_supplemental_attribute!(mgr, component, attr2) + IS.add_supplemental_attribute!(mgr, component, attr2) + end, + ) + @test length(mgr.data) == 0 + @test IS.get_num_attributes(mgr.associations) == 0 +end + +@testset "Test bulk addition of supplemental attributes with error, existing attrs" begin + mgr = IS.SupplementalAttributeManager() + attr1 = IS.TestSupplemental(; value = 1.0) + attr2 = IS.GeographicInfo(; geo_json = Dict("x" => 2.0)) + component = IS.TestComponent("component1", 1) + IS.begin_supplemental_attributes_update(mgr) do + IS.add_supplemental_attribute!(mgr, component, attr1) + IS.add_supplemental_attribute!(mgr, component, attr2) + end + + attr3 = IS.TestSupplemental(; value = 3.0) + attr4 = IS.GeographicInfo(; geo_json = Dict("x" => 3.0)) + @test_throws( + ArgumentError, + IS.begin_supplemental_attributes_update(mgr) do + IS.add_supplemental_attribute!(mgr, component, attr3) + IS.add_supplemental_attribute!(mgr, component, attr4) + IS.add_supplemental_attribute!(mgr, component, attr4) + end, + ) + @test length(mgr.data) == 2 + @test length(mgr.data[IS.TestSupplemental]) == 1 + @test length(mgr.data[IS.GeographicInfo]) == 1 + @test IS.get_num_attributes(mgr.associations) == 2 +end + +@testset "Test bulk removal of supplemental attributes with error" begin + mgr = IS.SupplementalAttributeManager() + attr1 = IS.TestSupplemental(; value = 1.0) + attr2 = IS.TestSupplemental(; value = 2.0) + attr3 = IS.GeographicInfo(; geo_json = Dict("x" => 3.0)) + component = IS.TestComponent("component1", 1) + IS.begin_supplemental_attributes_update(mgr) do + IS.add_supplemental_attribute!(mgr, component, attr1) + IS.add_supplemental_attribute!(mgr, component, attr2) + IS.add_supplemental_attribute!(mgr, component, attr3) + end + + @test_throws( + ArgumentError, + IS.begin_supplemental_attributes_update(mgr) do + IS.remove_supplemental_attribute!(mgr, component, attr2) + IS.remove_supplemental_attribute!(mgr, component, attr3) + IS.remove_supplemental_attribute!(mgr, component, attr3) + end, + ) + @test length(mgr.data) == 2 + @test length(mgr.data[IS.TestSupplemental]) == 2 + @test length(mgr.data[IS.GeographicInfo]) == 1 + @test IS.get_num_attributes(mgr.associations) == 3 +end + @testset "Test clear_supplemental_attributes" begin data = IS.SystemData(; time_series_in_memory = true) geo_supplemental_attribute = IS.GeographicInfo()