1 files changed, 31 insertions, 37 deletions

diff --git a/src/mailman/database/alembic/env.py b/src/mailman/database/alembic/env.py
index 402c4f7da..d1caec58d 100644
--- a/src/mailman/database/alembic/env.py
+++ b/src/mailman/database/alembic/env.py
@@ -1,4 +1,4 @@
-# Copyright (C) 2006-2014 by the Free Software Foundation, Inc.
+# Copyright (C) 2014 by the Free Software Foundation, Inc.
 #
 # This file is part of GNU Mailman.
 #
@@ -15,40 +15,45 @@
 # You should have received a copy of the GNU General Public License along with
 # GNU Mailman.  If not, see <http://www.gnu.org/licenses/>.
 
-from __future__ import with_statement
+"""Alembic migration environment."""
+
+from __future__ import absolute_import, print_function, unicode_literals
+
+__metaclass__ = type
+__all__ = [
+    'run_migrations_offline',
+    'run_migrations_online',
+    ]
+
 
 from alembic import context
-from alembic.config import Config
-from sqlalchemy import create_engine, pool
+from contextlib import closing
+from logging.config import fileConfig
+from sqlalchemy import create_engine
 
 from mailman.core import initialize
 from mailman.config import config
-from mailman.utilities.string import expand
+from mailman.database.alembic import alembic_cfg
 from mailman.database.model import Model
+from mailman.utilities.modules import expand_path
+from mailman.utilities.string import expand
 
-target_metadata = Model.metadata
 
+fileConfig(alembic_cfg.config_file_name)
 
+
 def run_migrations_offline():
     """Run migrations in 'offline' mode.
 
-    This configures the context with just a URL
-    and not an Engine, though an Engine is acceptable
-    here as well.  By skipping the Engine creation
-    we don't even need a DBAPI to be available.
-
-    Calls to context.execute() here emit the given string to the
-    script output.
+    This configures the context with just a URL and not an Engine,
+    though an Engine is acceptable here as well.  By skipping the Engine
+    creation we don't even need a DBAPI to be available.
 
+    Calls to context.execute() here emit the given string to the script
+    output.
     """
-    if not config.initialized:
-        initialize.initialize_1(context.config.config_file_name)
-    alembic_cfg= Config()
-    alembic_cfg.set_main_option(
-        "script_location", config.alembic['script_location'])
     url = expand(config.database.url, config.paths)
-    context.configure(url=url, target_metadata=target_metadata)
-
+    context.configure(url=url, target_metadata=Model.metadata)
     with context.begin_transaction():
         context.run_migrations()
 
@@ -56,30 +61,19 @@ def run_migrations_offline():
 def run_migrations_online():
     """Run migrations in 'online' mode.
 
-    In this scenario we need to create an Engine
-    and associate a connection with the context.
-
+    In this scenario we need to create an Engine and associate a
+    connection with the context.
     """
-    if not config.initialized:
-        initialize.initialize_1(context.config.config_file_name)
-    alembic_cfg= Config()
-    alembic_cfg.set_main_option(
-        "script_location", config.alembic['script_location'])
-    alembic_cfg.set_section_option('logger_alembic' ,'level' , 'ERROR')
     url = expand(config.database.url, config.paths)
     engine = create_engine(url)
 
     connection = engine.connect()
-    context.configure(
-        connection=connection,
-        target_metadata=target_metadata
-    )
-
-    try:
+    with closing(connection):
+        context.configure(
+            connection=connection, target_metadata=Model.metadata)
         with context.begin_transaction():
             context.run_migrations()
-    finally:
-        connection.close()
+
 
 if context.is_offline_mode():
     run_migrations_offline()