Refactor: replace Flask-Restful with own, simple implementation

mwdb/core/service.py

@@ -1,97 +1,117 @@
+import sys
 import textwrap
-from functools import partial
 
 from apispec import APISpec
 from apispec.ext.marshmallow import MarshmallowPlugin
-from flask_restful import Api
+from apispec_webframeworks.flask import FlaskPlugin
+from flask import Blueprint, Flask, jsonify, request
+from flask.typing import ResponseReturnValue
 from sqlalchemy.exc import OperationalError
-from werkzeug.exceptions import HTTPException, ServiceUnavailable
+from werkzeug.exceptions import (
+    HTTPException,
+    InternalServerError,
+    MethodNotAllowed,
+    ServiceUnavailable,
+)
+from werkzeug.wrappers import Response
 
 from mwdb.version import app_version
 
-from . import log
-from .apispec_utils import ApispecFlaskRestful
+from .log import getLogger
 
+SUPPORTED_METHODS = ["head", "get", "post", "put", "delete", "patch"]
 
-class Service(Api):
-    def __init__(self, flask_app, *args, **kwargs):
-        self.spec = self._create_spec()
-        self.flask_app = flask_app
-        super().__init__(*args, **kwargs)
+logger = getLogger()
 
-    def _init_app(self, app):
-        # I want to log exceptions on my own
-        def dont_log(*_, **__):
-            pass
 
-        app.log_exception = dont_log
-        if (
-            isinstance(app.handle_exception, partial)
-            and app.handle_exception.func is self.error_router
-        ):
-            # Prevent double-initialization
-            return
-        super()._init_app(app)
+class Resource:
+    def __init__(self):
+        self.available_methods = [
+            method.upper() for method in SUPPORTED_METHODS if hasattr(self, method)
+        ]
+
+    def __call__(self, *args, **kwargs):
+        """
+        Acts as view function, calling appropriate method and
+        jsonifying response
+        """
+        if request.method not in self.available_methods:
+            raise MethodNotAllowed(
+                valid_methods=self.available_methods,
+                description="Method is not allowed for this endpoint",
+            )
+        method = request.method.lower()
+        response = getattr(self, method)(*args, **kwargs)
+        if isinstance(response, Response):
+            return response
+        return jsonify(response)
+
+    def get_methods(self):
+        """
+        Returns available methods for this resource
+        """
+        return [getattr(self, method) for method in self.available_methods]
 
-    def _create_spec(self):
-        spec = APISpec(
-            title="MWDB",
+
+class Service:
+    def __init__(self, app: Flask, blueprint: Blueprint) -> None:
+        self.app = app
+        self.blueprint = blueprint
+        self.blueprint.register_error_handler(Exception, self.error_handler)
+        self.spec = APISpec(
+            title="MWDB Core",
             version=app_version,
             openapi_version="3.0.2",
-            plugins=[ApispecFlaskRestful(), MarshmallowPlugin()],
-        )
+            plugins=[FlaskPlugin(), MarshmallowPlugin()],
+            info={
+                "description": textwrap.dedent(
+                    """
+                    MWDB API documentation.
 
-        spec.components.security_scheme(
+                    If you want to automate things, we recommend using
+                    <a href="http://github.com/CERT-Polska/mwdblib">
+                        mwdblib library
+                    </a>
+                    """
+                )
+            },
+            servers=[
+                {
+                    "url": "{scheme}://{host}",
+                    "description": "MWDB API endpoint",
+                    "variables": {
+                        "scheme": {"enum": ["http", "https"], "default": "https"},
+                        "host": {"default": "mwdb.cert.pl"},
+                    },
+                }
+            ],
+        )
+        self.spec.components.security_scheme(
             "bearerAuth", {"type": "http", "scheme": "bearer", "bearerFormat": "JWT"}
         )
-        spec.options["info"] = {
-            "description": textwrap.dedent(
-                """
-                MWDB API documentation.
 
-                If you want to automate things, we recommend using
-                <a href="http://github.com/CERT-Polska/mwdblib">mwdblib library</a>"""
-            )
-        }
-        spec.options["servers"] = [
-            {
-                "url": "{scheme}://{host}",
-                "description": "MWDB API endpoint",
-                "variables": {
-                    "scheme": {"enum": ["http", "https"], "default": "https"},
-                    "host": {"default": "mwdb.cert.pl"},
-                },
-            }
-        ]
-        return spec
+    def _make_error_response(self, exc: HTTPException) -> ResponseReturnValue:
+        return jsonify({"message": exc.description}), exc.code
 
-    def error_router(self, original_handler, e):
-        logger = log.getLogger()
-        if isinstance(e, HTTPException):
-            logger.error(str(e))
-        elif isinstance(e, OperationalError):
-            logger.error(str(e))
-            raise ServiceUnavailable("Request canceled due to statement timeout")
+    def error_handler(self, exc: Exception) -> ResponseReturnValue:
+        if isinstance(exc, HTTPException):
+            return self._make_error_response(exc)
+        elif isinstance(exc, OperationalError):
+            return self._make_error_response(
+                ServiceUnavailable("Request canceled due to statement timeout")
+            )
         else:
