Flask login with Flask-JWT-Extended and SQLAlchemy and MongoDb token storages

Aug 25, 2022 – 8 min read

Intro

In this article I will speak about how to implement user authorization workflow with Flask-JWT-Extended that includes registering a new user, login, logout, refresh token, revoke a list of tokens as well as accessing pages that require authorization.

• Flask-JWT-Extended
• flask_sqlalchemy OR pymongo: Used to store the sessions and the user. Using a storage for storing the session is not required by Flask-JWT-Extended but we’ll use it in order to provide additional features – get the current user tokens and revoke a list of tokens.

How does Flask-JWT-Extended work?

Once the user is logged in, each subsequent request will include the JWT, allowing the user to access routes, services, and resources that are permitted with that token.

Local storage

We’ll use SQLAlchemy and MongoDb to store the uses and the user’s sessions.

If you want to store the users and the sessions in MongoDb, please install Mongo Compass on your local machine to access the database.

Code and folder structure

The Github repository of the project is located here. Download it on your machine and let’s get started.

The structure of the project is the following:

auth/
    models/
        mongodbModels.py 
        sqlaModels.py
    appAuth.py
    validators.py
.env.example
.gitignore
app.py
appExample.py
install.py
README.md
requirements.txt

The responsibilities of the files/folders is the following:

• auth – contains all files responsible for registering a new user, login, logout, refresh token and revoke a list of tokens. That includes API endpoints, validation of data received from the user, SQLAlchemy and MongoDb models.
• .env – environment variables required by the application
• app.py – setting variables required by the application, registering blueprints (API endpoints), initialising services and running the Flask app.
• appExample.py – contains two API endpoints. One that requires user authorization and one that does not.
• install.py – we run that script only once to build SQlAlchemy database.
• requirements.txt – contains python packages required by the application

API endpoints

The application supports the following endpoints:

POST '/v1/auth/register'
    --header 'Content-Type: application/json'
    --data-raw '{
        'email': 'EMAIL_ADDRESS',
        'password': 'PASSWORD'
    }'

POST '/v1/auth/login'
    --header 'Content-Type: application/json'
    --data-raw '{
        'email_or_username': 'EMAIL_OR_USERNAME',
        'password': 'PASSWORD'
    }'

POST '/v1/auth/logout'
    --header 'Authorization: Bearer TOKEN'

POST '/v1/auth/refresh'
    --header 'Authorization: Bearer TOKEN'

DELETE '/v1/auth/revoke-tokens'
    --header 'Authorization: Bearer TOKEN'
    --header 'Content-Type: application/json'
    --data-raw '{
        'tokens': [
            'TOKEN 1',
            'TOKEN 2',
            ...,
            'TOKEN N'
        ]
    }'

GET '/'

GET '/profile'
    --header 'Authorization: Bearer TOKEN'

Code review

Let’s review the code and give detailed explanation of what it does.

install.py

from app import app
from auth.sqlaModels import db

with app.app_context():
    db.create_all()

This code loads the Flask app and SQLAlchemy model and builds the SQLAlchemy database – tables user and sessions. After the build is completed you will find database.db file in the main folder of the project.

app.py

import os
from dotenv import load_dotenv
from flask import Flask, jsonify

Load libraries required to initialise a Flask application, load and access environment variables from .env.

from appExample import appExample

from auth.appAuth import appAuth, jwt

# SQLAlchemy
from auth.sqlaModels import db, bcrypt
# END SQLAlchemy

# MongoDb
from auth.mongodbModels import bcrypt
# END MongoDb

load_dotenv()

app = Flask(__name__)
app.config['JWT_SECRET_KEY'] = os.environ.get('SECRET_KEY')
app.config['JWT_COOKIE_SECURE'] = os.environ.get('COOKIE_SECURE')
app.config['JWT_ACCESS_TOKEN_EXPIRES'] = int(os.environ.get('ACCESS_TOKEN_LIFETIME'))
app.config['JWT_REFRESH_TOKEN_EXPIRES'] = int(os.environ.get('ACCESS_TOKEN_LIFETIME'))

• JWT_SECRET_KEY. A secret key required by JWT to encode the token sent to the Front-end.
• JWT_COOKIE_SECURE. A configuration telling JWT to accept requests via SSL or not.
• JWT_ACCESS_TOKEN_EXPIRES. Expiration time of the access token – in seconds.
• JWT_REFRESH_TOKEN_EXPIRES. Expiration time of the refresh token – in seconds.

app.register_blueprint(appAuth)
app.register_blueprint(appExample)

# SQLAlchemy
db.init_app(app)
# END SQLAlchemy

jwt.init_app(app)
bcrypt.init_app(app)

@app.errorhandler(404)
@app.errorhandler(405)
def uri_not_found(e) -> tuple:
    return jsonify([]), 404

if __name__ == '__main__':
    app.run(
        host = '0.0.0.0',
        port = os.getenv('PORT'),
        debug = bool(os.getenv('DEBUG'))
    )

auth/models/sqlaModels.py

from datetime import datetime
from dotenv import load_dotenv

