def _get_all_database_entity_field_names(entity: DatabaseEntity,
entity_field_type: EntityFieldType,
direction_checker):
"""Returns a set of field_names that correspond to any set fields on the
provided DatabaseEntity |entity| that match the provided
|entity_field_type|.
"""
back_edges = set()
forward_edges = set()
flat_fields = set()
foreign_keys = set()
for relationship_field_name in entity.get_relationship_property_names():
if direction_checker.is_back_edge(entity, relationship_field_name):
back_edges.add(relationship_field_name)
else:
forward_edges.add(relationship_field_name)
for foreign_key_name in entity.get_foreign_key_names():
foreign_keys.add(foreign_key_name)
for column_field_name in entity.get_column_property_names():
if column_field_name not in foreign_keys:
flat_fields.add(column_field_name)
if entity_field_type is EntityFieldType.FLAT_FIELD:
return flat_fields
if entity_field_type is EntityFieldType.FOREIGN_KEYS:
return foreign_keys
if entity_field_type is EntityFieldType.FORWARD_EDGE:
return forward_edges
if entity_field_type is EntityFieldType.BACK_EDGE:
return back_edges
if entity_field_type is EntityFieldType.ALL:
return flat_fields | foreign_keys | forward_edges | back_edges
raise ValueError(
f"Unrecognized EntityFieldType [{entity_field_type}] on entity [{entity}]"
)
def _get_related_entities(self,
entity: DatabaseEntity) -> List[DatabaseEntity]:
"""Returns list of all entities related to |entity|"""
related_entities = []
for relationship_name in entity.get_relationship_property_names():
# TODO(#1145): For County schema, fix direction checker to gracefully
# handle the fact that SentenceRelationship exists in the schema
# but not in the entity layer.
if self.get_system_level() == SystemLevel.STATE:
# Skip back edges
direction_checker = SchemaEdgeDirectionChecker.state_direction_checker(
)
if direction_checker.is_back_edge(entity, relationship_name):
continue
related = getattr(entity, relationship_name)
# Relationship can return either a list or a single item
if isinstance(related, list):
related_entities.extend(related)
elif related is not None:
related_entities.append(related)
return related_entities