1"""Celko's "Nested Sets" Tree Structure.
2
3http://www.intelligententerprise.com/001020/celko.jhtml
4
5"""
6
7from sqlalchemy import (create_engine, Column, Integer, String, select, case,
8    func)
9from sqlalchemy.orm import Session, aliased
10from sqlalchemy.ext.declarative import declarative_base
11from sqlalchemy import event
12
13Base = declarative_base()
14
15class Employee(Base):
16    __tablename__ = 'personnel'
17    __mapper_args__ = {
18        'batch': False  # allows extension to fire for each
19                        # instance before going to the next.
20    }
21
22    parent = None
23
24    emp = Column(String, primary_key=True)
25
26    left = Column("lft", Integer, nullable=False)
27    right = Column("rgt", Integer, nullable=False)
28
29    def __repr__(self):
30        return "Employee(%s, %d, %d)" % (self.emp, self.left, self.right)
31
32@event.listens_for(Employee, "before_insert")
33def before_insert(mapper, connection, instance):
34    if not instance.parent:
35        instance.left = 1
36        instance.right = 2
37    else:
38        personnel = mapper.mapped_table
39        right_most_sibling = connection.scalar(
40            select([personnel.c.rgt]).
41                where(personnel.c.emp == instance.parent.emp)
42        )
43
44        connection.execute(
45            personnel.update(
46                personnel.c.rgt >= right_most_sibling).values(
47                    lft=case(
48                        [(personnel.c.lft > right_most_sibling,
49                            personnel.c.lft + 2)],
50                        else_=personnel.c.lft
51                    ),
52                    rgt=case(
53                        [(personnel.c.rgt >= right_most_sibling,
54                                personnel.c.rgt + 2)],
55                            else_=personnel.c.rgt
56                      )
57            )
58        )
59        instance.left = right_most_sibling
60        instance.right = right_most_sibling + 1
61
62    # before_update() would be needed to support moving of nodes
63    # after_delete() would be needed to support removal of nodes.
64
65engine = create_engine('sqlite://', echo=True)
66
67Base.metadata.create_all(engine)
68
69session = Session(bind=engine)
70
71albert = Employee(emp='Albert')
72bert = Employee(emp='Bert')
73chuck = Employee(emp='Chuck')
74donna = Employee(emp='Donna')
75eddie = Employee(emp='Eddie')
76fred = Employee(emp='Fred')
77
78bert.parent = albert
79chuck.parent = albert
80donna.parent = chuck
81eddie.parent = chuck
82fred.parent = chuck
83
84# the order of "add" is important here.  elements must be added in
85# the order in which they should be INSERTed.
86session.add_all([albert, bert, chuck, donna, eddie, fred])
87session.commit()
88
89print(session.query(Employee).all())
90
91# 1. Find an employee and all their supervisors, no matter how deep the tree.
92ealias = aliased(Employee)
93print(session.query(Employee).\
94            filter(ealias.left.between(Employee.left, Employee.right)).\
95            filter(ealias.emp == 'Eddie').all())
96
97#2. Find the employee and all their subordinates.
98# (This query has a nice symmetry with the first query.)
99print(session.query(Employee).\
100    filter(Employee.left.between(ealias.left, ealias.right)).\
101    filter(ealias.emp == 'Chuck').all())
102
103#3. Find the level of each node, so you can print the tree
104# as an indented listing.
105for indentation, employee in session.query(
106            func.count(Employee.emp).label('indentation') - 1, ealias).\
107    filter(ealias.left.between(Employee.left, Employee.right)).\
108    group_by(ealias.emp).\
109        order_by(ealias.left):
110    print("    " * indentation + str(employee))
111
112