support both sqlalchemy engines and connections Fixes pandas-dev#7877

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:
@@ -62,7 +62,7 @@ def compile_big_int_sqlite(type_, compiler, **kw):
 
     if _SQLALCHEMY_INSTALLED:
         import sqlalchemy
-        return isinstance(con, sqlalchemy.engine.Engine)
+        return isinstance(con, sqlalchemy.engine.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':
@@ -637,7 +637,8 @@ def exists(self):
 
     def sql_schema(self):
         from sqlalchemy.schema import CreateTable
-        return str(CreateTable(self.table).compile(self.pd_sql.engine))
+        engine = self.pd_sql.connectable.engine
+        return str(CreateTable(self.table).compile(engine))
 
     def _execute_create(self):
         # Inserting table into database, add to MetaData object
@@ -993,7 +994,7 @@ class SQLDatabase(PandasSQL):
 
     Parameters
     ----------
-    engine : SQLAlchemy engine
+    engine : SQLAlchemy Connectable (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
@@ -1007,19 +1008,35 @@ class SQLDatabase(PandasSQL):
     """
 
     def __init__(self, engine, schema=None, meta=None):
-        self.engine = engine
+        self.connectable = engine
         if not meta:
             from sqlalchemy.schema import MetaData
-            meta = MetaData(self.engine, schema=schema)
+            meta = MetaData(self.connectable, schema=schema)
 
         self.meta = meta
 
+    class _RunTransaction(object):
+        def __init__(self, connectable):
+            tx = connectable.begin()
+            if hasattr(tx, 'execute'):
+                self.connectable = tx
+            else:
+                self.connectable = connectable
+
+        def __enter__(self, *args, **kwargs):
+            self.connectable
+
+    @contextmanager
     def run_transaction(self):
-        return self.engine.begin()
+        with self.connectable.begin() as tx:
+            if hasattr(tx, 'execute'):
+                yield tx
+            else:
+                yield self.connectable
 
     def execute(self, *args, **kwargs):
         """Simple passthrough to SQLAlchemy engine"""
-        return self.engine.execute(*args, **kwargs)
+        return self.connectable.execute(*args, **kwargs)
 
     def read_table(self, table_name, index_col=None, coerce_float=True,
                    parse_dates=None, columns=None, schema=None,
@@ -1187,7 +1204,8 @@ def to_sql(self, frame, name, if_exists='fail', index=True,
         table.create()
         table.insert(chunksize)
         # check for potentially case sensitivity issues (GH7815)
-        if name not in self.engine.table_names(schema=schema or self.meta.schema):
+        engine = self.connectable.engine
+        if name not in engine.table_names(schema=schema or self.meta.schema):
             warnings.warn("The provided table name '{0}' is not found exactly "
                           "as such in the database after writing the table, "
                           "possibly due to case sensitivity issues. Consider "
@@ -1198,7 +1216,8 @@ def tables(self):
         return self.meta.tables
 
     def has_table(self, name, schema=None):
-        return self.engine.has_table(name, schema or self.meta.schema)
+        engine = self.connectable.engine
+        return engine.has_table(name, schema or self.meta.schema)
 
     def get_table(self, table_name, schema=None):
         schema = schema or self.meta.schema
@@ -1217,7 +1236,8 @@ def get_table(self, table_name, schema=None):
 
     def drop_table(self, table_name, schema=None):
         schema = schema or self.meta.schema
-        if self.engine.has_table(table_name, schema):
+        engine = self.connectable.engine
+        if engine.has_table(table_name, schema):
             self.meta.reflect(only=[table_name], schema=schema)
             self.get_table(table_name, schema).drop()
             self.meta.clear()

pandas/io/tests/test_sql.py

@@ -9,7 +9,8 @@
     - `TestSQLiteFallbackApi`: test the public API with a sqlite DBAPI connection
 - Tests for the different SQL flavors (flavor specific type conversions)
     - Tests for the sqlalchemy mode: `_TestSQLAlchemy` is the base class with
-      common methods, the different tested flavors (sqlite3, MySQL, PostgreSQL)
+      common methods, `_TestSQLAlchemyConn` tests the API with a SQLAlchemy
+      Connection object. The different tested flavors (sqlite3, MySQL, PostgreSQL)
       derive from the base class
     - Tests for the fallback mode (`TestSQLiteFallback` and `TestMySQLLegacy`)
 
@@ -854,6 +855,30 @@ def test_sqlalchemy_type_mapping(self):
         self.assertTrue(isinstance(table.table.c['time'].type, sqltypes.DateTime))
 
 
+class _EngineToConnMixin(object):
+    """
+    A mixin that causes setup_connect to create a conn rather than an engine.
+    """
+
+    def setup_connect(self):
+        try:
+            engine = self.connect()  # connect lies, it returns an engine
+            conn = engine.connect()
+            self.addCleanup(conn.close)
+            tx = conn.begin()
+            self.addCleanup(tx.rollback)
+            self.pandasSQL = sql.SQLDatabase(conn)
+            self.conn = conn
+        except sqlalchemy.exc.OperationalError:
+            raise nose.SkipTest("Can't connect to {0} server".format(
+                self.flavor
+            ))
+
+
+class TestSQLApiConn(_EngineToConnMixin, TestSQLApi):
+    pass
+
+
 class TestSQLiteFallbackApi(_TestSQLApi):
     """
     Test the public sqlite connection fallback API
@@ -1347,8 +1372,11 @@ def test_double_precision(self):
         self.assertTrue(isinstance(col_dict['i64'].type, sqltypes.BigInteger))
 
 
+class _TestSQLAlchemyConn(_EngineToConnMixin, _TestSQLAlchemy):
+    pass
+
 
-class TestSQLiteAlchemy(_TestSQLAlchemy):
+class _TestSQLiteAlchemy(object):
     """
     Test the sqlalchemy backend against an in-memory sqlite database.
 
@@ -1404,7 +1432,7 @@ def test_bigint_warning(self):
             self.assertEqual(len(w), 0, "Warning triggered for other table")
 
 
-class TestMySQLAlchemy(_TestSQLAlchemy):
+class _TestMySQLAlchemy(object):
     """
     Test the sqlalchemy backend against an MySQL database.
 
@@ -1478,7 +1506,7 @@ def test_read_procedure(self):
         tm.assert_frame_equal(df, res2)
 
 
-class TestPostgreSQLAlchemy(_TestSQLAlchemy):
+class _TestPostgreSQLAlchemy(object):
     """
     Test the sqlalchemy backend against an PostgreSQL database.
 
@@ -1574,6 +1602,31 @@ def test_datetime_with_time_zone(self):
         # "2000-06-01 00:00:00-07:00" should convert to "2000-06-01 07:00:00"
         self.assertEqual(df.DateColWithTz[1], Timestamp('2000-06-01 07:00:00'))
 
+
+class TestMySQLAlchemy(_TestMySQLAlchemy, _TestSQLAlchemy):
+    pass
+
+
+class TestMySQLAlchemyConn(_TestMySQLAlchemy, _TestSQLAlchemyConn): 
+    pass
+
+
+class TestPostgreSQLAlchemy(_TestPostgreSQLAlchemy, _TestSQLAlchemy):
+    pass
+
+
+class TestPostgreSQLAlchemyConn(_TestPostgreSQLAlchemy, _TestSQLAlchemyConn):
+    pass
+
+
+class TestSQLiteAlchemy(_TestSQLiteAlchemy, _TestSQLAlchemy):
+    pass
+
+
+class TestSQLiteAlchemyConn(_TestSQLiteAlchemy, _TestSQLAlchemyConn):
+    pass
+
+
 #------------------------------------------------------------------------------
 #--- Test Sqlite / MySQL fallback