From 4c8523efc057421e9f19f0302faa92ca82065849 Mon Sep 17 00:00:00 2001
From: yxjames <zhugeyx@163.com>
Date: Fri, 17 Jun 2016 08:14:26 -0700
Subject: [PATCH] Prevent potential db upgrading problem (#628)

* fix general db upgrade problem

* add comments
---
 .../c3a8f8611885_materializing_permission.py  | 21 ++++++++++++---
 ...74f7aad_add_new_field_is_restricted_to_.py | 27 +++++++++++++++----
 2 files changed, 39 insertions(+), 9 deletions(-)

diff --git a/caravel/migrations/versions/c3a8f8611885_materializing_permission.py b/caravel/migrations/versions/c3a8f8611885_materializing_permission.py
index 1d65a172a4fa6..08db27b3c0f3c 100644
--- a/caravel/migrations/versions/c3a8f8611885_materializing_permission.py
+++ b/caravel/migrations/versions/c3a8f8611885_materializing_permission.py
@@ -13,15 +13,28 @@
 from alembic import op
 import sqlalchemy as sa
 from caravel import db
-from caravel import models
-
-
+from sqlalchemy.ext.declarative import declarative_base
+from sqlalchemy import (
+    Column, Integer, String, ForeignKey)
+
+Base = declarative_base()
+
+class Slice(Base):
+    """Declarative class to do query in upgrade"""
+    __tablename__ = 'slices'
+    id = Column(Integer, primary_key=True)
+    slice_name = Column(String(250))
+    druid_datasource_id = Column(Integer, ForeignKey('datasources.id'))
+    table_id = Column(Integer, ForeignKey('tables.id'))
+    perm = Column(String(2000))
+    
 def upgrade():
     bind = op.get_bind()
     op.add_column('slices', sa.Column('perm', sa.String(length=2000), nullable=True))
     session = db.Session(bind=bind)
 
-    for slc in session.query(models.Slice).all():
+    # Use Slice class defined here instead of models.Slice
+    for slc in session.query(Slice).all():
         if slc.datasource:
             slc.perm = slc.datasource.perm
             session.merge(slc)
diff --git a/caravel/migrations/versions/d8bc074f7aad_add_new_field_is_restricted_to_.py b/caravel/migrations/versions/d8bc074f7aad_add_new_field_is_restricted_to_.py
index 127ee27e6845c..3af362892f50c 100644
--- a/caravel/migrations/versions/d8bc074f7aad_add_new_field_is_restricted_to_.py
+++ b/caravel/migrations/versions/d8bc074f7aad_add_new_field_is_restricted_to_.py
@@ -13,9 +13,24 @@
 from alembic import op
 import sqlalchemy as sa
 from caravel import db
-from caravel import models
-
-
+from sqlalchemy.ext.declarative import declarative_base
+from sqlalchemy import (
+    Column, Integer, Boolean)
+
+Base = declarative_base()
+
+class DruidMetric(Base):
+    """Declarative class used to do query in upgrade"""
+    __tablename__ = 'metrics'
+    id = Column(Integer, primary_key=True)
+    is_restricted = Column(Boolean, default=False, nullable=True)
+
+class SqlMetric(Base):
+    """Declarative class used to do query in upgrade"""
+    __tablename__ = 'sql_metrics'
+    id = Column(Integer, primary_key=True)
+    is_restricted = Column(Boolean, default=False, nullable=True)
+    
 def upgrade():
     op.add_column('metrics', sa.Column('is_restricted', sa.Boolean(), nullable=True))
     op.add_column('sql_metrics', sa.Column('is_restricted', sa.Boolean(), nullable=True))
@@ -23,10 +38,12 @@ def upgrade():
     bind = op.get_bind()
     session = db.Session(bind=bind)
 
-    for obj in session.query(models.DruidMetric).all():
+    # don't use models.DruidMetric 
+    # because it assumes the context is consistent with the application
+    for obj in session.query(DruidMetric).all():
         obj.is_restricted = False
 
-    for obj in session.query(models.SqlMetric).all():
+    for obj in session.query(SqlMetric).all():
         obj.is_restricted = False
 
     session.commit()