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 pymongo
from bson import ObjectId
from bson.errors import InvalidId
from flask import flash
from jinja2 import contextfunction
......@@ -241,6 +243,13 @@ class ModelView(BaseModelView):
return count, results
def _get_valid_id(self, id):
try:
return ObjectId(id)
except InvalidId:
return id
def get_one(self, id):
"""
Return single model instance by ID
......@@ -248,8 +257,7 @@ class ModelView(BaseModelView):
:param id:
Model ID
"""
# TODO: Validate if it is valid ID
return self.coll.find_one({'_id': id})
return self.coll.find_one({'_id': self._get_valid_id(id)})
def edit_form(self, obj):
"""
......@@ -336,7 +344,7 @@ class ModelView(BaseModelView):
# TODO: Optimize me
for pk in ids:
self.coll.remove({'_id': pk})
self.coll.remove({'_id': self._get_valid_id(pk)})
count += 1
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