From 8ea697c93e9e2e039549ff514626c23e469eeb32 Mon Sep 17 00:00:00 2001
From: Mike Bayer <mike_mp@zzzcomputing.com>
Date: Fri, 23 Jan 2015 20:05:02 -0500
Subject: - 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.

---
 alembic/config.py                | 46 +++++++++++++++++++++++++++++++++++++---
 alembic/templates/generic/env.py | 15 ++++++-------
 alembic/templates/pylons/env.py  | 17 ++++-----------
 alembic/util.py                  |  2 +-
 4 files changed, 54 insertions(+), 26 deletions(-)

(limited to 'alembic')

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
 
-- 
cgit v1.2.1