3
3
import asyncio
4
4
import inspect
5
5
import warnings
6
+ import weakref
6
7
from contextvars import ContextVar
7
- from typing import Optional
8
+ from typing import Optional , Callable , Any
8
9
9
10
from sqlalchemy import text
10
11
from sqlalchemy .engine import Engine
23
24
from sqlalchemy .sql import ClauseElement , WARN_LINTING
24
25
from sqlalchemy .util .concurrency import greenlet_spawn
25
26
27
+ from .loader import Loader , LoaderResult
28
+
26
29
27
30
async def create_engine (
28
31
url , * arg , isolation_level = None , min_size = 1 , max_size = None , ** kw
@@ -212,6 +215,18 @@ async def _execute_sa10(self, object_, multiparams, params):
212
215
return await asyncio .wait_for (coro , timeout )
213
216
return await coro
214
217
218
+ def _load_result (self , result ):
219
+ options = result .context .execution_options
220
+ loader = options .get ("loader" )
221
+ model = options .get ("model" )
222
+ if loader is None and model is not None :
223
+ if isinstance (model , weakref .ref ):
224
+ model = model ()
225
+ loader = Loader .get (model )
226
+ if loader is not None and options .get ("return_model" , True ):
227
+ result = LoaderResult (result , loader )
228
+ return result
229
+
215
230
async def _execute (self , object_ , params_20style ):
216
231
if isinstance (object_ , str ):
217
232
return await self .exec_driver_sql (object_ , params_20style )
@@ -281,6 +296,7 @@ async def all(self, clause, *multiparams, **params):
281
296
282
297
"""
283
298
result = await self ._execute_sa10 (clause , multiparams , params )
299
+ result = self ._load_result (result )
284
300
return result .all ()
285
301
286
302
async def first (self , clause , * multiparams , ** params ):
@@ -293,6 +309,7 @@ async def first(self, clause, *multiparams, **params):
293
309
294
310
"""
295
311
result = await self ._execute_sa10 (clause , multiparams , params )
312
+ result = self ._load_result (result )
296
313
try :
297
314
return result .first ()
298
315
except ResourceClosedError as e :
@@ -313,6 +330,7 @@ async def one_or_none(self, clause, *multiparams, **params):
313
330
314
331
"""
315
332
result = await self ._execute_sa10 (clause , multiparams , params )
333
+ result = self ._load_result (result )
316
334
return result .one_or_none ()
317
335
318
336
async def one (self , clause , * multiparams , ** params ):
@@ -328,6 +346,7 @@ async def one(self, clause, *multiparams, **params):
328
346
329
347
"""
330
348
result = await self ._execute_sa10 (clause , multiparams , params )
349
+ result = self ._load_result (result )
331
350
return result .one ()
332
351
333
352
async def scalar (self , clause , * multiparams , ** params ):
@@ -352,7 +371,7 @@ async def status(self, clause, *multiparams, **params):
352
371
353
372
"""
354
373
result = await self ._execute_sa10 (clause , multiparams , params )
355
- return f"SELECT { result .rowcount } " , result . all ()
374
+ return result .context
356
375
357
376
class _IterateResult (StartableContext ):
358
377
def __init__ (self , conn , * args ):
@@ -633,6 +652,10 @@ async def status(self, clause, *multiparams, **params):
633
652
async with self .acquire (reuse = True ) as conn :
634
653
return await conn .status (clause , * multiparams , ** params )
635
654
655
+ async def run_sync (self , fn : Callable , * arg , ** kw ) -> Any :
656
+ async with self .acquire (reuse = True ) as conn :
657
+ return await conn .run_sync (fn , * arg , ** kw )
658
+
636
659
class _CompileConnection :
637
660
def __init__ (self , dialect ):
638
661
self .dialect = dialect
@@ -737,7 +760,7 @@ def repr(self, color=False):
737
760
return repr (self )
738
761
739
762
def __repr__ (self ):
740
- return f' { self .__class__ .__name__ } <{ self .sync_engine .pool .status ()} >'
763
+ return f" { self .__class__ .__name__ } <{ self .sync_engine .pool .status ()} >"
741
764
742
765
743
766
class AsyncOptionEngine (OptionEngineMixin , GinoEngine ):
0 commit comments