Skip to content

Instantly share code, notes, and snippets.

@lanbugs
Created April 8, 2024 19:59
Show Gist options
  • Save lanbugs/f0f777a7ff32da00f1dec48825614088 to your computer and use it in GitHub Desktop.
Save lanbugs/f0f777a7ff32da00f1dec48825614088 to your computer and use it in GitHub Desktop.
APIFlask and MongoDB use of Marshmallow Schemas to read and write directly to database
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)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment