def read_db_entity_trees_of_cls_to_merge(
session: Session,
state_code: str,
schema_cls: Type[StateBase]
) -> List[List[EntityTree]]:
"""
Returns a list of lists of EntityTree where each inner list is a group
of EntityTrees with entities of class |schema_cls| that need to be merged
because their entities have the same external_id.
Will assert if schema_cls does not have a person_id or external_id field.
"""
external_ids = dao.read_external_ids_of_cls_with_external_id_match(
session, state_code, schema_cls)
people = dao.read_people_by_cls_external_ids(
session, state_code, schema_cls, external_ids)
all_cls_trees = get_all_entity_trees_of_cls(people, schema_cls)
external_ids_map: Dict[str, List[EntityTree]] = defaultdict(list)
for tree in all_cls_trees:
if not isinstance(tree.entity, schema_cls):
raise ValueError(f'Unexpected entity type [{type(tree.entity)}]')
if tree.entity.external_id in external_ids:
external_ids_map[tree.entity.external_id].append(tree)
return [tree_list for _, tree_list in external_ids_map.items()]
def test_readObjsWithExternalIdMatch(self) -> None:
person_1 = schema.StatePerson(person_id=1, state_code=_STATE_CODE)
sentence_group_1 = schema.StateSentenceGroup(
sentence_group_id=1,
external_id=_EXTERNAL_ID,
status=StateSentenceStatus.PRESENT_WITHOUT_INFO.value,
state_code=_STATE_CODE,
person=person_1,
)
person_1.sentence_groups = [sentence_group_1]
person_2 = schema.StatePerson(person_id=2, state_code=_STATE_CODE)
sentence_group_1_dup = schema.StateSentenceGroup(
sentence_group_id=2,
external_id=_EXTERNAL_ID,
status=StateSentenceStatus.PRESENT_WITHOUT_INFO.value,
state_code=_STATE_CODE,
person=person_2,
)
sentence_group_2 = schema.StateSentenceGroup(
sentence_group_id=3,
external_id=_EXTERNAL_ID2,
status=StateSentenceStatus.PRESENT_WITHOUT_INFO.value,
state_code=_STATE_CODE,
person=person_2,
)
placeholder_sentence_group = schema.StateSentenceGroup(
sentence_group_id=4,
status=StateSentenceStatus.PRESENT_WITHOUT_INFO.value,
state_code=_STATE_CODE,
person=person_2,
)
person_2.sentence_groups = [
sentence_group_1_dup,
sentence_group_2,
placeholder_sentence_group,
]
with SessionFactory.using_database(self.database_key,
autocommit=False) as session:
session.add(person_1)
session.add(person_2)
session.flush()
# Act
external_ids = dao.read_external_ids_of_cls_with_external_id_match(
session, _STATE_CODE, schema.StateSentenceGroup)
# Assert
self.assertEqual(external_ids, [_EXTERNAL_ID])
def test_readObjsWithExternalIdMatch(self):
person_1 = schema.StatePerson(person_id=1)
sentence_group_1 = schema.StateSentenceGroup(
sentence_group_id=1,
external_id=_EXTERNAL_ID,
status=StateSentenceStatus.PRESENT_WITHOUT_INFO.value,
state_code=_STATE_CODE,
person=person_1)
person_1.sentence_groups = [sentence_group_1]
person_2 = schema.StatePerson(person_id=2)
sentence_group_1_dup = schema.StateSentenceGroup(
sentence_group_id=2,
external_id=_EXTERNAL_ID,
status=StateSentenceStatus.PRESENT_WITHOUT_INFO.value,
state_code=_STATE_CODE,
person=person_2)
sentence_group_2 = schema.StateSentenceGroup(
sentence_group_id=3,
external_id=_EXTERNAL_ID2,
status=StateSentenceStatus.PRESENT_WITHOUT_INFO.value,
state_code=_STATE_CODE,
person=person_2)
placeholder_sentence_group = schema.StateSentenceGroup(
sentence_group_id=4,
status=StateSentenceStatus.PRESENT_WITHOUT_INFO.value,
state_code=_STATE_CODE,
person=person_2)
person_2.sentence_groups = [
sentence_group_1_dup, sentence_group_2, placeholder_sentence_group
]
session = SessionFactory.for_schema_base(StateBase)
session.add(person_1)
session.add(person_2)
session.commit()
# Act
external_ids = dao.read_external_ids_of_cls_with_external_id_match(
session, _STATE_CODE, schema.StateSentenceGroup)
# Assert
self.assertEqual(external_ids, [_EXTERNAL_ID])