pr comments

getredash · Apr 29, 2024 · f30a674 · f30a674
1 parent 92070cc
commit f30a674
Show file tree

Hide file tree

Showing 9 changed files with 55 additions and 55 deletions.
diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -98,7 +98,7 @@ cassandra-driver = "3.21.0"
 certifi = ">=2019.9.11"
 cmem-cmempy = "21.2.3"
 databend-driver = "0.17.1"
-databend-sqlalchemy = "0.4.5"
+databend-sqlalchemy = "0.4.6"
 google-api-python-client = "1.7.11"
 gspread = "5.11.2"
 impyla = "0.16.0"

diff --git a/redash/handlers/organization.py b/redash/handlers/organization.py
@@ -1,6 +1,5 @@
 from flask_login import current_user, login_required
 from sqlalchemy import distinct, func
-from sqlalchemy.sql.expression import select
 
 from redash import models
 from redash.authentication import current_org
@@ -32,8 +31,11 @@ def organization_status(org_slug=None):
             )
         ),
         "dashboards": models.db.session.scalar(
-            select(func.count(models.Dashboard.id)).where(
-                models.Dashboard.org == current_org, models.Dashboard.is_archived.is_(False)
+            models.Dashboard.all(
+                current_org,
+                [],
+                None,
+                columns=[func.count(distinct(models.Dashboard.id))],
             )
         ),
     }

diff --git a/redash/models/__init__.py b/redash/models/__init__.py
@@ -1074,10 +1074,19 @@ def name_as_slug(self):
 
     @classmethod
     def all(cls, org, group_ids, user_id, columns=None, distinct=None):
+        conditions = [
+            cls.is_archived.is_(False),
+            cls.org == org,
+        ]
         if columns is None:
             columns = [cls, User.id, User.name, User.details, User.email]
         if distinct is None:
             distinct = [func.lower(cls.name), cls.created_at, cls.slug]
+        if len(group_ids) > 0 or user_id is not None:
+            conditions = conditions + [
+                or_(DataSourceGroup.group_id.in_(group_ids), cls.user_id == user_id),
+                or_(cls.user_id == user_id, cls.is_draft.is_(False)),
+            ]
         query = (
             select(*columns)
             .join(User)
@@ -1086,12 +1095,7 @@ def all(cls, org, group_ids, user_id, columns=None, distinct=None):
             .outerjoin(Visualization)
             .outerjoin(Query)
             .outerjoin(DataSourceGroup, Query.data_source_id == DataSourceGroup.data_source_id)
-            .where(
-                cls.is_archived.is_(False),
-                cls.org == org,
-                or_(DataSourceGroup.group_id.in_(group_ids), cls.user_id == user_id),
-                or_(cls.user_id == user_id, cls.is_draft.is_(False)),
-            )
+            .where(*conditions)
         )
 
         return query

diff --git a/redash/models/users.py b/redash/models/users.py
@@ -358,14 +358,8 @@ def grant(cls, obj, access_type, grantee, grantor):
 
     @classmethod
     def revoke(cls, obj, grantee, access_type=None):
-        q = delete(cls).where(cls.object_id == obj.id, cls.object_type == obj.__tablename__)
-
-        if access_type:
-            q = q.where(AccessPermission.access_type == access_type)
-
-        if grantee:
-            q = q.where(AccessPermission.grantee == grantee)
-
+        conditions = AccessPermission._query_condition(obj, access_type, grantee, None)
+        q = delete(cls).where(*conditions)
         return db.session.execute(q).rowcount
 
     @classmethod
@@ -377,18 +371,20 @@ def exists(cls, obj, access_type, grantee):
         return len(cls.find(obj, access_type, grantee)) > 0
 
     @classmethod
-    def _query(cls, obj, access_type=None, grantee=None, grantor=None):
-        q = select(cls).where(cls.object_id == obj.id, cls.object_type == obj.__tablename__)
-
+    def _query_condition(cls, obj, access_type=None, grantee=None, grantor=None):
+        conditions = [cls.object_id == obj.id, cls.object_type == obj.__tablename__]
         if access_type:
-            q = q.where(AccessPermission.access_type == access_type)
-
+            conditions.append(AccessPermission.access_type == access_type)
         if grantee:
-            q = q.where(AccessPermission.grantee == grantee)
-
+            conditions.append(AccessPermission.grantee == grantee)
         if grantor:
