From 7327f6bb16cfcf0b7ddf12f5403c81aa3757211e Mon Sep 17 00:00:00 2001
From: Thomas Grainger <tom@yplanapp.com>
Date: Wed, 27 May 2015 12:45:06 +0100
Subject: [PATCH] support both sqlalchemy engines and connections Fixes #7877

update pymysql to 0.6.3 to avoid cursor bugs

Add documentation and tests for SQLAlchemy connectables

explicit reference to connection/engine in docs

Temporary table test

pass compile the connectable
---
 ci/requirements-2.7.txt         |   2 +-
 doc/source/io.rst               |   9 ++-
 doc/source/whatsnew/v0.17.0.txt |   1 +
 pandas/io/sql.py                |  83 +++++++++++--------
 pandas/io/tests/test_sql.py     | 138 ++++++++++++++++++++++++++++----
 5 files changed, 180 insertions(+), 53 deletions(-)

diff --git a/ci/requirements-2.7.txt b/ci/requirements-2.7.txt
index 0d515f300f5a7..951c8798bef15 100644
--- a/ci/requirements-2.7.txt
+++ b/ci/requirements-2.7.txt
@@ -17,7 +17,7 @@ boto=2.36.0
 bottleneck=0.8.0
 psycopg2=2.5.2
 patsy
-pymysql=0.6.1
+pymysql=0.6.3
 html5lib=1.0b2
 beautiful-soup=4.2.1
 httplib2=0.8
diff --git a/doc/source/io.rst b/doc/source/io.rst
index 185deb4b9cae8..7a4318fb02cfc 100644
--- a/doc/source/io.rst
+++ b/doc/source/io.rst
@@ -3555,9 +3555,16 @@ below and the SQLAlchemy `documentation <http://docs.sqlalchemy.org/en/rel_0_9/c
 .. ipython:: python
 
    from sqlalchemy import create_engine
-   # Create your connection.
+   # Create your engine.
    engine = create_engine('sqlite:///:memory:')
 
+If you want to manage your own connections you can pass one of those instead:
+
+.. ipython:: python
+
+   with engine.connect() as conn, conn.begin():
+       data = pd.read_sql_table('data', conn)
+
 Writing DataFrames
 ~~~~~~~~~~~~~~~~~~
 
diff --git a/doc/source/whatsnew/v0.17.0.txt b/doc/source/whatsnew/v0.17.0.txt
index cd41c4fc82146..da16734dc873b 100644
--- a/doc/source/whatsnew/v0.17.0.txt
+++ b/doc/source/whatsnew/v0.17.0.txt
@@ -26,6 +26,7 @@ Check the :ref:`API Changes <whatsnew_0170.api>` and :ref:`deprecations <whatsne
 New features
 ~~~~~~~~~~~~
 
+- SQL io functions now accept a SQLAlchemy connectable. (:issue:`7877`)
 
 .. _whatsnew_0170.enhancements.other:
 
diff --git a/pandas/io/sql.py b/pandas/io/sql.py
index 8d8768c08fe02..ef8360f0ff459 100644
--- a/pandas/io/sql.py
+++ b/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
 
@@ -139,7 +139,7 @@ def execute(sql, con, cur=None, params=None):
     ----------
     sql : string
         Query to be executed
-    con : SQLAlchemy engine or sqlite3 DBAPI2 connection
+    con : SQLAlchemy connectable(engine/connection) or sqlite3 DBAPI2 connection
         Using SQLAlchemy makes it possible to use any DB supported by that
         library.
         If a DBAPI2 object, only sqlite3 is supported.
