sync sqlalchemy migrate with latest changes

rhodecode/lib/dbmigrate/migrate/versioning/schema.py

@@ -1,220 +1,221
 """
    Database schema version management.
 """
 import sys
 import logging
 
 from sqlalchemy import (Table, Column, MetaData, String, Text, Integer,
     create_engine)
 from sqlalchemy.sql import and_
-from sqlalchemy import exceptions as sa_exceptions
+from sqlalchemy import exc as sa_exceptions
 from sqlalchemy.sql import bindparam
 
 from rhodecode.lib.dbmigrate.migrate import exceptions
 from rhodecode.lib.dbmigrate.migrate.changeset import SQLA_07
 from rhodecode.lib.dbmigrate.migrate.versioning import genmodel, schemadiff
 from rhodecode.lib.dbmigrate.migrate.versioning.repository import Repository
 from rhodecode.lib.dbmigrate.migrate.versioning.util import load_model
 from rhodecode.lib.dbmigrate.migrate.versioning.version import VerNum
 
 
 log = logging.getLogger(__name__)
 
+
 class ControlledSchema(object):
     """A database under version control"""
 
     def __init__(self, engine, repository):
         if isinstance(repository, basestring):
             repository = Repository(repository)
         self.engine = engine
         self.repository = repository
         self.meta = MetaData(engine)
         self.load()
 
     def __eq__(self, other):
         """Compare two schemas by repositories and versions"""
         return (self.repository is other.repository \
             and self.version == other.version)
 
     def load(self):
         """Load controlled schema version info from DB"""
         tname = self.repository.version_table
         try:
             if not hasattr(self, 'table') or self.table is None:
                     self.table = Table(tname, self.meta, autoload=True)
 
             result = self.engine.execute(self.table.select(
                 self.table.c.repository_id == str(self.repository.id)))
 
             data = list(result)[0]
         except:
             cls, exc, tb = sys.exc_info()
             raise exceptions.DatabaseNotControlledError, exc.__str__(), tb
 
         self.version = data['version']
         return data
 
     def drop(self):
         """
         Remove version control from a database.
         """
         if SQLA_07:
             try:
                 self.table.drop()
             except sa_exceptions.DatabaseError:
                 raise exceptions.DatabaseNotControlledError(str(self.table))
         else:
             try:
                 self.table.drop()
             except (sa_exceptions.SQLError):
                 raise exceptions.DatabaseNotControlledError(str(self.table))
 
     def changeset(self, version=None):
         """API to Changeset creation.
 
         Uses self.version for start version and engine.name
         to get database name.
         """
         database = self.engine.name
         start_ver = self.version
         changeset = self.repository.changeset(database, start_ver, version)
         return changeset
 
     def runchange(self, ver, change, step):
         startver = ver
         endver = ver + step
         # Current database version must be correct! Don't run if corrupt!
         if self.version != startver:
             raise exceptions.InvalidVersionError("%s is not %s" % \
                                                      (self.version, startver))
         # Run the change
         change.run(self.engine, step)
 
         # Update/refresh database version
         self.update_repository_table(startver, endver)
         self.load()
 
     def update_repository_table(self, startver, endver):
         """Update version_table with new information"""
         update = self.table.update(and_(self.table.c.version == int(startver),
              self.table.c.repository_id == str(self.repository.id)))
         self.engine.execute(update, version=int(endver))
 
     def upgrade(self, version=None):
         """
         Upgrade (or downgrade) to a specified version, or latest version.
         """
         changeset = self.changeset(version)
         for ver, change in changeset:
             self.runchange(ver, change, changeset.step)
 
     def update_db_from_model(self, model):
         """
         Modify the database to match the structure of the current Python model.
         """
         model = load_model(model)
 
         diff = schemadiff.getDiffOfModelAgainstDatabase(
             model, self.engine, excludeTables=[self.repository.version_table]
             )
         genmodel.ModelGenerator(diff,self.engine).runB2A()
 
         self.update_repository_table(self.version, int(self.repository.latest))
 
         self.load()
 
     @classmethod
     def create(cls, engine, repository, version=None):
         """
         Declare a database to be under a repository's version control.
 
         :raises: :exc:`DatabaseAlreadyControlledError`
         :returns: :class:`ControlledSchema`
         """
         # Confirm that the version # is valid: positive, integer,
         # exists in repos
         if isinstance(repository, basestring):
             repository = Repository(repository)
         version = cls._validate_version(repository, version)
         table = cls._create_table_version(engine, repository, version)
         # TODO: history table
         # Load repository information and return
         return cls(engine, repository)
 
     @classmethod
     def _validate_version(cls, repository, version):
         """
         Ensures this is a valid version number for this repository.
 
         :raises: :exc:`InvalidVersionError` if invalid
         :return: valid version number
         """
         if version is None:
             version = 0
         try:
             version = VerNum(version) # raises valueerror
             if version < 0 or version > repository.latest:
                 raise ValueError()
         except ValueError:
             raise exceptions.InvalidVersionError(version)
         return version
 
     @classmethod
     def _create_table_version(cls, engine, repository, version):
         """
         Creates the versioning table in a database.
 
         :raises: :exc:`DatabaseAlreadyControlledError`
         """
         # Create tables
         tname = repository.version_table
         meta = MetaData(engine)
 
         table = Table(
             tname, meta,
             Column('repository_id', String(250), primary_key=True),
             Column('repository_path', Text),
             Column('version', Integer), )
 
         # there can be multiple repositories/schemas in the same db
         if not table.exists():
             table.create()
 
         # test for existing repository_id
         s = table.select(table.c.repository_id == bindparam("repository_id"))
         result = engine.execute(s, repository_id=repository.id)
         if result.fetchone():
             raise exceptions.DatabaseAlreadyControlledError
 
         # Insert data
         engine.execute(table.insert().values(
                            repository_id=repository.id,
                            repository_path=repository.path,
                            version=int(version)))
         return table
 
     @classmethod
     def compare_model_to_db(cls, engine, model, repository):
         """
         Compare the current model against the current database.
         """
         if isinstance(repository, basestring):
             repository = Repository(repository)
         model = load_model(model)
 
         diff = schemadiff.getDiffOfModelAgainstDatabase(
             model, engine, excludeTables=[repository.version_table])
         return diff
 
     @classmethod
     def create_model(cls, engine, repository, declarative=False):
         """
         Dump the current database as a Python model.
         """
         if isinstance(repository, basestring):
             repository = Repository(repository)
 
         diff = schemadiff.getDiffOfModelAgainstDatabase(
             MetaData(), engine, excludeTables=[repository.version_table]
             )
         return genmodel.ModelGenerator(diff, engine, declarative).genBDefinition()