-            q = q.where(AccessPermission.grantor == grantor)
+            conditions.append(AccessPermission.grantor == grantor)
+        return conditions
 
+    @classmethod
+    def _query(cls, obj, access_type=None, grantee=None, grantor=None):
+        conditions = AccessPermission._query_condition(obj, access_type, grantee, grantor)
+        q = select(cls).where(*conditions)
         return db.session.scalars(q).all()
 
     def to_dict(self):

diff --git a/redash/query_runner/databend.py b/redash/query_runner/databend.py
@@ -80,11 +80,7 @@ def run_query(self, query, user):
         try:
             cursor = connection.execute(text(query)).cursor
             columns = self.fetch_columns([(i[0], self._define_column_type(i[1])) for i in cursor.description])
-            rows = []
-            for row in cursor:
-                if row is None:
-                    break
-                rows.append(dict(zip((column["name"] for column in columns), row)))
+            rows = [dict(zip((column["name"] for column in columns), row)) for row in cursor]
 
             data = {"columns": columns, "rows": rows}
             error = None

diff --git a/tests/models/test_data_sources.py b/tests/models/test_data_sources.py
@@ -1,5 +1,6 @@
 import mock
 from mock import patch
+from sqlalchemy import func
 from sqlalchemy.sql.expression import select
 
 from redash.models import DataSource, Query, QueryResult, db
@@ -166,7 +167,7 @@ def test_deletes_child_models(self):
         data_source.delete()
         self.assertIsNone(db.session.get(DataSource, data_source.id))
         self.assertEqual(
-            0, len(db.session.scalars(select(QueryResult).where(QueryResult.data_source == data_source)).all())
+            0, db.session.scalar(select(func.count(QueryResult.id)).where(QueryResult.data_source == data_source))
         )
 
     @patch("redash.redis_connection.delete")

diff --git a/tests/test_cli.py b/tests/test_cli.py
@@ -2,6 +2,7 @@
 
 import mock
 from click.testing import CliRunner
+from sqlalchemy import func
 from sqlalchemy.sql.expression import select
 
 from redash.cli import manager
@@ -22,7 +23,7 @@ def test_interactive_new(self):
         )
         self.assertFalse(result.exception)
         self.assertEqual(result.exit_code, 0)
-        self.assertEqual(len(db.session.scalars(select(DataSource)).all()), 1)
+        self.assertEqual(db.session.scalar(select(func.count(DataSource.id))), 1)
         ds = db.session.scalar(select(DataSource))
         self.assertEqual(ds.name, "test")
         self.assertEqual(ds.type, "pg")
@@ -44,7 +45,7 @@ def test_options_new(self):
         )
         self.assertFalse(result.exception)
         self.assertEqual(result.exit_code, 0)
-        self.assertEqual(len(db.session.scalars(select(DataSource)).all()), 1)
+        self.assertEqual(db.session.scalar(select(func.count(DataSource.id))), 1)
         ds = db.session.scalar(select(DataSource))
         self.assertEqual(ds.name, "test")
         self.assertEqual(ds.type, "pg")
@@ -57,7 +58,7 @@ def test_bad_type_new(self):
         self.assertTrue(result.exception)
         self.assertEqual(result.exit_code, 1)
         self.assertIn("not supported", result.output)
-        self.assertEqual(len(db.session.scalars(select(DataSource)).all()), 0)
+        self.assertEqual(db.session.scalar(select(func.count(DataSource.id))), 0)
 
     def test_bad_options_new(self):
         runner = CliRunner()
@@ -76,7 +77,7 @@ def test_bad_options_new(self):
         self.assertTrue(result.exception)
         self.assertEqual(result.exit_code, 1)
         self.assertIn("invalid configuration", result.output)
