generate_schema_config.py

from formats import message_header
from formats.mission_plan import mission_plan_schema
from formats.mission_plan_encoded import mission_plan_encoded_schema
from formats.observation import observation_schema
from formats.observation_encoded import observation_encoded_schema
from formats.planning_configuration import planning_configuration_schema
from formats.platform_status import platform_status_schema
from formats.platform_status_encoded import platform_status_encoded_schema
from formats.survey import survey_schema
from formats.survey_encoded import survey_encoded_schema
from formats.acknowledgement import acknowledgement_schema
from formats.alert import alert_schema

from flasgger import Swagger
from flask import Flask

import argparse
import json
import os
import re
import requests


# Enable running on domain sub-path
URL_PREFIX = os.getenv("URL_PREFIX", "")
# Allow env override of default host
FLASK_HOST = os.getenv("FLASK_HOST", "localhost")
# Allow env override of default port
FLASK_PORT = os.getenv("FLASK_PORT", 5000)
# Switch on debug mode if env var is truthy
FLASK_DEBUG = os.getenv("FLASK_DEBUG", "False").lower() in ("true", "1", "t")

swagger_config = {
    "openapi": "3.0.2",
    "swagger_ui": True,
    "specs_route": "/",
    "info": {
        "title": "SoAR Backbone Message Formats",
        "version": "1.0",
        "description": "SoAR message protocol in schemas",
    },
    "specs": [
        {
            "endpoint": "swagger",
            "route": "/soar_protocol.json",
        }
    ],
    "url_prefix": URL_PREFIX,
    "paths": {},
    "components": {
        "schemas": {
            "MESSAGE": {
                "type": "object",
                "description": "Full message definition with"
                + " message-metadata in `header` and different"
                + " message type schemas under `payload`",
                "properties": {
                    "header": {
                        "$ref": "#/components/schemas/header",
                    },
                    "payload": {"$ref": "#/components/schemas/payload"},
                },
                "required": ["header", "payload"],
            },
            "payload": {
                "discriminator": {
                    "propertyName": "message_type",
                    "mapping": {
                        "alert": "#/components/schemas/alert",
                        "mission_plan": "#/components/schemas/mission_plan",
                        "mission_plan_encoded": "#/components/schemas/"
                        + "mission_plan_encoded",
                        "observation": "#/components/schemas/observation",
                        "observation_encoded": "#/components/schemas/"
                        + "observation_encoded",
                        "planning_configuration": "#/components/schemas/"
                        + "planning_configuration",
                        "platform_status": "#/components/schemas/platform_status",
                        "platform_status_encoded": "#/components/schemas/"
                        + "platform_status_encoded",
                        "acknowledgement": "#/components/schemas/acknowledgement",
                        "survey": "#/components/schemas/survey",
                        "survey_encoded": "#/components/schemas/" + "survey_encoded",
                    },
                },
                "oneOf": [
                    {"$ref": "#/components/schemas/alert"},
                    {"$ref": "#/components/schemas/acknowledgement"},
                    {"$ref": "#/components/schemas/mission_plan"},
                    {"$ref": "#/components/schemas/mission_plan_encoded"},
                    {"$ref": "#/components/schemas/observation"},
                    {"$ref": "#/components/schemas/observation_encoded"},
                    {"$ref": "#/components/schemas/planning_configuration"},
                    {"$ref": "#/components/schemas/platform_status"},
                    {"$ref": "#/components/schemas/platform_status_encoded"},
                    {"$ref": "#/components/schemas/survey"},
                    {"$ref": "#/components/schemas/survey_encoded"},
                ],
            },
            "header": message_header,
            "mission_plan": mission_plan_schema,
            "mission_plan_encoded": mission_plan_encoded_schema,
            "observation": observation_schema,
            "observation_encoded": observation_encoded_schema,
            "planning_configuration": planning_configuration_schema,
            "platform_status": platform_status_schema,
            "platform_status_encoded": platform_status_encoded_schema,
            "survey": survey_schema,
            "survey_encoded": survey_encoded_schema,
            "acknowledgement": acknowledgement_schema,
            "alert": alert_schema,
        }
    },
}


def resolve_ref(ref):
    """
    Get schema URL, parse JSON
    Return None if either fails
    """
    try:
        res = requests.get(ref)
        if res.status_code == 200:
            return res.json()
        else:
            return None
    except (json.JSONDecodeError, ValueError):
        return None