@@ -282,14 +282,14 @@ def read_sql_table(table_name, con, schema=None, index_col=None,
                    chunksize=None):
     """Read SQL database table into a DataFrame.
 
-    Given a table name and an SQLAlchemy engine, returns a DataFrame.
+    Given a table name and an SQLAlchemy connectable, returns a DataFrame.
     This function does not support DBAPI connections.
 
     Parameters
     ----------
     table_name : string
         Name of SQL table in database
-    con : SQLAlchemy engine
+    con : SQLAlchemy connectable
         Sqlite DBAPI connection mode not supported
     schema : string, default None
         Name of SQL schema in database to query (if database flavor
@@ -328,9 +328,9 @@ 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.")
+                                  "SQLAlchemy connectable.")
     import sqlalchemy
     from sqlalchemy.schema import MetaData
     meta = MetaData(con, schema=schema)
@@ -362,7 +362,7 @@ def read_sql_query(sql, con, index_col=None, coerce_float=True, params=None,
     ----------
     sql : string
         SQL query to be executed
-    con : SQLAlchemy engine or sqlite3 DBAPI2 connection
+    con : SQLAlchemy connectable(engine/connection) or sqlite3 DBAPI2 connection
         Using SQLAlchemy makes it possible to use any DB supported by that
         library.
         If a DBAPI2 object, only sqlite3 is supported.
@@ -420,7 +420,7 @@ def read_sql(sql, con, index_col=None, coerce_float=True, params=None,
     ----------
     sql : string
         SQL query to be executed or database table name.
-    con : SQLAlchemy engine or DBAPI2 connection (fallback mode)
+    con : SQLAlchemy connectable(engine/connection) or DBAPI2 connection (fallback mode)
         Using SQLAlchemy makes it possible to use any DB supported by that
         library.
         If a DBAPI2 object, only sqlite3 is supported.
@@ -504,14 +504,14 @@ def to_sql(frame, name, con, flavor='sqlite', schema=None, if_exists='fail',
     frame : DataFrame
     name : string
         Name of SQL table
-    con : SQLAlchemy engine or sqlite3 DBAPI2 connection
+    con : SQLAlchemy connectable(engine/connection) or sqlite3 DBAPI2 connection
         Using SQLAlchemy makes it possible to use any DB supported by that
         library.
         If a DBAPI2 object, only sqlite3 is supported.
     flavor : {'sqlite', 'mysql'}, default 'sqlite'
-        The flavor of SQL to use. Ignored when using SQLAlchemy engine.
+        The flavor of SQL to use. Ignored when using SQLAlchemy connectable.
         'mysql' is deprecated and will be removed in future versions, but it
-        will be further supported through SQLAlchemy engines.
+        will be further supported through SQLAlchemy connectables.
     schema : string, default None
         Name of SQL schema in database to write to (if database flavor
         supports this). If None, use default schema (default).
@@ -557,14 +557,14 @@ def has_table(table_name, con, flavor='sqlite', schema=None):
     ----------
     table_name: string
         Name of SQL table
-    con: SQLAlchemy engine or sqlite3 DBAPI2 connection
+    con: SQLAlchemy connectable(engine/connection) or sqlite3 DBAPI2 connection
         Using SQLAlchemy makes it possible to use any DB supported by that
         library.
         If a DBAPI2 object, only sqlite3 is supported.
     flavor: {'sqlite', 'mysql'}, default 'sqlite'
-        The flavor of SQL to use. Ignored when using SQLAlchemy engine.
+        The flavor of SQL to use. Ignored when using SQLAlchemy connectable.
         'mysql' is deprecated and will be removed in future versions, but it
-        will be further supported through SQLAlchemy engines.
+        will be further supported through SQLAlchemy connectables.
     schema : string, default None
         Name of SQL schema in database to write to (if database flavor supports
         this). If None, use default schema (default).
@@ -581,7 +581,7 @@ def has_table(table_name, con, flavor='sqlite', schema=None):
 
 _MYSQL_WARNING = ("The 'mysql' flavor with DBAPI connection is deprecated "
                   "and will be removed in future versions. "
-                  "MySQL will be further supported with SQLAlchemy engines.")
+                  "MySQL will be further supported with SQLAlchemy connectables.")
 
 
 def pandasSQL_builder(con, flavor=None, schema=None, meta=None,
@@ -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,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.connectable))
 
     def _execute_create(self):
         # Inserting table into database, add to MetaData object
@@ -982,11 +982,11 @@ class PandasSQL(PandasObject):
     """
 
     def read_sql(self, *args, **kwargs):
