- Added a new feature :attr:`.Config.attributes`, to help with the use

case of sharing state such as engines and connections on the outside
with a series of Alembic API calls; also added a new cookbook section
to describe this simple but pretty important use case.

9 files changed, 187 insertions, 39 deletions

diff --git a/alembic/config.py b/alembic/config.py
index 27bb31a..7f813d2 100644
--- a/alembic/config.py
+++ b/alembic/config.py
@@ -40,6 +40,13 @@ class Config(object):
         alembic_cfg.set_main_option("url", "postgresql://foo/bar")
         alembic_cfg.set_section_option("mysection", "foo", "bar")
 
+    For passing non-string values to environments, such as connections and
+    engines, use the :attr:`.Config.attributes` dictionary::
+
+        with engine.begin() as connection:
+            alembic_cfg.attributes['connection'] = connection
+            command.upgrade(alembic_cfg, "head")
+
     :param file_: name of the .ini file to open.
     :param ini_section: name of the main Alembic section within the
      .ini file
@@ -49,7 +56,7 @@ class Config(object):
     :param stdout: buffer where the "print" output of commands will be sent.
      Defaults to ``sys.stdout``.
 
-     ..versionadded:: 0.4
+     .. versionadded:: 0.4
 
     :param config_args: A dictionary of keys and values that will be used
      for substitution in the alembic config file.  The dictionary as given
@@ -59,13 +66,22 @@ class Config(object):
      dictionary before the dictionary is passed to ``SafeConfigParser()``
      to parse the .ini file.
 
-     ..versionadded:: 0.7.0
+     .. versionadded:: 0.7.0
+
+    :param attributes: optional dictionary of arbitrary Python keys/values,
+     which will be populated into the :attr:`.Config.attributes` dictionary.
+
+     .. versionadded:: 0.7.5
+
+     .. seealso::
+
+        :ref:`connection_sharing`
 
     """
 
     def __init__(self, file_=None, ini_section='alembic', output_buffer=None,
                  stdout=sys.stdout, cmd_opts=None,
-                 config_args=util.immutabledict()):
+                 config_args=util.immutabledict(), attributes=None):
         """Construct a new :class:`.Config`
 
         """
@@ -75,6 +91,8 @@ class Config(object):
         self.stdout = stdout
         self.cmd_opts = cmd_opts
         self.config_args = dict(config_args)
+        if attributes:
+            self.attributes.update(attributes)
 
     cmd_opts = None
     """The command-line options passed to the ``alembic`` script.
@@ -101,6 +119,28 @@ class Config(object):
 
     """
 
+    @util.memoized_property
+    def attributes(self):
+        """A Python dictionary for storage of additional state.
+
+
+        This is a utility dictionary which can include not just strings but
+        engines, connections, schema objects, or anything else.
+        Use this to pass objects into an env.py script, such as passing
+        a :class:`.Connection` when calling
+        commands from :mod:`alembic.command` programmatically.
+
+        .. versionadded:: 0.7.5
+
+        .. seealso::
+
+            :ref:`connection_sharing`
+
+            :paramref:`.Config.attributes`
+
+        """
+        return {}
+
     def print_stdout(self, text, *arg):
         """Render a message to standard out."""
 
diff --git a/alembic/templates/generic/env.py b/alembic/templates/generic/env.py
index fccd445..280006d 100644
--- a/alembic/templates/generic/env.py
+++ b/alembic/templates/generic/env.py
@@ -49,22 +49,19 @@ def run_migrations_online():
     and associate a connection with the context.
 
     """