rhodecode/lib/dbmigrate/migrate/versioning/schemadiff.py

@@ -1,293 +1,295
 """
    Schema differencing support.
 """
 
 import logging
 import sqlalchemy
 
 from rhodecode.lib.dbmigrate.migrate.changeset import SQLA_06
 from sqlalchemy.types import Float
 
 log = logging.getLogger(__name__)
 
+
 def getDiffOfModelAgainstDatabase(metadata, engine, excludeTables=None):
     """
     Return differences of model against database.
 
     :return: object which will evaluate to :keyword:`True` if there \
       are differences else :keyword:`False`.
     """
-    db_metadata = sqlalchemy.MetaData(engine, reflect=True)
+    db_metadata = sqlalchemy.MetaData(engine)
+    db_metadata.reflect()
 
     # sqlite will include a dynamically generated 'sqlite_sequence' table if
     # there are autoincrement sequences in the database; this should not be
     # compared.
     if engine.dialect.name == 'sqlite':
         if 'sqlite_sequence' in db_metadata.tables:
             db_metadata.remove(db_metadata.tables['sqlite_sequence'])
 
     return SchemaDiff(metadata, db_metadata,
                       labelA='model',
                       labelB='database',
                       excludeTables=excludeTables)
 
 
 def getDiffOfModelAgainstModel(metadataA, metadataB, excludeTables=None):
     """
     Return differences of model against another model.
 
     :return: object which will evaluate to :keyword:`True` if there \
       are differences else :keyword:`False`.
     """
     return SchemaDiff(metadataA, metadataB, excludeTables=excludeTables)
 
 
 class ColDiff(object):
     """
     Container for differences in one :class:`~sqlalchemy.schema.Column`
     between two :class:`~sqlalchemy.schema.Table` instances, ``A``
     and ``B``.
 
     .. attribute:: col_A
 
       The :class:`~sqlalchemy.schema.Column` object for A.
 
     .. attribute:: col_B
 
       The :class:`~sqlalchemy.schema.Column` object for B.
 
     .. attribute:: type_A
 
       The most generic type of the :class:`~sqlalchemy.schema.Column`
       object in A.
 
     .. attribute:: type_B
 
       The most generic type of the :class:`~sqlalchemy.schema.Column`
       object in A.
 
     """
 
     diff = False
 
     def __init__(self,col_A,col_B):
         self.col_A = col_A
         self.col_B = col_B
 
         self.type_A = col_A.type
         self.type_B = col_B.type
 
         self.affinity_A = self.type_A._type_affinity
         self.affinity_B = self.type_B._type_affinity
 
         if self.affinity_A is not self.affinity_B:
             self.diff = True
             return
 
         if isinstance(self.type_A,Float) or isinstance(self.type_B,Float):
             if not (isinstance(self.type_A,Float) and isinstance(self.type_B,Float)):
                 self.diff=True
                 return
 
         for attr in ('precision','scale','length'):
             A = getattr(self.type_A,attr,None)
             B = getattr(self.type_B,attr,None)
             if not (A is None or B is None) and A!=B:
                 self.diff=True
                 return
 
     def __nonzero__(self):
         return self.diff
 
 class TableDiff(object):
     """
     Container for differences in one :class:`~sqlalchemy.schema.Table`
     between two :class:`~sqlalchemy.schema.MetaData` instances, ``A``
     and ``B``.
 
     .. attribute:: columns_missing_from_A
 
       A sequence of column names that were found in B but weren't in
       A.
 
     .. attribute:: columns_missing_from_B
 
       A sequence of column names that were found in A but weren't in
       B.
 
     .. attribute:: columns_different
 
       A dictionary containing information about columns that were
       found to be different.
       It maps column names to a :class:`ColDiff` objects describing the
       differences found.
     """
     __slots__ = (
         'columns_missing_from_A',
         'columns_missing_from_B',
         'columns_different',
         )
 
     def __nonzero__(self):
         return bool(
             self.columns_missing_from_A or
             self.columns_missing_from_B or
             self.columns_different
             )
 
 class SchemaDiff(object):
     """
     Compute the difference between two :class:`~sqlalchemy.schema.MetaData`
     objects.
 
     The string representation of a :class:`SchemaDiff` will summarise
     the changes found between the two
     :class:`~sqlalchemy.schema.MetaData` objects.
 
     The length of a :class:`SchemaDiff` will give the number of
     changes found, enabling it to be used much like a boolean in
     expressions.
 
     :param metadataA:
       First :class:`~sqlalchemy.schema.MetaData` to compare.
 
     :param metadataB:
       Second :class:`~sqlalchemy.schema.MetaData` to compare.
 
     :param labelA:
       The label to use in messages about the first
       :class:`~sqlalchemy.schema.MetaData`.
 
     :param labelB:
       The label to use in messages about the second
       :class:`~sqlalchemy.schema.MetaData`.
 
     :param excludeTables:
       A sequence of table names to exclude.
 
     .. attribute:: tables_missing_from_A
 
       A sequence of table names that were found in B but weren't in
       A.
 
     .. attribute:: tables_missing_from_B
 
       A sequence of table names that were found in A but weren't in
       B.
 
     .. attribute:: tables_different
 
       A dictionary containing information about tables that were found
       to be different.
       It maps table names to a :class:`TableDiff` objects describing the
       differences found.
     """
 
     def __init__(self,
                  metadataA, metadataB,
                  labelA='metadataA',
                  labelB='metadataB',
                  excludeTables=None):
 
         self.metadataA, self.metadataB = metadataA, metadataB
         self.labelA, self.labelB = labelA, labelB
         self.label_width = max(len(labelA),len(labelB))
         excludeTables = set(excludeTables or [])
 
         A_table_names = set(metadataA.tables.keys())
         B_table_names = set(metadataB.tables.keys())
 
         self.tables_missing_from_A = sorted(
             B_table_names - A_table_names - excludeTables
             )
         self.tables_missing_from_B = sorted(
             A_table_names - B_table_names - excludeTables
             )
 
         self.tables_different = {}
         for table_name in A_table_names.intersection(B_table_names):
 
             td = TableDiff()
 
             A_table = metadataA.tables[table_name]
             B_table = metadataB.tables[table_name]
 
             A_column_names = set(A_table.columns.keys())
             B_column_names = set(B_table.columns.keys())
 
             td.columns_missing_from_A = sorted(
                 B_column_names - A_column_names
                 )
 
             td.columns_missing_from_B = sorted(
                 A_column_names - B_column_names
                 )
 
             td.columns_different = {}
 
             for col_name in A_column_names.intersection(B_column_names):
 
                 cd = ColDiff(
                     A_table.columns.get(col_name),
                     B_table.columns.get(col_name)
                     )
 
                 if cd:
                     td.columns_different[col_name]=cd
 
             # XXX - index and constraint differences should
             #       be checked for here
 
             if td:
                 self.tables_different[table_name]=td
 
     def __str__(self):
         ''' Summarize differences. '''
         out = []
         column_template ='      %%%is: %%r' % self.label_width
 
         for names,label in (
             (self.tables_missing_from_A,self.labelA),
             (self.tables_missing_from_B,self.labelB),
             ):
             if names:
                 out.append(
                     '  tables missing from %s: %s' % (
                         label,', '.join(sorted(names))
                         )
                     )
 
         for name,td in sorted(self.tables_different.items()):
             out.append(
                '  table with differences: %s' % name
                )
             for names,label in (
                 (td.columns_missing_from_A,self.labelA),
                 (td.columns_missing_from_B,self.labelB),
                 ):
                 if names:
                     out.append(
                         '    %s missing these columns: %s' % (
                             label,', '.join(sorted(names))
                             )
                         )
             for name,cd in td.columns_different.items():
                 out.append('    column with differences: %s' % name)
                 out.append(column_template % (self.labelA,cd.col_A))
                 out.append(column_template % (self.labelB,cd.col_B))
 
         if out:
             out.insert(0, 'Schema diffs:')
             return '\n'.join(out)
         else:
             return 'No schema diffs'
 
     def __len__(self):
         """
         Used in bool evaluation, return of 0 means no diffs.
         """
         return (
             len(self.tables_missing_from_A) +
             len(self.tables_missing_from_B) +
             len(self.tables_different)
             )