415 lines
13 KiB
Python

"""
Schema module providing common schema operations.
"""
import sqlalchemy
from migrate.changeset.databases.visitor import (get_engine_visitor,
run_single_visitor)
from migrate.changeset.exceptions import *
__all__ = [
'create_column',
'drop_column',
'alter_column',
'rename_table',
'rename_index',
'ChangesetTable',
'ChangesetColumn',
'ChangesetIndex',
]
def create_column(column, table=None, *p, **k):
"""Create a column, given the table
API to :meth:`column.create`
"""
if table is not None:
return table.create_column(column, *p, **k)
return column.create(*p, **k)
def drop_column(column, table=None, *p, **k):
"""Drop a column, given the table
API to :meth:`column.drop`
"""
if table is not None:
return table.drop_column(column, *p, **k)
return column.drop(*p, **k)
def rename_table(table, name, engine=None):
"""Rename a table.
If Table instance is given, engine is not used.
API to :meth:`table.rename`
:param table: Table to be renamed
:param name: new name
:param engine: Engine instance
:type table: string or Table instance
:type name: string
:type engine: obj
"""
table = _to_table(table, engine)
table.rename(name)
def rename_index(index, name, table=None, engine=None):
"""Rename an index.
If Index and Table object instances are given,
table and engine are not used.
API to :meth:`index.rename`
:param index: Index to be renamed
:param name: new name
:param table: Table to which Index is reffered
:param engine: Engine instance
:type index: string or Index instance
:type name: string
:type table: string or Table instance
:type engine: obj
"""
index = _to_index(index, table, engine)
index.rename(name)
def alter_column(*p, **k):
"""Alter a column.
Parameters: column name, table name, an engine, and the properties
of that column to change
API to :meth:`column.alter`
"""
if len(p) and isinstance(p[0], sqlalchemy.Column):
col = p[0]
else:
col = None
if 'table' not in k:
k['table'] = col.table
if 'engine' not in k:
k['engine'] = k['table'].bind
engine = k['engine']
delta = _ColumnDelta(*p, **k)
delta.result_column.delta = delta
delta.result_column.table = delta.table
visitorcallable = get_engine_visitor(engine, 'schemachanger')
engine._run_visitor(visitorcallable, delta.result_column)
# Update column
if col is not None:
# Special case: change column key on rename, if key not
# explicit
#
# Used by SA : table.c.[key]
#
# This fails if the key was explit AND equal to the column
# name. (It changes the key name when it shouldn't.)
#
# Not much we can do about it.
if 'name' in delta.keys():
if (col.name == col.key):
newname = delta['name']
del col.table.c[col.key]
setattr(col, 'key', newname)
col.table.c[col.key] = col
# Change all other attrs
for key, val in delta.iteritems():
setattr(col, key, val)
def _to_table(table, engine=None):
"""Return if instance of Table, else construct new with metadata"""
if isinstance(table, sqlalchemy.Table):
return table
# Given: table name, maybe an engine
meta = sqlalchemy.MetaData()
if engine is not None:
meta.bind = engine
return sqlalchemy.Table(table, meta)
def _to_index(index, table=None, engine=None):
"""Return if instance of Index, else construct new with metadata"""
if isinstance(index, sqlalchemy.Index):
return index
# Given: index name; table name required
table = _to_table(table, engine)
ret = sqlalchemy.Index(index)
ret.table = table
return ret
class _ColumnDelta(dict):
"""Extracts the differences between two columns/column-parameters"""
def __init__(self, *p, **k):
"""Extract ALTER-able differences from two columns.
May receive parameters arranged in several different ways:
* old_column_object,new_column_object,*parameters Identifies
attributes that differ between the two columns.
Parameters specified outside of either column are always
executed and override column differences.
* column_object,[current_name,]*parameters Parameters
specified are changed; table name is extracted from column
object. Name is changed to column_object.name from
current_name, if current_name is specified. If not
specified, name is unchanged.
* current_name,table,*parameters 'table' may be either an
object or a name
"""
# Things are initialized differently depending on how many column
# parameters are given. Figure out how many and call the appropriate
# method.
if len(p) >= 1 and isinstance(p[0], sqlalchemy.Column):
# At least one column specified
if len(p) >= 2 and isinstance(p[1], sqlalchemy.Column):
# Two columns specified
func = self._init_2col
else:
# Exactly one column specified
func = self._init_1col
else:
# Zero columns specified
func = self._init_0col
diffs = func(*p, **k)
self._set_diffs(diffs)
# Column attributes that can be altered
diff_keys = ('name',
'type',
'nullable',
'default',
'server_default',
'primary_key',
'foreign_key')
@property
def table(self):
if isinstance(self._table, sqlalchemy.Table):
return self._table
def _init_0col(self, current_name, *p, **k):
p, k = self._init_normalize_params(p, k)
table = k.pop('table')
self.current_name = current_name
self._table = table
self.result_column = table.c.get(current_name, None)
return k
def _init_1col(self, col, *p, **k):
p, k = self._init_normalize_params(p, k)
self._table = k.pop('table', None) or col.table
self.result_column = col.copy()
if 'current_name' in k:
# Renamed
self.current_name = k.pop('current_name')
k.setdefault('name', col.name)
else:
self.current_name = col.name
return k
def _init_2col(self, start_col, end_col, *p, **k):
p, k = self._init_normalize_params(p, k)
self.result_column = start_col.copy()
self._table = k.pop('table', None) or start_col.table \
or end_col.table
self.current_name = start_col.name
for key in ('name', 'nullable', 'default', 'server_default',
'primary_key', 'foreign_key'):
val = getattr(end_col, key, None)
if getattr(start_col, key, None) != val:
k.setdefault(key, val)
if not self.column_types_eq(start_col.type, end_col.type):
k.setdefault('type', end_col.type)
return k
def _init_normalize_params(self, p, k):
p = list(p)
if len(p):
k.setdefault('name', p.pop(0))
if len(p):
k.setdefault('type', p.pop(0))
# TODO: sequences? FKs?
return p, k
def _set_diffs(self, diffs):
for key in self.diff_keys:
if key in diffs:
self[key] = diffs[key]
if getattr(self, 'result_column', None) is not None:
setattr(self.result_column, key, diffs[key])
def column_types_eq(self, this, that):
ret = isinstance(this, that.__class__)
ret = ret or isinstance(that, this.__class__)
# String length is a special case
if ret and isinstance(that, sqlalchemy.types.String):
ret = (getattr(this, 'length', None) == \
getattr(that, 'length', None))
return ret
class ChangesetTable(object):
"""Changeset extensions to SQLAlchemy tables."""
def create_column(self, column):
"""Creates a column.
The column parameter may be a column definition or the name of
a column in this table.
"""
if not isinstance(column, sqlalchemy.Column):
# It's a column name
column = getattr(self.c, str(column))
column.create(table=self)
def drop_column(self, column):
"""Drop a column, given its name or definition."""
if not isinstance(column, sqlalchemy.Column):
# It's a column name
try:
column = getattr(self.c, str(column))
except AttributeError:
# That column isn't part of the table. We don't need
# its entire definition to drop the column, just its
# name, so create a dummy column with the same name.
column = sqlalchemy.Column(str(column))
column.drop(table=self)
def rename(self, name, *args, **kwargs):
"""Rename this table.
This changes both the database name and the name of this
Python object
"""
engine = self.bind
self.new_name = name
visitorcallable = get_engine_visitor(engine, 'schemachanger')
run_single_visitor(engine, visitorcallable, self, *args, **kwargs)
# Fix metadata registration
self.name = name
self.deregister()
self._set_parent(self.metadata)
def _meta_key(self):
return sqlalchemy.schema._get_table_key(self.name, self.schema)
def deregister(self):
"""Remove this table from its metadata"""
key = self._meta_key()
meta = self.metadata
if key in meta.tables:
del meta.tables[key]
class ChangesetColumn(object):
"""Changeset extensions to SQLAlchemy columns"""
def alter(self, *p, **k):
"""Alter a column's definition: ``ALTER TABLE ALTER COLUMN``.
May supply a new column object, or a list of properties to
change.
For example; the following are equivalent:
col.alter(Column('myint', Integer, nullable=False))
col.alter('myint', Integer, nullable=False)
col.alter(name='myint', type=Integer, nullable=False)
Column name, type, default, and nullable may be changed
here. Note that for column defaults, only PassiveDefaults are
managed by the database - changing others doesn't make sense.
:param table: Table to be altered
:param engine: Engine to be used
"""
if 'table' not in k:
k['table'] = self.table
if 'engine' not in k:
k['engine'] = k['table'].bind
return alter_column(self, *p, **k)
def create(self, table=None, index_name=None, unique_name=None,
primary_key_name=None, *args, **kwargs):
"""Create this column in the database.
Assumes the given table exists. ``ALTER TABLE ADD COLUMN``,
for most databases.
"""
self.index_name = index_name
self.unique_name = unique_name
self.primary_key_name = primary_key_name
for cons in ('index_name', 'unique_name', 'primary_key_name'):
self._check_sanity_constraints(cons)
self.add_to_table(table)
engine = self.table.bind
visitorcallable = get_engine_visitor(engine, 'columngenerator')
engine._run_visitor(visitorcallable, self, *args, **kwargs)
return self
def drop(self, table=None, *args, **kwargs):
"""Drop this column from the database, leaving its table intact.
``ALTER TABLE DROP COLUMN``, for most databases.
"""
if table is not None:
self.table = table
self.remove_from_table(self.table)
engine = self.table.bind
visitorcallable = get_engine_visitor(engine, 'columndropper')
engine._run_visitor(visitorcallable, self, *args, **kwargs)
return self
def add_to_table(self, table):
if table and not self.table:
self._set_parent(table)
def remove_from_table(self, table):
# TODO: remove indexes, primary keys, constraints, etc
if table.c.contains_column(self):
table.c.remove(self)
def _check_sanity_constraints(self, name):
obj = getattr(self, name)
if (getattr(self, name[:-5]) and not obj):
raise InvalidConstraintError("Column.create() accepts index_name,"
" primary_key_name and unique_name to generate constraints")
if not isinstance(obj, basestring) and obj is not None:
raise InvalidConstraintError(
"%s argument for column must be constraint name" % name)
class ChangesetIndex(object):
"""Changeset extensions to SQLAlchemy Indexes."""
__visit_name__ = 'index'
def rename(self, name, *args, **kwargs):
"""Change the name of an index.
This changes both the Python object name and the database
name.
"""
engine = self.table.bind
self.new_name = name
visitorcallable = get_engine_visitor(engine, 'schemachanger')
engine._run_visitor(visitorcallable, self, *args, **kwargs)
self.name = name