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