Skip to content

Commit

Permalink
Made Relationshiploader utilize the new and improved DataLoader imple…
Browse files Browse the repository at this point in the history
…mentation housed inside graphene, if possible (graphene >=3.1.1)
  • Loading branch information
flipbit03 committed Sep 12, 2022
1 parent b3657b0 commit 5b2c075
Showing 1 changed file with 15 additions and 16 deletions.
31 changes: 15 additions & 16 deletions graphene_sqlalchemy/batching.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,15 +2,28 @@
from asyncio import get_event_loop
from typing import Any, Dict

import aiodataloader
import sqlalchemy
from sqlalchemy.orm import Session, strategies
from sqlalchemy.orm.query import QueryContext

from .utils import is_graphene_version_less_than, is_sqlalchemy_version_less_than


class RelationshipLoader(aiodataloader.DataLoader):
def get_data_loader_impl() -> Any: # pragma: no cover
"""Graphene >= 3.1.1 ships a copy of aiodataloader with minor fixes. To preserve backward-compatibility,
aiodataloader is used in conjunction with older versions of graphene"""
if is_graphene_version_less_than("3.1.1"):
from aiodataloader import DataLoader
else:
from graphene.utils.dataloader import DataLoader

return DataLoader


DataLoader = get_data_loader_impl()


class RelationshipLoader(DataLoader):
cache = False

def __init__(self, relationship_prop, selectin_loader):
Expand Down Expand Up @@ -92,20 +105,6 @@ async def batch_load_fn(self, parents):
] = {}


def get_data_loader_impl() -> Any: # pragma: no cover
"""Graphene >= 3.1.1 ships a copy of aiodataloader with minor fixes. To preserve backward-compatibility,
aiodataloader is used in conjunction with older versions of graphene"""
if is_graphene_version_less_than("3.1.1"):
from aiodataloader import DataLoader
else:
from graphene.utils.dataloader import DataLoader

return DataLoader


DataLoader = get_data_loader_impl()


def get_batch_resolver(relationship_prop):
"""Get the resolve function for the given relationship."""

Expand Down

0 comments on commit 5b2c075

Please sign in to comment.