support both sqlalchemy engines and connections Fixes pandas-dev#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()