def rename_ref(ref):
    """
    Convert remote ref URL into a name that can
    be used for a local ref in the schema
    Remote the URL scheme and replace / with .
    """
    # remove url scheme
    deschemed = re.sub(r"^[htps]*\:*[/]{2}", "", ref)
    # replace / with . since the name will be in a path
    return re.sub(r"[/]", ".", deschemed)


def nested_replace(source, key, value, replace_with):
    """
    Find all instances of a key value pair in a nested
    dictionary and replace the value with replace_with
    """
    for k, v in source.items():
        if k == key and v == value:
            source[k] = replace_with
        elif type(v) is list:
            for item in v:
                if type(item) is dict:
                    nested_replace(item, key, value, replace_with)
        if type(v) is dict:
            nested_replace(v, key, value, replace_with)


def inject_schema(schema, remote_ref):
    """
    Given a parent schema and a remote ref

    1. get the remote ref schema
    2. create a local reference name (without path separators)
    3. insert into components.schemas
    4. replace remote references with local references


    returns True if resolved and injected
    """
    local_name = rename_ref(remote_ref)
    local_ref = f"#/components/schemas/{local_name}"
    ref_schema = resolve_ref(remote_ref)
    if ref_schema is not None:
        nested_replace(schema, "$ref", remote_ref, local_ref)
        schema["components"]["schemas"][local_name] = ref_schema
        return True
    else:
        return False


def import_remote_refs():
    """
    inject the following remote refs into the schema
    and replace the remote refs with local refs

    returns True if all schemas resolved and injected
    """
    ref_imports = [
        "https://geojson.org/schema/Feature.json",
        "https://geojson.org/schema/FeatureCollection.json",
        "https://geojson.org/schema/LineString.json",
        "https://geojson.org/schema/Point.json",
        "https://geojson.org/schema/Polygon.json",
    ]

    return all([inject_schema(swagger_config, ref) for ref in ref_imports])


def configure_flask(swagger_config):
    """
    Setup a flask app, load flasgger

    and then patch to remove invalid
    definitions:{} object
    """
    app = Flask(__name__)
    Swagger(app, config=swagger_config, merge=True)

    # Replace schema route to remove invalid
    # definitions: {}
    # Should be fixed if Flassger 0.9.7 is released
    #
    # The last release of flasgger was Aug 2020
    # This bug was fixed in Nov 2021
    # There is a pre-release from May 2023
    # Until the fix gets released we have to
    # remove the invalid definitions object
    # from the spec
    @app.after_request
    def after_request_decorator(response):
        """
        I didn't want to mess with flasgger so
        this blunt workaround that runs on every
        route and then checks whether it's required
        """
        is_response = type(response).__name__ == "Response"
        is_json = is_response and response.content_type == "application/json"
        if is_json:
            parsed = response.json
            if "definitions" in parsed:
                del parsed["definitions"]
            response.data = json.dumps(parsed)

        return response

    return app


def serve(swagger_config):
    """
    Run as local flask app on FLASK_PORT|5000
    """
    app = configure_flask(swagger_config)
    app.run(debug=FLASK_DEBUG, host=FLASK_HOST, port=FLASK_PORT)


def compile_schema(swagger_config):
    """Extract the output schema from flasgger

    The only way I have found to do this is to
    use a test client to make the GET request
    for the page

    The function that returns the definition
    can't be called outside the flask app context
    """
    app = configure_flask(swagger_config)
    route = swagger_config["specs"][0]["route"]
    client = app.test_client()
    response = client.get(route)
    spec = response.json
    return spec


def write_schema(swagger_config, file_path):
    """
    Dump schema to specified file
    """
    spec = compile_schema(swagger_config)
    json_schema = json.dumps(spec, indent=2)

    with open(file_path, "w") as f:
        f.write(json_schema)


def get_options():
    """
    Parse script arguments
    """
    parser = argparse.ArgumentParser(
        description="Generate the schema",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument(
        "-s",
        "--serve",
        dest="run_flask",
        action="store_true",
        help="Run flask app",
        default=False,
    )
    parser.add_argument(
        "-f",
        "--file",
        dest="output_file",
        action="store_true",
        help="Save output to schema file",
        default=False,
    )
    parser.add_argument("filename", nargs="?", default="project/soar/swagger.json")
    args = parser.parse_args()
    config = vars(args)
    # If no flag is specified default to running the flask server
    if all(v is False for v in config.values()):
        config["run_flask"] = True
    return config


if __name__ == "__main__":
    import_remote_refs()
    
    # Parse script args
    config = get_options()

    # Output compiled schema
    if config.get("output_file"):
        write_schema(swagger_config, config.get("filename"))

    # Run flask app
    if config.get("run_flask"):
        serve(swagger_config)