- [bug] Default prefix for autogenerate

  directives is "op.", matching the
  mako templates. [#18]

- [feature] Add alembic_module_prefix argument
  to configure() to complement
  sqlalchemy_module_prefix. [#18]

diff --git a/CHANGES b/CHANGES
index 0263e247bc18e52f5238ebba0f88313e5865a2bb..3dffe90e1c535faa4b1aae3896e01324f2e76a9c 100644 (file)
--- a/CHANGES
+++ b/CHANGES
@@ -21,6 +21,14 @@
   handle self-referential foreign key
   correctly [#17]
 
+- [bug] Default prefix for autogenerate
+  directives is "op.", matching the 
+  mako templates. [#18]
+
+- [feature] Add alembic_module_prefix argument
+  to configure() to complement 
+  sqlalchemy_module_prefix. [#18]
+
 0.1.0
 =====
 - Initial release.  Status of features:

diff --git a/alembic/autogenerate.py b/alembic/autogenerate.py
index 5c8c6f2bcd196f59aeaffbf80e807ba593fef5cb..7e4045b5d753f2efad8ec04ca3211e78210c7b87 100644 (file)
--- a/alembic/autogenerate.py
+++ b/alembic/autogenerate.py
@@ -302,8 +302,9 @@ def _invoke_modify_command(updown, args, autogen_context):
 # render python
 
 def _add_table(table, autogen_context):
-    return "create_table(%(tablename)r,\n%(args)s\n)" % {
+    return "%(prefix)screate_table(%(tablename)r,\n%(args)s\n)" % {
         'tablename':table.name,
+        'prefix':_alembic_autogenerate_prefix(),
         'args':',\n'.join(
             [_render_column(col, autogen_context) for col in table.c] +
             sorted([rcons for rcons in 
@@ -315,15 +316,24 @@ def _add_table(table, autogen_context):
     }
 
 def _drop_table(table, autogen_context):
-    return "drop_table(%r)" % table.name
+    return "%(prefix)sdrop_table(%(tname)r)" % {
+            "prefix":_alembic_autogenerate_prefix(),
+            "tname":table.name
+        }
 
 def _add_column(tname, column, autogen_context):
-    return "add_column(%r, %s)" % (
-            tname, 
-            _render_column(column, autogen_context))
+    return "%(prefix)sadd_column(%(tname)r, %(column)s)" % {
+            "prefix":_alembic_autogenerate_prefix(),
+            "tname":tname,
+            "column":_render_column(column, autogen_context)
+            }
 
 def _drop_column(tname, column, autogen_context):
-    return "drop_column(%r, %r)" % (tname, column.name)
+    return "%(prefix)sdrop_column(%(tname)r, %(cname)r)" % {
+            "prefix":_alembic_autogenerate_prefix(),
+            "tname":tname,
+            "cname":column.name
+            }
 
 def _modify_col(tname, cname, 
                 autogen_context,
@@ -333,16 +343,19 @@ def _modify_col(tname, cname,
                 existing_type=None,
                 existing_nullable=None,
                 existing_server_default=False):
-    prefix = _autogenerate_prefix()
+    sqla_prefix = _sqlalchemy_autogenerate_prefix()
     indent = " " * 11
-    text = "alter_column(%r, %r" % (tname, cname)
+    text = "%(prefix)salter_column(%(tname)r, %(cname)r" % {
+                            'prefix':_alembic_autogenerate_prefix(), 
+                            'tname':tname, 
+                            'cname':cname}
     text += ", \n%sexisting_type=%s" % (indent, 
-                    _repr_type(prefix, existing_type, autogen_context))
+                    _repr_type(sqla_prefix, existing_type, autogen_context))
     if server_default is not False:
         text += ", \n%sserver_default=%s" % (indent, 
                         _render_server_default(server_default, autogen_context),)
     if type_ is not None:
-        text += ", \n%stype_=%s" % (indent, _repr_type(prefix, type_, autogen_context))
+        text += ", \n%stype_=%s" % (indent, _repr_type(sqla_prefix, type_, autogen_context))
     if nullable is not None:
         text += ", \n%snullable=%r" % (
                         indent, nullable,)
@@ -359,9 +372,12 @@ def _modify_col(tname, cname,
     text += ")"
     return text
 
-def _autogenerate_prefix():
+def _sqlalchemy_autogenerate_prefix():
     return _context_opts['sqlalchemy_module_prefix'] or ''
 
+def _alembic_autogenerate_prefix():
+    return _context_opts['alembic_module_prefix'] or ''
+
 def _render_column(column, autogen_context):
     opts = []
     if column.server_default:
@@ -372,9 +388,9 @@ def _render_column(column, autogen_context):
 
     # TODO: for non-ascii colname, assign a "key"
     return "%(prefix)sColumn(%(name)r, %(type)s, %(kw)s)" % {
-        'prefix':_autogenerate_prefix(),
+        'prefix':_sqlalchemy_autogenerate_prefix(),
         'name':column.name,
-        'type':_repr_type(_autogenerate_prefix(), column.type, autogen_context),
+        'type':_repr_type(_sqlalchemy_autogenerate_prefix(), column.type, autogen_context),
         'kw':", ".join(["%s=%s" % (kwname, val) for kwname, val in opts])
     }
 
@@ -416,7 +432,7 @@ def _render_primary_key(constraint):
     if constraint.name:
         opts.append(("name", repr(constraint.name)))
     return "%(prefix)sPrimaryKeyConstraint(%(args)s)" % {
-        "prefix":_autogenerate_prefix(),
+        "prefix":_sqlalchemy_autogenerate_prefix(),
         "args":", ".join(
             [repr(c.key) for c in constraint.columns] +
             ["%s=%s" % (kwname, val) for kwname, val in opts]
@@ -429,7 +445,7 @@ def _render_foreign_key(constraint):
         opts.append(("name", repr(constraint.name)))
     # TODO: deferrable, initially, etc.
     return "%(prefix)sForeignKeyConstraint([%(cols)s], [%(refcols)s], %(args)s)" % {
-        "prefix":_autogenerate_prefix(),
+        "prefix":_sqlalchemy_autogenerate_prefix(),
         "cols":", ".join(f.parent.key for f in constraint.elements),
         "refcols":", ".join(repr(f._get_colspec()) for f in constraint.elements),
         "args":", ".join(
@@ -442,7 +458,7 @@ def _render_check_constraint(constraint):
     if constraint.name:
         opts.append(("name", repr(constraint.name)))
     return "%(prefix)sCheckConstraint('TODO')" % {
-            "prefix":_autogenerate_prefix()
+            "prefix":_sqlalchemy_autogenerate_prefix()
         }
 
 _constraint_renderers = {

diff --git a/alembic/context.py b/alembic/context.py
index 60081ba8d125337b7b45017858a891da5a207aa4..a49ffa46d71f1f1ea06d604e0b4a4427a5e1c2da 100644 (file)
--- a/alembic/context.py
+++ b/alembic/context.py
@@ -320,6 +320,7 @@ def configure(
         compare_server_default=False,
         upgrade_token="upgrades",
         downgrade_token="downgrades",
+        alembic_module_prefix="op.",
         sqlalchemy_module_prefix="sa.",
         **kw
     ):
@@ -427,6 +428,11 @@ def configure(
      candidate downgrade operations will be present in this
      template variable when ``script.py.mako`` is rendered.  Defaults to 
      ``downgrades``.
+
+    :param alembic_module_prefix: When autogenerate refers to Alembic 
+     :module:`alembic.op` constructs, this prefix will be used
+     (i.e. ``op.create_table``)  Defaults to "``op.``".
+     Can be ``None`` to indicate no prefix.  
      
     :param sqlalchemy_module_prefix: When autogenerate refers to SQLAlchemy 
      :class:`~sqlalchemy.schema.Column` or type classes, this prefix will be used
@@ -475,6 +481,7 @@ def configure(
     opts['upgrade_token'] = upgrade_token
     opts['downgrade_token'] = downgrade_token
     opts['sqlalchemy_module_prefix'] = sqlalchemy_module_prefix
+    opts['alembic_module_prefix'] = alembic_module_prefix
     opts.update(kw)
 
     _context = Context(

diff --git a/tests/test_autogenerate.py b/tests/test_autogenerate.py
index 2beed5888f7efec898ef911293f9e3f97a8d073a..5a84dcd7398e2f85dffbbf05fc0403fdd1d04cb9 100644 (file)
--- a/tests/test_autogenerate.py
+++ b/tests/test_autogenerate.py
@@ -165,52 +165,52 @@ class AutogenerateDiffTest(TestCase):
         autogenerate.produce_migration_diffs(template_args, self.autogen_context)
         eq_(re.sub(r"u'", "'", template_args['upgrades']),
 """### commands auto generated by Alembic - please adjust! ###
-    create_table('item',
+    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.ForeignKeyConstraint([order_id], ['order.order_id'], ),
     sa.PrimaryKeyConstraint('id')
     )
-    drop_table('extra')
-    add_column('address', sa.Column('street', sa.String(length=50), nullable=True))
-    add_column('order', sa.Column('user_id', sa.Integer(), nullable=True))
-    alter_column('order', 'amount', 
+    op.drop_table('extra')
+    op.add_column('address', sa.Column('street', sa.String(length=50), nullable=True))
+    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='0')
-    drop_column('user', 'pw')
-    alter_column('user', 'a1', 
+    op.drop_column('user', 'pw')
+    op.alter_column('user', 'a1', 
                existing_type=sa.TEXT(), 
                server_default='x', 
                existing_nullable=True)
-    alter_column('user', 'name', 
+    op.alter_column('user', 'name', 
                existing_type=sa.VARCHAR(length=50), 
                nullable=False)
     ### end Alembic commands ###""")
         eq_(re.sub(r"u'", "'", template_args['downgrades']),
 """### commands auto generated by Alembic - please adjust! ###
-    drop_table('item')
-    create_table('extra',
+    op.drop_table('item')
+    op.create_table('extra',
     sa.Column('x', sa.CHAR(), nullable=True),
     sa.Column('uid', sa.INTEGER(), nullable=True),
     sa.ForeignKeyConstraint([uid], ['user.id'], ),
     sa.PrimaryKeyConstraint()
     )
-    drop_column('address', 'street')
-    drop_column('order', 'user_id')
-    alter_column('order', 'amount', 
+    op.drop_column('address', 'street')
+    op.drop_column('order', 'user_id')
+    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='0')
-    add_column('user', sa.Column('pw', sa.VARCHAR(length=50), nullable=True))
-    alter_column('user', 'a1', 
+    op.add_column('user', sa.Column('pw', sa.VARCHAR(length=50), nullable=True))
+    op.alter_column('user', 'a1', 
                existing_type=sa.TEXT(), 
                server_default=None, 
                existing_nullable=True)
-    alter_column('user', 'name', 
+    op.alter_column('user', 'name', 
                existing_type=sa.VARCHAR(length=50), 
                nullable=True)
     ### end Alembic commands ###""")
@@ -274,6 +274,7 @@ class AutogenRenderTest(TestCase):
     @requires_07
     def setup_class(cls):
         context._context_opts['sqlalchemy_module_prefix'] = 'sa.'
+        context._context_opts['alembic_module_prefix'] = 'op.'
 
     def test_render_table_upgrade(self):
         m = MetaData()
@@ -285,7 +286,7 @@ class AutogenRenderTest(TestCase):
         )
         eq_ignore_whitespace(
             autogenerate._add_table(t, {}),
-            "create_table('test',"
+            "op.create_table('test',"
             "sa.Column('id', sa.Integer(), nullable=False),"
             "sa.Column('address_id', sa.Integer(), nullable=True),"
             "sa.Column('timestamp', sa.DATETIME(), "
@@ -300,14 +301,14 @@ class AutogenRenderTest(TestCase):
     def test_render_drop_table(self):
         eq_(
             autogenerate._drop_table(Table("sometable", MetaData()), {}),
-            "drop_table('sometable')"
+            "op.drop_table('sometable')"
         )
 
     def test_render_add_column(self):
         eq_(
             autogenerate._add_column(
                     "foo", Column("x", Integer, server_default="5"), {}),
-            "add_column('foo', sa.Column('x', sa.Integer(), "
+            "op.add_column('foo', sa.Column('x', sa.Integer(), "
                 "server_default='5', nullable=True))"
         )
 
@@ -316,7 +317,7 @@ class AutogenRenderTest(TestCase):
             autogenerate._drop_column(
                     "foo", Column("x", Integer, server_default="5"), {}),
 
-            "drop_column('foo', 'x')"
+            "op.drop_column('foo', 'x')"
         )
 
     def test_render_modify_type(self):
@@ -325,7 +326,7 @@ class AutogenRenderTest(TestCase):
                         "sometable", "somecolumn", 
                         {},
                         type_=CHAR(10), existing_type=CHAR(20)),
-            "alter_column('sometable', 'somecolumn', "
+            "op.alter_column('sometable', 'somecolumn', "
                 "existing_type=sa.CHAR(length=20), type_=sa.CHAR(length=10))"
         )
 
@@ -336,7 +337,7 @@ class AutogenRenderTest(TestCase):
                         {},
                         existing_type=Integer(),
                         nullable=True),
-            "alter_column('sometable', 'somecolumn', "
+            "op.alter_column('sometable', 'somecolumn', "
             "existing_type=sa.Integer(), nullable=True)"
         )
 
@@ -348,7 +349,7 @@ class AutogenRenderTest(TestCase):
                         existing_type=Integer(),
                         existing_server_default="5",
                         nullable=True),
-            "alter_column('sometable', 'somecolumn', "
+            "op.alter_column('sometable', 'somecolumn', "
             "existing_type=sa.Integer(), nullable=True, "
             "existing_server_default='5')"
         )