Coverage for postrfp/model/questionnaire/renumber.py: 94%

1""" 

2Section and Question Numbering Module 

3 

4This module provides efficient functions for renumbering sections and questions 

5in hierarchical trees. It uses a base36-encoded materialized path approach where 

6each two characters in a node's b36_number represent a level in the hierarchy. 

7 

8The algorithm works as follows: 

91. Fetch all relevant sections and questions as lightweight tuples 

102. Build a parent-to-children lookup dictionary to represent the hierarchy 

113. Traverse the hierarchy to generate new node numbers based on position 

124. Apply all updates to the database in a single bulk operation 

13 

14Features: 

15- Supports renumbering from any section in the hierarchy (not just the root) 

16- Uses efficient queries with materialized path pattern matching 

17- Performs bulk updates for better performance with large trees 

18 

19Example: 

20 To renumber a section and all its descendants: 

21 

22 ```python 

23 from postrfp.model.questionnaire.renumber import renumber_tree 

24 from sqlalchemy.orm import Session 

25 

26 session = Session() 

27 section = session.get(Section, section_id) 

28 renumber_tree(session, section) 

29 ``` 

30""" 

31 

32from sqlalchemy.orm import Session 

33from sqlalchemy import update, text, select, func, Integer 

34 

35from postrfp.model import Section, QuestionInstance 

36from postrfp.model.questionnaire.b36 import materialised_paths 

37 

38 

39def fetch_section_tuples(session: Session, section): 

40 """ 

41 Fetch the section and all its descendants as lightweight tuples, 

42 using recursive query based on parent_id relationship instead of b36_number. 

43 

44 Args: 

45 session: SQLAlchemy session 

46 section: The section to start renumbering from 

47 """ 

48 

49 # Define the base table for sections 

50 sections = Section.__table__ 

51 

52 # Base case: the section itself 

53 base_query = ( 

54 select( 

55 sections.c.id.cast(Integer).label("id"), 

56 sections.c.parent_id.cast(Integer).label("parent_id"), 

57 sections.c.position.cast(Integer).label("position"), 

58 ) 

59 .where(sections.c.id == section.id) 

60 .where(sections.c.project_id == section.project_id) 

61 ) 

62 

63 # Create the recursive CTE 

64 descendants_cte = base_query.cte(name="descendants", recursive=True) 

65 

66 # Recursive case: all children of sections already in the result 

67 recursive_query = ( 

68 select( 

69 sections.c.id.cast(Integer).label("id"), 

70 sections.c.parent_id.cast(Integer).label("parent_id"), 

71 sections.c.position.cast(Integer).label("position"), 

72 ) 

73 .select_from( 

74 sections.join(descendants_cte, sections.c.parent_id == descendants_cte.c.id) 

75 ) 

76 .where(sections.c.project_id == section.project_id) 

77 ) 

78 

79 # Union the base case and recursive case 

80 descendants_cte = descendants_cte.union_all(recursive_query) 

81 

82 # Final query to select from the CTE 

83 final_query = select( 

84 descendants_cte.c.id, descendants_cte.c.parent_id, descendants_cte.c.position 

85 ).order_by(descendants_cte.c.parent_id, descendants_cte.c.position) 

86 

87 result = session.execute(final_query) 

88 return result.all() 

89 

90 

91def fetch_all_section_tuples(session: Session, root_section: Section): 

92 """ 

93 Fast path function to fetch all sections in a project when starting from root section. 

94 Avoids the recursive CTE when we know we need the entire project. 

95 

96 Args: 

97 session: SQLAlchemy session 

98 root_section: The root section (with parent_id=None) 

99 """ 

100 return session.execute( 

101 select(Section.id, Section.parent_id, Section.position) 

102 .where(Section.project_id == root_section.project_id) 

103 .order_by(Section.parent_id, Section.position) 

104 ).all() 

105 

106 

107def save_section_numbers(session, section_numbers): 

108 """Apply the generated node numbers to the database using ORM updates to ensure proper type conversion""" 

109 

110 # Get all the sections that need updating 

111 sections_to_update = ( 

112 session.query(Section).filter(Section.id.in_(section_numbers.keys())).all() 

113 ) 

114 

115 # Update each section through the ORM to ensure type conversion 

116 for section in sections_to_update: 

117 section.b36_number = section_numbers[section.id] 

118 

119 session.flush() 

120 

121 # Verify the values were saved correctly 

122 for section in sections_to_update: 

123 session.refresh(section) 

124 

125 

126def update_qinstance_numbers(session: Session, section_ids): 

127 """ 

128 Derive "b36_number" values for all QuestionInstances based on the 

129 b36_number of parent section b36_number 

130 """ 

131 # Create a subquery to get the section numbers 

132 section_query = ( 

133 select(Section.id, Section.b36_number) 

134 .where(Section.id.in_(section_ids)) 

135 .subquery("section_numbers") 

136 ) 

137 stmt = ( 

138 update(QuestionInstance) 

139 .where(QuestionInstance.section_id.in_(section_ids)) 

140 .values( 

141 b36_number=func.concat( 

142 section_query.c.b36_number, 

143 func.lpad(func.conv(QuestionInstance.position, 10, 36), 2, "0"), 

144 ) 

145 ) 

146 .where(QuestionInstance.section_id == section_query.c.id) 

147 ) 

148 session.execute(stmt) 

149 

150 

151def update_all_project_questions(session: Session, project_id): 

152 """ 

153 Fast path function to update all question numbers in a project. 

154 Used when renumbering from root section. 

155 

156 Args: 

157 session: SQLAlchemy session 

158 project_id: The project ID 

159 """ 

160 # Join approach is more efficient when updating all questions in a project 

161 stmt = text(""" 

162 UPDATE question_instances qi 

163 JOIN sections s ON qi.section_id = s.id 

164 SET qi.b36_number = CONCAT( 

165 s.b36_number, 

166 LPAD(CONV(qi.position, 10, 36), 2, '0') 

167 ) 

168 WHERE s.project_id = :project_id 

169 """) 

170 

171 session.execute(stmt, {"project_id": project_id}) 

172 

173 

174def renumber_tree(session: Session, root_section: Section) -> None: 

175 """ 

176 Efficiently renumber a section tree and all its questions. 

177 Works with any section at any level of the hierarchy. 

178 

179 Args: 

180 session: The SQLAlchemy session 

181 section: The Section object to renumber from (can be at any level) 

182 """ 

183 # Fetch all sections 

184 sec_nodes = fetch_section_tuples(session, root_section) 

185 # build a dict of {id: b36_number} (path) - this is the main algorithm 

186 sec_paths = materialised_paths(sec_nodes, root_section) 

187 # save the new b36_number values to Sections by id 

188 save_section_numbers(session, sec_paths) 

189 # Use the section b36_numbers to derive the question b36_numbers (in sql) 

190 update_qinstance_numbers(session, sec_paths.keys())