Commit 085c6e77 authored by Mostafa Rokooie's avatar Mostafa Rokooie

Fix ObjectId casting bug in get_one and delete_action

parent d6cd59d0
import logging import logging
import pymongo import pymongo
from bson import ObjectId
from bson.errors import InvalidId
from flask import flash from flask import flash
from jinja2 import contextfunction from jinja2 import contextfunction
...@@ -241,6 +243,13 @@ class ModelView(BaseModelView): ...@@ -241,6 +243,13 @@ class ModelView(BaseModelView):
return count, results return count, results
def _get_valid_id(self, id):
try:
return ObjectId(id)
except InvalidId:
return id
def get_one(self, id): def get_one(self, id):
""" """
Return single model instance by ID Return single model instance by ID
...@@ -248,8 +257,7 @@ class ModelView(BaseModelView): ...@@ -248,8 +257,7 @@ class ModelView(BaseModelView):
:param id: :param id:
Model ID Model ID
""" """
# TODO: Validate if it is valid ID return self.coll.find_one({'_id': self._get_valid_id(id)})
return self.coll.find_one({'_id': id})
def edit_form(self, obj): def edit_form(self, obj):
""" """
...@@ -336,7 +344,7 @@ class ModelView(BaseModelView): ...@@ -336,7 +344,7 @@ class ModelView(BaseModelView):
# TODO: Optimize me # TODO: Optimize me
for pk in ids: for pk in ids:
self.coll.remove({'_id': pk}) self.coll.remove({'_id': self._get_valid_id(pk)})
count += 1 count += 1
flash(ngettext('Model was successfully deleted.', flash(ngettext('Model was successfully deleted.',
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment