From ab0a448c8799ab7a90b557b54b64ac37dd813399 Mon Sep 17 00:00:00 2001 From: Vladislav Denisov Date: Tue, 10 Jan 2017 15:08:36 +0300 Subject: [PATCH] Add get_by_id & get_by_name methods for Query and DataSource classes --- redash/models.py | 13 ++++++++++++- redash/query_runner/python.py | 8 ++++---- 2 files changed, 16 insertions(+), 5 deletions(-) diff --git a/redash/models.py b/redash/models.py index f45d204650..0858cd9e67 100644 --- a/redash/models.py +++ b/redash/models.py @@ -514,6 +514,14 @@ def update_group_permission(self, group, view_only): 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 @@ def recent(cls, group_ids, user_id=None, limit=20): 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 - diff --git a/redash/query_runner/python.py b/redash/query_runner/python.py index 6e89734e3e..cc8848e93a 100644 --- a/redash/query_runner/python.py +++ b/redash/query_runner/python.py @@ -155,9 +155,9 @@ def execute_query(self, data_source_name_or_id, query): 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 @@ def get_query_result(self, query_id): """ 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: