APIFlask and MongoDB use of Marshmallow Schemas to read and write directly to database

app.py

from apiflask import APIFlask
from apiflask import Schema
from apiflask import fields
from bson import ObjectId
from pymongo import MongoClient

class DB():
    def __init__(self, app=None):
        if app is not None:
            self.init_app(app)

    def init_app(self, app):
        self.config = app.config

        self.client = MongoClient(self.config['MONGODB_URI'])
        self.db = self.client[self.config['MONGODB_DB']]

    def save_to_db(self, collection, data):
        return self.db[collection].insert_one(data)

    def load_from_db(self, collection, query):
        return self.db[collection].find_one(query)


db = DB()
app = APIFlask(__name__)
app.config['MONGODB_URI'] = "mongodb://root:example@localhost/"
app.config['MONGODB_DB'] = "test_database"
db.init_app(app)

class ObjectIdField(fields.Field):
    def _serialize(self, value, attr, obj):
        if value is None:
            return ''
        return str(value)

    def _deserialize(self, value, attr, data):
        try:
            return ObjectId(value)
        except (TypeError, ValueError):
            self.fail('invalid')


class UserSchemaOut(Schema):
    _id = ObjectIdField(attribute='_id')
    name = fields.String()
    email = fields.String()


class UserSchemaIn(Schema):
    name = fields.String()
    email = fields.String()


@app.get('/users/<string:user_id>')
@app.output(UserSchemaOut)
def get_user(user_id: str):
    user = db.load_from_db("users", {'_id': ObjectId(user_id)})
    if user:
        return user
    else:
        return {'message': 'User not found'}, 404

@app.post('/users')
@app.input(UserSchemaIn)
def create_user(json_data):
    id = str(db.save_to_db("users", json_data).inserted_id)
    return {'message': 'User created', 'id': id}, 201

if __name__ == '__main__':
    app.run(debug=True, port=5555)