support both sqlalchemy engines and connections Fixes #7877

Author: Thomas Grainger

pandas/io/sql.py

@@ -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
 
@@ -637,7 +637,7 @@ def exists(self):
 
     def sql_schema(self):
         from sqlalchemy.schema import CreateTable
-        return str(CreateTable(self.table).compile(self.pd_sql.engine))
+        return str(CreateTable(self.table).compile(self.pd_sql.connection))
 
     def _execute_create(self):
         # Inserting table into database, add to MetaData object
@@ -1006,20 +1006,31 @@ class SQLDatabase(PandasSQL):
 
     """
 
-    def __init__(self, engine, schema=None, meta=None):
-        self.engine = engine
+    def __init__(self, connection, schema=None, meta=None):
+        import sqlalchemy.engine
+        if isinstance(connection, sqlalchemy.engine.Engine):
+            self.connection = connection.connect()
+        else:
+            self.connection = connection
         if not meta:
             from sqlalchemy.schema import MetaData
-            meta = MetaData(self.engine, schema=schema)
+            meta = MetaData(self.connection, schema=schema)
 
         self.meta = meta
 
+    @contextmanager
     def run_transaction(self):
-        return self.engine.begin()
+        trans = self.connection.begin()
+        try:
+            yield self.connection
+            trans.commit()
+        except:
+            trans.rollback()
+            raise
 
     def execute(self, *args, **kwargs):
         """Simple passthrough to SQLAlchemy engine"""
-        return self.engine.execute(*args, **kwargs)
+        return self.connection.execute(*args, **kwargs)
 
     def read_table(self, table_name, index_col=None, coerce_float=True,
                    parse_dates=None, columns=None, schema=None,
@@ -1187,7 +1198,7 @@ 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):
+        if name not in self.connection.engine.table_names(schema=schema or self.meta.schema, connection=self.connection):
             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 +1209,7 @@ 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)
+        return self.connection.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 +1228,7 @@ 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):
+        if self.connection.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,18 @@ def test_sqlalchemy_type_mapping(self):
         self.assertTrue(isinstance(table.table.c['time'].type, sqltypes.DateTime))
 
 
+class TestSQLApiConn(TestSQLApi):
+    def setUp(self):
+        super(TestSQLApiConn, self).setUp()
+        conn = self.conn
+        conn.connect()
+        self.addCleanup(conn.close)
+        tx = conn.begin()
+        self.addCleanup(tx.rollback)
+
+        self.conn = conn
+
+
 class TestSQLiteFallbackApi(_TestSQLApi):
     """
     Test the public sqlite connection fallback API
@@ -1347,8 +1360,18 @@ def test_double_precision(self):
         self.assertTrue(isinstance(col_dict['i64'].type, sqltypes.BigInteger))
 
 
+class _TestSQLAlchemyConn(_TestSQLAlchemy):
+    def setUp(self):
+        super(_TestSQLAlchemyConn, self).setUp()
+        conn = self.conn.connect()
+        self.addCleanup(conn.close)
+        tx = conn.begin()
+        self.addCleanup(tx.rollback)
+
+        self.conn = conn
+
 
-class TestSQLiteAlchemy(_TestSQLAlchemy):
+class _TestSQLiteAlchemy(object):
     """
     Test the sqlalchemy backend against an in-memory sqlite database.
 
@@ -1404,7 +1427,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 +1501,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 +1597,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