Add get_by_id & get_by_name methods for Query and DataSource classes

redash/models.py

@@ -514,6 +514,14 @@ class DataSource(BelongsToOrgMixin, db.Model):
     def query_runner(self):
         return get_query_runner(self.type, self.options)
 
+    @classmethod
+    def get_by_id(cls, _id):
+        return cls.query.filter(cls.id == _id).one()
+
+    @classmethod
+    def get_by_name(cls, name):
+        return cls.query.filter(cls.name == name).one()
+
     @classmethod
     def all(cls, org, group_ids=None):
         data_sources = cls.query.filter(cls.org == org).order_by(cls.id.asc())
@@ -830,6 +838,10 @@ class Query(ChangeTrackingMixin, TimestampMixin, BelongsToOrgMixin, db.Model):
 
         return query
 
+    @classmethod
+    def get_by_id(cls, _id):
+        return cls.query.filter(cls.id == _id).one()
+
     def fork(self, user):
         forked_list = ['org', 'data_source', 'latest_query_data', 'description',
                        'query_text', 'query_hash']
@@ -1497,4 +1509,3 @@ def init_db():
     #XXX remove after fixing User.group_ids
     db.session.commit()
     return default_org, admin_group, default_group
-

redash/query_runner/python.py

@@ -155,9 +155,9 @@ class Python(BaseQueryRunner):
             if type(data_source_name_or_id) == int:
                 data_source = models.DataSource.get_by_id(data_source_name_or_id)
             else:
-                data_source = models.DataSource.get(models.DataSource.name==data_source_name_or_id)
-        except models.DataSource.DoesNotExist:
-            raise Exception("Wrong data source name/id: %s." % data_source_name_or_id)
+                data_source = models.DataSource.get_by_name(data_source_name_or_id)
+        except models.NoResultFound:
+                raise Exception("Wrong data source name/id: %s." % data_source_name_or_id)
 
         # TODO: pass the user here...
         data, error = data_source.query_runner.run_query(query, None)
@@ -175,7 +175,7 @@ class Python(BaseQueryRunner):
         """
         try:
             query = models.Query.get_by_id(query_id)
-        except models.Query.DoesNotExist:
+        except models.NoResultFound:
             raise Exception("Query id %s does not exist." % query_id)
 
         if query.latest_query_data is None: