Speed up accessors for time series and supplemental attributes

src/supplemental_attribute_associations.jl

@@ -23,16 +23,21 @@
 
 mutable struct SupplementalAttributeAssociations
     db::SQLite.DB
+    # If we don't cache SQL statements, there is a cost of 3-4 us on every query.
+    cached_statements::Dict{String, SQLite.Stmt}
     # If you add any fields, ensure they are managed in deepcopy_internal below.
 end
 
 """
 Construct a new SupplementalAttributeAssociations with an in-memory database.
 """
-function SupplementalAttributeAssociations()
-    associations = SupplementalAttributeAssociations(SQLite.DB())
+function SupplementalAttributeAssociations(; create_indexes = true)
+    associations =
+        SupplementalAttributeAssociations(SQLite.DB(), Dict{String, SQLite.Stmt}())
     _create_attribute_associations_table!(associations)
-    _create_indexes!(associations)
+    if create_indexes
+        _create_indexes!(associations)
+    end
     @debug "Initialized new supplemental attributes association table" _group =
         LOG_GROUP_SUPPLEMENTAL_ATTRIBUTES
     return associations
@@ -61,18 +66,23 @@
     SQLite.createindex!(
         associations.db,
         SUPPLEMENTAL_ATTRIBUTE_TABLE_NAME,
-        "by_attribute_and_component",
+        "by_attribute",
         [
             "attribute_uuid",
             "component_uuid",
+            "component_type",
         ];
         unique = false,
     )
     SQLite.createindex!(
         associations.db,
         SUPPLEMENTAL_ATTRIBUTE_TABLE_NAME,
         "by_component",
-        "component_uuid";
+        [
+            "component_uuid",
+            "attribute_uuid",
+            "attribute_type",
+        ];
         unique = false,
     )
     return
@@ -84,7 +94,8 @@
     end
     new_db = SQLite.DB()
     backup(new_db, val.db)
-    new_associations = SupplementalAttributeAssociations(new_db)
+    new_associations =
+        SupplementalAttributeAssociations(new_db, Dict{String, SQLite.Stmt}())
     dict[val] = new_associations
     return new_associations
 end
@@ -106,8 +117,8 @@
             string(nameof(typeof(component))),
         )
         params = chop(repeat("?,", length(row)))
