support both sqlalchemy engines and connections

Author: Thomas Grainger

pandas/io/sql.py

@@ -38,7 +38,7 @@ class DatabaseError(IOError):
 _SQLALCHEMY_INSTALLED = None
 
 
-def _is_sqlalchemy_engine(con):
+def _is_sqlalchemy_connectable(con):
     global _SQLALCHEMY_INSTALLED
     if _SQLALCHEMY_INSTALLED is None:
         try:
@@ -61,8 +61,8 @@ def compile_big_int_sqlite(type_, compiler, **kw):
             _SQLALCHEMY_INSTALLED = False
 
     if _SQLALCHEMY_INSTALLED:
-        import sqlalchemy
-        return isinstance(con, sqlalchemy.engine.Engine)
+        from sqlalchemy.engine import interfaces
+        return isinstance(con, interfaces.Connectable)
     else:
         return False
 
@@ -328,7 +328,7 @@ def read_sql_table(table_name, con, schema=None, index_col=None,
     read_sql
 
     """
-    if not _is_sqlalchemy_engine(con):
+    if not _is_sqlalchemy_connectable(con):
         raise NotImplementedError("read_sql_table only supported for "
                                   "SQLAlchemy engines.")
     import sqlalchemy
@@ -592,7 +592,7 @@ def pandasSQL_builder(con, flavor=None, schema=None, meta=None,
     """
     # When support for DBAPI connections is removed,
     # is_cursor should not be necessary.
-    if _is_sqlalchemy_engine(con):
+    if _is_sqlalchemy_connectable(con):
         return SQLDatabase(con, schema=schema, meta=meta)
     else:
         if flavor == 'mysql':
@@ -993,7 +993,7 @@ class SQLDatabase(PandasSQL):
 
     Parameters
     ----------
-    engine : SQLAlchemy engine
+    engine : SQLAlchemy engine or connection
         Engine to connect with the database. Using SQLAlchemy makes it
         possible to use any DB supported by that library.
     schema : string, default None
@@ -1014,8 +1014,28 @@ def __init__(self, engine, schema=None, meta=None):
 
         self.meta = meta
 
+
+    class SQLAlchemyTransaction(object):
+        """
+        Context manager for sql alchemy transactions
+        that returns the relevent connectable.
+        """
+        def __init__(connectable):
+            self.connectable = connectable
+
+        def __enter__(self):
+            self.tx = self.connectable.begin()
+            if _is_sqlalchemy_connectable(self.tx):
+                return self.tx
+            else:
+                return self.connectable
+
+        def __exit__(self, *args, **kwargs):
+            return self.tx.__exit__(*args, **kwargs)
+
+
     def run_transaction(self):
-        return self.engine.begin()
+        return self.SQLAlchemyTransaction(self.engine)
 
     def execute(self, *args, **kwargs):
         """Simple passthrough to SQLAlchemy engine"""