Fix bulk update of supplemental attributes

NREL-Sienna · Jan 10, 2025 · 4e2deb8 · 4e2deb8
1 parent 6887658
commit 4e2deb8
Show file tree

Hide file tree

Showing 3 changed files with 109 additions and 11 deletions.
diff --git 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
@@ -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
@@ -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()