import re

from alembic import autogenerate
from alembic.migration import MigrationContext
from alembic.testing import TestBase
from alembic.testing import eq_

from ._autogen_fixtures import AutogenTest, ModelOne, _default_include_object


class AutogenerateDiffTest(ModelOne, AutogenTest, TestBase):
    __only_on__ = 'sqlite'

    def test_render_nothing(self):
        context = MigrationContext.configure(
            connection=self.bind.connect(),
            opts={
                'compare_type': True,
                'compare_server_default': True,
                'target_metadata': self.m1,
                'upgrade_token': "upgrades",
                'downgrade_token': "downgrades",
            }
        )
        template_args = {}
        autogenerate._render_migration_diffs(context, template_args)

        eq_(re.sub(r"u'", "'", template_args['upgrades']),
            """### commands auto generated by Alembic - please adjust! ###
    pass
    ### end Alembic commands ###""")
        eq_(re.sub(r"u'", "'", template_args['downgrades']),
            """### commands auto generated by Alembic - please adjust! ###
    pass
    ### end Alembic commands ###""")

    def test_render_nothing_batch(self):
        context = MigrationContext.configure(
            connection=self.bind.connect(),
            opts={
                'compare_type': True,
                'compare_server_default': True,
                'target_metadata': self.m1,
                'upgrade_token': "upgrades",
                'downgrade_token': "downgrades",
                'alembic_module_prefix': 'op.',
                'sqlalchemy_module_prefix': 'sa.',
                'render_as_batch': True,
                'include_symbol': lambda name, schema: False
            }
        )
        template_args = {}
        autogenerate._render_migration_diffs(context, template_args)

        eq_(re.sub(r"u'", "'", template_args['upgrades']),
            """### commands auto generated by Alembic - please adjust! ###
    pass
    ### end Alembic commands ###""")
        eq_(re.sub(r"u'", "'", template_args['downgrades']),
            """### commands auto generated by Alembic - please adjust! ###
    pass
    ### end Alembic commands ###""")

    def test_render_diffs_standard(self):
        """test a full render including indentation"""

        template_args = {}
        autogenerate._render_migration_diffs(self.context, template_args)
        eq_(re.sub(r"u'", "'", template_args['upgrades']),
            """### commands auto generated by Alembic - please adjust! ###
    op.create_table('item',
    sa.Column('id', sa.Integer(), nullable=False),
    sa.Column('description', sa.String(length=100), nullable=True),
    sa.Column('order_id', sa.Integer(), nullable=True),
    sa.CheckConstraint('len(description) > 5'),
    sa.ForeignKeyConstraint(['order_id'], ['order.order_id'], ),
    sa.PrimaryKeyConstraint('id')
    )
    op.drop_table('extra')
    op.add_column('address', sa.Column('street', sa.String(length=50), \
nullable=True))
    op.create_unique_constraint('uq_email', 'address', ['email_address'])
    op.add_column('order', sa.Column('user_id', sa.Integer(), nullable=True))
    op.alter_column('order', 'amount',
               existing_type=sa.NUMERIC(precision=8, scale=2),
               type_=sa.Numeric(precision=10, scale=2),
               nullable=True,
               existing_server_default=sa.text('0'))
    op.create_foreign_key(None, 'order', 'user', ['user_id'], ['id'])
    op.alter_column('user', 'a1',
               existing_type=sa.TEXT(),
               server_default='x',
               existing_nullable=True)
    op.alter_column('user', 'name',
               existing_type=sa.VARCHAR(length=50),
               nullable=False)
    op.drop_index('pw_idx', table_name='user')
    op.drop_column('user', 'pw')
    ### end Alembic commands ###""")

        eq_(re.sub(r"u'", "'", template_args['downgrades']),
            """### commands auto generated by Alembic - please adjust! ###
    op.add_column('user', sa.Column('pw', sa.VARCHAR(length=50), \
nullable=True))
    op.create_index('pw_idx', 'user', ['pw'], unique=False)
    op.alter_column('user', 'name',
               existing_type=sa.VARCHAR(length=50),
               nullable=True)
    op.alter_column('user', 'a1',
               existing_type=sa.TEXT(),
               server_default=None,
               existing_nullable=True)
    op.drop_constraint(None, 'order', type_='foreignkey')
    op.alter_column('order', 'amount',
               existing_type=sa.Numeric(precision=10, scale=2),
               type_=sa.NUMERIC(precision=8, scale=2),
               nullable=False,
               existing_server_default=sa.text('0'))
    op.drop_column('order', 'user_id')
    op.drop_constraint('uq_email', 'address', type_='unique')
    op.drop_column('address', 'street')
    op.create_table('extra',
    sa.Column('x', sa.CHAR(), nullable=True),
    sa.Column('uid', sa.INTEGER(), nullable=True),
    sa.ForeignKeyConstraint(['uid'], ['user.id'], )
    )
    op.drop_table('item')
    ### end Alembic commands ###""")

    def test_render_diffs_batch(self):
        """test a full render in batch mode including indentation"""

        template_args = {}
        self.context.opts['render_as_batch'] = True
        autogenerate._render_migration_diffs(self.context, template_args)

        eq_(re.sub(r"u'", "'", template_args['upgrades']),
            """### commands auto generated by Alembic - please adjust! ###
    op.create_table('item',
    sa.Column('id', sa.Integer(), nullable=False),
    sa.Column('description', sa.String(length=100), nullable=True),
    sa.Column('order_id', sa.Integer(), nullable=True),
    sa.CheckConstraint('len(description) > 5'),
    sa.ForeignKeyConstraint(['order_id'], ['order.order_id'], ),
    sa.PrimaryKeyConstraint('id')
    )
    op.drop_table('extra')
    with op.batch_alter_table('address', schema=None) as batch_op:
        batch_op.add_column(sa.Column('street', sa.String(length=50), nullable=True))
        batch_op.create_unique_constraint('uq_email', ['email_address'])

    with op.batch_alter_table('order', schema=None) as batch_op:
        batch_op.add_column(sa.Column('user_id', sa.Integer(), nullable=True))
        batch_op.alter_column('amount',
               existing_type=sa.NUMERIC(precision=8, scale=2),
               type_=sa.Numeric(precision=10, scale=2),
               nullable=True,
               existing_server_default=sa.text('0'))
        batch_op.create_foreign_key(None, 'user', ['user_id'], ['id'])

    with op.batch_alter_table('user', schema=None) as batch_op:
        batch_op.alter_column('a1',
               existing_type=sa.TEXT(),
               server_default='x',
               existing_nullable=True)
        batch_op.alter_column('name',
               existing_type=sa.VARCHAR(length=50),
               nullable=False)
        batch_op.drop_index('pw_idx')
        batch_op.drop_column('pw')

    ### end Alembic commands ###""")

        eq_(re.sub(r"u'", "'", template_args['downgrades']),
            """### commands auto generated by Alembic - please adjust! ###
    with op.batch_alter_table('user', schema=None) as batch_op:
        batch_op.add_column(sa.Column('pw', sa.VARCHAR(length=50), nullable=True))
        batch_op.create_index('pw_idx', ['pw'], unique=False)
        batch_op.alter_column('name',
               existing_type=sa.VARCHAR(length=50),
               nullable=True)
        batch_op.alter_column('a1',
               existing_type=sa.TEXT(),
               server_default=None,
               existing_nullable=True)

    with op.batch_alter_table('order', schema=None) as batch_op:
        batch_op.drop_constraint(None, type_='foreignkey')
        batch_op.alter_column('amount',
               existing_type=sa.Numeric(precision=10, scale=2),
               type_=sa.NUMERIC(precision=8, scale=2),
               nullable=False,
               existing_server_default=sa.text('0'))
        batch_op.drop_column('user_id')

    with op.batch_alter_table('address', schema=None) as batch_op:
        batch_op.drop_constraint('uq_email', type_='unique')
        batch_op.drop_column('street')

    op.create_table('extra',
    sa.Column('x', sa.CHAR(), nullable=True),
    sa.Column('uid', sa.INTEGER(), nullable=True),
    sa.ForeignKeyConstraint(['uid'], ['user.id'], )
    )
    op.drop_table('item')
    ### end Alembic commands ###""")

    def test_imports_maintined(self):
        template_args = {}
        self.context.opts['render_as_batch'] = True

        def render_item(type_, col, autogen_context):
            autogen_context.imports.add(
                "from mypackage import my_special_import"
            )
            autogen_context.imports.add(
                "from foobar import bat"
            )

        self.context.opts["render_item"] = render_item
        autogenerate._render_migration_diffs(self.context, template_args)
        eq_(
            set(
                template_args['imports'].split("\n")
            ),
            set([
                "from foobar import bat",
                "from mypackage import my_special_import"
            ])
        )


