Improve crumb collection support for inventories

This commit is contained in:
Ben 2019-06-07 00:38:31 +01:00
parent 1c4ee8d2ff
commit d1d2b004c6

View File

@ -28,7 +28,7 @@ class BaseCrumbsCollection(dict):
try: try:
return self[k] return self[k]
except KeyError as e: except KeyError as e:
query = self._model.load(parent=self._inventory_model).where( query = self._inventory_model.load(parent=self._model).where(
(self._inventory_key_column == self._inventory_id) & (self._inventory_value_column == k) (self._inventory_key_column == self._inventory_id) & (self._inventory_value_column == k)
) if self._is_inventory else self._model.query.where(self._model_key_column == k) ) if self._is_inventory else self._model.query.where(self._model_key_column == k)
result = await query.gino.first() result = await query.gino.first()
@ -38,25 +38,36 @@ class BaseCrumbsCollection(dict):
raise e raise e
async def set(self, k=None, **kwargs): async def set(self, k=None, **kwargs):
if self._is_inventory and k: if self._is_inventory:
kwargs = {self._inventory_key: self._inventory_id, self._inventory_value: k} kwargs = {self._inventory_key: self._inventory_id, self._inventory_value: k}
self[k] = await self._inventory_model.create(**kwargs) model_instance = await self._inventory_model.create(**kwargs)
k = getattr(model_instance, self._inventory_value)
self[k] = model_instance
else: else:
model_instance = await self._model.create(**kwargs) model_instance = await self._model.create(**kwargs)
k = getattr(model_instance, self._key) k = getattr(model_instance, self._key)
self[k] = model_instance self[k] = model_instance
return self[k] return self[k]
async def delete(self, k):
query = self._inventory_model.delete.where(
(self._inventory_key_column == self._inventory_id) & (self._inventory_value_column == k)
) if self._is_inventory else self._model.delete.where(self._model_key_column == k)
await query.gino.status()
if k in self:
del self[k]
async def __collect(self): async def __collect(self):
query = self._model.load(parent=self._inventory_model).where( query = self._inventory_model.load(parent=self._model).where(
self._inventory_key_column == self._inventory_id self._inventory_key_column == self._inventory_id
) if self._is_inventory else self._model.query ) if self._is_inventory else self._model.query
async with db.transaction(): async with db.transaction():
collected = query.gino.iterate() collected = query.gino.iterate()
self.update( self.update({
{getattr(model_instance, self._key): model_instance async for model_instance in collected} getattr(model_instance, self._inventory_value if self._is_inventory else self._key): model_instance
) async for model_instance in collected
})
@classmethod @classmethod
async def get_collection(cls, *args, **kwargs): async def get_collection(cls, *args, **kwargs):