ENH: add schema support to sql functions

Author: JoergRittinger

pandas/io/sql.py

@@ -38,7 +38,7 @@ def _is_sqlalchemy_engine(con):
         try:
             import sqlalchemy
             _SQLALCHEMY_INSTALLED = True
-            
+
             from distutils.version import LooseVersion
             ver = LooseVersion(sqlalchemy.__version__)
             # For sqlalchemy versions < 0.8.2, the BIGINT type is recognized
@@ -47,7 +47,7 @@ def _is_sqlalchemy_engine(con):
             if ver < '0.8.2':
                 from sqlalchemy import BigInteger
                 from sqlalchemy.ext.compiler import compiles
-        
+
                 @compiles(BigInteger, 'sqlite')
                 def compile_big_int_sqlite(type_, compiler, **kw):
                     return 'INTEGER'
@@ -145,7 +145,7 @@ def _safe_fetch(cur):
         if not isinstance(result, list):
             result = list(result)
         return result
-    except Exception as e: # pragma: no cover
+    except Exception as e:  # pragma: no cover
         excName = e.__class__.__name__
         if excName == 'OperationalError':
             return []
@@ -187,7 +187,7 @@ def tquery(sql, con=None, cur=None, retry=True):
             con.commit()
         except Exception as e:
             excName = e.__class__.__name__
-            if excName == 'OperationalError': # pragma: no cover
+            if excName == 'OperationalError':  # pragma: no cover
                 print('Failed to commit, may need to restart interpreter')
             else:
                 raise
@@ -199,7 +199,7 @@ def tquery(sql, con=None, cur=None, retry=True):
     if result and len(result[0]) == 1:
         # python 3 compat
         result = list(lzip(*result)[0])
-    elif result is None: # pragma: no cover
+    elif result is None:  # pragma: no cover
         result = []
 
     return result
@@ -253,7 +253,7 @@ def uquery(sql, con=None, cur=None, retry=True, params=None):
 #------------------------------------------------------------------------------
 #--- Read and write to DataFrames
 
