From 5cadd6794cfb6c70f54e968b9e82b2c991f5ea5b Mon Sep 17 00:00:00 2001
From: Maxime Beauchemin <maximebeauchemin@gmail.com>
Date: Wed, 13 Apr 2016 17:28:12 -0700
Subject: [PATCH] Allowing to specify schema for tables (#330)

* Allowing to specify schema for tables

* Installed postgres and tested/fixed
---
 .../bb51420eaf83_add_schema_to_table_model.py | 22 +++++++++++++++++++
 caravel/models.py                             | 18 ++++++++++-----
 caravel/views.py                              |  8 +++++--
 dev-reqs.txt                                  |  1 +
 4 files changed, 41 insertions(+), 8 deletions(-)
 create mode 100644 caravel/migrations/versions/bb51420eaf83_add_schema_to_table_model.py

diff --git a/caravel/migrations/versions/bb51420eaf83_add_schema_to_table_model.py b/caravel/migrations/versions/bb51420eaf83_add_schema_to_table_model.py
new file mode 100644
index 000000000000..f86e332300ff
--- /dev/null
+++ b/caravel/migrations/versions/bb51420eaf83_add_schema_to_table_model.py
@@ -0,0 +1,22 @@
+"""add schema to table model
+
+Revision ID: bb51420eaf83
+Revises: 867bf4f117f9
+Create Date: 2016-04-11 22:41:06.185955
+
+"""
+
+# revision identifiers, used by Alembic.
+revision = 'bb51420eaf83'
+down_revision = '867bf4f117f9'
+
+from alembic import op
+import sqlalchemy as sa
+
+
+def upgrade():
+    op.add_column('tables', sa.Column('schema', sa.String(length=256), nullable=True))
+
+
+def downgrade():
+    op.drop_column('tables', 'schema')
diff --git a/caravel/models.py b/caravel/models.py
index 3588d13e1c07..1e572ed92641 100644
--- a/caravel/models.py
+++ b/caravel/models.py
@@ -369,11 +369,12 @@ def get_extra(self):
                 logging.error(e)
         return extra
 
-    def get_table(self, table_name):
+    def get_table(self, table_name, schema=None):
         extra = self.get_extra()
         meta = MetaData(**extra.get('metadata_params', {}))
         return Table(
             table_name, meta,
+            schema=schema or None,
             autoload=True,
             autoload_with=self.get_sqla_engine())
 
@@ -417,6 +418,7 @@ class SqlaTable(Model, Queryable, AuditMixinNullable):
         'Database', backref='tables', foreign_keys=[database_id])
     offset = Column(Integer, default=0)
     cache_timeout = Column(Integer)
+    schema = Column(String(256))
 
     baselink = "tablemodelview"
 
@@ -592,7 +594,11 @@ def query(  # sqla
 
         select_exprs += metrics_exprs
         qry = select(select_exprs)
-        from_clause = table(self.table_name)
+
+        tbl = table(self.table_name)
+        if self.schema:
+            tbl.schema = self.schema
+
         if not columns:
             qry = qry.group_by(*groupby_exprs)
 
@@ -625,7 +631,7 @@ def query(  # sqla
 
         if timeseries_limit and groupby:
             subq = select(inner_select_exprs)
-            subq = subq.select_from(table(self.table_name))
+            subq = subq.select_from(tbl)
             subq = subq.where(and_(*(where_clause_and + inner_time_filter)))
             subq = subq.group_by(*inner_groupby_exprs)
             subq = subq.order_by(desc(main_metric_expr))
@@ -635,9 +641,9 @@ def query(  # sqla
                 on_clause.append(
                     groupby_exprs[i] == column("__" + gb))
 
-            from_clause = from_clause.join(subq.alias(), and_(*on_clause))
+            tbl = tbl.join(subq.alias(), and_(*on_clause))
 
-        qry = qry.select_from(from_clause)
+        qry = qry.select_from(tbl)
 
         engine = self.database.get_sqla_engine()
         sql = "{}".format(
@@ -653,7 +659,7 @@ def query(  # sqla
     def fetch_metadata(self):
         """Fetches the metadata for the table and merges it in"""
         try:
-            table = self.database.get_table(self.table_name)
+            table = self.database.get_table(self.table_name, schema=self.schema)
         except Exception as e:
             flash(str(e))
             flash(
diff --git a/caravel/views.py b/caravel/views.py
index f902982339be..ff499b04f800 100644
--- a/caravel/views.py
+++ b/caravel/views.py
@@ -185,14 +185,18 @@ class TableModelView(CaravelModelView, DeleteMixin):  # noqa
         'table_link', 'database', 'sql_link', 'is_featured',
         'changed_by_', 'changed_on']
     add_columns = [
-        'table_name', 'database', 'default_endpoint', 'offset', 'cache_timeout']
+        'table_name', 'database', 'schema',
+        'default_endpoint', 'offset', 'cache_timeout']
     edit_columns = [
-        'table_name', 'is_featured', 'database', 'description', 'owner',
+        'table_name', 'is_featured', 'database', 'schema', 'description', 'owner',
         'main_dttm_col', 'default_endpoint', 'offset', 'cache_timeout']
     related_views = [TableColumnInlineView, SqlMetricInlineView]
     base_order = ('changed_on', 'desc')
     description_columns = {
         'offset': "Timezone offset (in hours) for this datasource",
+        'schema': (
+            "Schema, as used only in some databases like Postgres, Redshift "
+            "and DB2"),
         'description': Markup(
             "Supports <a href='https://daringfireball.net/projects/markdown/'>"
             "markdown</a>"),
diff --git a/dev-reqs.txt b/dev-reqs.txt
index 783c93c2f789..73cfa730c57f 100644
--- a/dev-reqs.txt
+++ b/dev-reqs.txt
@@ -1,4 +1,5 @@
 coveralls
+nose
 sphinx
 sphinx_bootstrap_theme
 sphinxcontrib.youtube