Coverage for /Users/davegaeddert/Developer/dropseed/plain/plain-models/plain/models/backends/ddl_references.py: 50%

137 statements  

« prev     ^ index     » next       coverage.py v7.6.9, created at 2024-12-23 11:16 -0600

1""" 

2Helpers to manipulate deferred DDL statements that might need to be adjusted or 

3discarded within when executing a migration. 

4""" 

5 

6from copy import deepcopy 

7 

8 

9class Reference: 

10 """Base class that defines the reference interface.""" 

11 

12 def references_table(self, table): 

13 """ 

14 Return whether or not this instance references the specified table. 

15 """ 

16 return False 

17 

18 def references_column(self, table, column): 

19 """ 

20 Return whether or not this instance references the specified column. 

21 """ 

22 return False 

23 

24 def rename_table_references(self, old_table, new_table): 

25 """ 

26 Rename all references to the old_name to the new_table. 

27 """ 

28 pass 

29 

30 def rename_column_references(self, table, old_column, new_column): 

31 """ 

32 Rename all references to the old_column to the new_column. 

33 """ 

34 pass 

35 

36 def __repr__(self): 

37 return f"<{self.__class__.__name__} {str(self)!r}>" 

38 

39 def __str__(self): 

40 raise NotImplementedError( 

41 "Subclasses must define how they should be converted to string." 

42 ) 

43 

44 

45class Table(Reference): 

46 """Hold a reference to a table.""" 

47 

48 def __init__(self, table, quote_name): 

49 self.table = table 

50 self.quote_name = quote_name 

51 

52 def references_table(self, table): 

53 return self.table == table 

54 

55 def rename_table_references(self, old_table, new_table): 

56 if self.table == old_table: 

57 self.table = new_table 

58 

59 def __str__(self): 

60 return self.quote_name(self.table) 

61 

62 

63class TableColumns(Table): 

64 """Base class for references to multiple columns of a table.""" 

65 

66 def __init__(self, table, columns): 

67 self.table = table 

68 self.columns = columns 

69 

70 def references_column(self, table, column): 

71 return self.table == table and column in self.columns 

72 

73 def rename_column_references(self, table, old_column, new_column): 

74 if self.table == table: 

75 for index, column in enumerate(self.columns): 

76 if column == old_column: 

77 self.columns[index] = new_column 

78 

79 

80class Columns(TableColumns): 

81 """Hold a reference to one or many columns.""" 

82 

83 def __init__(self, table, columns, quote_name, col_suffixes=()): 

84 self.quote_name = quote_name 

85 self.col_suffixes = col_suffixes 

86 super().__init__(table, columns) 

87 

88 def __str__(self): 

89 def col_str(column, idx): 

90 col = self.quote_name(column) 

91 try: 

92 suffix = self.col_suffixes[idx] 

93 if suffix: 

94 col = f"{col} {suffix}" 

95 except IndexError: 

96 pass 

97 return col 

98 

99 return ", ".join( 

100 col_str(column, idx) for idx, column in enumerate(self.columns) 

101 ) 

102 

103 

104class IndexName(TableColumns): 

105 """Hold a reference to an index name.""" 

106 

107 def __init__(self, table, columns, suffix, create_index_name): 

108 self.suffix = suffix 

109 self.create_index_name = create_index_name 

110 super().__init__(table, columns) 

111 

112 def __str__(self): 

113 return self.create_index_name(self.table, self.columns, self.suffix) 

114 

115 

116class IndexColumns(Columns): 

117 def __init__(self, table, columns, quote_name, col_suffixes=(), opclasses=()): 

118 self.opclasses = opclasses 

119 super().__init__(table, columns, quote_name, col_suffixes) 

120 

121 def __str__(self): 

122 def col_str(column, idx): 

123 # Index.__init__() guarantees that self.opclasses is the same 

124 # length as self.columns. 

125 col = f"{self.quote_name(column)} {self.opclasses[idx]}" 

126 try: 

127 suffix = self.col_suffixes[idx] 

128 if suffix: 

129 col = f"{col} {suffix}" 