-        SQLite.DBInterface.execute(
-            associations.db,
+        _execute_cached(
+            associations,
             "INSERT INTO $SUPPLEMENTAL_ATTRIBUTE_TABLE_NAME VALUES($params)",
             row,
         )
@@ -189,69 +200,102 @@
     return _execute_count(associations, query)
 end
 
+const _HAS_ASSOCIATION_BY_ATTRIBUTE = """
+    SELECT attribute_uuid
+    FROM $SUPPLEMENTAL_ATTRIBUTE_TABLE_NAME
+    WHERE attribute_uuid = ?
+    LIMIT 1
+"""
+
 """
 Return true if there is at least one association matching the inputs.
 """
 function has_association(
     associations::SupplementalAttributeAssociations,
     attribute::SupplementalAttribute,
 )
-    return _has_association(associations, attribute, "attribute_uuid")
+    # Note: Unlike the other has_association methods, this is not covered by an index.
+    params = (string(get_uuid(attribute)),)
+    return !isempty(
+        Tables.rowtable(
+            _execute_cached(associations, _HAS_ASSOCIATION_BY_ATTRIBUTE, params),
+        ),
+    )
 end
 
+const _HAS_ASSOCIATION_BY_COMPONENT_ATTRIBUTE = """
+    SELECT attribute_uuid
+    FROM $SUPPLEMENTAL_ATTRIBUTE_TABLE_NAME
+    WHERE attribute_uuid = ? AND component_uuid = ?
+    LIMIT 1
+"""
 function has_association(
     associations::SupplementalAttributeAssociations,
     component::InfrastructureSystemsComponent,
     attribute::SupplementalAttribute,
 )
     a_uuid = get_uuid(attribute)
     c_uuid = get_uuid(component)
-    query = """
-        SELECT attribute_uuid
-        FROM $SUPPLEMENTAL_ATTRIBUTE_TABLE_NAME
-        WHERE attribute_uuid = ? AND component_uuid = ?
-        LIMIT 1
-    """
-    params = [string(a_uuid), string(c_uuid)]
-    return !isempty(Tables.rowtable(_execute(associations, query, params)))
+    params = (string(a_uuid), string(c_uuid))
+    return !isempty(
+        _execute_cached(associations, _HAS_ASSOCIATION_BY_COMPONENT_ATTRIBUTE, params),
+    )
 end
 
+const _HAS_ASSOCIATION_BY_COMPONENT = """
+    SELECT attribute_uuid
+    FROM $SUPPLEMENTAL_ATTRIBUTE_TABLE_NAME
+    WHERE component_uuid = ?
+    LIMIT 1
+"""
 function has_association(
     associations::SupplementalAttributeAssociations,
     component::InfrastructureSystemsComponent,
-    attribute_type::Union{Nothing, Type{<:SupplementalAttribute}} = nothing,
 )
-    c_str = "component_uuid = ?"
-    params = [string(get_uuid(component))]
-    where_clause = if isnothing(attribute_type)
-        c_str
-    else
-        a_str = _get_attribute_type_string!(params, attribute_type)
-        "$c_str AND $a_str"
-    end
+    params = (string(get_uuid(component)),)
+    return !isempty(
+        Tables.rowtable(
+            _execute_cached(associations, _HAS_ASSOCIATION_BY_COMPONENT, params),
+        ),
+    )
+end
 
-    query = """
-        SELECT attribute_uuid
-        FROM $SUPPLEMENTAL_ATTRIBUTE_TABLE_NAME
-        WHERE $where_clause
-        LIMIT 1
-    """
-    return !isempty(Tables.rowtable(_execute(associations, query, params)))
+const _HAS_ASSOCIATION_BY_COMP_ATTR_TYPE = """
+    SELECT attribute_uuid
+    FROM $SUPPLEMENTAL_ATTRIBUTE_TABLE_NAME
+    WHERE component_uuid = ? AND attribute_type = ?
+    LIMIT 1
+"""
+function has_association(
+    associations::SupplementalAttributeAssociations,
+    component::InfrastructureSystemsComponent,
+    attribute_type::Type{<:SupplementalAttribute},
+)
+    params = (string(get_uuid(component)), string(nameof(attribute_type)))
+    return !isempty(
+        Tables.rowtable(
+            _execute_cached(associations, _HAS_ASSOCIATION_BY_COMP_ATTR_TYPE, params),
+        ),
+    )
 end
 
+const _LIST_ASSOCIATED_COMP_UUIDS = """
+    SELECT component_uuid
+    FROM $SUPPLEMENTAL_ATTRIBUTE_TABLE_NAME
+    WHERE attribute_uuid = ?
+"""
+
 """
 Return the component UUIDs associated with the attribute.
 """
 function list_associated_component_uuids(
     associations::SupplementalAttributeAssociations,
     attribute::SupplementalAttribute,
 )
-    query = """
-        SELECT component_uuid
-        FROM $SUPPLEMENTAL_ATTRIBUTE_TABLE_NAME
-        WHERE attribute_uuid = '$(get_uuid(attribute))'
-    """
-    table = Tables.columntable(_execute(associations, query))
+    params = (string(get_uuid(attribute)),)
+    table = Tables.columntable(
+        _execute_cached(associations, _LIST_ASSOCIATED_COMP_UUIDS, params),
+    )
     return Base.UUID.(table.component_uuid)
 end
 
@@ -276,7 +320,7 @@
         FROM $SUPPLEMENTAL_ATTRIBUTE_TABLE_NAME
         WHERE $where_clause
     """
-    table = Tables.columntable(_execute(associations, query, params))
+    table = Tables.columntable(_execute_cached(associations, query, params))
     return Base.UUID.(table.attribute_uuid)
 end
 
@@ -288,10 +332,9 @@
     component::InfrastructureSystemsComponent,
     attribute::SupplementalAttribute,
 )
-    a_uuid = get_uuid(attribute)
-    c_uuid = get_uuid(component)
-    where_clause = "WHERE attribute_uuid = '$a_uuid' AND component_uuid = '$c_uuid'"
-    num_deleted = _remove_associations!(associations, where_clause)
+    where_clause = "WHERE attribute_uuid = ? AND component_uuid = ?"
+    params = (string(get_uuid(attribute)), string(get_uuid(component)))
+    num_deleted = _remove_associations!(associations, where_clause, params)
     if num_deleted != 1
         error("Bug: unexpected number of deletions: $num_deleted. Should have been 1.")
     end
@@ -304,13 +347,20 @@
     associations::SupplementalAttributeAssociations,
     type::Type{<:SupplementalAttribute},
 )
-    where_clause = "WHERE attribute_type = '$(nameof(type))'"
-    num_deleted = _remove_associations!(associations, where_clause)
+    where_clause = "WHERE attribute_type = ?"
+    params = (string(nameof(type)),)
+    num_deleted = _remove_associations!(associations, where_clause, params)
     @debug "Deleted $num_deleted supplemental attribute associations" _group =
         LOG_GROUP_SUPPLEMENTAL_ATTRIBUTES
     return
 end
 
+const _REPLACE_COMP_UUID_SA = """
+    UPDATE $SUPPLEMENTAL_ATTRIBUTE_TABLE_NAME
+    SET component_uuid = ?
+    WHERE component_uuid = ?
+"""
+
 """
 Replace the component UUID in the table.
 """
@@ -319,12 +369,8 @@
     old_uuid::Base.UUID,
     new_uuid::Base.UUID,
 )
-    query = """
-        UPDATE $SUPPLEMENTAL_ATTRIBUTE_TABLE_NAME
-        SET component_uuid = '$new_uuid'
-        WHERE component_uuid = '$old_uuid'
-    """
-    _execute(associations, query)
+    params = (string(new_uuid), string(old_uuid))
+    _execute_cached(associations, _REPLACE_COMP_UUID_SA, params)
     return
 end
 
@@ -350,8 +396,9 @@
 """
 Add records to the database. Expects output from [`to_records`](@ref).
 """
-function load_records!(associations::SupplementalAttributeAssociations, records)
-    isempty(records) && return
+function from_records(::Type{SupplementalAttributeAssociations}, records)
+    associations = SupplementalAttributeAssociations(; create_indexes = false)
+    isempty(records) && return associations
 
     columns = ("attribute_uuid", "attribute_type", "component_uuid", "component_type")
     num_rows = length(records)
@@ -368,31 +415,19 @@
         "INSERT INTO $SUPPLEMENTAL_ATTRIBUTE_TABLE_NAME VALUES($params)",
         NamedTuple(Symbol(k) => v for (k, v) in data),
     )
-    return
-end
-
-function _has_association(
-    associations::SupplementalAttributeAssociations,
-    val::Union{InfrastructureSystemsComponent, SupplementalAttribute},
-    column::String,
-)
-    uuid = get_uuid(val)
-    query = """
-        SELECT attribute_uuid
-        FROM $SUPPLEMENTAL_ATTRIBUTE_TABLE_NAME
-        WHERE $column = '$uuid'
-        LIMIT 1
-    """
-    return !isempty(Tables.rowtable(_execute(associations, query)))
+    _create_indexes!(associations)
+    return associations
 end
 
 function _remove_associations!(
     associations::SupplementalAttributeAssociations,
     where_clause::AbstractString,
+    params,
 )
-    _execute(
+    _execute_cached(
         associations,
         "DELETE FROM $SUPPLEMENTAL_ATTRIBUTE_TABLE_NAME $where_clause",
+        params,
     )
     table = Tables.rowtable(_execute(associations, "SELECT CHANGES() AS changes"))
     @assert_op length(table) == 1
@@ -420,6 +455,16 @@
     return match_fn(table_x, table_y)
 end
 
+function _make_stmt(associations::SupplementalAttributeAssociations, query::String)
+    return get!(
+        () -> SQLite.Stmt(associations.db, query),
+        associations.cached_statements,
+        query,
+    )
+end
+
+_execute_cached(s::SupplementalAttributeAssociations, q, p = nothing) =
+    execute(_make_stmt(s, q), p, LOG_GROUP_TIME_SERIES)
 _execute(s::SupplementalAttributeAssociations, q, p = nothing) =
     execute(s.db, q, p, LOG_GROUP_SUPPLEMENTAL_ATTRIBUTES)
 _execute_count(s::SupplementalAttributeAssociations, q, p = nothing) =

src/supplemental_attribute_manager.jl

@@ -1,21 +1,45 @@
 const SupplementalAttributesByType =
     Dict{DataType, Dict{Base.UUID, <:SupplementalAttribute}}
 
-struct SupplementalAttributeManager <: InfrastructureSystemsContainer
+mutable struct SupplementalAttributeManager <: InfrastructureSystemsContainer
     data::SupplementalAttributesByType
     associations::SupplementalAttributeAssociations
 end
 
-function SupplementalAttributeManager(data::SupplementalAttributesByType)
-    return SupplementalAttributeManager(data, SupplementalAttributeAssociations())
-end
-
 function SupplementalAttributeManager()
-    return SupplementalAttributeManager(SupplementalAttributesByType())
+    return SupplementalAttributeManager(
+        SupplementalAttributesByType(),
+        SupplementalAttributeAssociations(),
+    )
 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,
@@ -244,7 +268,9 @@ function deserialize(
         @debug "Deserialized $(summary(attr))" _group = LOG_GROUP_SERIALIZATION
     end
 
-    mgr = SupplementalAttributeManager(SupplementalAttributesByType(attributes))
-    load_records!(mgr.associations, data["associations"])
+    mgr = SupplementalAttributeManager(
+        SupplementalAttributesByType(attributes),
+        from_records(SupplementalAttributeAssociations, data["associations"]),
+    )
     return mgr
 end