|
2 | 2 | from asyncio import get_event_loop |
3 | 3 | from typing import Any, Dict |
4 | 4 |
|
5 | | -import aiodataloader |
6 | 5 | import sqlalchemy |
7 | 6 | from sqlalchemy.orm import Session, strategies |
8 | 7 | from sqlalchemy.orm.query import QueryContext |
9 | 8 |
|
10 | 9 | from .utils import is_graphene_version_less_than, is_sqlalchemy_version_less_than |
11 | 10 |
|
12 | 11 |
|
13 | | -class RelationshipLoader(aiodataloader.DataLoader): |
| 12 | +def get_data_loader_impl() -> Any: # pragma: no cover |
| 13 | + """Graphene >= 3.1.1 ships a copy of aiodataloader with minor fixes. To preserve backward-compatibility, |
| 14 | + aiodataloader is used in conjunction with older versions of graphene""" |
| 15 | + if is_graphene_version_less_than("3.1.1"): |
| 16 | + from aiodataloader import DataLoader |
| 17 | + else: |
| 18 | + from graphene.utils.dataloader import DataLoader |
| 19 | + |
| 20 | + return DataLoader |
| 21 | + |
| 22 | + |
| 23 | +DataLoader = get_data_loader_impl() |
| 24 | + |
| 25 | + |
| 26 | +class RelationshipLoader(DataLoader): |
14 | 27 | cache = False |
15 | 28 |
|
16 | 29 | def __init__(self, relationship_prop, selectin_loader): |
@@ -92,20 +105,6 @@ async def batch_load_fn(self, parents): |
92 | 105 | ] = {} |
93 | 106 |
|
94 | 107 |
|
95 | | -def get_data_loader_impl() -> Any: # pragma: no cover |
96 | | - """Graphene >= 3.1.1 ships a copy of aiodataloader with minor fixes. To preserve backward-compatibility, |
97 | | - aiodataloader is used in conjunction with older versions of graphene""" |
98 | | - if is_graphene_version_less_than("3.1.1"): |
99 | | - from aiodataloader import DataLoader |
100 | | - else: |
101 | | - from graphene.utils.dataloader import DataLoader |
102 | | - |
103 | | - return DataLoader |
104 | | - |
105 | | - |
106 | | -DataLoader = get_data_loader_impl() |
107 | | - |
108 | | - |
109 | 108 | def get_batch_resolver(relationship_prop): |
110 | 109 | """Get the resolve function for the given relationship.""" |
111 | 110 |
|
|
0 commit comments