-def read_sql_table(table_name, con, index_col=None, coerce_float=True,
+def read_sql_table(table_name, con, schema=None, index_col=None, coerce_float=True,
                    parse_dates=None, columns=None):
     """Read SQL database table into a DataFrame.
 
@@ -266,6 +266,7 @@ def read_sql_table(table_name, con, index_col=None, coerce_float=True,
         Name of SQL table in database
     con : SQLAlchemy engine
         Sqlite DBAPI connection mode not supported
+    schema : Name of SQL schema in database.
     index_col : string, optional
         Column to set as index
     coerce_float : boolean, default True
@@ -300,14 +301,14 @@ def read_sql_table(table_name, con, index_col=None, coerce_float=True,
     from sqlalchemy.schema import MetaData
     meta = MetaData(con)
     try:
-        meta.reflect(only=[table_name])
+        meta.reflect(only=[table_name], schema=schema)
     except sqlalchemy.exc.InvalidRequestError:
         raise ValueError("Table %s not found" % table_name)
 
     pandas_sql = PandasSQLAlchemy(con, meta=meta)
     table = pandas_sql.read_table(
         table_name, index_col=index_col, coerce_float=coerce_float,
-        parse_dates=parse_dates, columns=columns)
+        parse_dates=parse_dates, columns=columns, schema=schema)
 
     if table is not None:
         return table
@@ -437,8 +438,8 @@ def read_sql(sql, con, index_col=None, coerce_float=True, params=None,
             coerce_float=coerce_float, parse_dates=parse_dates)
 
 
-def to_sql(frame, name, con, flavor='sqlite', if_exists='fail', index=True,
-           index_label=None, chunksize=None):
+def to_sql(frame, name, con, flavor='sqlite', schema=None, if_exists='fail',
+           index=True, index_label=None, chunksize=None):
     """
     Write records stored in a DataFrame to a SQL database.
 
@@ -455,6 +456,7 @@ def to_sql(frame, name, con, flavor='sqlite', if_exists='fail', index=True,
         The flavor of SQL to use. Ignored when using SQLAlchemy engine.
         'mysql' is deprecated and will be removed in future versions, but it
         will be further supported through SQLAlchemy engines.
+    schema : Name of SQL schema in database.
     if_exists : {'fail', 'replace', 'append'}, default 'fail'
         - fail: If table exists, do nothing.
         - replace: If table exists, drop it, recreate it, and insert data.
@@ -481,10 +483,11 @@ def to_sql(frame, name, con, flavor='sqlite', if_exists='fail', index=True,
         raise NotImplementedError
 
     pandas_sql.to_sql(frame, name, if_exists=if_exists, index=index,
-                      index_label=index_label, chunksize=chunksize)
+                      index_label=index_label, schema=schema,
+                      chunksize=chunksize)
 
 
-def has_table(table_name, con, flavor='sqlite'):
+def has_table(table_name, con, flavor='sqlite', schema=None):
     """
     Check if DataBase has named table.
 
@@ -500,13 +503,14 @@ def has_table(table_name, con, flavor='sqlite'):
         The flavor of SQL to use. Ignored when using SQLAlchemy engine.
         'mysql' is deprecated and will be removed in future versions, but it
         will be further supported through SQLAlchemy engines.
+    schema : Name of SQL schema in database.
 
     Returns
     -------
     boolean
     """
     pandas_sql = pandasSQL_builder(con, flavor=flavor)
-    return pandas_sql.has_table(table_name)
+    return pandas_sql.has_table(table_name, schema=schema)
 
 table_exists = has_table
 
@@ -540,24 +544,26 @@ class PandasSQLTable(PandasObject):
     """
     # TODO: support for multiIndex
     def __init__(self, name, pandas_sql_engine, frame=None, index=True,
-                 if_exists='fail', prefix='pandas', index_label=None):
+                 if_exists='fail', prefix='pandas', index_label=None,
+                 schema=None):
         self.name = name
         self.pd_sql = pandas_sql_engine
         self.prefix = prefix
         self.frame = frame
         self.index = self._index_name(index, index_label)
+        self.schema = schema
 
         if frame is not None:
             # We want to write a frame
-            if self.pd_sql.has_table(self.name):
+            if self.pd_sql.has_table(self.name, self.schema):
                 if if_exists == 'fail':
                     raise ValueError("Table '%s' already exists." % name)
                 elif if_exists == 'replace':
-                    self.pd_sql.drop_table(self.name)
+                    self.pd_sql.drop_table(self.name, self.schema)
                     self.table = self._create_table_statement()
                     self.create()
                 elif if_exists == 'append':
-                    self.table = self.pd_sql.get_table(self.name)
+                    self.table = self.pd_sql.get_table(self.name, self.schema)
                     if self.table is None:
                         self.table = self._create_table_statement()
                 else:
@@ -568,13 +574,13 @@ def __init__(self, name, pandas_sql_engine, frame=None, index=True,
                 self.create()
         else:
             # no data provided, read-only mode
-            self.table = self.pd_sql.get_table(self.name)
+            self.table = self.pd_sql.get_table(self.name, self.schema)
 
         if self.table is None:
             raise ValueError("Could not init table '%s'" % name)
 
     def exists(self):
-        return self.pd_sql.has_table(self.name)
+        return self.pd_sql.has_table(self.name, self.schema)
 
     def sql_schema(self):
         from sqlalchemy.schema import CreateTable
@@ -709,7 +715,7 @@ def _create_table_statement(self):
         columns = [Column(name, typ)
                    for name, typ in column_names_and_types]
 
-        return Table(self.name, self.pd_sql.meta, *columns)
+        return Table(self.name, self.pd_sql.meta, *columns, schema=self.schema)
 
     def _harmonize_columns(self, parse_dates=None):
         """ Make a data_frame's column type align with an sql_table
@@ -843,9 +849,10 @@ def execute(self, *args, **kwargs):
         return self.engine.execute(*args, **kwargs)
 
     def read_table(self, table_name, index_col=None, coerce_float=True,
-                   parse_dates=None, columns=None):
+                   parse_dates=None, columns=None, schema=None):
 
-        table = PandasSQLTable(table_name, self, index=index_col)
+        table = PandasSQLTable(
+            table_name, self, index=index_col, schema=schema)
         return table.read(coerce_float=coerce_float,
                           parse_dates=parse_dates, columns=columns)
 
@@ -868,26 +875,29 @@ def read_sql(self, sql, index_col=None, coerce_float=True,
         return data_frame
 
     def to_sql(self, frame, name, if_exists='fail', index=True,
-               index_label=None, chunksize=None):
+               index_label=None, schema=None, chunksize=None):
         table = PandasSQLTable(
             name, self, frame=frame, index=index, if_exists=if_exists,
-            index_label=index_label)
+            index_label=index_label, schema=schema)
         table.insert(chunksize)
 
     @property
     def tables(self):
         return self.meta.tables
 
-    def has_table(self, name):
-        return self.engine.has_table(name)
+    def has_table(self, name, schema=None):
+        return self.engine.has_table(name, schema)
 
-    def get_table(self, table_name):
-        return self.meta.tables.get(table_name)
+    def get_table(self, table_name, schema=None):
+        if schema:
+            return self.meta.tables.get('.'.join([schema, table_name]))
+        else:
+            return self.meta.tables.get(table_name)
 
-    def drop_table(self, table_name):
-        if self.engine.has_table(table_name):
-            self.meta.reflect(only=[table_name])
-            self.get_table(table_name).drop()
+    def drop_table(self, table_name, schema=None):
+        if self.engine.has_table(table_name, schema):
+            self.meta.reflect(only=[table_name], schema=schema)
+            self.get_table(table_name, schema).drop()
             self.meta.clear()
 
     def _create_sql_schema(self, frame, table_name):