From 91d951ac422a6411e59b2ec195011ea4dc44b0b1 Mon Sep 17 00:00:00 2001
From: "Rich @ RadICS" <RichRadics@users.noreply.github.com>
Date: Fri, 12 May 2017 18:37:24 +0200
Subject: [PATCH] Change hardcoded references to 'User' security model to allow
 custom class override (#2728)

---
 superset/connectors/druid/models.py | 2 +-
 superset/connectors/sqla/models.py  | 4 ++--
 superset/models/core.py             | 6 +++---
 superset/models/sql_lab.py          | 5 +++--
 4 files changed, 9 insertions(+), 8 deletions(-)

diff --git a/superset/connectors/druid/models.py b/superset/connectors/druid/models.py
index 4ad6a6f951914..73d7368cbf7e8 100644
--- a/superset/connectors/druid/models.py
+++ b/superset/connectors/druid/models.py
@@ -323,7 +323,7 @@ class DruidDatasource(Model, BaseDatasource):
         'DruidCluster', backref='datasources', foreign_keys=[cluster_name])
     user_id = Column(Integer, ForeignKey('ab_user.id'))
     owner = relationship(
-        'User',
+        sm.user_model,
         backref=backref('datasources', cascade='all, delete-orphan'),
         foreign_keys=[user_id])
 
diff --git a/superset/connectors/sqla/models.py b/superset/connectors/sqla/models.py
index a5e445ff0b443..4f8c2af341699 100644
--- a/superset/connectors/sqla/models.py
+++ b/superset/connectors/sqla/models.py
@@ -20,7 +20,7 @@
 from flask_appbuilder import Model
 from flask_babel import lazy_gettext as _
 
-from superset import db, utils, import_util
+from superset import db, utils, import_util, sm
 from superset.connectors.base import BaseDatasource, BaseColumn, BaseMetric
 from superset.utils import DTTM_ALIAS, QueryStatus
 from superset.models.helpers import QueryResult
@@ -169,7 +169,7 @@ class SqlaTable(Model, BaseDatasource):
     fetch_values_predicate = Column(String(1000))
     user_id = Column(Integer, ForeignKey('ab_user.id'))
     owner = relationship(
-        'User',
+        sm.user_model,
         backref='tables',
         foreign_keys=[user_id])
     database = relationship(
diff --git a/superset/models/core.py b/superset/models/core.py
index f7412fbc1173e..ef66671dfdfc5 100644
--- a/superset/models/core.py
+++ b/superset/models/core.py
@@ -103,7 +103,7 @@ class Slice(Model, AuditMixinNullable, ImportMixin):
     description = Column(Text)
     cache_timeout = Column(Integer)
     perm = Column(String(1000))
-    owners = relationship("User", secondary=slice_user)
+    owners = relationship(sm.user_model, secondary=slice_user)
 
     export_fields = ('slice_name', 'datasource_type', 'datasource_name',
                      'viz_type', 'params', 'cache_timeout')
@@ -304,7 +304,7 @@ class Dashboard(Model, AuditMixinNullable, ImportMixin):
     slug = Column(String(255), unique=True)
     slices = relationship(
         'Slice', secondary=dashboard_slices, backref='dashboards')
-    owners = relationship("User", secondary=dashboard_user)
+    owners = relationship(sm.user_model, secondary=dashboard_user)
 
     export_fields = ('dashboard_title', 'position_json', 'json_metadata',
                      'description', 'css', 'slug')
@@ -728,7 +728,7 @@ class Log(Model):
     dashboard_id = Column(Integer)
     slice_id = Column(Integer)
     json = Column(Text)
-    user = relationship('User', backref='logs', foreign_keys=[user_id])
+    user = relationship(sm.user_model, backref='logs', foreign_keys=[user_id])
     dttm = Column(DateTime, default=datetime.utcnow)
     dt = Column(Date, default=date.today())
     duration_ms = Column(Integer)
diff --git a/superset/models/sql_lab.py b/superset/models/sql_lab.py
index 18c67f4561249..84964b9cf57c5 100644
--- a/superset/models/sql_lab.py
+++ b/superset/models/sql_lab.py
@@ -20,6 +20,7 @@
 )
 from sqlalchemy.orm import backref, relationship
 
+from superset import sm
 from superset.utils import QueryStatus
 from superset.models.helpers import AuditMixinNullable
 
@@ -79,7 +80,7 @@ class Query(Model):
         backref=backref('queries', cascade='all, delete-orphan')
     )
     user = relationship(
-        'User',
+        sm.user_model,
         foreign_keys=[user_id])
 
     __table_args__ = (
@@ -147,7 +148,7 @@ class SavedQuery(Model, AuditMixinNullable):
     description = Column(Text)
     sql = Column(Text)
     user = relationship(
-        'User',
+        sm.user_model,
         backref=backref('saved_queries', cascade='all, delete-orphan'),
         foreign_keys=[user_id])
     database = relationship(