from flask_bcrypt import Bcrypt
from flask_sqlalchemy import SQLAlchemy

from sqlalchemy import (
    String,
    Column,
    Integer,
    DateTime,
    ForeignKey
)
from sqlalchemy.orm.exc import NoResultFound
from sqlalchemy.ext.declarative import declared_attr

load_dotenv()
db = SQLAlchemy()
bcrypt = Bcrypt()

class User(db.Model):
    id = Column(Integer, primary_key = True)

    email = Column(String(250), nullable = False, unique = True)
    username = Column(String(250), nullable = False, unique = True)
    password = Column(String(250), unique = False)

• get_by_id. Get the user by id.
• register. register a new user.
• get_by_email. Get the user by email
• get_by_username. Get the user by username. In our case the username is identical to the email but it is up to you to change that depending on your needs.
• get_verified. Get a user by email or username and password. If the user does not exist return None. This method is used when logging in the user.
• get_sessions. Get a list of user’s sessions.

class Session(db.Model):
    @declared_attr
    def __tablename__(cls):
        return 'session'

    id = Column(Integer, primary_key = True)
    user = db.relationship(User)
    user_id = Column(Integer, ForeignKey(User.id))

    jti_access = Column(String(36), nullable = False, index = True)
    jti_refresh = Column(String(36), nullable = False, index = True)

    created_at = Column(DateTime, default = datetime.now(), nullable = False)

• add. Add a new session.
• get_by_jti_access. Get the session associated with jti_access.
• get_by_jti_refresh. Get the session associated with jti_refresh.
• refresh. Refresh a session.
• delete. Delete a session.
• get_user_sessions. Get the sessions of a user.

auth/models/mongodbModels.py

db_path = os.environ.get('MONGODB_PATH')
db_name = str(os.environ.get('MONGODB_NAME'))
client = MongoClient(db_path)
db = client[db_name]
bcrypt = Bcrypt()

auth/validators.py

auth/appAuth.py

from flask_jwt_extended import (
    get_jwt,
    JWTManager,
    decode_token,
    jwt_required,
    current_user,
    get_jwt_identity,
    create_access_token,
    create_refresh_token,
)

from flask import Blueprint, request, jsonify

from .validators import (
    login_validator,
    register_validator,
    revoke_tokens_validator
)

# SQLAlchemy
from .sqlaModels import User, Session
# END SQLAlchemy

# MongoDb
from .mongodbModels import User, Session
# END MongoDb

jwt = JWTManager()

appAuth = Blueprint('appAuth', __name__)

@jwt.user_identity_loader
def user_identity_lookup(sub):
    id = sub['id'] if type(sub) is dict else sub
    return id

@jwt.user_lookup_loader
def user_lookup_callback(jwt_header: dict, jwt_data: dict) -> User:
    sub = jwt_data['sub']
    id = sub['id'] if type(sub) is dict else sub
    return User.get_by_id(id)

@jwt.token_in_blocklist_loader
def check_if_token_revoked(jwt_header: dict, jwt_payload: dict) -> bool:
    jti = jwt_payload['jti']

    if jwt_payload['type'] == 'access':
        return Session.get_by_jti_access(jti) is None

    if jwt_payload['type'] == 'refresh':
        return Session.get_by_jti_refresh(jti) is None

@jwt.expired_token_loader
def my_expired_token_callback(jwt_header: dict, jwt_payload: dict) -> tuple:
    return jsonify({'msg': 'Access denied!'}), 401

• register. Handle the process of registering a new user, creating new access and refresh tokens, stores them in the storage and returns them to the Front-end.
• login. Handle the process of logging in an existing user, creating new access and refresh tokens, stores them in the storage and returns them to the Front-end.
• logout. Handle the process of logging out an existing user. It deletes the token from the storage.
• refresh. Refresh a token.
• revoke_tokens. Revokes a list of tokens.

appExample.py

from flask import Blueprint, jsonify
from flask_jwt_extended import jwt_required, current_user

# SQLAlchemy
from auth.sqlaModels import User
# END SQLAlchemy

# MongoDb
from auth.mongodbModels import User
# END MongoDb

appExample = Blueprint('appExample', __name__)

@appExample.route('/')
def home() -> tuple:
    return jsonify([]), 200

@appExample.route('/profile')
@jwt_required()
def profile() -> tuple:
    sessions = User.get_sessions(current_user.id)
    return jsonify(
        {
            'username': current_user.username,
            'email': current_user.email,
            'sessions': [
                {
                    'access_token': session.jti_access,
                    'created_at': session.created_at
                }
                for session in sessions
            ]
        }
    ), 200

Configuration Of The Project

Installation

python3 -m venv env
source env/bin/activate
pip3 install -r requirements.txt

python3 install.py

Run

source env/bin/activate
python3 app.py

Check The Storage

sqlite3 database.db

.tables

select * from TABLE_NAME;

.exit

Stop

deactivate

Additional readings

• Flask-Login With SQLAlchemy And MongoDb Token Storages

If you need assistance with your  Flask application or you have suggestions for improving the content, don’t hesitate to contact me.