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,24 @@ 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
 
+    @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 +1193,13 @@ 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
+        with self.connectable.connect() as conn:
+            table_names = engine.table_names(
+                schema=schema or self.meta.schema,
+                connection=conn,
+            )
+        if name not in table_names:
             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 +1210,11 @@ 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.connectable.run_callable(
+            self.connectable.dialect.has_table,
+            name,
+            schema or self.meta.schema,
+        )
 
     def get_table(self, table_name, schema=None):
         schema = schema or self.meta.schema
@@ -1217,7 +1233,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.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,31 @@ 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(self):
+        super(_EngineToConnMixin, self).setUp()
+        engine = self.conn
+        conn = engine.connect()
+        self.__tx = conn.begin()
+        self.pandasSQL = sql.SQLDatabase(conn)
+        self.__engine = engine
+        self.conn = conn
+
+    def tearDown(self):
+        self.__tx.rollback()
+        self.conn.close()
+        self.conn = self.__engine
+        self.pandasSQL = sql.SQLDatabase(self.__engine)
+
+
+class TestSQLApiConn(_EngineToConnMixin, TestSQLApi):
+    pass
+
+
 class TestSQLiteFallbackApi(_TestSQLApi):
     """
     Test the public sqlite connection fallback API
@@ -990,9 +1016,6 @@ def setup_connect(self):
         except sqlalchemy.exc.OperationalError:
             raise nose.SkipTest("Can't connect to {0} server".format(self.flavor))
 
-    def tearDown(self):
-        raise NotImplementedError()
-
     def test_aread_sql(self):
         self._read_sql_iris()
 
@@ -1347,8 +1370,12 @@ def test_double_precision(self):
         self.assertTrue(isinstance(col_dict['i64'].type, sqltypes.BigInteger))
 
 
+class _TestSQLAlchemyConn(_EngineToConnMixin, _TestSQLAlchemy):
+    def test_transactions(self):
+        raise nose.SkipTest("Nested transactions rollbacks don't work with Pandas")
+
 
-class TestSQLiteAlchemy(_TestSQLAlchemy):
+class _TestSQLiteAlchemy(object):
     """
     Test the sqlalchemy backend against an in-memory sqlite database.
 
@@ -1365,8 +1392,8 @@ def setup_driver(cls):
         cls.driver = None
 
     def tearDown(self):
+        super(_TestSQLiteAlchemy, self).tearDown()
         # in memory so tables should not be removed explicitly
-        pass
 
     def test_default_type_conversion(self):
         df = sql.read_sql_table("types_test_data", self.conn)
@@ -1404,7 +1431,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.
 
@@ -1425,6 +1452,7 @@ def setup_driver(cls):
             raise nose.SkipTest('pymysql not installed')
 
     def tearDown(self):
+        super(_TestMySQLAlchemy, self).tearDown()
         c = self.conn.execute('SHOW TABLES')
         for table in c.fetchall():
             self.conn.execute('DROP TABLE %s' % table[0])
@@ -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.
 
@@ -1499,6 +1527,7 @@ def setup_driver(cls):
             raise nose.SkipTest('psycopg2 not installed')
 
     def tearDown(self):
+        super(_TestPostgreSQLAlchemy, self).tearDown()
         c = self.conn.execute(
             "SELECT table_name FROM information_schema.tables"
             " WHERE table_schema = 'public'")
@@ -1550,15 +1579,18 @@ def test_schema_support(self):
 
         ## specifying schema in user-provided meta
 
-        engine2 = self.connect()
-        meta = sqlalchemy.MetaData(engine2, schema='other')
-        pdsql = sql.SQLDatabase(engine2, meta=meta)
-        pdsql.to_sql(df, 'test_schema_other2', index=False)
-        pdsql.to_sql(df, 'test_schema_other2', index=False, if_exists='replace')
-        pdsql.to_sql(df, 'test_schema_other2', index=False, if_exists='append')
-        res1 = sql.read_sql_table('test_schema_other2', self.conn, schema='other')
-        res2 = pdsql.read_table('test_schema_other2')
-        tm.assert_frame_equal(res1, res2)
+        # The schema won't be applied on another Connection
+        # because of transactional schemas
+        if isinstance(self.conn, sqlalchemy.engine.Engine):
+            engine2 = self.connect()
+            meta = sqlalchemy.MetaData(engine2, schema='other')
+            pdsql = sql.SQLDatabase(engine2, meta=meta)
+            pdsql.to_sql(df, 'test_schema_other2', index=False)
+            pdsql.to_sql(df, 'test_schema_other2', index=False, if_exists='replace')
+            pdsql.to_sql(df, 'test_schema_other2', index=False, if_exists='append')
+            res1 = sql.read_sql_table('test_schema_other2', self.conn, schema='other')
+            res2 = pdsql.read_table('test_schema_other2')
+            tm.assert_frame_equal(res1, res2)
 
     def test_datetime_with_time_zone(self):
         # Test to see if we read the date column with timezones that
@@ -1574,6 +1606,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