-            logger.exception("Unhandled exception occurred")
-
-        # Handle all exceptions using handle_error, not only for owned routes
-        try:
-            return self.handle_error(e)
-        except Exception:
-            logger.exception("Exception from handle_error occurred")
-            pass
-        # If something went wrong - fallback to original behavior
-        return super().error_router(original_handler, e)
+            # Unknown exception, return ISE 500
+            logger.exception("Internal server error", exc_info=sys.exc_info())
+            return self._make_error_response(
+                InternalServerError("Internal server error")
+            )
 
-    def add_resource(self, resource, *urls, undocumented=False, **kwargs):
-        super().add_resource(resource, *urls, **kwargs)
+    def add_resource(
+        self, resource: Resource, *urls: str, undocumented: bool = False
+    ) -> None:
+        for url in urls:
+            endpoint = f"{self.blueprint.name}.{resource.__name__.lower()}"
+            self.blueprint.add_url_rule(rule=url, endpoint=endpoint, view_func=resource)
         if not undocumented:
-            self.spec.path(resource=resource, api=self, app=self.flask_app)
-
-    def relative_url_for(self, resource, **values):
-        path = self.url_for(resource, **values)
-        return path[len(self.blueprint.url_prefix) :]
-
-    def endpoint_for(self, resource):
-        return f"{self.blueprint.name}.{resource}"
+            self.spec.path(view=resource, app=self.app)

mwdb/core/util.py

@@ -22,8 +22,6 @@ def token_hex(nbytes=None):
     InstanceMetadataFetcher,
     InstanceMetadataProvider,
 )
-from flask_restful import abort
-from flask_sqlalchemy import Pagination
 
 
 def config_dhash(obj):
@@ -130,21 +128,6 @@ def calc_crc32(stream):
     return "{:08x}".format(csum)
 
 
-def paginate_fast(q, page, per_page):
-    if page < 1:
-        abort(404)
-
-    if per_page < 0:
-        abort(404)
-
-    items = q.limit(per_page).offset((page - 1) * per_page).all()
-
-    if not items and page != 1:
-        abort(404)
-
-    return Pagination(q, page, per_page, 0, items)
-
-
 def is_true(flag):
     # "True", "true", "1"
     if isinstance(flag, str) and flag and flag.lower() in ["true", "1"]:

mwdb/resources/api_key.py

@@ -2,11 +2,11 @@
 from datetime import datetime
 
 from flask import g, request
-from flask_restful import Resource
 from sqlalchemy.orm.exc import NoResultFound
 from werkzeug.exceptions import Forbidden, NotFound
 
 from mwdb.core.capabilities import Capabilities
+from mwdb.core.service import Resource
 from mwdb.model import APIKey, User, db
 from mwdb.schema.api_key import (
     APIKeyIdentifierBase,

mwdb/resources/attribute.py

@@ -1,9 +1,9 @@
 from flask import g, request
-from flask_restful import Resource
 from werkzeug.exceptions import BadRequest, Conflict, Forbidden, NotFound
 
 from mwdb.core.capabilities import Capabilities
 from mwdb.core.plugins import hooks
+from mwdb.core.service import Resource
 from mwdb.model import Attribute, AttributeDefinition, AttributePermission, Group, db
 from mwdb.schema.attribute import (
     AttributeDefinitionCreateRequestSchema,

mwdb/resources/auth.py

@@ -2,7 +2,6 @@
 
 import requests
 from flask import g, request
-from flask_restful import Resource
 from sqlalchemy import exists, func
 from sqlalchemy.orm import joinedload
 from sqlalchemy.orm.exc import NoResultFound
@@ -13,6 +12,7 @@
 from mwdb.core.config import app_config
 from mwdb.core.mail import MailError, send_email_notification
 from mwdb.core.plugins import hooks
+from mwdb.core.service import Resource
 from mwdb.model import Group, Member, User, db
 from mwdb.schema.auth import (
     AuthLoginRequestSchema,

mwdb/resources/comment.py

@@ -1,9 +1,9 @@
 from flask import g, request
-from flask_restful import Resource
 from werkzeug.exceptions import NotFound
 
 from mwdb.core.capabilities import Capabilities
 from mwdb.core.plugins import hooks
+from mwdb.core.service import Resource
 from mwdb.model import Comment, db
 from mwdb.schema.comment import CommentItemResponseSchema, CommentRequestSchema
 

mwdb/resources/config.py

@@ -1,13 +1,13 @@
 from datetime import datetime, timedelta
 
 from flask import g, request
-from flask_restful import Resource
 from sqlalchemy import func
 from werkzeug.exceptions import BadRequest, Conflict, Forbidden, NotFound
 
 from mwdb.core.capabilities import Capabilities
 from mwdb.core.deprecated import DeprecatedFeature, deprecated_endpoint
 from mwdb.core.plugins import hooks
+from mwdb.core.service import Resource
 from mwdb.model import Config, TextBlob, db
 from mwdb.model.object import ObjectTypeConflictError
 from mwdb.schema.blob import BlobCreateSpecSchema

mwdb/resources/download.py

@@ -1,9 +1,9 @@
 from flask import Response
-from flask_restful import Resource
 from werkzeug.exceptions import Forbidden, NotFound
 
 from mwdb.core.app import api
 from mwdb.core.deprecated import DeprecatedFeature, deprecated_endpoint
+from mwdb.core.service import Resource
 from mwdb.model import File
 from mwdb.schema.download import DownloadURLResponseSchema