diff --git a/src/plexosdb/db.py b/src/plexosdb/db.py index 7052793..aa363d7 100644 --- a/src/plexosdb/db.py +++ b/src/plexosdb/db.py @@ -352,7 +352,7 @@ def add_category(self, class_enum: ClassEnum, /, name: str) -> int: return self.get_category_id(class_enum, name) class_id = self.get_class_id(class_enum) - rank = self.get_category_max_id(class_enum) or 1 # Default to max rank of 1 if no category exists + rank = self.get_category_max_id(class_enum) + 1 or 1 # Default to max rank of 1 if no category exists params = (class_id, name, rank) placeholders = ", ".join("?" * len(params)) query = f"INSERT INTO {Schema.Categories.name}(class_id, name, rank) VALUES({placeholders})" diff --git a/src/plexosdb/enums.py b/src/plexosdb/enums.py index 308d878..d4042a4 100644 --- a/src/plexosdb/enums.py +++ b/src/plexosdb/enums.py @@ -145,9 +145,14 @@ def str2enum(string: str, schema_enum: type[Enum] = Schema) -> Schema | None: def get_default_collection(class_enum: ClassEnum) -> CollectionEnum: """Return default collection for class.""" - # Special case for Data File - if class_enum == ClassEnum.DataFile: - return CollectionEnum.DataFiles + # Special cases for Data File and Battery objects + special_cases = { + ClassEnum.DataFile: CollectionEnum.DataFiles, + ClassEnum.Battery: CollectionEnum.Batteries, + } + + if class_enum in special_cases: + return special_cases[class_enum] collection_name = f"{class_enum}s" if collection_name not in CollectionEnum.__members__: