houdini/Houdini/Data/__init__.py

47 lines
1.4 KiB
Python

from gino import Gino
db = Gino()
class BaseCrumbsCollection(dict):
def __init__(self, *args, **kwargs):
super().__init__()
self._db = db
self._model = kwargs.get('model')
self._key = kwargs.get('key')
self._inventory_model = kwargs.get('inventory_model')
self._inventory_id = kwargs.get('inventory_id')
self._is_inventory = self._inventory_model is not None and self._inventory_id is not None
async def get(self, k):
if self._is_inventory:
return self[k]
try:
return self[k]
except KeyError as e:
result = await self._model.query.where(getattr(self._model, self._key) == k).gino.first()
if result:
self[k] = result
return result
raise e
async def __collect(self):
query = self._model.load(parent=self._inventory_model).where(
self._inventory_model.PenguinID == self._inventory_id
) if self._is_inventory else self._model.query
async with db.transaction():
collected = query.gino.iterate()
self.update(
{getattr(model_instance, self._key): model_instance async for model_instance in collected}
)
@classmethod
async def get_collection(cls, *args, **kwargs):
cc = cls(*args, **kwargs)
await cc.__collect()
return cc