flask "get_or_404" like function but with another status code

Question

What I know:

We all know that flask has a useful query.get_or_404 we can call it to any class object and return the object or raise a 404 error if the object is not found.

The problem:

I have a very large application and using that function for my queries, but it becomes now a bit confusing when I'm sending data to the front end and other application consuming my APIs. When an object they are querying is not found it returns 404 and the same behavior happens also when a page is not found.

What I want to achieve:

I would like to have a meaningful response that is different for every object that is not found, and different for the normal 404 error message.

Example :

if I have this query:

user = User.query.get_or_404(id)

if the user is not found I want to raise an HTTP error and return a message like user not found

What I have tried So far:

    user = User.query.get(id)
    if user:
        #do something
    else 
        return {'status':'01', 'description': 'user not found'} 
        # or raise a http error 

The problem with this approach I can't maintain the application I'm working on, it will require me to change everywhere I was using get_or_404 with that code.

What I'm thinking about :

Creating a function like query.get_or_404but with another status message like query.get_or_415 for example and add an error handler for 415 HTTP code so that if an object is not found it can return {'status':'0message:ge : 'object of the class is not found'}

How can I achieve it?

I have checked for that function in Flask but was unable to find it

Anyone with suggestions?

Solution 1

As noted by bereal, you should indeed use flask-SQLAlchemy's BaseQuery and add your own functionality from there. I'll check to see if I can add some messaging system to this approach as well.

import json
from flask import Flask, abort, Response
from flask_sqlalchemy import SQLAlchemy, BaseQuery

class CustomBaseQuery(BaseQuery):
    def get_or_415(self, ident):
        model_class_name = ''
        try:
            model_class_name = self._mapper_zero().class_.__name__
        except Exception as e:
            print(e)

        rv = self.get(ident)
        if not rv:
            error_message = json.dumps({'message': model_class_name + ' ' + str(ident) + ' not found'})
            abort(Response(error_message, 415))
        return rv

app = Flask(__name__)
db = SQLAlchemy(app, query_class=CustomBaseQuery)

class User(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    username = db.Column(db.String(80), unique=True, nullable=False)

db.create_all()
user = User(username='foo')
db.session.add(user)
db.session.commit()

@app.route('/')
def index():
    User.query.get_or_415(1)
    User.query.get_or_415(2)
    return ''

When going to the index, it returns:

{"message": "User 2 not found"}

Solution 2

Based on this answer

You can write like this:

from sqlalchemy.orm import Query

class MyQuery(Query):

  def get_or_415(self, pk):
    instance = self.get(pk)
    if not instance:
      raise HttpException(code=415)
    return instance

Then you can use flask error handling to handle your exception and response the way you want