130 except IndexError: 

131 pass 

132 return col 

133 

134 return ", ".join( 

135 col_str(column, idx) for idx, column in enumerate(self.columns) 

136 ) 

137 

138 

139class ForeignKeyName(TableColumns): 

140 """Hold a reference to a foreign key name.""" 

141 

142 def __init__( 

143 self, 

144 from_table, 

145 from_columns, 

146 to_table, 

147 to_columns, 

148 suffix_template, 

149 create_fk_name, 

150 ): 

151 self.to_reference = TableColumns(to_table, to_columns) 

152 self.suffix_template = suffix_template 

153 self.create_fk_name = create_fk_name 

154 super().__init__( 

155 from_table, 

156 from_columns, 

157 ) 

158 

159 def references_table(self, table): 

160 return super().references_table(table) or self.to_reference.references_table( 

161 table 

162 ) 

163 

164 def references_column(self, table, column): 

165 return super().references_column( 

166 table, column 

167 ) or self.to_reference.references_column(table, column) 

168 

169 def rename_table_references(self, old_table, new_table): 

170 super().rename_table_references(old_table, new_table) 

171 self.to_reference.rename_table_references(old_table, new_table) 

172 

173 def rename_column_references(self, table, old_column, new_column): 

174 super().rename_column_references(table, old_column, new_column) 

175 self.to_reference.rename_column_references(table, old_column, new_column) 

176 

177 def __str__(self): 

178 suffix = self.suffix_template % { 

179 "to_table": self.to_reference.table, 

180 "to_column": self.to_reference.columns[0], 

181 } 

182 return self.create_fk_name(self.table, self.columns, suffix) 

183 

184 

185class Statement(Reference): 

186 """ 

187 Statement template and formatting parameters container. 

188 

189 Allows keeping a reference to a statement without interpolating identifiers 

190 that might have to be adjusted if they're referencing a table or column 

191 that is removed 

192 """ 

193 

194 def __init__(self, template, **parts): 

195 self.template = template 

196 self.parts = parts 

197 

198 def references_table(self, table): 

199 return any( 

200 hasattr(part, "references_table") and part.references_table(table) 

201 for part in self.parts.values() 

202 ) 

203 

204 def references_column(self, table, column): 

205 return any( 

206 hasattr(part, "references_column") and part.references_column(table, column) 

207 for part in self.parts.values() 

208 ) 

209 

210 def rename_table_references(self, old_table, new_table): 

211 for part in self.parts.values(): 

212 if hasattr(part, "rename_table_references"): 

213 part.rename_table_references(old_table, new_table) 

214 

215 def rename_column_references(self, table, old_column, new_column): 

216 for part in self.parts.values(): 

217 if hasattr(part, "rename_column_references"): 

218 part.rename_column_references(table, old_column, new_column) 

219 

220 def __str__(self): 

221 return self.template % self.parts 

222 

223 

224class Expressions(TableColumns): 

225 def __init__(self, table, expressions, compiler, quote_value): 

226 self.compiler = compiler 

227 self.expressions = expressions 

228 self.quote_value = quote_value 

229 columns = [ 

230 col.target.column 

231 for col in self.compiler.query._gen_cols([self.expressions]) 

232 ] 

233 super().__init__(table, columns) 

234 

235 def rename_table_references(self, old_table, new_table): 

236 if self.table != old_table: 

237 return 

238 self.expressions = self.expressions.relabeled_clone({old_table: new_table}) 

239 super().rename_table_references(old_table, new_table) 

240 

241 def rename_column_references(self, table, old_column, new_column): 

242 if self.table != table: 

243 return 

244 expressions = deepcopy(self.expressions) 

245 self.columns = [] 

246 for col in self.compiler.query._gen_cols([expressions]): 

247 if col.target.column == old_column: 

248 col.target.column = new_column 

249 self.columns.append(col.target.column) 

250 self.expressions = expressions 

251 

252 def __str__(self): 

253 sql, params = self.compiler.compile(self.expressions) 

254 params = map(self.quote_value, params) 

255 return sql % tuple(params)