added 0.6 TODO, all api now uses engine.dispose() to handle pool correctly

3 files changed, 63 insertions, 21 deletions

diff --git a/migrate/versioning/api.py b/migrate/versioning/api.py
index 287f749..1237721 100644
--- a/migrate/versioning/api.py
+++ b/migrate/versioning/api.py
@@ -6,7 +6,7 @@
     changed order of positional arguments so all accept `url` and `repository`
     as first arguments.
 
-   .. versionchanged:: 0.5.4 
+   .. 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.
@@ -32,7 +32,7 @@ import logging
 
 from migrate.versioning import (exceptions, repository, schema, version,
     script as script_) # command name conflict
-from migrate.versioning.util import catch_known_errors, construct_engine
+from migrate.versioning.util import catch_known_errors, with_engine
 
 
 log = logging.getLogger(__name__)
@@ -134,6 +134,7 @@ def version(repository, **opts):
     return repo.latest
 
 
+@with_engine
 def db_version(url, repository, **opts):
     """%prog db_version URL REPOSITORY_PATH
 
@@ -143,7 +144,7 @@ def db_version(url, repository, **opts):
 
     The url should be any valid SQLAlchemy connection string.
     """
-    engine = construct_engine(url, **opts)
+    engine = opts.pop('engine')
     schema = ControlledSchema(engine, repository)
     return schema.version
 
@@ -199,7 +200,8 @@ def downgrade(url, repository, version, **opts):
     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]
 
@@ -208,7 +210,7 @@ def test(url, repository, **opts):
     bad state. You should therefore better run the test on a copy of
     your database.
     """
-    engine = construct_engine(url, **opts)
+    engine = opts.pop('engine')
     repos = Repository(repository)
     script = repos.version(None).script()
 
@@ -223,6 +225,7 @@ def test(url, repository, **opts):
     log.info("Success")
 
 
+@with_engine
 def version_control(url, repository, version=None, **opts):
     """%prog version_control URL REPOSITORY_PATH [VERSION]
 
@@ -242,16 +245,17 @@ def version_control(url, repository, version=None, **opts):
     identical to what it would be if the database were created from
     scratch.
     """
-    engine = construct_engine(url, **opts)
+    engine = opts.pop('engine')
     ControlledSchema.create(engine, repository, version)
 
 
+@with_engine
 def drop_version_control(url, repository, **opts):
     """%prog drop_version_control URL REPOSITORY_PATH
 
     Removes version control from a database.
     """
-    engine = construct_engine(url, **opts)
+    engine = opts.pop('engine')
     schema = ControlledSchema(engine, repository)
     schema.drop()
 
@@ -275,6 +279,7 @@ def manage(file, **opts):
     Repository.create_manage_file(file, **opts)
 
 
+@with_engine
 def compare_model_to_db(url, repository, model, **opts):
     """%prog compare_model_to_db URL REPOSITORY_PATH MODEL
 
@@ -283,10 +288,11 @@ def compare_model_to_db(url, repository, model, **opts):
 
     NOTE: This is EXPERIMENTAL.
     """  # TODO: get rid of EXPERIMENTAL label
-    engine = construct_engine(url, **opts)
+    engine = opts.pop('engine')
     return ControlledSchema.compare_model_to_db(engine, model, repository)
 
 
+@with_engine
 def create_model(url, repository, **opts):
     """%prog create_model URL REPOSITORY_PATH [DECLERATIVE=True]
 
@@ -294,12 +300,13 @@ def create_model(url, repository, **opts):
 
     NOTE: This is EXPERIMENTAL.
     """  # TODO: get rid of EXPERIMENTAL label
-    engine = construct_engine(url, **opts)
+    engine = opts.pop('engine')
     declarative = opts.get('declarative', False)
     return ControlledSchema.create_model(engine, repository, declarative)
 
 
 @catch_known_errors
+@with_engine
 def make_update_script_for_model(url, repository, oldmodel, model, **opts):
     """%prog make_update_script_for_model URL OLDMODEL MODEL REPOSITORY_PATH
 
@@ -308,11 +315,12 @@ def make_update_script_for_model(url, repository, oldmodel, model, **opts):
 
     NOTE: This is EXPERIMENTAL.
     """  # TODO: get rid of EXPERIMENTAL label
-    engine = construct_engine(url, **opts)
+    engine = opts.pop('engine')
     return PythonScript.make_update_script_for_model(
         engine, oldmodel, model, repository, **opts)
 
 
+@with_engine
 def update_db_from_model(url, repository, model, **opts):
     """%prog update_db_from_model URL REPOSITORY_PATH MODEL
 
