From 92a771f5ed301e491774714a2183f43c15c55a97 Mon Sep 17 00:00:00 2001 From: Curtis Castrapel Date: Tue, 13 Nov 2018 09:14:21 -0800 Subject: [PATCH] More accurate db count functionality --- lemur/database.py | 30 +++++++++++++++++++++++++++--- 1 file changed, 27 insertions(+), 3 deletions(-) diff --git a/lemur/database.py b/lemur/database.py index 5cfe340a..82fb0423 100644 --- a/lemur/database.py +++ b/lemur/database.py @@ -10,8 +10,8 @@ .. moduleauthor:: Kevin Glisson """ from inflection import underscore -from sqlalchemy import exc, func -from sqlalchemy.orm import make_transient +from sqlalchemy import exc, func, distinct +from sqlalchemy.orm import make_transient, lazyload from sqlalchemy.sql import and_, or_ from lemur.exceptions import AttrNotFound, DuplicateError @@ -273,7 +273,31 @@ def get_count(q): :param q: :return: """ - count_q = q.statement.with_only_columns([func.count()]).group_by(None).order_by(None) + disable_group_by = False + if len(q._entities) > 1: + # currently support only one entity + raise Exception('only one entity is supported for get_count, got: %s' % q) + entity = q._entities[0] + if hasattr(entity, 'column'): + # _ColumnEntity has column attr - on case: query(Model.column)... + col = entity.column + if q._group_by and q._distinct: + # which query can have both? + raise NotImplementedError + if q._group_by or q._distinct: + col = distinct(col) + if q._group_by: + # need to disable group_by and enable distinct - we can do this because we have only 1 entity + disable_group_by = True + count_func = func.count(col) + else: + # _MapperEntity doesn't have column attr - on case: query(Model)... + count_func = func.count() + if q._group_by and not disable_group_by: + count_func = count_func.over(None) + count_q = q.options(lazyload('*')).statement.with_only_columns([count_func]).order_by(None) + if disable_group_by: + count_q = count_q.group_by(None) count = q.session.execute(count_q).scalar() return count