Search for dbapi's DatabaseError class on the root driver module

Psycopg2 2.6+ uses a connection class in the package psycopg2.extensions, from
where ``DatabaseError`` cannot be imported.

Fixes issue #7

2 files changed, 7 insertions, 8 deletions

diff --git a/yoyo/connections.py b/yoyo/connections.py
index 608fc6b..27d22fc 100755
--- a/yoyo/connections.py
+++ b/yoyo/connections.py
@@ -1,5 +1,9 @@
+from __future__ import absolute_import
+
 from functools import wraps
 
+from . import exceptions
+
 _schemes = {}
 
 drivers = {
@@ -30,6 +34,7 @@ def connection_for(scheme):
         @wraps(func)
         def with_driver(*args, **kwargs):
             driver = __import__(drivers[scheme], globals(), locals())
+            exceptions.register(driver.DatabaseError)
             return func(driver, *args, **kwargs)
         _schemes[scheme] = with_driver
 
diff --git a/yoyo/migrations.py b/yoyo/migrations.py
index e7273d7..c800b03 100755
--- a/yoyo/migrations.py
+++ b/yoyo/migrations.py
@@ -423,15 +423,9 @@ def create_migrations_table(conn, tablename):
 
 def initialize_connection(conn, tablename):
     """
-    Initialize the DBAPI connection for use.
-
-    - Registers the connection's DatabaseError class in ``yoyo.exceptions``
-
-    - Creates the migrations table if not already existing
-
+    Initialize the connection for use by creating the migrations table if
+    it does not already exist.
     """
-    module = inspect.getmodule(type(conn))
-    exceptions.register(module.DatabaseError)
     create_migrations_table(conn, tablename)