From 6e037b8523734b52f86a91a094880b2e8b2e4ac7 Mon Sep 17 00:00:00 2001 From: pesap Date: Wed, 11 Feb 2026 15:12:03 -0700 Subject: [PATCH 1/2] fix: Properly resolve system name and add system_name for bulk insertion Fixes: #94, #95 --- src/plexosdb/db.py | 40 ++++-- src/plexosdb/utils.py | 174 +++++++++++++++++++++++++-- tests/test_plexosdb_delete_object.py | 33 +++++ tests/test_plexosdb_from_records.py | 21 ++++ 4 files changed, 248 insertions(+), 20 deletions(-) diff --git a/src/plexosdb/db.py b/src/plexosdb/db.py index 86d48a4..04bf7e5 100644 --- a/src/plexosdb/db.py +++ b/src/plexosdb/db.py @@ -32,11 +32,13 @@ from .utils import ( apply_scenario_tags, create_membership_record, + get_system_object_name, insert_property_texts, insert_property_values, no_space, normalize_names, plan_property_inserts, + resolve_membership_id, ) from .xml_handler import XMLHandler @@ -695,7 +697,8 @@ def add_object( if not collection_enum: collection_enum = get_default_collection(class_enum) - _ = self.add_membership(ClassEnum.System, class_enum, "System", name, collection_enum) + system_name = get_system_object_name(self) + _ = self.add_membership(ClassEnum.System, class_enum, system_name, name, collection_enum) return object_id def add_objects( @@ -756,7 +759,8 @@ def add_objects( collection_enum = get_default_collection(class_enum) object_ids = self.get_objects_id(names, class_enum=class_enum) parent_class_id = self.get_class_id(ClassEnum.System) - parent_object_id = self.get_object_id(ClassEnum.System, "System") + system_name = get_system_object_name(self) + parent_object_id = self.get_object_id(ClassEnum.System, system_name) collection_id = self.get_collection_id( collection_enum, parent_class_enum=ClassEnum.System, child_class_enum=class_enum ) @@ -1013,7 +1017,8 @@ def add_property( Class enumeration of the parent object. If None, defaults to ClassEnum.System, by default None parent_object_name : str | None, optional - Name of the parent object. If None, defaults to "System", by default None + Name of the parent object. If None, membership is resolved from + `parent_class_enum`, `collection_enum`, and `object_name`. Returns ------- @@ -1079,7 +1084,14 @@ def add_property( child_class_enum=object_class_enum, collection_enum=collection_enum, ) - membership_id = self.get_membership_id(parent_object_name or "System", object_name, collection_enum) + membership_id = resolve_membership_id( + self, + object_name, + object_class=object_class_enum, + collection=collection_enum, + parent_class=parent_class_enum, + parent_object_name=parent_object_name, + ) query = f"INSERT INTO {Schema.Data.name}(membership_id, property_id, value) values (?, ?, ?)" _ = self._db.execute(query, (membership_id, property_id, value)) @@ -1848,8 +1860,9 @@ def copy_object( ) system_collection = get_default_collection(object_class) - old_sys_id = self.get_membership_id("System", original_object_name, system_collection) - new_sys_id = self.get_membership_id("System", new_object_name, system_collection) + system_name = get_system_object_name(self) + old_sys_id = self.get_membership_id(system_name, original_object_name, system_collection) + new_sys_id = self.get_membership_id(system_name, new_object_name, system_collection) membership_mapping[old_sys_id] = new_sys_id if not copy_properties: @@ -2119,6 +2132,7 @@ def delete_property( property_name: str, collection: CollectionEnum | None = None, parent_class: ClassEnum | None = None, + parent_object_name: str | None = None, scenario: str | None = None, ) -> None: """Delete a property from an object. @@ -2141,6 +2155,9 @@ def delete_property( parent_class : ClassEnum | None, optional Parent class enumeration for the property. If None, defaults to ClassEnum.System, by default None + parent_object_name : str | None, optional + Parent object name. If None, membership is resolved by parent class, + child class, collection, and child object name. scenario : str | None, optional Name of the scenario to filter by. If specified, only deletes property data associated with this scenario, by default None @@ -2194,9 +2211,14 @@ def delete_property( collection_enum=collection, ) - # For parent object name, default to "System" if not specified - parent_object_name = "System" # This matches the pattern used in add_property - membership_id = self.get_membership_id(parent_object_name, object_name, collection) + membership_id = resolve_membership_id( + self, + object_name, + object_class=object_class, + collection=collection, + parent_class=parent_class, + parent_object_name=parent_object_name, + ) # Build the delete query if scenario is not None: diff --git a/src/plexosdb/utils.py b/src/plexosdb/utils.py index d6b2e9f..b0f087c 100644 --- a/src/plexosdb/utils.py +++ b/src/plexosdb/utils.py @@ -12,10 +12,11 @@ from loguru import logger +from .enums import ClassEnum from .exceptions import NotFoundError if TYPE_CHECKING: - from plexosdb import ClassEnum, CollectionEnum, PlexosDB + from plexosdb import CollectionEnum, PlexosDB from plexosdb.db_manager import SQLiteManager @@ -282,7 +283,13 @@ def plan_property_inserts( collection, parent_class_enum=parent_class, child_class_enum=object_class ) collection_properties = _fetch_collection_properties(db, collection_id=collection_id) - name_to_membership = _resolve_membership_map(db, normalized_records, object_class=object_class) + name_to_membership = _resolve_membership_map( + db, + normalized_records, + object_class=object_class, + parent_class=parent_class, + collection=collection, + ) property_id_map = {prop: pid for prop, pid in collection_properties} params, metadata_map = _build_property_rows( @@ -304,17 +311,46 @@ def _resolve_membership_map( normalized_records: list[dict[str, Any]], *, object_class: ClassEnum, + parent_class: ClassEnum, + collection: CollectionEnum, ) -> dict[str, int]: """Resolve membership ids for each object name.""" component_names = tuple({d["name"] for d in normalized_records if d.get("name") is not None}) - try: - memberships = db.get_memberships_system(component_names, object_class=object_class) - except Exception as exc: - missing = ", ".join(sorted(name for name in component_names if name)) - raise NotFoundError( - f"Objects not found: {missing}. Add them with `add_object` or `add_objects` before " - "adding properties." - ) from exc + if not component_names: + return {} + + memberships: list[tuple[str, int]] | list[dict[str, Any]] + if parent_class == ClassEnum.System: + try: + memberships = db.get_memberships_system( + component_names, + object_class=object_class, + collection=collection, + ) + except Exception as exc: + missing = ", ".join(sorted(name for name in component_names if name)) + raise NotFoundError( + f"Objects not found: {missing}. Add them with `add_object` or `add_objects` before " + "adding properties." + ) from exc + else: + collection_id = db.get_collection_id( + collection, parent_class_enum=parent_class, child_class_enum=object_class + ) + parent_class_id = db.get_class_id(parent_class) + child_class_id = db.get_class_id(object_class) + placeholders = ",".join("?" for _ in component_names) + query = f""" + SELECT child_object.name, mem.membership_id + FROM t_membership AS mem + INNER JOIN t_object AS child_object ON child_object.object_id = mem.child_object_id + WHERE mem.parent_class_id = ? + AND mem.child_class_id = ? + AND mem.collection_id = ? + AND child_object.name IN ({placeholders}) + """ + params: tuple[Any, ...] = (parent_class_id, child_class_id, collection_id, *component_names) + memberships = db._db.fetchall(query, params) if not memberships: missing = ", ".join(sorted(name for name in component_names if name)) @@ -323,7 +359,123 @@ def _resolve_membership_map( "adding properties." ) - return {membership["name"]: membership["membership_id"] for membership in memberships} + name_to_membership: dict[str, int] = {} + ambiguous_objects: set[str] = set() + for membership in memberships: + if isinstance(membership, dict): + object_name = membership["name"] + membership_id = membership["membership_id"] + else: + object_name = membership[0] + membership_id = membership[1] + existing_membership_id = name_to_membership.get(object_name) + if existing_membership_id is not None and existing_membership_id != membership_id: + ambiguous_objects.add(object_name) + continue + name_to_membership[object_name] = membership_id + + if ambiguous_objects: + ambiguous_names = ", ".join(sorted(ambiguous_objects)) + raise ValueError( + "Multiple memberships found for objects in the same parent class/collection: " + f"{ambiguous_names}. Resolve membership ambiguity before bulk insert." + ) + + logger.trace("Resolved {} memberships for collection {}", len(name_to_membership), collection.value) + return name_to_membership + + +def get_system_object_name(db: PlexosDB) -> str: + """Return the canonical System object name for this model. + + Parameters + ---------- + db : PlexosDB + Database instance to query. + + Returns + ------- + str + The name of the System object. + + Raises + ------ + NotFoundError + If no System object exists in the model. + ValueError + If multiple System objects exist and none is named "System". + """ + system_objects = db.list_objects_by_class(ClassEnum.System) + if not system_objects: + raise NotFoundError("No System object found in the model.") + if len(system_objects) == 1: + logger.trace("Resolved System object: {}", system_objects[0]) + return system_objects[0] + if "System" in system_objects: + logger.trace("Multiple System objects found, defaulting to 'System'") + return "System" + raise ValueError( + "Multiple System objects found and no default could be inferred. " + "Pass an explicit `parent_object_name`." + ) + + +def resolve_membership_id( + db: PlexosDB, + object_name: str, + *, + object_class: ClassEnum, + collection: CollectionEnum, + parent_class: ClassEnum, + parent_object_name: str | None = None, +) -> int: + """Resolve a single membership ID for property operations. + + Parameters + ---------- + db : PlexosDB + Database instance to query. + object_name : str + Name of the child object. + object_class : ClassEnum + Class of the child object. + collection : CollectionEnum + Collection defining the membership relationship. + parent_class : ClassEnum + Class of the parent object. + parent_object_name : str | None, optional + Explicit parent object name. If provided, uses direct lookup. + + Returns + ------- + int + The membership ID. + + Raises + ------ + NotFoundError + If no matching membership is found. + ValueError + If multiple memberships exist (ambiguous). + """ + if parent_object_name is not None: + logger.trace("Resolving membership via explicit parent: {}", parent_object_name) + return db.get_membership_id(parent_object_name, object_name, collection) + + mapping = _resolve_membership_map( + db, + [{"name": object_name}], + object_class=object_class, + parent_class=parent_class, + collection=collection, + ) + if object_name not in mapping: + raise NotFoundError( + f"No membership found for '{object_name}' in collection '{collection.value}' " + f"with parent class '{parent_class.value}'." + ) + logger.trace("Resolved membership_id={} for {}", mapping[object_name], object_name) + return mapping[object_name] def _build_property_rows( diff --git a/tests/test_plexosdb_delete_object.py b/tests/test_plexosdb_delete_object.py index 2152698..8826576 100644 --- a/tests/test_plexosdb_delete_object.py +++ b/tests/test_plexosdb_delete_object.py @@ -83,6 +83,39 @@ def test_delete_property_with_scenario(db_base: PlexosDB): assert properties[0]["scenario_name"] is None +def test_delete_property_with_renamed_system_object(db_base: PlexosDB): + from plexosdb.enums import ClassEnum, CollectionEnum + + db = db_base + object_name = "BW01" + property_name = "Max Capacity" + + system_name = db.list_objects_by_class(ClassEnum.System)[0] + db.add_object(ClassEnum.Generator, object_name) + db.add_property( + ClassEnum.Generator, + object_name, + property_name, + 150.0, + collection_enum=CollectionEnum.Generators, + parent_class_enum=ClassEnum.System, + parent_object_name=system_name, + ) + + if system_name == "System": + db.update_object(ClassEnum.System, "System", new_name="NEM") + + db.delete_property( + ClassEnum.Generator, + object_name, + property_name=property_name, + collection=CollectionEnum.Generators, + parent_class=ClassEnum.System, + ) + + assert not db.has_properties(ClassEnum.Generator, object_name) + + def test_delete_property_fails_with_nonexistent_object(db_base: PlexosDB): from plexosdb.enums import ClassEnum from plexosdb.exceptions import NotFoundError diff --git a/tests/test_plexosdb_from_records.py b/tests/test_plexosdb_from_records.py index 3e4c702..7469375 100644 --- a/tests/test_plexosdb_from_records.py +++ b/tests/test_plexosdb_from_records.py @@ -224,3 +224,24 @@ def test_add_properties_from_records_unknown_property(db_instance_with_schema: P ) assert db._db.fetchone("SELECT COUNT(*) FROM t_data")[0] == 0 + + +def test_add_properties_from_records_respects_parent_membership(db_with_topology: PlexosDB): + from plexosdb import ClassEnum, CollectionEnum + from plexosdb.exceptions import NotFoundError + + db = db_with_topology + system_name = db.list_objects_by_class(ClassEnum.System)[0] + system_membership_id = db.get_membership_id(system_name, "node-01", CollectionEnum.Nodes) + db._db.execute("DELETE FROM t_membership WHERE membership_id = ?", (system_membership_id,)) + + with pytest.raises(NotFoundError, match="Objects not found"): + db.add_properties_from_records( + [{"name": "node-01", "property": "Load", "value": 123.0}], + object_class=ClassEnum.Node, + parent_class=ClassEnum.System, + collection=CollectionEnum.Nodes, + scenario="Base Case", + ) + + assert db._db.fetchone("SELECT COUNT(*) FROM t_data")[0] == 0 From ed48c9ed4fe0cbe1a8db76982f2f7768743295a7 Mon Sep 17 00:00:00 2001 From: pesap Date: Wed, 11 Feb 2026 15:20:03 -0700 Subject: [PATCH 2/2] test: Fixing resource warning and improving coverage --- tests/conftest.py | 3 ++ tests/test_db_manager_iterators.py | 5 +++ tests/test_plexosdb_from_records.py | 27 ++++++++++++++ tests/test_utils.py | 56 ++++++++++++++++++++++++++++- 4 files changed, 90 insertions(+), 1 deletion(-) diff --git a/tests/conftest.py b/tests/conftest.py index 28dffc7..32d6492 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -53,6 +53,7 @@ def db_instance(): """Create a base DB instance that lasts the entire test session.""" db = PlexosDB() yield db + db._db.close() @pytest.fixture() @@ -63,6 +64,7 @@ def db_instance_with_xml(data_folder, tmp_path): shutil.copy(xml_fname, xml_copy) db = PlexosDB.from_xml(xml_path=xml_copy) yield db + db._db.close() xml_copy.unlink() @@ -152,6 +154,7 @@ def db_instance_with_schema() -> PlexosDB: "INSERT INTO t_property_report(property_id, collection_id, name) VALUES (1, 1, 'Units')" ) yield db + db._db.close() @pytest.fixture(scope="function") diff --git a/tests/test_db_manager_iterators.py b/tests/test_db_manager_iterators.py index 45f410c..2552c52 100644 --- a/tests/test_db_manager_iterators.py +++ b/tests/test_db_manager_iterators.py @@ -303,6 +303,7 @@ def test_iter_dicts_reraises_sqlite_error() -> None: from plexosdb.db_manager import SQLiteManager db = SQLiteManager() + original_conn = db._con mock_conn = MagicMock() mock_cursor = MagicMock() mock_cursor.execute.side_effect = sqlite3.Error("Query failed") @@ -313,12 +314,15 @@ def test_iter_dicts_reraises_sqlite_error() -> None: with pytest.raises(sqlite3.Error): list(db.iter_dicts("SELECT * FROM nonexistent")) + original_conn.close() + def test_iter_dicts_cursor_cleanup_on_error() -> None: """Test that cursor is closed even when error occurs during iteration.""" from plexosdb.db_manager import SQLiteManager db = SQLiteManager() + original_conn = db._con mock_conn = MagicMock() mock_cursor = MagicMock() mock_cursor.execute.side_effect = sqlite3.Error("Query failed") @@ -330,6 +334,7 @@ def test_iter_dicts_cursor_cleanup_on_error() -> None: list(db.iter_dicts("SELECT * FROM test")) mock_cursor.close.assert_called_once() + original_conn.close() def test_fetchmany_happy_path(db_with_large_dataset: SQLiteManager) -> None: diff --git a/tests/test_plexosdb_from_records.py b/tests/test_plexosdb_from_records.py index 7469375..3bd7a65 100644 --- a/tests/test_plexosdb_from_records.py +++ b/tests/test_plexosdb_from_records.py @@ -245,3 +245,30 @@ def test_add_properties_from_records_respects_parent_membership(db_with_topology ) assert db._db.fetchone("SELECT COUNT(*) FROM t_data")[0] == 0 + + +def test_add_properties_from_records_non_system_parent(db_with_topology: PlexosDB): + from plexosdb import ClassEnum, CollectionEnum + + db = db_with_topology + + db.add_object(ClassEnum.Reserve, "TestReserve") + db.add_membership( + ClassEnum.Reserve, + ClassEnum.Region, + "TestReserve", + "region-01", + CollectionEnum.Regions, + ) + + records = [{"name": "region-01", "property": "Load Risk", "value": 5.0}] + db.add_properties_from_records( + records, + object_class=ClassEnum.Region, + parent_class=ClassEnum.Reserve, + collection=CollectionEnum.Regions, + scenario="Base", + ) + + data_count = db._db.fetchone("SELECT COUNT(*) FROM t_data")[0] + assert data_count == 1 diff --git a/tests/test_utils.py b/tests/test_utils.py index 7cba188..f99049a 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -1,6 +1,7 @@ import pytest -from plexosdb.utils import batched, get_sql_query +from plexosdb.utils import batched, get_sql_query, get_system_object_name, resolve_membership_id +from plexosdb.exceptions import NotFoundError @pytest.mark.parametrize( @@ -17,3 +18,56 @@ def test_batched(): for element in batched(test_list, 2): assert len(element) == 2 + + +def test_get_system_object_name_no_system_objects(db_instance_with_schema): + db = db_instance_with_schema + db._db.execute("DELETE FROM t_object WHERE class_id = 1") + + with pytest.raises(NotFoundError, match="No System object found"): + get_system_object_name(db) + + +def test_get_system_object_name_multiple_with_default(db_instance_with_schema): + import uuid + + db = db_instance_with_schema + db._db.execute( + "INSERT INTO t_object(name, class_id, GUID) VALUES (?, 1, ?)", + ("AnotherSystem", str(uuid.uuid4())), + ) + + result = get_system_object_name(db) + assert result == "System" + + +def test_get_system_object_name_multiple_no_default(db_instance_with_schema): + import uuid + + db = db_instance_with_schema + db._db.execute("UPDATE t_object SET name = 'NEM' WHERE name = 'System'") + db._db.execute( + "INSERT INTO t_object(name, class_id, GUID) VALUES (?, 1, ?)", + ("AnotherSystem", str(uuid.uuid4())), + ) + + with pytest.raises(ValueError, match="Multiple System objects found"): + get_system_object_name(db) + + +def test_resolve_membership_id_not_found(db_instance_with_schema): + from plexosdb.enums import ClassEnum, CollectionEnum + + db = db_instance_with_schema + db.add_object(ClassEnum.Generator, "TestGen") + + db._db.execute("DELETE FROM t_membership") + + with pytest.raises(NotFoundError, match="Objects not found"): + resolve_membership_id( + db, + "TestGen", + object_class=ClassEnum.Generator, + collection=CollectionEnum.Generators, + parent_class=ClassEnum.System, + )