Skip to content

Commit

Permalink
Fix: Count Performance
Browse files Browse the repository at this point in the history
  • Loading branch information
arunsureshkumar committed Apr 12, 2023
1 parent f171ddd commit f4a2037
Show file tree
Hide file tree
Showing 4 changed files with 74 additions and 25 deletions.
47 changes: 35 additions & 12 deletions graphene_mongo/fields.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,8 @@
from graphene.types.utils import get_type
from graphene.utils.str_converters import to_snake_case
from graphql import GraphQLResolveInfo
from graphql_relay import from_global_id, cursor_to_offset
from graphql_relay import from_global_id
from graphql_relay.connection.arrayconnection import cursor_to_offset
from mongoengine import QuerySet
from mongoengine.base import get_document
from promise import Promise
Expand Down Expand Up @@ -355,28 +356,44 @@ def default_resolver(self, _root, info, required_fields=None, resolved=None, **a
if before:
before = cursor_to_offset(before)

queryset = None

if resolved is not None:
items = resolved

if isinstance(items, QuerySet):
queryset = items.clone()
try:
count = items.count(with_limit_and_skip=True)
if last is not None and after is not None:
count = items.count(with_limit_and_skip=False)
else:
count = None
except OperationFailure:
count = len(items)
else:
queryset = None
count = len(items)

skip, limit, reverse = find_skip_and_limit(first=first, last=last, after=after, before=before,
count=count)

if limit:
if reverse:
items = items[::-1][skip:skip + limit]
else:
items = items[skip:skip + limit]
elif skip:
items = items[skip:]
iterables = items
if isinstance(items, QuerySet):
if limit:
if reverse:
items = items.order_by("-pk").skip(skip).limit(limit)
else:
items = items.skip(skip).limit(limit)
elif skip:
items = items.skip(skip)
else:
if limit:
if reverse:
items = items[::-1][skip:skip + limit]
else:
items = items[skip:skip + limit]
elif skip:
items = items[skip:]
iterables = list(items)
list_length = len(iterables)

elif callable(getattr(self.model, "objects", None)):
Expand Down Expand Up @@ -454,7 +471,13 @@ def default_resolver(self, _root, info, required_fields=None, resolved=None, **a
iterables = items
list_length = len(iterables)

has_next_page = True if (0 if limit is None else limit) + (0 if skip is None else skip) < count else False
if count:
has_next_page = True if (0 if limit is None else limit) + (0 if skip is None else skip) < count else False
else:
if queryset:
has_next_page = bool(queryset(pk__gt=iterables[-1].pk).limit(1).first())
else:
has_next_page = False
has_previous_page = True if skip else False
if reverse:
iterables = list(iterables)
Expand Down Expand Up @@ -570,7 +593,7 @@ def connection_resolver(cls, resolver, connection_type, root, info, **args):

return on_resolve(iterable)

def wrap_resolve(self, parent_resolver):
def get_resolver(self, parent_resolver):
super_resolver = self.resolver or parent_resolver
resolver = partial(
self.chained_resolver, super_resolver, isinstance(super_resolver, partial)
Expand Down
44 changes: 34 additions & 10 deletions graphene_mongo/fields_async.py
Original file line number Diff line number Diff line change
Expand Up @@ -92,13 +92,18 @@ async def default_resolver(self, _root, info, required_fields=None, resolved=Non
if before:
before = cursor_to_offset(before)

queryset = None

if resolved is not None:
items = resolved

if isinstance(items, QuerySet):
try:
count = await sync_to_async(items.count, thread_sensitive=False,
executor=ThreadPoolExecutor())(with_limit_and_skip=True)
if last is not None and after is not None:
count = await sync_to_async(items.count, thread_sensitive=False,
executor=ThreadPoolExecutor())(with_limit_and_skip=False)
else:
count = None
except OperationFailure:
count = len(items)
else:
Expand All @@ -107,13 +112,24 @@ async def default_resolver(self, _root, info, required_fields=None, resolved=Non
skip, limit, reverse = find_skip_and_limit(first=first, last=last, after=after, before=before,
count=count)

if limit:
if reverse:
items = items[::-1][skip:skip + limit]
else:
items = items[skip:skip + limit]
elif skip:
items = items[skip:]
if isinstance(items, QuerySet):
queryset = items.clone()
if limit:
if reverse:
items = items.order_by("-pk").skip(skip).limit(limit)
else:
items = items.skip(skip).limit(limit)
elif skip:
items = items.skip(skip)
else:
queryset = None
if limit:
if reverse:
items = items[::-1][skip:skip + limit]
else:
items = items[skip:skip + limit]
elif skip:
items = items[skip:]
iterables = await sync_to_async(list, thread_sensitive=False, executor=ThreadPoolExecutor())(items)
list_length = len(iterables)

Expand Down Expand Up @@ -198,7 +214,15 @@ async def default_resolver(self, _root, info, required_fields=None, resolved=Non
iterables)
list_length = len(iterables)

has_next_page = True if (0 if limit is None else limit) + (0 if skip is None else skip) < count else False
if count:
has_next_page = True if (0 if limit is None else limit) + (0 if skip is None else skip) < count else False
else:
if queryset:
has_next_page = bool(await sync_to_async(queryset(pk__gt=iterables[-1].pk).limit(1).first,
thread_sensitive=False,
executor=ThreadPoolExecutor())())
else:
has_next_page = False
has_previous_page = True if skip else False
if reverse:
iterables = await sync_to_async(list, thread_sensitive=False, executor=ThreadPoolExecutor())(
Expand Down
4 changes: 2 additions & 2 deletions graphene_mongo/tests/test_fields.py
Original file line number Diff line number Diff line change
Expand Up @@ -48,15 +48,15 @@ def test_field_args_with_unconverted_field():


@pytest.mark.asyncio
async def test_default_resolver_with_colliding_objects_field():
async def test_default_resolver_with_colliding_objects_field_async():
field = AsyncMongoengineConnectionField(nodes_async.ErroneousModelAsyncNode)

connection = await field.default_resolver(None, {})
assert 0 == len(connection.iterable)


@pytest.mark.asyncio
async def test_default_resolver_connection_list_length(fixtures):
async def test_default_resolver_connection_list_length_async(fixtures):
field = AsyncMongoengineConnectionField(nodes_async.ArticleAsyncNode)

connection = await field.default_resolver(None, {}, **{"first": 1})
Expand Down
4 changes: 3 additions & 1 deletion graphene_mongo/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -194,7 +194,7 @@ def ast_to_dict(node, include_loc=False):
return node


def find_skip_and_limit(first, last, after, before, count):
def find_skip_and_limit(first, last, after, before, count=None):
reverse = False
skip = 0
limit = None
Expand All @@ -217,6 +217,8 @@ def find_skip_and_limit(first, last, after, before, count):
limit = last
skip = before - last
elif last is not None and after is not None:
if not count:
raise ValueError("Count Missing")
reverse = True
if last + after < count:
limit = last
Expand Down

0 comments on commit f4a2037

Please sign in to comment.