-        raise ValueError("PandasSQL must be created with an SQLAlchemy engine"
+        raise ValueError("PandasSQL must be created with an SQLAlchemy connectable"
                          " or connection+sql flavor")
 
     def to_sql(self, *args, **kwargs):
-        raise ValueError("PandasSQL must be created with an SQLAlchemy engine"
+        raise ValueError("PandasSQL must be created with an SQLAlchemy connectable"
                          " or connection+sql flavor")
 
 
@@ -997,8 +997,8 @@ class SQLDatabase(PandasSQL):
 
     Parameters
     ----------
-    engine : SQLAlchemy engine
-        Engine to connect with the database. Using SQLAlchemy makes it
+    engine : SQLAlchemy connectable
+        Connectable to connect with the database. Using SQLAlchemy makes it
         possible to use any DB supported by that library.
     schema : string, default None
         Name of SQL schema in database to write to (if database flavor
@@ -1011,19 +1011,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)
+        """Simple passthrough to SQLAlchemy connectable"""
+        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,
@@ -1191,7 +1196,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 "
@@ -1202,7 +1213,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
@@ -1221,7 +1236,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()
@@ -1610,12 +1625,12 @@ def get_schema(frame, name, flavor='sqlite', keys=None, con=None, dtype=None):
     name : string
         name of SQL table
     flavor : {'sqlite', 'mysql'}, default 'sqlite'
-        The flavor of SQL to use. Ignored when using SQLAlchemy engine.
+        The flavor of SQL to use. Ignored when using SQLAlchemy connectable.
         'mysql' is deprecated and will be removed in future versions, but it
         will be further supported through SQLAlchemy engines.
     keys : string or sequence
         columns to use a primary key
-    con: an open SQL database connection object or an SQLAlchemy engine
+    con: an open SQL database connection object or a SQLAlchemy connectable
         Using SQLAlchemy makes it possible to use any DB supported by that
         library.
         If a DBAPI2 object, only sqlite3 is supported.
@@ -1673,8 +1688,8 @@ def write_frame(frame, name, con, flavor='sqlite', if_exists='fail', **kwargs):
 
     - With ``to_sql`` the index is written to the sql database by default. To
       keep the behaviour this function you need to specify ``index=False``.
-    - The new ``to_sql`` function supports sqlalchemy engines to work with
-      different sql flavors.
+    - The new ``to_sql`` function supports sqlalchemy connectables to work
+      with different sql flavors.
 
     See also
     --------
diff --git a/pandas/io/tests/test_sql.py b/pandas/io/tests/test_sql.py
index d8bc3c61f68f0..ba951c7cb513d 100644
--- a/pandas/io/tests/test_sql.py
+++ b/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`)
 
@@ -43,6 +44,8 @@
     import sqlalchemy
     import sqlalchemy.schema
     import sqlalchemy.sql.sqltypes as sqltypes
+    from sqlalchemy.ext import declarative
+    from sqlalchemy.orm import session as sa_session
     SQLALCHEMY_INSTALLED = True
 except ImportError:
     SQLALCHEMY_INSTALLED = False
@@ -867,6 +870,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
@@ -1003,9 +1031,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()
 
@@ -1359,9 +1384,58 @@ def test_double_precision(self):
         self.assertTrue(isinstance(col_dict['i32'].type, sqltypes.Integer))
         self.assertTrue(isinstance(col_dict['i64'].type, sqltypes.BigInteger))
 
+    def test_connectable_issue_example(self):
+        # This tests the example raised in issue
+        # https://github.com/pydata/pandas/issues/10104
+
+        def foo(connection):
+            query = 'SELECT test_foo_data FROM test_foo_data'
+            return sql.read_sql_query(query, con=connection)
+
+        def bar(connection, data):
+            data.to_sql(name='test_foo_data', con=connection, if_exists='append')
+
+        def main(connectable):
+            with connectable.connect() as conn:
+                with conn.begin():
+                    foo_data = conn.run_callable(foo)
+                    conn.run_callable(bar, foo_data)
+
+        DataFrame({'test_foo_data': [0, 1, 2]}).to_sql('test_foo_data', self.conn)
+        main(self.conn)
+
+    def test_temporary_table(self):
+        test_data = u'Hello, World!'
+        expected = DataFrame({'spam': [test_data]})
+        Base = declarative.declarative_base()
+
+        class Temporary(Base):
+            __tablename__ = 'temp_test'
+            __table_args__ = {'prefixes': ['TEMPORARY']}
+            id = sqlalchemy.Column(sqlalchemy.Integer, primary_key=True)
+            spam = sqlalchemy.Column(sqlalchemy.Unicode(30), nullable=False)
+
+        Session = sa_session.sessionmaker(bind=self.conn)
+        session = Session()
+        with session.transaction:
+            conn = session.connection()
+            Temporary.__table__.create(conn)
+            session.add(Temporary(spam=test_data))
+            session.flush()
+            df = sql.read_sql_query(
+                sql=sqlalchemy.select([Temporary.spam]),
+                con=conn,
+            )
+
+        tm.assert_frame_equal(df, expected)
+
+
+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.
 
@@ -1378,8 +1452,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)
@@ -1417,7 +1491,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.
 
@@ -1438,6 +1512,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])
@@ -1491,7 +1566,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.
 
@@ -1512,6 +1587,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'")
@@ -1563,15 +1639,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
@@ -1587,6 +1666,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