class AutogenerateDiffTestWSchema(ModelOne, AutogenTest, TestBase):
    __only_on__ = 'postgresql'
    schema = "test_schema"

    def test_render_nothing(self):
        context = MigrationContext.configure(
            connection=self.bind.connect(),
            opts={
                'compare_type': True,
                'compare_server_default': True,
                'target_metadata': self.m1,
                'upgrade_token': "upgrades",
                'downgrade_token': "downgrades",
                'alembic_module_prefix': 'op.',
                'sqlalchemy_module_prefix': 'sa.',
                'include_symbol': lambda name, schema: False
            }
        )
        template_args = {}
        autogenerate._render_migration_diffs(context, template_args)

        eq_(re.sub(r"u'", "'", template_args['upgrades']),
            """### commands auto generated by Alembic - please adjust! ###
    pass
    ### end Alembic commands ###""")
        eq_(re.sub(r"u'", "'", template_args['downgrades']),
            """### commands auto generated by Alembic - please adjust! ###
    pass
    ### end Alembic commands ###""")

    def test_render_diffs_extras(self):
        """test a full render including indentation (include and schema)"""

        template_args = {}
        self.context.opts.update({
            'include_object': _default_include_object,
            'include_schemas': True
        })
        autogenerate._render_migration_diffs(self.context, template_args)

        eq_(re.sub(r"u'", "'", template_args['upgrades']),
            """### commands auto generated by Alembic - please adjust! ###
    op.create_table('item',
    sa.Column('id', sa.Integer(), nullable=False),
    sa.Column('description', sa.String(length=100), nullable=True),
    sa.Column('order_id', sa.Integer(), nullable=True),
    sa.CheckConstraint('len(description) > 5'),
    sa.ForeignKeyConstraint(['order_id'], ['%(schema)s.order.order_id'], ),
    sa.PrimaryKeyConstraint('id'),
    schema='%(schema)s'
    )
    op.drop_table('extra', schema='%(schema)s')
    op.add_column('address', sa.Column('street', sa.String(length=50), \
nullable=True), schema='%(schema)s')
    op.create_unique_constraint('uq_email', 'address', ['email_address'], \
schema='test_schema')
    op.add_column('order', sa.Column('user_id', sa.Integer(), nullable=True), \
schema='%(schema)s')
    op.alter_column('order', 'amount',
               existing_type=sa.NUMERIC(precision=8, scale=2),
               type_=sa.Numeric(precision=10, scale=2),
               nullable=True,
               existing_server_default=sa.text('0'),
               schema='%(schema)s')
    op.create_foreign_key(None, 'order', 'user', ['user_id'], ['id'], \
source_schema='%(schema)s', referent_schema='%(schema)s')
    op.alter_column('user', 'a1',
               existing_type=sa.TEXT(),
               server_default='x',
               existing_nullable=True,
               schema='%(schema)s')
    op.alter_column('user', 'name',
               existing_type=sa.VARCHAR(length=50),
               nullable=False,
               schema='%(schema)s')
    op.drop_index('pw_idx', table_name='user', schema='test_schema')
    op.drop_column('user', 'pw', schema='%(schema)s')
    ### end Alembic commands ###""" % {"schema": self.schema})

        eq_(re.sub(r"u'", "'", template_args['downgrades']),
            """### commands auto generated by Alembic - please adjust! ###
    op.add_column('user', sa.Column('pw', sa.VARCHAR(length=50), \
autoincrement=False, nullable=True), schema='%(schema)s')
    op.create_index('pw_idx', 'user', ['pw'], unique=False, schema='%(schema)s')
    op.alter_column('user', 'name',
               existing_type=sa.VARCHAR(length=50),
               nullable=True,
               schema='%(schema)s')
    op.alter_column('user', 'a1',
               existing_type=sa.TEXT(),
               server_default=None,
               existing_nullable=True,
               schema='%(schema)s')
    op.drop_constraint(None, 'order', schema='%(schema)s', type_='foreignkey')
    op.alter_column('order', 'amount',
               existing_type=sa.Numeric(precision=10, scale=2),
               type_=sa.NUMERIC(precision=8, scale=2),
               nullable=False,
               existing_server_default=sa.text('0'),
               schema='%(schema)s')
    op.drop_column('order', 'user_id', schema='%(schema)s')
    op.drop_constraint('uq_email', 'address', schema='test_schema', type_='unique')
    op.drop_column('address', 'street', schema='%(schema)s')
    op.create_table('extra',
    sa.Column('x', sa.CHAR(length=1), autoincrement=False, nullable=True),
    sa.Column('uid', sa.INTEGER(), autoincrement=False, nullable=True),
    sa.ForeignKeyConstraint(['uid'], ['%(schema)s.user.id'], \
name='extra_uid_fkey'),
    schema='%(schema)s'
    )
    op.drop_table('item', schema='%(schema)s')
    ### end Alembic commands ###""" % {"schema": self.schema})

