"""

   SQLAlchemy migrate provides two APIs :mod:`migrate.versioning` for

   database schema version and repository management and

   :mod:`migrate.changeset` that allows to define database schema changes

   using Python.

"""

from rhodecode.lib.dbmigrate.migrate.versioning import *

from rhodecode.lib.dbmigrate.migrate.changeset import *

"""

   This module extends SQLAlchemy and provides additional DDL [#]_

   support.

   .. [#] SQL Data Definition Language

"""

import re

import warnings

import sqlalchemy

from sqlalchemy import __version__ as _sa_version

warnings.simplefilter('always', DeprecationWarning)

_sa_version = tuple(int(re.match("\d+", x).group(0))

                    for x in _sa_version.split("."))

SQLA_06 = _sa_version >= (0, 6)

del re

del _sa_version

from rhodecode.lib.dbmigrate.migrate.changeset.schema import *

from rhodecode.lib.dbmigrate.migrate.changeset.constraint import *

sqlalchemy.schema.Table.__bases__ += (ChangesetTable,)

sqlalchemy.schema.Column.__bases__ += (ChangesetColumn,)

sqlalchemy.schema.Index.__bases__ += (ChangesetIndex,)

sqlalchemy.schema.DefaultClause.__bases__ += (ChangesetDefaultClause,)

"""

   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 rhodecode.lib.dbmigrate.migrate.exceptions import *

from rhodecode.lib.dbmigrate.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',

]

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.

    :returns: A :class:`ColumnDelta` instance representing the change.

    """

    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']

    # enough tests seem to break when metadata is always altered

    # that this crutch has to be left in until they can be sorted

    # out

    k['alter_metadata']=True

    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 metadata: A :class:`MetaData` instance to store

                         reflected table names

        :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()

                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 connection: reuse connection istead of creating new one.

        :type connection: :class:`sqlalchemy.engine.base.Connection` instance

        """

        engine = self.bind

        self.new_name = name

        visitorcallable = get_engine_visitor(engine, 'schemachanger')

        run_single_visitor(engine, visitorcallable, self, connection, **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):

        """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 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 populate_default: bool

        :type connection: :class:`sqlalchemy.engine.base.Connection` instance

        :returns: self

        """

        self.populate_default = populate_default

        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, 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 connection: reuse connection istead of creating new one.

        :type connection: :class:`sqlalchemy.engine.base.Connection` instance

        """

        if table is not None:

            self.table = table

        engine = self.table.bind

        visitorcallable = get_engine_visitor(engine, 'columndropper')

        engine._run_visitor(visitorcallable, self, connection, **kwargs)

        self.remove_from_table(self.table, unset_table=False)

        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 connection: reuse connection istead of creating new one.

        :type connection: :class:`sqlalchemy.engine.base.Connection` instance

        """

        engine = self.table.bind

        self.new_name = name

        visitorcallable = get_engine_visitor(engine, 'schemachanger')

        engine._run_visitor(visitorcallable, self, connection, **kwargs)

        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)

"""

   Provide exception classes for :mod:`migrate`

"""

class Error(Exception):

    """Error base class."""

class ApiError(Error):

    """Base class for API errors."""

class KnownError(ApiError):

    """A known error condition."""

class UsageError(ApiError):

    """A known error condition where help should be displayed."""

class ControlledSchemaError(Error):

    """Base class for controlled schema errors."""

class InvalidVersionError(ControlledSchemaError):

    """Invalid version number."""

class DatabaseNotControlledError(ControlledSchemaError):

    """Database should be under version control, but it's not."""

class DatabaseAlreadyControlledError(ControlledSchemaError):

    """Database shouldn't be under version control, but it is"""

class WrongRepositoryError(ControlledSchemaError):

    """This database is under version control by another repository."""

class NoSuchTableError(ControlledSchemaError):

    """The table does not exist."""

class PathError(Error):

    """Base class for path errors."""

class PathNotFoundError(PathError):

    """A path with no file was required; found a file."""

class PathFoundError(PathError):

    """A path with a file was required; found no file."""

class RepositoryError(Error):

    """Base class for repository errors."""

class InvalidRepositoryError(RepositoryError):

    """Invalid repository error."""

class ScriptError(Error):

    """Base class for script errors."""

class InvalidScriptError(ScriptError):

    """Invalid script error."""

class NotSupportedError(Error):

    """Not supported error"""

class InvalidConstraintError(Error):

    """Invalid constraint error"""

class MigrateDeprecationWarning(DeprecationWarning):

    """Warning for deprecated features in Migrate"""

"""

   This module provides an external API to the versioning system.

   .. versionchanged:: 0.6.0

    :func:`migrate.versioning.api.test` and schema diff functions

    changed order of positional arguments so all accept `url` and `repository`

    as first arguments.

   .. versionchanged:: 0.5.4

    ``--preview_sql`` displays source file when using SQL scripts.

    If Python script is used, it runs the action with mocked engine and

    returns captured SQL statements.

   .. versionchanged:: 0.5.4

    Deprecated ``--echo`` parameter in favour of new

    :func:`migrate.versioning.util.construct_engine` behavior.

"""

# Dear migrate developers,

#

# please do not comment this module using sphinx syntax because its

# docstrings are presented as user help and most users cannot

# interpret sphinx annotated ReStructuredText.

#

# Thanks,

# Jan Dittberner

import sys

import inspect

import logging

from rhodecode.lib.dbmigrate.migrate import exceptions