-        self.assertEqual(len(db.session.scalars(select(DataSource)).all()), 0)
+        self.assertEqual(db.session.scalar(select(func.count(DataSource.id))), 0)
 
     def test_list(self):
         self.factory.create_data_source(
@@ -152,7 +153,7 @@ def test_connection_delete(self):
         self.assertFalse(result.exception)
         self.assertEqual(result.exit_code, 0)
         self.assertIn("Deleting", result.output)
-        self.assertEqual(len(db.session.scalars(select(DataSource)).all()), 0)
+        self.assertEqual(db.session.scalar(select(func.count(DataSource.id))), 0)
 
     def test_connection_bad_delete(self):
         self.factory.create_data_source(
@@ -165,7 +166,7 @@ def test_connection_bad_delete(self):
         self.assertTrue(result.exception)
         self.assertEqual(result.exit_code, 1)
         self.assertIn("Couldn't find", result.output)
-        self.assertEqual(len(db.session.scalars(select(DataSource)).all()), 1)
+        self.assertEqual(db.session.scalar(select(func.count(DataSource.id))), 1)
 
     def test_options_edit(self):
         self.factory.create_data_source(
@@ -190,7 +191,7 @@ def test_options_edit(self):
         )
         self.assertFalse(result.exception)
         self.assertEqual(result.exit_code, 0)
-        self.assertEqual(len(db.session.scalars(select(DataSource)).all()), 1)
+        self.assertEqual(db.session.scalar(select(func.count(DataSource.id))), 1)
         ds = db.session.scalar(select(DataSource))
         self.assertEqual(ds.name, "test2")
         self.assertEqual(ds.type, "pg")
@@ -240,13 +241,13 @@ def test_bad_options_edit(self):
 
 class GroupCommandTests(BaseTestCase):
     def test_create(self):
-        gcount = len(db.session.scalars(select(Group)).all())
+        gcount = db.session.scalar(select(func.count(Group.id)))
         perms = ["create_query", "edit_query", "view_query"]
         runner = CliRunner()
         result = runner.invoke(manager, ["groups", "create", "test", "--permissions", ",".join(perms)])
         self.assertFalse(result.exception)
         self.assertEqual(result.exit_code, 0)
-        self.assertEqual(len(db.session.scalars(select(Group)).all()), gcount + 1)
+        self.assertEqual(db.session.scalar(select(func.count(Group.id))), gcount + 1)
         g = db.session.scalar(select(Group).order_by(Group.id.desc()))
         db.session.add(self.factory.org)
         self.assertEqual(g.org_id, self.factory.org.id)
@@ -360,7 +361,7 @@ def test_create(self):
         self.assertFalse(result.exception)
         self.assertEqual(result.exit_code, 0)
 
-        ucount = len(db.session.scalars(select(Organization)).all())
+        ucount = db.session.scalar(select(func.count(Organization.id)))
 
         self.assertEqual(ucount, 2)
 
@@ -455,20 +456,20 @@ def test_create_bad(self):
 
     def test_delete(self):
         self.factory.create_user(email="foobar@example.com")
-        ucount = len(db.session.scalars(select(User)).all())
+        ucount = db.session.scalar(select(func.count(User.id)))
         runner = CliRunner()
         result = runner.invoke(manager, ["users", "delete", "foobar@example.com"])
         self.assertFalse(result.exception)
         self.assertEqual(result.exit_code, 0)
-        self.assertEqual(len(db.session.scalars(select(User).where(User.email == "foobar@example.com")).all()), 0)
-        self.assertEqual(len(db.session.scalars(select(User)).all()), ucount - 1)
+        self.assertEqual(db.session.scalar(select(func.count(User.id)).where(User.email == "foobar@example.com")), 0)
+        self.assertEqual(db.session.scalar(select(func.count(User.id))), ucount - 1)
 
     def test_delete_bad(self):
-        ucount = len(db.session.scalars(select(User)).all())
+        ucount = db.session.scalar(select(func.count(User.id)))
         runner = CliRunner()
         result = runner.invoke(manager, ["users", "delete", "foobar@example.com"])
         self.assertIn("Deleted 0 users", result.output)
-        self.assertEqual(len(db.session.scalars(select(User)).all()), ucount)
+        self.assertEqual(db.session.scalar(select(func.count(User.id))), ucount)
 
     def test_password(self):
         self.factory.create_user(email="foobar@example.com")

diff --git a/tests/test_handlers.py b/tests/test_handlers.py
@@ -1,6 +1,7 @@
 from flask_login import current_user
 from funcy import project
 from mock import patch
+from sqlalchemy import func
 from sqlalchemy.sql.expression import select
 
 from redash import models, settings
@@ -320,4 +321,4 @@ def test_delete(self):
         models.db.session.add(qs)
         models.db.session.commit()
         self.make_request("delete", "/api/query_snippets/1", user=self.factory.user)
-        self.assertEqual(len(models.db.session.scalars(select(models.QuerySnippet)).all()), 0)
+        self.assertEqual(models.db.session.scalar(select(func.count(models.QuerySnippet.id))), 0)