Mercurial > kallithea

comparison rhodecode/lib/dbmigrate/migrate/changeset/ansisql.py @ 833:9753e0907827 beta

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
added dbmigrate package, added model changes moved out upgrade db command to that package
author Marcin Kuzminski <marcin@python-works.com>
date Sat, 11 Dec 2010 01:54:12 +0100
parents
children 08d2dcd71666
comparison
equal deleted inserted replaced
832:634596f81cfd 833:9753e0907827
1 """
2 Extensions to SQLAlchemy for altering existing tables.
3
4 At the moment, this isn't so much based off of ANSI as much as
5 things that just happen to work with multiple databases.
6 """
7 import StringIO
8
9 import sqlalchemy as sa
10 from sqlalchemy.schema import SchemaVisitor
11 from sqlalchemy.engine.default import DefaultDialect
12 from sqlalchemy.sql import ClauseElement
13 from sqlalchemy.schema import (ForeignKeyConstraint,
14 PrimaryKeyConstraint,
15 CheckConstraint,
16 UniqueConstraint,
17 Index)
18
19 from migrate import exceptions
20 from migrate.changeset import constraint, SQLA_06
21
22 if not SQLA_06:
23 from sqlalchemy.sql.compiler import SchemaGenerator, SchemaDropper
24 else:
25 from sqlalchemy.schema import AddConstraint, DropConstraint
26 from sqlalchemy.sql.compiler import DDLCompiler
27 SchemaGenerator = SchemaDropper = DDLCompiler
28
29
30 class AlterTableVisitor(SchemaVisitor):
31 """Common operations for ``ALTER TABLE`` statements."""
32
33 if SQLA_06:
34 # engine.Compiler looks for .statement
35 # when it spawns off a new compiler
36 statement = ClauseElement()
37
38 def append(self, s):
39 """Append content to the SchemaIterator's query buffer."""
40
41 self.buffer.write(s)
42
43 def execute(self):
44 """Execute the contents of the SchemaIterator's buffer."""
45 try:
46 return self.connection.execute(self.buffer.getvalue())
47 finally:
48 self.buffer.truncate(0)
49
50 def __init__(self, dialect, connection, **kw):
51 self.connection = connection
52 self.buffer = StringIO.StringIO()
53 self.preparer = dialect.identifier_preparer
54 self.dialect = dialect
55
56 def traverse_single(self, elem):
57 ret = super(AlterTableVisitor, self).traverse_single(elem)
58 if ret:
59 # adapt to 0.6 which uses a string-returning
60 # object
61 self.append(" %s" % ret)
62
63 def _to_table(self, param):
64 """Returns the table object for the given param object."""
65 if isinstance(param, (sa.Column, sa.Index, sa.schema.Constraint)):
66 ret = param.table
67 else:
68 ret = param
69 return ret
70
71 def start_alter_table(self, param):
72 """Returns the start of an ``ALTER TABLE`` SQL-Statement.
73
74 Use the param object to determine the table name and use it
75 for building the SQL statement.
76
77 :param param: object to determine the table from
78 :type param: :class:`sqlalchemy.Column`, :class:`sqlalchemy.Index`,
79 :class:`sqlalchemy.schema.Constraint`, :class:`sqlalchemy.Table`,
80 or string (table name)
81 """
82 table = self._to_table(param)
83 self.append('\nALTER TABLE %s ' % self.preparer.format_table(table))
84 return table
85
86
87 class ANSIColumnGenerator(AlterTableVisitor, SchemaGenerator):
88 """Extends ansisql generator for column creation (alter table add col)"""
89
90 def visit_column(self, column):
91 """Create a column (table already exists).
92
93 :param column: column object
94 :type column: :class:`sqlalchemy.Column` instance
95 """
96 if column.default is not None:
97 self.traverse_single(column.default)
98
99 table = self.start_alter_table(column)
100 self.append("ADD ")
101 self.append(self.get_column_specification(column))
102
103 for cons in column.constraints:
104 self.traverse_single(cons)
105 self.execute()
106
107 # ALTER TABLE STATEMENTS
108
109 # add indexes and unique constraints
110 if column.index_name:
111 Index(column.index_name,column).create()
112 elif column.unique_name:
113 constraint.UniqueConstraint(column,
114 name=column.unique_name).create()
115
116 # SA bounds FK constraints to table, add manually
117 for fk in column.foreign_keys:
118 self.add_foreignkey(fk.constraint)
119
120 # add primary key constraint if needed
121 if column.primary_key_name:
122 cons = constraint.PrimaryKeyConstraint(column,
123 name=column.primary_key_name)
124 cons.create()
125
126 if SQLA_06:
127 def add_foreignkey(self, fk):
128 self.connection.execute(AddConstraint(fk))
129
130 class ANSIColumnDropper(AlterTableVisitor, SchemaDropper):
131 """Extends ANSI SQL dropper for column dropping (``ALTER TABLE
132 DROP COLUMN``).
133 """
134
135 def visit_column(self, column):
136 """Drop a column from its table.
137
138 :param column: the column object
139 :type column: :class:`sqlalchemy.Column`
140 """
141 table = self.start_alter_table(column)
142 self.append('DROP COLUMN %s' % self.preparer.format_column(column))
143 self.execute()
144
145
146 class ANSISchemaChanger(AlterTableVisitor, SchemaGenerator):
147 """Manages changes to existing schema elements.
148
149 Note that columns are schema elements; ``ALTER TABLE ADD COLUMN``
150 is in SchemaGenerator.
151
152 All items may be renamed. Columns can also have many of their properties -
153 type, for example - changed.
154
155 Each function is passed a tuple, containing (object, name); where
156 object is a type of object you'd expect for that function
157 (ie. table for visit_table) and name is the object's new
158 name. NONE means the name is unchanged.
159 """
160
161 def visit_table(self, table):
162 """Rename a table. Other ops aren't supported."""
163 self.start_alter_table(table)
164 self.append("RENAME TO %s" % self.preparer.quote(table.new_name,
165 table.quote))
166 self.execute()
167
168 def visit_index(self, index):
169 """Rename an index"""
170 if hasattr(self, '_validate_identifier'):
171 # SA <= 0.6.3
172 self.append("ALTER INDEX %s RENAME TO %s" % (
173 self.preparer.quote(
174 self._validate_identifier(
175 index.name, True), index.quote),
176 self.preparer.quote(
177 self._validate_identifier(
178 index.new_name, True), index.quote)))
179 else:
180 # SA >= 0.6.5
181 self.append("ALTER INDEX %s RENAME TO %s" % (
182 self.preparer.quote(
183 self._index_identifier(
184 index.name), index.quote),
185 self.preparer.quote(
186 self._index_identifier(
187 index.new_name), index.quote)))
188 self.execute()
189
190 def visit_column(self, delta):
191 """Rename/change a column."""
192 # ALTER COLUMN is implemented as several ALTER statements
193 keys = delta.keys()
194 if 'type' in keys:
195 self._run_subvisit(delta, self._visit_column_type)
196 if 'nullable' in keys:
197 self._run_subvisit(delta, self._visit_column_nullable)
198 if 'server_default' in keys:
199 # Skip 'default': only handle server-side defaults, others
200 # are managed by the app, not the db.
201 self._run_subvisit(delta, self._visit_column_default)
202 if 'name' in keys:
203 self._run_subvisit(delta, self._visit_column_name, start_alter=False)
204
205 def _run_subvisit(self, delta, func, start_alter=True):
206 """Runs visit method based on what needs to be changed on column"""
207 table = self._to_table(delta.table)
208 col_name = delta.current_name
209 if start_alter:
210 self.start_alter_column(table, col_name)
211 ret = func(table, delta.result_column, delta)
212 self.execute()
213
214 def start_alter_column(self, table, col_name):
215 """Starts ALTER COLUMN"""
216 self.start_alter_table(table)
217 self.append("ALTER COLUMN %s " % self.preparer.quote(col_name, table.quote))
218
219 def _visit_column_nullable(self, table, column, delta):
220 nullable = delta['nullable']
221 if nullable:
222 self.append("DROP NOT NULL")
223 else:
224 self.append("SET NOT NULL")
225
226 def _visit_column_default(self, table, column, delta):
227 default_text = self.get_column_default_string(column)
228 if default_text is not None:
229 self.append("SET DEFAULT %s" % default_text)
230 else:
231 self.append("DROP DEFAULT")
232
233 def _visit_column_type(self, table, column, delta):
234 type_ = delta['type']
235 if SQLA_06:
236 type_text = str(type_.compile(dialect=self.dialect))
237 else:
238 type_text = type_.dialect_impl(self.dialect).get_col_spec()
239 self.append("TYPE %s" % type_text)
240
241 def _visit_column_name(self, table, column, delta):
242 self.start_alter_table(table)
243 col_name = self.preparer.quote(delta.current_name, table.quote)
244 new_name = self.preparer.format_column(delta.result_column)
245 self.append('RENAME COLUMN %s TO %s' % (col_name, new_name))
246
247
248 class ANSIConstraintCommon(AlterTableVisitor):
249 """
250 Migrate's constraints require a separate creation function from
251 SA's: Migrate's constraints are created independently of a table;
252 SA's are created at the same time as the table.
253 """
254
255 def get_constraint_name(self, cons):
256 """Gets a name for the given constraint.
257
258 If the name is already set it will be used otherwise the
259 constraint's :meth:`autoname <migrate.changeset.constraint.ConstraintChangeset.autoname>`
260 method is used.
261
262 :param cons: constraint object
263 """
264 if cons.name is not None:
265 ret = cons.name
266 else:
267 ret = cons.name = cons.autoname()
268 return self.preparer.quote(ret, cons.quote)
269
270 def visit_migrate_primary_key_constraint(self, *p, **k):
271 self._visit_constraint(*p, **k)
272
273 def visit_migrate_foreign_key_constraint(self, *p, **k):
274 self._visit_constraint(*p, **k)
275
276 def visit_migrate_check_constraint(self, *p, **k):
277 self._visit_constraint(*p, **k)
278
279 def visit_migrate_unique_constraint(self, *p, **k):
280 self._visit_constraint(*p, **k)
281
282 if SQLA_06:
283 class ANSIConstraintGenerator(ANSIConstraintCommon, SchemaGenerator):
284 def _visit_constraint(self, constraint):
285 constraint.name = self.get_constraint_name(constraint)
286 self.append(self.process(AddConstraint(constraint)))
287 self.execute()
288
289 class ANSIConstraintDropper(ANSIConstraintCommon, SchemaDropper):
290 def _visit_constraint(self, constraint):
291 constraint.name = self.get_constraint_name(constraint)
292 self.append(self.process(DropConstraint(constraint, cascade=constraint.cascade)))
293 self.execute()
294
295 else:
296 class ANSIConstraintGenerator(ANSIConstraintCommon, SchemaGenerator):
297
298 def get_constraint_specification(self, cons, **kwargs):
299 """Constaint SQL generators.
300
301 We cannot use SA visitors because they append comma.
302 """
303
304 if isinstance(cons, PrimaryKeyConstraint):
305 if cons.name is not None:
306 self.append("CONSTRAINT %s " % self.preparer.format_constraint(cons))
307 self.append("PRIMARY KEY ")
308 self.append("(%s)" % ', '.join(self.preparer.quote(c.name, c.quote)
309 for c in cons))
310 self.define_constraint_deferrability(cons)
311 elif isinstance(cons, ForeignKeyConstraint):
312 self.define_foreign_key(cons)
313 elif isinstance(cons, CheckConstraint):
314 if cons.name is not None:
315 self.append("CONSTRAINT %s " %
316 self.preparer.format_constraint(cons))
317 self.append("CHECK (%s)" % cons.sqltext)
318 self.define_constraint_deferrability(cons)
319 elif isinstance(cons, UniqueConstraint):
320 if cons.name is not None:
321 self.append("CONSTRAINT %s " %
322 self.preparer.format_constraint(cons))
323 self.append("UNIQUE (%s)" % \
324 (', '.join(self.preparer.quote(c.name, c.quote) for c in cons)))
325 self.define_constraint_deferrability(cons)
326 else:
327 raise exceptions.InvalidConstraintError(cons)
328
329 def _visit_constraint(self, constraint):
330
331 table = self.start_alter_table(constraint)
332 constraint.name = self.get_constraint_name(constraint)
333 self.append("ADD ")
334 self.get_constraint_specification(constraint)
335 self.execute()
336
337
338 class ANSIConstraintDropper(ANSIConstraintCommon, SchemaDropper):
339
340 def _visit_constraint(self, constraint):
341 self.start_alter_table(constraint)
342 self.append("DROP CONSTRAINT ")
343 constraint.name = self.get_constraint_name(constraint)
344 self.append(self.preparer.format_constraint(constraint))
345 if constraint.cascade:
346 self.cascade_constraint(constraint)
347 self.execute()
348
349 def cascade_constraint(self, constraint):
350 self.append(" CASCADE")
351
352
353 class ANSIDialect(DefaultDialect):
354 columngenerator = ANSIColumnGenerator
355 columndropper = ANSIColumnDropper
356 schemachanger = ANSISchemaChanger
357 constraintgenerator = ANSIConstraintGenerator
358 constraintdropper = ANSIConstraintDropper