from rhodecode.lib.dbmigrate.migrate.versioning import repository, schema, version, \

    script as script_ # command name conflict

from rhodecode.lib.dbmigrate.migrate.versioning.util import catch_known_errors, with_engine

log = logging.getLogger(__name__)

command_desc = {

    'help': 'displays help on a given command',

    'create': 'create an empty repository at the specified path',

    'script': 'create an empty change Python script',

    'script_sql': 'create empty change SQL scripts for given database',

    'version': 'display the latest version available in a repository',

    'db_version': 'show the current version of the repository under version control',

    'source': 'display the Python code for a particular version in this repository',

    'version_control': 'mark a database as under this repository\'s version control',

    'upgrade': 'upgrade a database to a later version',

    'downgrade': 'downgrade a database to an earlier version',

    'drop_version_control': 'removes version control from a database',

    'manage': 'creates a Python script that runs Migrate with a set of default values',

    'test': 'performs the upgrade and downgrade command on the given database',

    'compare_model_to_db': 'compare MetaData against the current database state',

    'create_model': 'dump the current database as a Python model to stdout',

    'make_update_script_for_model': 'create a script changing the old MetaData to the new (current) MetaData',

    'update_db_from_model': 'modify the database to match the structure of the current MetaData',

}

__all__ = command_desc.keys()

Repository = repository.Repository

ControlledSchema = schema.ControlledSchema

VerNum = version.VerNum

PythonScript = script_.PythonScript

SqlScript = script_.SqlScript

# deprecated

def help(cmd=None, **opts):

    """%prog help COMMAND

    Displays help on a given command.

    """

    if cmd is None:

        raise exceptions.UsageError(None)

    try:

        func = globals()[cmd]

    except:

        raise exceptions.UsageError(

            "'%s' isn't a valid command. Try 'help COMMAND'" % cmd)

    ret = func.__doc__

    if sys.argv[0]:

        ret = ret.replace('%prog', sys.argv[0])

    return ret

@catch_known_errors

def create(repository, name, **opts):

    """%prog create REPOSITORY_PATH NAME [--table=TABLE]

    Create an empty repository at the specified path.

    You can specify the version_table to be used; by default, it is

    'migrate_version'.  This table is created in all version-controlled

    databases.

    """

    repo_path = Repository.create(repository, name, **opts)

@catch_known_errors

def script(description, repository, **opts):

    """%prog script DESCRIPTION REPOSITORY_PATH

    Create an empty change script using the next unused version number

    appended with the given description.

    For instance, manage.py script "Add initial tables" creates:

    repository/versions/001_Add_initial_tables.py

    """

    repo = Repository(repository)

    repo.create_script(description, **opts)

@catch_known_errors

    Create empty change SQL scripts for given DATABASE, where DATABASE

    or generic ('default').

    """

    repo = Repository(repository)

def version(repository, **opts):

    """%prog version REPOSITORY_PATH

    Display the latest version available in a repository.

    """

    repo = Repository(repository)

    return repo.latest

@with_engine

def db_version(url, repository, **opts):

    """%prog db_version URL REPOSITORY_PATH

    Show the current version of the repository with the given

    connection string, under version control of the specified

    repository.

    The url should be any valid SQLAlchemy connection string.

    """

    engine = opts.pop('engine')

    schema = ControlledSchema(engine, repository)

    return schema.version

def source(version, dest=None, repository=None, **opts):

    """%prog source VERSION [DESTINATION] --repository=REPOSITORY_PATH

    Display the Python code for a particular version in this

    repository.  Save it to the file at DESTINATION or, if omitted,

    send to stdout.

    """

    if repository is None:

        raise exceptions.UsageError("A repository must be specified")

    repo = Repository(repository)

    ret = repo.version(version).script().source()

    if dest is not None:

        dest = open(dest, 'w')

        dest.write(ret)

        dest.close()

        ret = None

    return ret

def upgrade(url, repository, version=None, **opts):

    """%prog upgrade URL REPOSITORY_PATH [VERSION] [--preview_py|--preview_sql]

    Upgrade a database to a later version.

    This runs the upgrade() function defined in your change scripts.

    By default, the database is updated to the latest available

    version. You may specify a version instead, if you wish.

    You may preview the Python or SQL code to be executed, rather than

    actually executing it, using the appropriate 'preview' option.

    """

    err = "Cannot upgrade a database of version %s to version %s. "\

        "Try 'downgrade' instead."

    return _migrate(url, repository, version, upgrade=True, err=err, **opts)

def downgrade(url, repository, version, **opts):

    """%prog downgrade URL REPOSITORY_PATH VERSION [--preview_py|--preview_sql]

    Downgrade a database to an earlier version.

    This is the reverse of upgrade; this runs the downgrade() function

    defined in your change scripts.

    You may preview the Python or SQL code to be executed, rather than

    actually executing it, using the appropriate 'preview' option.

    """

    err = "Cannot downgrade a database of version %s to version %s. "\

        "Try 'upgrade' instead."

    return _migrate(url, repository, version, upgrade=False, err=err, **opts)

@with_engine

def test(url, repository, **opts):

    """%prog test URL REPOSITORY_PATH [VERSION]

    Performs the upgrade and downgrade option on the given

    database. This is not a real test and may leave the database in a

    bad state. You should therefore better run the test on a copy of

    your database.

    """

    engine = opts.pop('engine')

    repos = Repository(repository)

    script = repos.version(None).script()

    # Upgrade

    log.info("Upgrading...")

    script.run(engine, 1)