-    engine = engine_from_config(
+    connectable = engine_from_config(
         config.get_section(config.config_ini_section),
         prefix='sqlalchemy.',
         poolclass=pool.NullPool)
 
-    connection = engine.connect()
-    context.configure(
-        connection=connection,
-        target_metadata=target_metadata
-    )
+    with connectable.connect() as connection:
+        context.configure(
+            connection=connection,
+            target_metadata=target_metadata
+        )
 
-    try:
         with context.begin_transaction():
             context.run_migrations()
-    finally:
-        connection.close()
 
 if context.is_offline_mode():
     run_migrations_offline()
diff --git a/alembic/templates/pylons/env.py b/alembic/templates/pylons/env.py
index 3329428..70eea4e 100644
--- a/alembic/templates/pylons/env.py
+++ b/alembic/templates/pylons/env.py
@@ -62,23 +62,14 @@ def run_migrations_online():
     # engine = meta.engine
     raise NotImplementedError("Please specify engine connectivity here")
 
-    if isinstance(engine, Engine):
-        connection = engine.connect()
-    else:
-        raise Exception(
-            'Expected engine instance got %s instead' % type(engine)
+    with engine.connect() as connection:
+        context.configure(
+            connection=connection,
+            target_metadata=target_metadata
         )
 
-    context.configure(
-        connection=connection,
-        target_metadata=target_metadata
-    )
-
-    try:
         with context.begin_transaction():
             context.run_migrations()
-    finally:
-        connection.close()
 
 if context.is_offline_mode():
     run_migrations_offline()
diff --git a/alembic/util.py b/alembic/util.py
index d9ec1c8..87bc7b1 100644
--- a/alembic/util.py
+++ b/alembic/util.py
@@ -323,7 +323,7 @@ class memoized_property(object):
 
     def __get__(self, obj, cls):
         if obj is None:
-            return None
+            return self
         obj.__dict__[self.__name__] = result = self.fget(obj)
         return result
 
diff --git a/docs/build/api.rst b/docs/build/api.rst
index 48da805..fea4e14 100644
--- a/docs/build/api.rst
+++ b/docs/build/api.rst
@@ -87,6 +87,19 @@ object, as in::
     alembic_cfg = Config("/path/to/yourapp/alembic.ini")
     command.upgrade(alembic_cfg, "head")
 
+In many cases, and perhaps more often than not, an application will wish
+to call upon a series of Alembic commands and/or other features.  It is
+usually a good idea to link multiple commands along a single connection
+and transaction, if feasible.  This can be achieved using the
+:attr:`.Config.attributes` dictionary in order to share a connection::
+
+    with engine.begin() as connection:
+        alembic_cfg.attributes['connection'] = connection
+        command.upgrade(alembic_cfg, "head")
+
+This recipe requires that ``env.py`` consumes this connection argument;
+see the example in :ref:`connection_sharing` for details.
+
 To write small API functions that make direct use of database and script directory
 information, rather than just running one of the built-in commands,
 use the :class:`.ScriptDirectory` and :class:`.MigrationContext`
diff --git a/docs/build/autogenerate.rst b/docs/build/autogenerate.rst
index ee9ccb9..8ad79ed 100644
--- a/docs/build/autogenerate.rst
+++ b/docs/build/autogenerate.rst
@@ -35,19 +35,14 @@ we can see the directive passed to :meth:`.EnvironmentContext.configure`::
         engine = engine_from_config(
                     config.get_section(config.config_ini_section), prefix='sqlalchemy.')
 
-        connection = engine.connect()
-        context.configure(
-                    connection=connection,
-                    target_metadata=target_metadata
-                    )
-
-        trans = connection.begin()
-        try:
-            context.run_migrations()
-            trans.commit()
-        except:
-            trans.rollback()
-            raise
+        with engine.connect() as connection:
+            context.configure(
+                        connection=connection,
+                        target_metadata=target_metadata
+                        )
+
+            with context.begin_transaction():
+                context.run_migrations()
 
 We can then use the ``alembic revision`` command in conjunction with the
 ``--autogenerate`` option.  Suppose
diff --git a/docs/build/changelog.rst b/docs/build/changelog.rst
index f9c6145..225db44 100644
--- a/docs/build/changelog.rst
+++ b/docs/build/changelog.rst
@@ -7,6 +7,27 @@ Changelog
     :version: 0.7.5
 
     .. change::
+      :tags: feature, commands
+
+      Added a new feature :attr:`.Config.attributes`, to help with the use
+      case of sharing state such as engines and connections on the outside
+      with a series of Alembic API calls; also added a new cookbook section
+      to describe this simple but pretty important use case.
+
+      .. seealso::
+
+          :ref:`connection_sharing`
+
+    .. change::
+      :tags: feature, environment
+
+      The format of the default ``env.py`` script has been refined a bit;
+      it now uses context managers not only for the scope of the transaction,
+      but also for connectivity from the starting engine.  The engine is also
+      now called a "connectable" in support of the use case of an external
+      connection being passed in.
+
+    .. change::
       :tags: feature, versioning
       :tickets: 267
 
diff --git a/docs/build/cookbook.rst b/docs/build/cookbook.rst
index d24aab9..8c1e0d7 100644
--- a/docs/build/cookbook.rst
+++ b/docs/build/cookbook.rst
@@ -187,3 +187,77 @@ To invoke our migrations with data included, we use the ``-x`` flag::
 The :meth:`.EnvironmentContext.get_x_argument` is an easy way to support
 new commandline options within environment and migration scripts.
 
+.. _connection_sharing:
+
+Sharing a Connection with a Series of Migration Commands and Environments
+=========================================================================
+
+It is often the case that an application will need to call upon a series
+of commands within :mod:`alembic.command`, where it would be advantageous
+for all operations to proceed along a single transaction.   The connectivity
+for a migration is typically solely determined within the ``env.py`` script
+of a migration environment, which is called within the scope of a command.
+
+The steps to take here are:
+
+1. Produce the :class:`~sqlalchemy.engine.Connection` object to use.
+
+2. Place it somewhere that ``env.py`` will be able to access it.  This
+   can be either a. a module-level global somewhere, or b.
+   an attribute which we place into the :attr:`.Config.attributes`
+   dictionary (if we are on an older Alembic version, we may also attach
+   an attribute directly to the :class:`.Config` object).
+
+3. The ``env.py`` script is modified such that it looks for this
+   :class:`~sqlalchemy.engine.Connection` and makes use of it, in lieu
+   of building up its own :class:`~sqlalchemy.engine.Engine` instance.
+
+We illustrate using :attr:`.Config.attributes`::
+
+    from alembic import command, config
+
+    cfg = config.Config("/path/to/yourapp/alembic.ini")
+    with engine.begin() as connection:
+        cfg.attributes['connection'] = connection
+        command.upgrade(cfg, "head")
+
+Then in ``env.py``::
+
+    def run_migrations_online():
+        connectable = config.attributes.get('connection', None)
+
+        if connectable is None:
+            # only create Engine if we don't have a Connection
+            # from the outside
+            connectable = engine_from_config(
+                config.get_section(config.config_ini_section),
+                prefix='sqlalchemy.',
+                poolclass=pool.NullPool)
+
+        # when connectable is already a Connection object, calling
+        # connect() gives us a *branched connection*.
+
+        with connectable.connect() as connection:
+            context.configure(
+                connection=connection,
+                target_metadata=target_metadata
+            )
+
+            with context.begin_transaction():
+                context.run_migrations()
+
+.. topic:: Branched Connections
+
+    Note that we are calling the ``connect()`` method, **even if we are
+    using a** :class:`~sqlalchemy.engine.Connection` **object to start with**.
+    The effect this has when calling :meth:`~sqlalchemy.engine.Connection.connect`
+    is that SQLAlchemy passes us a **branch** of the original connection; it
+    is in every way the same as the :class:`~sqlalchemy.engine.Connection`
+    we started with, except it provides **nested scope**; the
+    context we have here as well as the
+    :meth:`~sqlalchemy.engine.Connection.close` method of this branched
+    connection doesn't actually close the outer connection, which stays
+    active for continued use.
+
+.. versionadded:: 0.7.5 Added :attr:`.Config.attributes`.
+
diff --git a/tests/test_config.py b/tests/test_config.py
index 2d8f964..db37456 100644
--- a/tests/test_config.py
+++ b/tests/test_config.py
@@ -67,6 +67,23 @@ class ConfigTest(TestBase):
             ScriptDirectory.from_config, cfg
         )
 
+    def test_attributes_attr(self):
+        m1 = Mock()
+        cfg = config.Config()
+        cfg.attributes['connection'] = m1
+        eq_(
+            cfg.attributes['connection'], m1
+        )
+
+    def test_attributes_construtor(self):
+        m1 = Mock()
+        m2 = Mock()
+        cfg = config.Config(attributes={'m1': m1})
+        cfg.attributes['connection'] = m2
+        eq_(
+            cfg.attributes, {'m1': m1, 'connection': m2}
+        )
+
 
 class StdoutOutputEncodingTest(TestBase):