bedrock.generators.generate-asyncapi-specification

  1"""
  2Generates an AsyncAPI 3.1.0 specification for the application's WebSocket API.
  3
  4Scans endpoints/ for classes using the @websocket_broadcast decorator
  5and builds a spec with lifecycle events and dynamic broadcast channels.
  6
  7Usage:
  8    python generate-asyncapi-specification.py BASE_PATH [APP_NAME] [VERSION] [ENVIRONMENTS]
  9
 10Arguments:
 11    BASE_PATH    - Base path to the application (required)
 12    APP_NAME     - Application name (default: "bedrock")
 13    VERSION      - API version (default: "1.0.0")
 14    ENVIRONMENTS - Optional comma-separated "name|host" pairs
 15                   (e.g. "testing|app-ws.testing.example.com,production|app-ws.example.com")
 16"""
 17
 18import copy
 19import re
 20import sys
 21
 22from sqlalchemy.orm import RelationshipProperty
 23
 24from bedrock._helpers.classes import find_all_attributes
 25from bedrock._helpers.string import snake_case_to_camelCase
 26from bedrock.config._modules import _get_modules
 27from bedrock.generators._helpers.spec_generation_helpers import dump_dict_to_file
 28from bedrock.meta.model import convert_model_to_schema
 29
 30BASE_PATH = None
 31ASYNCAPI_PATH = None
 32
 33# --- Default WebSocket lifecycle schemas ---
 34
 35DEFAULT_SCHEMAS = {
 36    "connectMessage": {
 37        "type": "object",
 38        "properties": {
 39            "action": {"type": "string", "const": "connect"},
 40        },
 41    },
 42    "connectResponse": {
 43        "type": "object",
 44        "properties": {
 45            "statusCode": {"type": "integer", "example": 200},
 46            "body": {"type": "string", "example": "Connected"},
 47        },
 48    },
 49    "authenticateMessage": {
 50        "type": "object",
 51        "properties": {
 52            "action": {"type": "string", "const": "authenticate"},
 53            "token": {"type": "string", "example": "jwt-token"},
 54        },
 55    },
 56    "authenticationResponse": {
 57        "type": "object",
 58        "properties": {
 59            "statusCode": {"type": "integer", "example": 200},
 60            "body": {"type": "string", "example": "Connection 12345 Authenticated"},
 61        },
 62    },
 63    "pingMessage": {
 64        "type": "object",
 65        "properties": {
 66            "action": {"type": "string", "const": "ping"},
 67        },
 68    },
 69    "pongResponse": {
 70        "type": "object",
 71        "properties": {
 72            "statusCode": {"type": "integer", "example": 200},
 73            "body": {"type": "string", "const": "pong"},
 74        },
 75    },
 76    "subscribeMessage": {
 77        "type": "object",
 78        "properties": {
 79            "action": {"type": "string", "const": "subscribe"},
 80            "topic": {"type": "string", "example": "example-topic"},
 81        },
 82    },
 83    "subscribeResponse": {
 84        "type": "object",
 85        "properties": {
 86            "statusCode": {"type": "integer", "example": 200},
 87            "body": {"type": "string", "example": "Subscribed"},
 88        },
 89    },
 90    "unsubscribeMessage": {
 91        "type": "object",
 92        "properties": {
 93            "action": {"type": "string", "const": "unsubscribe"},
 94            "topic": {"type": "string", "example": "example-topic"},
 95        },
 96    },
 97    "unsubscribeResponse": {
 98        "type": "object",
 99        "properties": {
100            "statusCode": {"type": "integer", "example": 200},
101            "body": {"type": "string", "example": "Unsubscribed"},
102        },
103    },
104    "disconnectMessage": {
105        "type": "object",
106        "properties": {
107            "action": {"type": "string", "const": "disconnect"},
108        },
109    },
110    "disconnectResponse": {
111        "type": "object",
112        "properties": {
113            "statusCode": {"type": "integer", "example": 200},
114            "body": {"type": "string", "example": "Disconnected"},
115        },
116    },
117    "sendMessageResponse": {
118        "type": "object",
119        "properties": {
120            "type": {"type": "string", "const": "send_message"},
121            "topic": {"type": "string", "example": "topic-name"},
122            "message": {"type": "string", "example": "Hello, World!"},
123        },
124    },
125}
126
127# --- Default lifecycle messages (refs to schemas) ---
128
129DEFAULT_MESSAGES = {
130    name: {"name": name, "payload": {"$ref": f"#/components/schemas/{name}"}}
131    for name in DEFAULT_SCHEMAS
132}
133
134# --- Default lifecycle channel ---
135
136LIFECYCLE_CHANNEL = {
137    "description": "WebSocket lifecycle events (connect, authenticate, ping/pong, subscribe, unsubscribe, disconnect).",
138    "messages": {
139        name: {"$ref": f"#/components/messages/{name}"}
140        for name in DEFAULT_SCHEMAS
141    },
142}
143
144# --- Default lifecycle operations (client -> server and server -> client) ---
145
146DEFAULT_OPERATIONS = {
147    "sendConnect": {
148        "action": "send",
149        "summary": "Send a connect message to establish a WebSocket connection.",
150        "channel": {"$ref": "#/channels/lifecycle"},
151        "messages": [{"$ref": "#/channels/lifecycle/messages/connectMessage"}],
152    },
153    "receiveConnectResponse": {
154        "action": "receive",
155        "summary": "Receive a response to the connect message.",
156        "channel": {"$ref": "#/channels/lifecycle"},
157        "messages": [{"$ref": "#/channels/lifecycle/messages/connectResponse"}],
158    },
159    "sendAuthenticate": {
160        "action": "send",
161        "summary": "Send an authenticate message to authenticate the WebSocket connection.",
162        "channel": {"$ref": "#/channels/lifecycle"},
163        "messages": [{"$ref": "#/channels/lifecycle/messages/authenticateMessage"}],
164    },
165    "receiveAuthenticationResponse": {
166        "action": "receive",
167        "summary": "Receive a response to the authenticate message.",
168        "channel": {"$ref": "#/channels/lifecycle"},
169        "messages": [{"$ref": "#/channels/lifecycle/messages/authenticationResponse"}],
170    },
171    "sendPing": {
172        "action": "send",
173        "summary": "Send a ping message to keep the connection alive.",
174        "channel": {"$ref": "#/channels/lifecycle"},
175        "messages": [{"$ref": "#/channels/lifecycle/messages/pingMessage"}],
176    },
177    "receivePongResponse": {
178        "action": "receive",
179        "summary": "Receive a pong response to the ping message.",
180        "channel": {"$ref": "#/channels/lifecycle"},
181        "messages": [{"$ref": "#/channels/lifecycle/messages/pongResponse"}],
182    },
183    "sendSubscribe": {
184        "action": "send",
185        "summary": "Send a subscribe message to subscribe to a topic.",
186        "channel": {"$ref": "#/channels/lifecycle"},
187        "messages": [{"$ref": "#/channels/lifecycle/messages/subscribeMessage"}],
188    },
189    "receiveSubscribeResponse": {
190        "action": "receive",
191        "summary": "Receive a response to the subscribe message.",
192        "channel": {"$ref": "#/channels/lifecycle"},
193        "messages": [{"$ref": "#/channels/lifecycle/messages/subscribeResponse"}],
194    },
195    "sendUnsubscribe": {
196        "action": "send",
197        "summary": "Send an unsubscribe message to unsubscribe from a topic.",
198        "channel": {"$ref": "#/channels/lifecycle"},
199        "messages": [{"$ref": "#/channels/lifecycle/messages/unsubscribeMessage"}],
200    },
201    "receiveUnsubscribeResponse": {
202        "action": "receive",
203        "summary": "Receive a response to the unsubscribe message.",
204        "channel": {"$ref": "#/channels/lifecycle"},
205        "messages": [{"$ref": "#/channels/lifecycle/messages/unsubscribeResponse"}],
206    },
207    "sendDisconnect": {
208        "action": "send",
209        "summary": "Send a disconnect message to disconnect the WebSocket connection.",
210        "channel": {"$ref": "#/channels/lifecycle"},
211        "messages": [{"$ref": "#/channels/lifecycle/messages/disconnectMessage"}],
212    },
213    "receiveDisconnectResponse": {
214        "action": "receive",
215        "summary": "Receive a response to the disconnect message.",
216        "channel": {"$ref": "#/channels/lifecycle"},
217        "messages": [{"$ref": "#/channels/lifecycle/messages/disconnectResponse"}],
218    },
219    "receiveSendMessageResponse": {
220        "action": "receive",
221        "summary": "Receive a message pushed to a subscribed topic.",
222        "channel": {"$ref": "#/channels/lifecycle"},
223        "messages": [{"$ref": "#/channels/lifecycle/messages/sendMessageResponse"}],
224    },
225}
226
227DEFAULT_SPEC = {
228    "asyncapi": "3.1.0",
229    "info": {
230        "title": "",
231        "version": "",
232        "description": "",
233    },
234    "servers": {},
235    "channels": {
236        "lifecycle": LIFECYCLE_CHANNEL,
237    },
238    "operations": DEFAULT_OPERATIONS,
239    "components": {
240        "messages": DEFAULT_MESSAGES,
241        "schemas": DEFAULT_SCHEMAS,
242    },
243}
244
245
246# --- Helper functions ---
247
248def _rewrite_schema_refs(schema):
249    """Rewrite file-based $ref values to inline AsyncAPI component refs."""
250    if isinstance(schema, dict):
251        return {
252            key: re.sub(r'^\.{1,2}/(schemas/)?[\w]+\.yml#/(.+)$', r'#/components/schemas/\2', value)
253            if key == "$ref" and isinstance(value, str)
254            else _rewrite_schema_refs(value)
255            for key, value in schema.items()
256        }
257    if isinstance(schema, list):
258        return [_rewrite_schema_refs(item) for item in schema]
259    return schema
260
261
262def _get_sub_models(model_cls, seen=None):
263    """Recursively get all models referenced via relationships in the given model class."""
264    if seen is None:
265        seen = set()
266    seen.add(model_cls)
267    sub_models = []
268    properties = [
269        getattr(model_cls, prop) for prop in find_all_attributes(model_cls)
270        if prop not in ['metadata', 'registry']
271    ]
272    for prop in properties:
273        if isinstance(prop.prop, RelationshipProperty):
274            related_model = prop.prop.entity.entity
275            if related_model not in seen:
276                sub_models.append(related_model)
277                sub_models.extend(_get_sub_models(related_model, seen))
278    return sub_models
279
280
281def _make_schema_name(module_name, topic):
282    """Build a clean camelCase schema name from module name and topic."""
283    raw = f"{module_name}_{topic.replace('-', '_')}_message"
284    return snake_case_to_camelCase(raw)
285
286
287def _add_model_schema(spec, schema_name, model_cls):
288    """Convert a model to a schema, rewrite its refs, and add it with a matching message."""
289    model_schema = convert_model_to_schema(model_cls)
290    spec["components"]["schemas"][schema_name] = _rewrite_schema_refs(model_schema[model_cls.__name__])
291    spec["components"]["messages"][schema_name] = {
292        "name": schema_name,
293        "payload": {"$ref": f"#/components/schemas/{schema_name}"},
294    }
295
296
297def _add_topic(spec, topic, schema_name, seen_topics):
298    """Create or update a topic channel and its receive operation."""
299    if topic not in seen_topics:
300        spec["channels"][topic] = {
301            "description": f"Channel for topic: {topic}",
302            "messages": {
303                schema_name: {"$ref": f"#/components/messages/{schema_name}"},
304            },
305        }
306        operation_name = f"receive{snake_case_to_camelCase(topic.replace('-', '_'), capitalise_first_letter=True)}"
307        spec["operations"][operation_name] = {
308            "action": "receive",
309            "summary": f"Receive {topic} messages.",
310            "channel": {"$ref": f"#/channels/{topic}"},
311            "messages": [
312                {"$ref": f"#/channels/{topic}/messages/{schema_name}"},
313            ],
314        }
315        seen_topics[topic] = operation_name
316    else:
317        spec["channels"][topic]["messages"][schema_name] = {
318            "$ref": f"#/components/messages/{schema_name}",
319        }
320        spec["operations"][seen_topics[topic]]["messages"].append(
321            {"$ref": f"#/channels/{topic}/messages/{schema_name}"},
322        )
323
324
325# --- Endpoint discovery ---
326
327def get_websocket_endpoints(module_directory, module_prefix):
328    """Discover endpoint classes that use the @websocket_broadcast decorator."""
329    modules = _get_modules(dir_path=module_directory, module_prefix=module_prefix, excludes=[])
330
331    websocket_endpoints = {}
332    for module_name, endpoint_cls in modules.items():
333        if hasattr(endpoint_cls, "__uses_websockets__"):
334            websocket_endpoints[module_name] = endpoint_cls()
335
336    return websocket_endpoints
337
338
339# --- Spec generation ---
340
341def generate_asyncapi_spec(app_name, version, environments=None):
342    """Generate a complete AsyncAPI 3.1.0 specification."""
343    app_name_formatted = app_name.replace("-", " ").title()
344    spec = copy.deepcopy(DEFAULT_SPEC)
345    spec["info"]["title"] = f"{app_name_formatted} WebSocket API"
346    spec["info"]["version"] = version
347    spec["info"]["description"] = f"AsyncAPI specification for {app_name_formatted} WebSocket API routes."
348
349    if environments:
350        for env in environments:
351            spec["servers"][env["name"]] = {
352                "host": env["host"],
353                "protocol": "wss",
354                "description": f"{env['name'].capitalize()} WebSocket server.",
355            }
356
357    endpoints = get_websocket_endpoints(f"{BASE_PATH}/app/endpoints", "endpoints.")
358    seen_topics = {}
359
360    for module_name, endpoint_instance in endpoints.items():
361        topics = getattr(endpoint_instance, "__websocket_topics__", [])
362        related_model = getattr(endpoint_instance, "related_model", None)
363
364        for topic in topics:
365            schema_name = _make_schema_name(module_name, topic)
366
367            if related_model:
368                _add_model_schema(spec, schema_name, related_model)
369
370                for sub_model in _get_sub_models(related_model):
371                    sub_model_name = sub_model.__name__
372                    if sub_model_name not in spec["components"]["schemas"]:
373                        _add_model_schema(spec, sub_model_name, sub_model)
374            else:
375                spec["components"]["messages"][schema_name] = {
376                    "name": schema_name,
377                    "payload": {"$ref": f"#/components/schemas/{schema_name}"},
378                }
379
380            _add_topic(spec, topic, schema_name, seen_topics)
381
382    return spec
383
384
385if __name__ == "__main__":
386    BASE_PATH = sys.argv[1]
387    app_name = sys.argv[2] if len(sys.argv) > 2 else "bedrock"
388    version = sys.argv[3] if len(sys.argv) > 3 else "1.0.0"
389
390    default_environments = [f"{env}|{app_name}-ws.{env}.keyholding.com" for env in ["testing", "staging"]]
391    default_environments.append(f"production|{app_name}-ws.keyholding.com")
392    environments_with_hosts = sys.argv[4].split(",") if len(sys.argv) > 4 else default_environments
393    environments = [{"name": env_with_host.split("|")[0], "host": env_with_host.split("|")[1]} for env_with_host in
394                    environments_with_hosts]
395
396    ASYNCAPI_PATH = f"{BASE_PATH}/asyncapi/"
397    spec = generate_asyncapi_spec(app_name, version, environments)
398    dump_dict_to_file(ASYNCAPI_PATH, "asyncapi.spec.yaml", spec)