from flask_sqlalchemy import SQLAlchemy
from flask_wtf.csrf import CSRFProtect
from flask_sqlalchemy.session import Session
from flask import g, has_request_context


class RequestSession(Session):
    def commit(self):
        # Only the request boundary may commit a mutation and its receipt.
        if has_request_context() and getattr(g, "mutation_active", False):
            self.flush()
        else:
            super().commit()

    def rollback(self):
        if has_request_context() and getattr(g, "mutation_active", False):
            g.mutation_aborted = True
        super().rollback()


db = SQLAlchemy(session_options={"class_": RequestSession})
csrf = CSRFProtect()