Untitled

mail@pastecode.io avatar
unknown
plain_text
a month ago
2.2 kB
1
Indexable
Never

class TreeChildModelMixin:
    def __init__(self, *args, **kwargs):
        self.child_fields = getattr(self.__class__, self.TreeMeta.child_field_names)
        self.child_field_db_names = self.child_fields.field.get_attname()
        super().__init__(*args, **kwargs)

    def get_descendants(self, max_depth: int = None):
        SQL_QUERY = """
                WITH RECURSIVE AnswerHierarchy AS (
                SELECT 
                    id,
                    question,
                    NULL AS parent_id,
                    'root' AS relation_type,
                    NULL AS parent_answer_id
                FROM yourapp_answer
                WHERE id = 1
                
                UNION ALL
                
                SELECT 
                    child.id,
                    child.question,
                    parent.id AS parent_id,
                    CASE WHEN parent.positive_answer_id = child.id THEN 'positive' ELSE 'negative' END AS relation_type,
                    parent.id AS parent_answer_id
                FROM yourapp_answer AS child
                JOIN AnswerHierarchy AS parent ON (parent.id = child.positive_answer_id OR parent.id = child.negative_answer_id)
            )
            SELECT * FROM AnswerHierarchy;

            """.format(
            db_table=self._meta.db_table,
            object_id=self.id,
            child_field_db_names=self.child_field_db_names,
            filter_max_depth='WHERE ct.depth < {max_depth}'.format(max_depth=max_depth) if max_depth else ''
        )
        objects = self._meta.model.objects.raw(SQL_QUERY)
        return objects

    def is_cycled(self):
        if not self.id:
            return False
        for child in self.child_fields:
            descendants = [descendant.id for descendant in child.get_descendants()]
            parent_id = getattr(self, self.child_field_db_names)
            if parent_id in descendants:
                return True
        return False

    def clean(self):
        if self.is_cycled():
            parent_id = getattr(self, self.child_field_db_names)
            raise ValidationError(message=_('parent {} create recursion').format(parent_id))

    class TreeMeta:
        parent_field_name = 'parent'
Leave a Comment