@@ -322,13 +330,14 @@ def update_db_from_model(url, repository, model, **opts):
 
     NOTE: This is EXPERIMENTAL.
     """  # TODO: get rid of EXPERIMENTAL label
-    engine = construct_engine(url, **opts)
+    engine = opts.pop('engine')
     schema = ControlledSchema(engine, repository)
     schema.update_db_from_model(model)
 
-
+@with_engine
 def _migrate(url, repository, version, upgrade, err, **opts):
-    engine = construct_engine(url, **opts)
+    engine = opts.pop('engine')
+    url = str(engine.url)
     schema = ControlledSchema(engine, repository)
     version = _migrate_version(schema, version, upgrade, err)
 
diff --git a/migrate/versioning/script/py.py b/migrate/versioning/script/py.py
index 74d4903..bf893ce 100644
--- a/migrate/versioning/script/py.py
+++ b/migrate/versioning/script/py.py
@@ -11,7 +11,7 @@ from migrate.versioning import exceptions, genmodel, schemadiff
 from migrate.versioning.config import operations
 from migrate.versioning.template import Template
 from migrate.versioning.script import base
-from migrate.versioning.util import import_path, load_model, construct_engine
+from migrate.versioning.util import import_path, load_model, with_engine
 
 
 log = logging.getLogger(__name__)
@@ -102,18 +102,21 @@ class PythonScript(base.BaseScript):
     def preview_sql(self, url, step, **args):
         """Mocks SQLAlchemy Engine to store all executed calls in a string 
         and runs :meth:`PythonScript.run <migrate.versioning.script.py.PythonScript.run>`
-        
+
         :returns: SQL file
         """
         buf = StringIO()
         args['engine_arg_strategy'] = 'mock'
         args['engine_arg_executor'] = lambda s, p = '': buf.write(str(s) + p)
-        engine = construct_engine(url, **args)
 
-        self.run(engine, step)
+        @with_engine
+        def go(url, step, **kw):
+            engine = kw.pop('engine')
+            self.run(engine, step)
+            return buf.getvalue()
+
+        return go(url, step, **args)
 
-        return buf.getvalue()
-            
     def run(self, engine, step):
         """Core method of Script file. 
         Exectues :func:`update` or :func:`downgrade` functions
diff --git a/migrate/versioning/util/__init__.py b/migrate/versioning/util/__init__.py
index 01612b1..8d4eb2d 100644
--- a/migrate/versioning/util/__init__.py
+++ b/migrate/versioning/util/__init__.py
@@ -1,18 +1,23 @@
 #!/usr/bin/env python
 # -*- coding: utf-8 -*-
+""".. currentmodule:: migrate.versioning.util"""
 
 import warnings
+import logging
 from decorator import decorator
 from pkg_resources import EntryPoint
 
 from sqlalchemy import create_engine
 from sqlalchemy.engine import Engine
+from sqlalchemy.pool import StaticPool
 
 from migrate.versioning import exceptions
 from migrate.versioning.util.keyedinstance import KeyedInstance
 from migrate.versioning.util.importpath import import_path
 
 
+log = logging.getLogger(__name__)
+
 def load_model(dotted_name):
     """Import module and use module-level variable".
 
@@ -123,14 +128,39 @@ def construct_engine(engine, **opts):
             'engine_arg_echo=True or engine_dict={"echo": True}',
             DeprecationWarning)
         kwargs['echo'] = echo
-    
+
     # parse keyword arguments
     for key, value in opts.iteritems():
         if key.startswith('engine_arg_'):
             kwargs[key[11:]] = guess_obj_type(value)
-    
+
+    log.debug('Constructing engine')
+    # TODO: return create_engine(engine, poolclass=StaticPool, **kwargs)
+    # seems like 0.5.x branch does not work with engine.dispose and staticpool
     return create_engine(engine, **kwargs)
 
+@decorator
+def with_engine(f, *a, **kw):
+    """Decorator for :mod:`migrate.versioning.api` functions
+    to safely close resources after function usage.
+
+    Passes engine parameters to :func:`construct_engine` and
+    resulting parameter is available as kw['engine'].
+
+    Engine is disposed after wrapped function is executed.
+
+    .. versionadded: 0.6.0
+    """
+    url = a[0]
+    engine = construct_engine(url, **kw)
+
+    try:
+        return f(*a, engine=engine, **kw)
+    finally:
+        if isinstance(engine, Engine):
+            log.debug('Disposing SQLAlchemy engine %s', engine)
+            engine.dispose()
+
 
 class Memoize:
     """Memoize(fn) - an instance which acts like fn but memoizes its arguments