diff rhodecode/lib/dbmigrate/migrate/changeset/schema.py @ 833:9753e0907827 beta

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
line wrap: on
line diff
--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/rhodecode/lib/dbmigrate/migrate/changeset/schema.py	Sat Dec 11 01:54:12 2010 +0100
@@ -0,0 +1,669 @@
+"""
+   Schema module providing common schema operations.
+"""
+import warnings
+
+from UserDict import DictMixin
+
+import sqlalchemy
+
+from sqlalchemy.schema import ForeignKeyConstraint
+from sqlalchemy.schema import UniqueConstraint
+
+from migrate.exceptions import *
+from migrate.changeset import SQLA_06
+from migrate.changeset.databases.visitor import (get_engine_visitor,
+                                                 run_single_visitor)
+
+
+__all__ = [
+    'create_column',
+    'drop_column',
+    'alter_column',
+    'rename_table',
+    'rename_index',
+    'ChangesetTable',
+    'ChangesetColumn',
+    'ChangesetIndex',
+    'ChangesetDefaultClause',
+    'ColumnDelta',
+]
+
+DEFAULT_ALTER_METADATA = True
+
+
+def create_column(column, table=None, *p, **kw):
+    """Create a column, given the table.
+    
+    API to :meth:`ChangesetColumn.create`.
+    """
+    if table is not None:
+        return table.create_column(column, *p, **kw)
+    return column.create(*p, **kw)
+
+
+def drop_column(column, table=None, *p, **kw):
+    """Drop a column, given the table.
+    
+    API to :meth:`ChangesetColumn.drop`.
+    """
+    if table is not None:
+        return table.drop_column(column, *p, **kw)
+    return column.drop(*p, **kw)
+
+
+def rename_table(table, name, engine=None, **kw):
+    """Rename a table.
+
+    If Table instance is given, engine is not used.
+
+    API to :meth:`ChangesetTable.rename`.
+
+    :param table: Table to be renamed.
+    :param name: New name for Table.
+    :param engine: Engine instance.
+    :type table: string or Table instance
+    :type name: string
+    :type engine: obj
+    """
+    table = _to_table(table, engine)
+    table.rename(name, **kw)
+
+
+def rename_index(index, name, table=None, engine=None, **kw):
+    """Rename an index.
+
+    If Index instance is given,
+    table and engine are not used.
+
+    API to :meth:`ChangesetIndex.rename`.
+
+    :param index: Index to be renamed.
+    :param name: New name for index.
+    :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, **kw)
+
+
+def alter_column(*p, **k):
+    """Alter a column.
+
+    This is a helper function that creates a :class:`ColumnDelta` and
+    runs it.
+
+    :argument column:
+      The name of the column to be altered or a
+      :class:`ChangesetColumn` column representing it.
+
+    :param table:
+      A :class:`~sqlalchemy.schema.Table` or table name to
+      for the table where the column will be changed.
+
+    :param engine:
+      The :class:`~sqlalchemy.engine.base.Engine` to use for table
+      reflection and schema alterations.
+    
+    :param alter_metadata:
+      If `True`, which is the default, the
+      :class:`~sqlalchemy.schema.Column` will also modified.
+      If `False`, the :class:`~sqlalchemy.schema.Column` will be left
+      as it was.
+    
+    :returns: A :class:`ColumnDelta` instance representing the change.
+
+    
+    """
+    
+    k.setdefault('alter_metadata', DEFAULT_ALTER_METADATA)
+
+    if 'table' not in k and isinstance(p[0], sqlalchemy.Column):
+        k['table'] = p[0].table
+    if 'engine' not in k:
+        k['engine'] = k['table'].bind
+
+    # deprecation
+    if len(p) >= 2 and isinstance(p[1], sqlalchemy.Column):
+        warnings.warn(
+            "Passing a Column object to alter_column is deprecated."
+            " Just pass in keyword parameters instead.",
+            MigrateDeprecationWarning
+            )
+    engine = k['engine']
+    delta = ColumnDelta(*p, **k)
+
+    visitorcallable = get_engine_visitor(engine, 'schemachanger')
+    engine._run_visitor(visitorcallable, delta)
+
+    return delta
+
+
+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(DictMixin, sqlalchemy.schema.SchemaItem):
+    """Extracts the differences between two columns/column-parameters
+
+        May receive parameters arranged in several different ways:
+
+        * **current_column, new_column, \*p, \*\*kw**
+            Additional parameters can be specified to override column
+            differences.
+
+        * **current_column, \*p, \*\*kw**
+            Additional parameters alter current_column. Table name is extracted
+            from current_column object.
+            Name is changed to current_column.name from current_name,
+            if current_name is specified.
+
+        * **current_col_name, \*p, \*\*kw**
+            Table kw must specified.
+
+        :param table: Table at which current Column should be bound to.\
+        If table name is given, reflection will be used.
+        :type table: string or Table instance
+        :param alter_metadata: If True, it will apply changes to metadata.
+        :type alter_metadata: bool
+        :param metadata: If `alter_metadata` is true, \
+        metadata is used to reflect table names into
+        :type metadata: :class:`MetaData` instance
+        :param engine: When reflecting tables, either engine or metadata must \
+        be specified to acquire engine object.
+        :type engine: :class:`Engine` instance
+        :returns: :class:`ColumnDelta` instance provides interface for altered attributes to \
+        `result_column` through :func:`dict` alike object.
+
+        * :class:`ColumnDelta`.result_column is altered column with new attributes
+
+        * :class:`ColumnDelta`.current_name is current name of column in db
+
+
+    """
+
+    # Column attributes that can be altered
+    diff_keys = ('name', 'type', 'primary_key', 'nullable',
+        'server_onupdate', 'server_default', 'autoincrement')
+    diffs = dict()
+    __visit_name__ = 'column'
+
+    def __init__(self, *p, **kw):
+        self.alter_metadata = kw.pop("alter_metadata", False)
+        self.meta = kw.pop("metadata", None)
+        self.engine = kw.pop("engine", None)
+
+        # 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
+                diffs = self.compare_2_columns(*p, **kw)
+            else:
+                # Exactly one column specified
+                diffs = self.compare_1_column(*p, **kw)
+        else:
+            # Zero columns specified
+            if not len(p) or not isinstance(p[0], basestring):
+                raise ValueError("First argument must be column name")
+            diffs = self.compare_parameters(*p, **kw)
+
+        self.apply_diffs(diffs)
+
+    def __repr__(self):
+        return '<ColumnDelta altermetadata=%r, %s>' % (self.alter_metadata,
+            super(ColumnDelta, self).__repr__())
+
+    def __getitem__(self, key):
+        if key not in self.keys():
+            raise KeyError("No such diff key, available: %s" % self.diffs)
+        return getattr(self.result_column, key)
+
+    def __setitem__(self, key, value):
+        if key not in self.keys():
+            raise KeyError("No such diff key, available: %s" % self.diffs)
+        setattr(self.result_column, key, value)
+
+    def __delitem__(self, key):
+        raise NotImplementedError
+
+    def keys(self):
+        return self.diffs.keys()
+
+    def compare_parameters(self, current_name, *p, **k):
+        """Compares Column objects with reflection"""
+        self.table = k.pop('table')
+        self.result_column = self._table.c.get(current_name)
+        if len(p):
+            k = self._extract_parameters(p, k, self.result_column)
+        return k
+
+    def compare_1_column(self, col, *p, **k):
+        """Compares one Column object"""
+        self.table = k.pop('table', None)
+        if self.table is None:
+            self.table = col.table
+        self.result_column = col
+        if len(p):
+            k = self._extract_parameters(p, k, self.result_column)
+        return k
+
+    def compare_2_columns(self, old_col, new_col, *p, **k):
+        """Compares two Column objects"""
+        self.process_column(new_col)
+        self.table = k.pop('table', None)
+        # we cannot use bool() on table in SA06 
+        if self.table is None:
+            self.table = old_col.table
+        if self.table is None:
+            new_col.table
+        self.result_column = old_col
+
+        # set differences
+        # leave out some stuff for later comp
+        for key in (set(self.diff_keys) - set(('type',))):
+            val = getattr(new_col, key, None)
+            if getattr(self.result_column, key, None) != val:
+                k.setdefault(key, val)
+
+        # inspect types
+        if not self.are_column_types_eq(self.result_column.type, new_col.type):
+            k.setdefault('type', new_col.type)
+
+        if len(p):
+            k = self._extract_parameters(p, k, self.result_column)
+        return k
+
+    def apply_diffs(self, diffs):
+        """Populate dict and column object with new values"""
+        self.diffs = diffs
+        for key in self.diff_keys:
+            if key in diffs:
+                setattr(self.result_column, key, diffs[key])
+
+        self.process_column(self.result_column)
+
+        # create an instance of class type if not yet
+        if 'type' in diffs and callable(self.result_column.type):
+            self.result_column.type = self.result_column.type()
+
+        # add column to the table
+        if self.table is not None and self.alter_metadata:
+            self.result_column.add_to_table(self.table)
+
+    def are_column_types_eq(self, old_type, new_type):
+        """Compares two types to be equal"""
+        ret = old_type.__class__ == new_type.__class__
+
+        # String length is a special case
+        if ret and isinstance(new_type, sqlalchemy.types.String):
+            ret = (getattr(old_type, 'length', None) == \
+                       getattr(new_type, 'length', None))
+        return ret
+
+    def _extract_parameters(self, p, k, column):
+        """Extracts data from p and modifies diffs"""
+        p = list(p)
+        while len(p):
+            if isinstance(p[0], basestring):
+                k.setdefault('name', p.pop(0))
+            elif isinstance(p[0], sqlalchemy.types.AbstractType):
+                k.setdefault('type', p.pop(0))
+            elif callable(p[0]):
+                p[0] = p[0]()
+            else:
+                break
+
+        if len(p):
+            new_col = column.copy_fixed()
+            new_col._init_items(*p)
+            k = self.compare_2_columns(column, new_col, **k)
+        return k
+
+    def process_column(self, column):
+        """Processes default values for column"""
+        # XXX: this is a snippet from SA processing of positional parameters
+        if not SQLA_06 and column.args:
+            toinit = list(column.args)
+        else:
+            toinit = list()
+
+        if column.server_default is not None:
+            if isinstance(column.server_default, sqlalchemy.FetchedValue):
+                toinit.append(column.server_default)
+            else:
+                toinit.append(sqlalchemy.DefaultClause(column.server_default))
+        if column.server_onupdate is not None:
+            if isinstance(column.server_onupdate, FetchedValue):
+                toinit.append(column.server_default)
+            else:
+                toinit.append(sqlalchemy.DefaultClause(column.server_onupdate,
+                                            for_update=True))
+        if toinit:
+            column._init_items(*toinit)
+            
+        if not SQLA_06:
+            column.args = []
+
+    def _get_table(self):
+        return getattr(self, '_table', None)
+
+    def _set_table(self, table):
+        if isinstance(table, basestring):
+            if self.alter_metadata:
+                if not self.meta:
+                    raise ValueError("metadata must be specified for table"
+                        " reflection when using alter_metadata")
+                meta = self.meta
+                if self.engine:
+                    meta.bind = self.engine
+            else:
+                if not self.engine and not self.meta:
+                    raise ValueError("engine or metadata must be specified"
+                        " to reflect tables")
+                if not self.engine:
+                    self.engine = self.meta.bind
+                meta = sqlalchemy.MetaData(bind=self.engine)
+            self._table = sqlalchemy.Table(table, meta, autoload=True)
+        elif isinstance(table, sqlalchemy.Table):
+            self._table = table
+            if not self.alter_metadata:
+                self._table.meta = sqlalchemy.MetaData(bind=self._table.bind)
+
+    def _get_result_column(self):
+        return getattr(self, '_result_column', None)
+
+    def _set_result_column(self, column):
+        """Set Column to Table based on alter_metadata evaluation."""
+        self.process_column(column)
+        if not hasattr(self, 'current_name'):
+            self.current_name = column.name
+        if self.alter_metadata:
+            self._result_column = column
+        else:
+            self._result_column = column.copy_fixed()
+
+    table = property(_get_table, _set_table)
+    result_column = property(_get_result_column, _set_result_column)
+
+
+class ChangesetTable(object):
+    """Changeset extensions to SQLAlchemy tables."""
+
+    def create_column(self, column, *p, **kw):
+        """Creates a column.
+
+        The column parameter may be a column definition or the name of
+        a column in this table.
+
+        API to :meth:`ChangesetColumn.create`
+
+        :param column: Column to be created
+        :type column: Column instance or string
+        """
+        if not isinstance(column, sqlalchemy.Column):
+            # It's a column name
+            column = getattr(self.c, str(column))
+        column.create(table=self, *p, **kw)
+
+    def drop_column(self, column, *p, **kw):
+        """Drop a column, given its name or definition.
+
+        API to :meth:`ChangesetColumn.drop`
+
+        :param column: Column to be droped
+        :type column: Column instance or string
+        """
+        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), sqlalchemy.Integer())
+        column.drop(table=self, *p, **kw)
+
+    def rename(self, name, connection=None, **kwargs):
+        """Rename this table.
+
+        :param name: New name of the table.
+        :type name: string
+        :param alter_metadata: If True, table will be removed from metadata
+        :type alter_metadata: bool
+        :param connection: reuse connection istead of creating new one.
+        :type connection: :class:`sqlalchemy.engine.base.Connection` instance
+        """
+        self.alter_metadata = kwargs.pop('alter_metadata', DEFAULT_ALTER_METADATA)
+        engine = self.bind
+        self.new_name = name
+        visitorcallable = get_engine_visitor(engine, 'schemachanger')
+        run_single_visitor(engine, visitorcallable, self, connection, **kwargs)
+
+        # Fix metadata registration
+        if self.alter_metadata:
+            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):
+        """Makes a call to :func:`alter_column` for the column this
+        method is called on. 
+        """
+        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, populate_default=True, connection=None, **kwargs):
+        """Create this column in the database.
+
+        Assumes the given table exists. ``ALTER TABLE ADD COLUMN``,
+        for most databases.
+
+        :param table: Table instance to create on.
+        :param index_name: Creates :class:`ChangesetIndex` on this column.
+        :param unique_name: Creates :class:\
+`~migrate.changeset.constraint.UniqueConstraint` on this column.
+        :param primary_key_name: Creates :class:\
+`~migrate.changeset.constraint.PrimaryKeyConstraint` on this column.
+        :param alter_metadata: If True, column will be added to table object.
+        :param populate_default: If True, created column will be \
+populated with defaults
+        :param connection: reuse connection istead of creating new one.
+        :type table: Table instance
+        :type index_name: string
+        :type unique_name: string
+        :type primary_key_name: string
+        :type alter_metadata: bool
+        :type populate_default: bool
+        :type connection: :class:`sqlalchemy.engine.base.Connection` instance
+
+        :returns: self
+        """
+        self.populate_default = populate_default
+        self.alter_metadata = kwargs.pop('alter_metadata', DEFAULT_ALTER_METADATA)
+        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)
+
+        if self.alter_metadata:
+            self.add_to_table(table)
+        engine = self.table.bind
+        visitorcallable = get_engine_visitor(engine, 'columngenerator')
+        engine._run_visitor(visitorcallable, self, connection, **kwargs)
+
+        # TODO: reuse existing connection
+        if self.populate_default and self.default is not None:
+            stmt = table.update().values({self: engine._execute_default(self.default)})
+            engine.execute(stmt)
+
+        return self
+
+    def drop(self, table=None, connection=None, **kwargs):
+        """Drop this column from the database, leaving its table intact.
+
+        ``ALTER TABLE DROP COLUMN``, for most databases.
+
+        :param alter_metadata: If True, column will be removed from table object.
+        :type alter_metadata: bool
+        :param connection: reuse connection istead of creating new one.
+        :type connection: :class:`sqlalchemy.engine.base.Connection` instance
+        """
+        self.alter_metadata = kwargs.pop('alter_metadata', DEFAULT_ALTER_METADATA)
+        if table is not None:
+            self.table = table
+        engine = self.table.bind
+        if self.alter_metadata:
+            self.remove_from_table(self.table, unset_table=False)
+        visitorcallable = get_engine_visitor(engine, 'columndropper')
+        engine._run_visitor(visitorcallable, self, connection, **kwargs)
+        if self.alter_metadata:
+            self.table = None
+        return self
+
+    def add_to_table(self, table):
+        if table is not None  and self.table is None:
+            self._set_parent(table)
+
+    def _col_name_in_constraint(self, cons, name):
+        return False
+    
+    def remove_from_table(self, table, unset_table=True):
+        # TODO: remove primary keys, constraints, etc
+        if unset_table:
+            self.table = None
+            
+        to_drop = set()
+        for index in table.indexes:
+            columns = []
+            for col in index.columns:
+                if col.name != self.name:
+                    columns.append(col)
+            if columns:
+                index.columns = columns
+            else:
+                to_drop.add(index)
+        table.indexes = table.indexes - to_drop
+        
+        to_drop = set()
+        for cons in table.constraints:
+            # TODO: deal with other types of constraint
+            if isinstance(cons, (ForeignKeyConstraint,
+                                UniqueConstraint)):
+                for col_name in cons.columns:
+                    if not isinstance(col_name, basestring):
+                        col_name = col_name.name
+                    if self.name == col_name:
+                        to_drop.add(cons)
+        table.constraints = table.constraints - to_drop
+        
+        if table.c.contains_column(self):
+            table.c.remove(self)
+
+    # TODO: this is fixed in 0.6
+    def copy_fixed(self, **kw):
+        """Create a copy of this ``Column``, with all attributes."""
+        return sqlalchemy.Column(self.name, self.type, self.default,
+            key=self.key,
+            primary_key=self.primary_key,
+            nullable=self.nullable,
+            quote=self.quote,
+            index=self.index,
+            unique=self.unique,
+            onupdate=self.onupdate,
+            autoincrement=self.autoincrement,
+            server_default=self.server_default,
+            server_onupdate=self.server_onupdate,
+            *[c.copy(**kw) for c in self.constraints])
+
+    def _check_sanity_constraints(self, name):
+        """Check if constraints names are correct"""
+        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, connection=None, **kwargs):
+        """Change the name of an index.
+
+        :param name: New name of the Index.
+        :type name: string
+        :param alter_metadata: If True, Index object will be altered.
+        :type alter_metadata: bool
+        :param connection: reuse connection istead of creating new one.
+        :type connection: :class:`sqlalchemy.engine.base.Connection` instance
+        """
+        self.alter_metadata = kwargs.pop('alter_metadata', DEFAULT_ALTER_METADATA)
+        engine = self.table.bind
+        self.new_name = name
+        visitorcallable = get_engine_visitor(engine, 'schemachanger')
+        engine._run_visitor(visitorcallable, self, connection, **kwargs)
+        if self.alter_metadata:
+            self.name = name
+
+
+class ChangesetDefaultClause(object):
+    """Implements comparison between :class:`DefaultClause` instances"""
+
+    def __eq__(self, other):
+        if isinstance(other, self.__class__):
+            if self.arg == other.arg:
+                return True
+
+    def __ne__(self, other):
+        return not self.__eq__(other)