Show More
ansisql.py
315 lines
| 11.3 KiB
| text/x-python
|
PythonLexer
r1 | """ | |||
Extensions to SQLAlchemy for altering existing tables. | ||||
At the moment, this isn't so much based off of ANSI as much as | ||||
things that just happen to work with multiple databases. | ||||
""" | ||||
import StringIO | ||||
import sqlalchemy as sa | ||||
from sqlalchemy.schema import SchemaVisitor | ||||
from sqlalchemy.engine.default import DefaultDialect | ||||
from sqlalchemy.sql import ClauseElement | ||||
from sqlalchemy.schema import (ForeignKeyConstraint, | ||||
PrimaryKeyConstraint, | ||||
CheckConstraint, | ||||
UniqueConstraint, | ||||
Index) | ||||
import sqlalchemy.sql.compiler | ||||
from rhodecode.lib.dbmigrate.migrate import exceptions | ||||
from rhodecode.lib.dbmigrate.migrate.changeset import constraint | ||||
from rhodecode.lib.dbmigrate.migrate.changeset import util | ||||
from sqlalchemy.schema import AddConstraint, DropConstraint | ||||
from sqlalchemy.sql.compiler import DDLCompiler | ||||
SchemaGenerator = SchemaDropper = DDLCompiler | ||||
class AlterTableVisitor(SchemaVisitor): | ||||
"""Common operations for ``ALTER TABLE`` statements.""" | ||||
# engine.Compiler looks for .statement | ||||
# when it spawns off a new compiler | ||||
statement = ClauseElement() | ||||
def append(self, s): | ||||
"""Append content to the SchemaIterator's query buffer.""" | ||||
self.buffer.write(s) | ||||
def execute(self): | ||||
"""Execute the contents of the SchemaIterator's buffer.""" | ||||
try: | ||||
return self.connection.execute(self.buffer.getvalue()) | ||||
finally: | ||||
self.buffer.seek(0) | ||||
self.buffer.truncate() | ||||
def __init__(self, dialect, connection, **kw): | ||||
self.connection = connection | ||||
self.buffer = StringIO.StringIO() | ||||
self.preparer = dialect.identifier_preparer | ||||
self.dialect = dialect | ||||
def traverse_single(self, elem): | ||||
ret = super(AlterTableVisitor, self).traverse_single(elem) | ||||
if ret: | ||||
# adapt to 0.6 which uses a string-returning | ||||
# object | ||||
self.append(" %s" % ret) | ||||
def _to_table(self, param): | ||||
"""Returns the table object for the given param object.""" | ||||
if isinstance(param, (sa.Column, sa.Index, sa.schema.Constraint)): | ||||
ret = param.table | ||||
else: | ||||
ret = param | ||||
return ret | ||||
def start_alter_table(self, param): | ||||
"""Returns the start of an ``ALTER TABLE`` SQL-Statement. | ||||
Use the param object to determine the table name and use it | ||||
for building the SQL statement. | ||||
:param param: object to determine the table from | ||||
:type param: :class:`sqlalchemy.Column`, :class:`sqlalchemy.Index`, | ||||
:class:`sqlalchemy.schema.Constraint`, :class:`sqlalchemy.Table`, | ||||
or string (table name) | ||||
""" | ||||
table = self._to_table(param) | ||||
self.append('\nALTER TABLE %s ' % self.preparer.format_table(table)) | ||||
return table | ||||
class ANSIColumnGenerator(AlterTableVisitor, SchemaGenerator): | ||||
"""Extends ansisql generator for column creation (alter table add col)""" | ||||
def visit_column(self, column): | ||||
"""Create a column (table already exists). | ||||
:param column: column object | ||||
:type column: :class:`sqlalchemy.Column` instance | ||||
""" | ||||
if column.default is not None: | ||||
self.traverse_single(column.default) | ||||
table = self.start_alter_table(column) | ||||
self.append("ADD ") | ||||
self.append(self.get_column_specification(column)) | ||||
for cons in column.constraints: | ||||
self.traverse_single(cons) | ||||
self.execute() | ||||
# ALTER TABLE STATEMENTS | ||||
# add indexes and unique constraints | ||||
if column.index_name: | ||||
Index(column.index_name,column).create() | ||||
elif column.unique_name: | ||||
constraint.UniqueConstraint(column, | ||||
name=column.unique_name).create() | ||||
# SA bounds FK constraints to table, add manually | ||||
for fk in column.foreign_keys: | ||||
self.add_foreignkey(fk.constraint) | ||||
# add primary key constraint if needed | ||||
if column.primary_key_name: | ||||
cons = constraint.PrimaryKeyConstraint(column, | ||||
name=column.primary_key_name) | ||||
cons.create() | ||||
def add_foreignkey(self, fk): | ||||
self.connection.execute(AddConstraint(fk)) | ||||
class ANSIColumnDropper(AlterTableVisitor, SchemaDropper): | ||||
"""Extends ANSI SQL dropper for column dropping (``ALTER TABLE | ||||
DROP COLUMN``). | ||||
""" | ||||
def visit_column(self, column): | ||||
"""Drop a column from its table. | ||||
:param column: the column object | ||||
:type column: :class:`sqlalchemy.Column` | ||||
""" | ||||
table = self.start_alter_table(column) | ||||
self.append('DROP COLUMN %s' % self.preparer.format_column(column)) | ||||
self.execute() | ||||
class ANSISchemaChanger(AlterTableVisitor, SchemaGenerator): | ||||
"""Manages changes to existing schema elements. | ||||
Note that columns are schema elements; ``ALTER TABLE ADD COLUMN`` | ||||
is in SchemaGenerator. | ||||
All items may be renamed. Columns can also have many of their properties - | ||||
type, for example - changed. | ||||
Each function is passed a tuple, containing (object, name); where | ||||
object is a type of object you'd expect for that function | ||||
(ie. table for visit_table) and name is the object's new | ||||
name. NONE means the name is unchanged. | ||||
""" | ||||
def visit_table(self, table): | ||||
"""Rename a table. Other ops aren't supported.""" | ||||
self.start_alter_table(table) | ||||
q = util.safe_quote(table) | ||||
self.append("RENAME TO %s" % self.preparer.quote(table.new_name, q)) | ||||
self.execute() | ||||
def visit_index(self, index): | ||||
"""Rename an index""" | ||||
if hasattr(self, '_validate_identifier'): | ||||
# SA <= 0.6.3 | ||||
self.append("ALTER INDEX %s RENAME TO %s" % ( | ||||
self.preparer.quote( | ||||
self._validate_identifier( | ||||
index.name, True), index.quote), | ||||
self.preparer.quote( | ||||
self._validate_identifier( | ||||
index.new_name, True), index.quote))) | ||||
elif hasattr(self, '_index_identifier'): | ||||
# SA >= 0.6.5, < 0.8 | ||||
self.append("ALTER INDEX %s RENAME TO %s" % ( | ||||
self.preparer.quote( | ||||
self._index_identifier( | ||||
index.name), index.quote), | ||||
self.preparer.quote( | ||||
self._index_identifier( | ||||
index.new_name), index.quote))) | ||||
else: | ||||
# SA >= 0.8 | ||||
class NewName(object): | ||||
"""Map obj.name -> obj.new_name""" | ||||
def __init__(self, index): | ||||
self.name = index.new_name | ||||
self._obj = index | ||||
def __getattr__(self, attr): | ||||
if attr == 'name': | ||||
return getattr(self, attr) | ||||
return getattr(self._obj, attr) | ||||
self.append("ALTER INDEX %s RENAME TO %s" % ( | ||||
self._prepared_index_name(index), | ||||
self._prepared_index_name(NewName(index)))) | ||||
self.execute() | ||||
def visit_column(self, delta): | ||||
"""Rename/change a column.""" | ||||
# ALTER COLUMN is implemented as several ALTER statements | ||||
keys = delta.keys() | ||||
if 'type' in keys: | ||||
self._run_subvisit(delta, self._visit_column_type) | ||||
if 'nullable' in keys: | ||||
self._run_subvisit(delta, self._visit_column_nullable) | ||||
if 'server_default' in keys: | ||||
# Skip 'default': only handle server-side defaults, others | ||||
# are managed by the app, not the db. | ||||
self._run_subvisit(delta, self._visit_column_default) | ||||
if 'name' in keys: | ||||
self._run_subvisit(delta, self._visit_column_name, start_alter=False) | ||||
def _run_subvisit(self, delta, func, start_alter=True): | ||||
"""Runs visit method based on what needs to be changed on column""" | ||||
table = self._to_table(delta.table) | ||||
col_name = delta.current_name | ||||
if start_alter: | ||||
self.start_alter_column(table, col_name) | ||||
ret = func(table, delta.result_column, delta) | ||||
self.execute() | ||||
def start_alter_column(self, table, col_name): | ||||
"""Starts ALTER COLUMN""" | ||||
self.start_alter_table(table) | ||||
q = util.safe_quote(table) | ||||
self.append("ALTER COLUMN %s " % self.preparer.quote(col_name, q)) | ||||
def _visit_column_nullable(self, table, column, delta): | ||||
nullable = delta['nullable'] | ||||
if nullable: | ||||
self.append("DROP NOT NULL") | ||||
else: | ||||
self.append("SET NOT NULL") | ||||
def _visit_column_default(self, table, column, delta): | ||||
default_text = self.get_column_default_string(column) | ||||
if default_text is not None: | ||||
self.append("SET DEFAULT %s" % default_text) | ||||
else: | ||||
self.append("DROP DEFAULT") | ||||
def _visit_column_type(self, table, column, delta): | ||||
type_ = delta['type'] | ||||
type_text = str(type_.compile(dialect=self.dialect)) | ||||
self.append("TYPE %s" % type_text) | ||||
def _visit_column_name(self, table, column, delta): | ||||
self.start_alter_table(table) | ||||
q = util.safe_quote(table) | ||||
col_name = self.preparer.quote(delta.current_name, q) | ||||
new_name = self.preparer.format_column(delta.result_column) | ||||
self.append('RENAME COLUMN %s TO %s' % (col_name, new_name)) | ||||
class ANSIConstraintCommon(AlterTableVisitor): | ||||
""" | ||||
Migrate's constraints require a separate creation function from | ||||
SA's: Migrate's constraints are created independently of a table; | ||||
SA's are created at the same time as the table. | ||||
""" | ||||
def get_constraint_name(self, cons): | ||||
"""Gets a name for the given constraint. | ||||
If the name is already set it will be used otherwise the | ||||
constraint's :meth:`autoname <migrate.changeset.constraint.ConstraintChangeset.autoname>` | ||||
method is used. | ||||
:param cons: constraint object | ||||
""" | ||||
if cons.name is not None: | ||||
ret = cons.name | ||||
else: | ||||
ret = cons.name = cons.autoname() | ||||
return self.preparer.quote(ret, cons.quote) | ||||
def visit_migrate_primary_key_constraint(self, *p, **k): | ||||
self._visit_constraint(*p, **k) | ||||
def visit_migrate_foreign_key_constraint(self, *p, **k): | ||||
self._visit_constraint(*p, **k) | ||||
def visit_migrate_check_constraint(self, *p, **k): | ||||
self._visit_constraint(*p, **k) | ||||
def visit_migrate_unique_constraint(self, *p, **k): | ||||
self._visit_constraint(*p, **k) | ||||
class ANSIConstraintGenerator(ANSIConstraintCommon, SchemaGenerator): | ||||
def _visit_constraint(self, constraint): | ||||
constraint.name = self.get_constraint_name(constraint) | ||||
self.append(self.process(AddConstraint(constraint))) | ||||
self.execute() | ||||
class ANSIConstraintDropper(ANSIConstraintCommon, SchemaDropper): | ||||
def _visit_constraint(self, constraint): | ||||
constraint.name = self.get_constraint_name(constraint) | ||||
self.append(self.process(DropConstraint(constraint, cascade=constraint.cascade))) | ||||
self.execute() | ||||
class ANSIDialect(DefaultDialect): | ||||
columngenerator = ANSIColumnGenerator | ||||
columndropper = ANSIColumnDropper | ||||
schemachanger = ANSISchemaChanger | ||||
constraintgenerator = ANSIConstraintGenerator | ||||
constraintdropper = ANSIConstraintDropper | ||||