latch-asgi 0.2.0__py3-none-any.whl → 1.0.0__py3-none-any.whl

latch_asgi/datadog_propagator.py

@@ -1,3 +1,5 @@
+from typing import Self
+
 from opentelemetry import trace
 from opentelemetry.context.context import Context
 from opentelemetry.propagators import textmap
@@ -6,7 +8,7 @@ from opentelemetry.propagators import textmap
 class DDTraceContextTextMapPropagator(textmap.TextMapPropagator):
     # https://github.com/open-telemetry/opentelemetry-python-contrib/blob/934af7ea4f9b1e0294ced6a014d6eefdda156b2b/exporter/opentelemetry-exporter-datadog/src/opentelemetry/exporter/datadog/propagator.py
     def extract(
-        self,
+        self: Self,
         carrier: textmap.CarrierT,
         context: Context | None = None,
         getter: textmap.Getter[textmap.CarrierT] = textmap.default_getter,
@@ -28,12 +30,16 @@ class DDTraceContextTextMapPropagator(textmap.TextMapPropagator):
         sampling_priority = getter.get(carrier, "x-datadog-sampling-priority")
 
         trace_flags = trace.TraceFlags()
-        if sampling_priority is not None and len(sampling_priority) > 0:
-            if int(sampling_priority[0]) in (
+        if (
+            sampling_priority is not None
+            and len(sampling_priority) > 0
+            and int(sampling_priority[0])
+            in (
                 1,  # auto keep
                 2,  # user keep
-            ):
-                trace_flags = trace.TraceFlags(trace.TraceFlags.SAMPLED)
+            )
+        ):
+            trace_flags = trace.TraceFlags(trace.TraceFlags.SAMPLED)
 
         dd_origin = getter.get(carrier, "x-datadog-origin")
 
@@ -51,7 +57,7 @@ class DDTraceContextTextMapPropagator(textmap.TextMapPropagator):
         return trace.set_span_in_context(trace.NonRecordingSpan(span_context), context)
 
     def inject(
-        self,
+        self: Self,
         carrier: textmap.CarrierT,
         context: Context | None = None,
         setter: textmap.Setter[textmap.CarrierT] = textmap.default_setter,
@@ -69,11 +75,7 @@ class DDTraceContextTextMapPropagator(textmap.TextMapPropagator):
             "x-datadog-trace-id",
             str(span_context.trace_id & 0xFFFF_FFFF_FFFF_FFFF),
         )
-        setter.set(
-            carrier,
-            "x-datadog-parent-id",
-            str(span_context.span_id),
-        )
+        setter.set(carrier, "x-datadog-parent-id", str(span_context.span_id))
         setter.set(
             carrier,
             "x-datadog-sampling-priority",
@@ -87,7 +89,7 @@ class DDTraceContextTextMapPropagator(textmap.TextMapPropagator):
             )
 
     @property
-    def fields(self) -> set[str]:
+    def fields(self: Self) -> set[str]:
         return {
             "x-datadog-trace-id",
             "x-datadog-parent-id",

latch_asgi/framework/common.py

@@ -5,3 +5,36 @@ from opentelemetry.trace import get_tracer
 Headers: TypeAlias = dict[str | bytes, str | bytes]
 
 tracer = get_tracer(__name__)
+
+otel_header_whitelist = {
+    "host",
+    "content-type",
+    "content-length",
+    "accept",
+    "accept-encoding",
+    "accept-language",
+    "user-agent",
+    "dnt",
+    "sec-fetch-dest",
+    "sec-fetch-mode",
+    "sec-fetch-site",
+    "sec-fetch-user",
+    "sec-gpc",
+    "te",
+    "upgrade-insecure-requests",
+    "device-memory",
+    "downlink",
+    "dpr",
+    "ect",
+    "rtt",
+    "sec-ch-prefers-color-scheme",
+    "sec-ch-prefers-reduced-motion",
+    "sec-ch-ua",
+    "sec-ch-ua-arch",
+    "sec-ch-ua-full-version",
+    "sec-ch-ua-mobile",
+    "sec-ch-ua-model",
+    "sec-ch-ua-platform",
+    "sec-ch-ua-platfrom-version",
+    "viewport-width",
+}

latch_asgi/framework/http.py

@@ -1,12 +1,11 @@
 from http import HTTPStatus
-from typing import Any, Literal, TypeAlias, TypeVar, cast
+from typing import Any, Literal, Self, TypeAlias, TypeVar, cast
 
-import opentelemetry.context as context
-import simdjson
+import orjson
 from latch_data_validation.data_validation import DataValidationError, validate
 from latch_o11y.o11y import trace_function, trace_function_with_span
+from opentelemetry import context
 from opentelemetry.trace.span import Span
-from orjson import dumps
 
 from ..asgi_iface import (
     HTTPReceiveCallable,
@@ -18,17 +17,9 @@ from .common import Headers, tracer
 
 T = TypeVar("T")
 
-HTTPMethod: TypeAlias = (
-    Literal["GET"]
-    | Literal["HEAD"]
-    | Literal["POST"]
-    | Literal["PUT"]
-    | Literal["DELETE"]
-    | Literal["CONNECT"]
-    | Literal["OPTIONS"]
-    | Literal["TRACE"]
-    | Literal["PATCH"]
-)
+HTTPMethod: TypeAlias = Literal[
+    "GET", "HEAD", "POST", "PUT", "DELETE", "CONNECT", "OPTIONS", "TRACE", "PATCH"
+]
 
 # >>> O11y
 
@@ -43,35 +34,68 @@ def current_http_request_span() -> Span:
 
 
 class HTTPErrorResponse(RuntimeError):
-    def __init__(self, status: HTTPStatus, data: Any, *, headers: Headers = {}):
-        super().__init__()
+    def __init__(
+        self: Self, status: HTTPStatus, data: Any, *, headers: Headers | None = None
+    ) -> None:
+        if headers is None:
+            headers = {}
+
         self.status = status
         self.data = data
         self.headers = headers
 
 
 class HTTPInternalServerError(HTTPErrorResponse):
-    def __init__(self, data: Any, *, headers: Headers = {}):
+    def __init__(self: Self, data: Any, *, headers: Headers | None = None) -> None:
         super().__init__(HTTPStatus.INTERNAL_SERVER_ERROR, data, headers=headers)
 
 
 class HTTPBadRequest(HTTPErrorResponse):
-    def __init__(self, data: Any, *, headers: Headers = {}):
+    def __init__(self: Self, data: Any, *, headers: Headers | None = None) -> None:
         super().__init__(HTTPStatus.BAD_REQUEST, data, headers=headers)
 
 
 class HTTPForbidden(HTTPErrorResponse):
-    def __init__(self, data: Any, *, headers: Headers = {}):
+    def __init__(self: Self, data: Any, *, headers: Headers | None = None) -> None:
         super().__init__(HTTPStatus.FORBIDDEN, data, headers=headers)
 
 
-class HTTPConnectionClosedError(RuntimeError): ...
+class HTTPConnectionClosedError(RuntimeError):
+    ...
 
 
-# >>> I/O
+# >>>
+# I/O
+# >>>
 
 # todo(maximsmol): add max body length limit by default
 
+# >>> Receive
+
+
+@trace_function(tracer)
+async def receive_data(receive: HTTPReceiveCallable) -> bytes:
+    res = b""
+    more_body = True
+    while more_body:
+        msg = await receive()
+        if msg["type"] == "http.disconnect":
+            raise HTTPConnectionClosedError
+
+        res += msg["body"]
+        more_body = msg["more_body"]
+
+    # todo(maximsmol): accumulate instead of overriding
+    # todo(maximsmol): probably use the content-length header if present?
+    current_http_request_span().set_attribute("http.request.body.size", len(res))
+
+    return res
+
+
+@trace_function(tracer)
+async def receive_json(receive: HTTPReceiveCallable) -> Any:
+    return orjson.loads(await receive_data(receive))
+
 
 async def receive_class_ext(
     receive: HTTPReceiveCallable, cls: type[T]
@@ -89,53 +113,25 @@ async def receive_class(receive: HTTPReceiveCallable, cls: type[T]) -> T:
     return (await receive_class_ext(receive, cls))[1]
 
 
-@trace_function(tracer)
-async def receive_json(receive: HTTPReceiveCallable) -> Any:
-    res = await receive_data(receive)
-
-    p = simdjson.Parser()
-    try:
-        return p.parse(res, True)
-    except ValueError as e:
-        raise HTTPBadRequest("Failed to parse JSON") from e
-
-
-async def receive_data(receive: HTTPReceiveCallable):
-    res = b""
-    more_body = True
-    while more_body:
-        with tracer.start_as_current_span("read chunk") as s:
-            msg = await receive()
-            if msg.type == "http.disconnect":
-                raise HTTPConnectionClosedError()
-
-            res += msg.body
-            more_body = msg.more_body
+# >>> Send
 
-            s.set_attributes({"size": len(msg.body), "more_body": more_body})
 
-    # todo(maximsmol): accumulate instead of overriding
-    # todo(maximsmol): probably use the content-length header if present?
-    current_http_request_span().set_attribute("http.request_content_length", len(res))
-
-    return res
-
-
-@trace_function_with_span(tracer)
+@trace_function(tracer)
 async def send_http_data(
-    s: Span,
     send: HTTPSendCallable,
     status: HTTPStatus,
     data: str | bytes,
     /,
     *,
-    content_type: str | bytes | None = "text/plain",
-    headers: Headers = {},
-):
+    content_type: str | bytes | None = b"text/plain",
+    headers: Headers | None = None,
+) -> None:
+    if headers is None:
+        headers = {}
+
     if isinstance(data, str):
         data = data.encode("utf-8")
 
-    s.set_attribute("size", len(data))
     headers_to_send: list[tuple[bytes, bytes]] = [
         (b"Content-Length", str(len(data)).encode("latin-1"))
     ]
@@ -160,8 +156,6 @@ async def send_http_data(
         HTTPResponseBodyEvent(type="http.response.body", body=data, more_body=False)
     )
 
-    current_http_request_span().set_attribute("http.response_content_length", len(data))
-
 
 @trace_function(tracer)
 async def send_json(
@@ -171,10 +165,10 @@ async def send_json(
     /,
     *,
     content_type: str = "application/json",
-    headers: Headers = {},
-):
+    headers: Headers | None = None,
+) -> None:
     return await send_http_data(
-        send, status, dumps(data), content_type=content_type, headers=headers
+        send, status, orjson.dumps(data), content_type=content_type, headers=headers
     )
 
 
@@ -185,9 +179,9 @@ async def send_auto(
     data: str | bytes | Any,
     /,
     *,
-    headers: Headers = {},
-):
-    if isinstance(data, str) or isinstance(data, bytes):
+    headers: Headers | None = None,
+) -> None:
+    if isinstance(data, str | bytes):
         return await send_http_data(send, status, data, headers=headers)
 
     return await send_json(send, status, data, headers=headers)

latch_asgi/framework/websocket.py

@@ -1,11 +1,12 @@
 from enum import Enum
-from typing import Any, TypeVar, cast
+from typing import Any, Self, TypeVar, cast
 
-import opentelemetry.context as context
-import simdjson
+import orjson
 from latch_data_validation.data_validation import DataValidationError, validate
 from latch_o11y.o11y import trace_function, trace_function_with_span
+from opentelemetry import context
 from opentelemetry.trace.span import Span
+from opentelemetry.util.types import AttributeValue
 
 from ..asgi_iface import (
     WebsocketAcceptEvent,
@@ -18,39 +19,35 @@ from .common import Headers, tracer
 
 T = TypeVar("T")
 
-websocket_request_span_key = context.create_key("websocket_request_span")
+websocket_session_span_key = context.create_key("websocket_message_span")
 
 
-def current_websocket_request_span() -> Span:
-    return cast(Span, context.get_value(websocket_request_span_key))
+def current_websocket_session_span() -> Span:
+    return cast(Span, context.get_value(websocket_session_span_key))
 
 
 # >>> Error classes
 
 
 class WebsocketStatus(int, Enum):
+    """https://www.rfc-editor.org/rfc/rfc6455.html#section-7.4.1"""
+
     normal = 1000
     """
     1000 indicates a normal closure, meaning that the purpose for
     which the connection was established has been fulfilled.
-
-    https://www.rfc-editor.org/rfc/rfc6455.html#section-7.4.1
     """
 
     going_away = 1001
     """
     1001 indicates that an endpoint is "going away", such as a server
     going down or a browser having navigated away from a page.
-
-    https://www.rfc-editor.org/rfc/rfc6455.html#section-7.4.1
     """
 
     protocol_error = 1002
     """
     1002 indicates that an endpoint is terminating the connection due
     to a protocol error.
-
-    https://www.rfc-editor.org/rfc/rfc6455.html#section-7.4.1
     """
 
     unsupported = 1003
@@ -59,15 +56,11 @@ class WebsocketStatus(int, Enum):
     because it has received a type of data it cannot accept (e.g., an
     endpoint that understands only text data MAY send this if it
     receives a binary message).
-
-    https://www.rfc-editor.org/rfc/rfc6455.html#section-7.4.1
     """
 
     reserved = 1004
     """
     Reserved.  The specific meaning might be defined in the future.
-
-    https://www.rfc-editor.org/rfc/rfc6455.html#section-7.4.1
     """
 
     no_status = 1005
@@ -76,8 +69,6 @@ class WebsocketStatus(int, Enum):
     Close control frame by an endpoint.  It is designated for use in
     applications expecting a status code to indicate that no status
     code was actually present.
-
-    https://www.rfc-editor.org/rfc/rfc6455.html#section-7.4.1
     """
 
     abnormal = 1006
@@ -87,8 +78,6 @@ class WebsocketStatus(int, Enum):
     applications expecting a status code to indicate that the
     connection was closed abnormally, e.g., without sending or
     receiving a Close control frame.
-
-    https://www.rfc-editor.org/rfc/rfc6455.html#section-7.4.1
     """
 
     unsupported_payload = 1007
@@ -97,8 +86,6 @@ class WebsocketStatus(int, Enum):
     because it has received data within a message that was not
     consistent with the type of the message (e.g., non-UTF-8 [RFC3629]
     data within a text message).
-
-    https://www.rfc-editor.org/rfc/rfc6455.html#section-7.4.1
     """
 
     policy_violation = 1008
@@ -108,8 +95,6 @@ class WebsocketStatus(int, Enum):
     is a generic status code that can be returned when there is no
     other more suitable status code (e.g., 1003 or 1009) or if there
     is a need to hide specific details about the policy.
-
-    https://www.rfc-editor.org/rfc/rfc6455.html#section-7.4.1
     """
 
     too_large = 1009
@@ -117,8 +102,6 @@ class WebsocketStatus(int, Enum):
     1009 indicates that an endpoint is terminating the connection
     because it has received a message that is too big for it to
     process.
-
-    https://www.rfc-editor.org/rfc/rfc6455.html#section-7.4.1
     """
 
     mandatory_extension = 1010
@@ -130,8 +113,6 @@ class WebsocketStatus(int, Enum):
     are needed SHOULD appear in the /reason/ part of the Close frame.
     Note that this status code is not used by the server, because it
     can fail the WebSocket handshake instead.
-
-    https://www.rfc-editor.org/rfc/rfc6455.html#section-7.4.1
     """
 
     server_error = 1011
@@ -139,8 +120,6 @@ class WebsocketStatus(int, Enum):
     1011 indicates that a server is terminating the connection because
     it encountered an unexpected condition that prevented it from
     fulfilling the request.
-
-    https://www.rfc-editor.org/rfc/rfc6455.html#section-7.4.1
     """
 
     tls_handshake_fail = 1015
@@ -150,79 +129,112 @@ class WebsocketStatus(int, Enum):
     applications expecting a status code to indicate that the
     connection was closed due to a failure to perform a TLS handshake
     (e.g., the server certificate can't be verified).
-
-    https://www.rfc-editor.org/rfc/rfc6455.html#section-7.4.1
     """
 
 
-class WebsocketError(RuntimeError):
-    def __init__(self, status: WebsocketStatus, data: Any, *, headers: Headers = {}):
-        super().__init__()
+class WebsocketErrorResponse(RuntimeError):
+    def __init__(self: Self, status: WebsocketStatus, data: Any) -> None:
         self.status = status
         self.data = data
-        self.headers = headers
 
 
-class WebsocketErrorResponse(WebsocketError): ...
+class WebsocketBadMessage(WebsocketErrorResponse):
+    def __init__(self: Self, data: Any) -> None:
+        super().__init__(WebsocketStatus.policy_violation, data)
 
 
-class WebsocketConnectionClosedError(RuntimeError): ...
+class WebsocketInternalServerError(WebsocketErrorResponse):
+    def __init__(self: Self, data: Any) -> None:
+        super().__init__(WebsocketStatus.server_error, data)
 
 
-class WebsocketBadMessage(WebsocketErrorResponse):
-    def __init__(self, data: Any, *, headers: Headers = {}):
-        super().__init__(WebsocketStatus.policy_violation, data, headers=headers)
+class WebsocketConnectionClosedError(RuntimeError):
+    def __init__(self: Self, code: WebsocketStatus) -> None:
+        self.code = code
 
 
-class WebsocketInternalServerError(WebsocketErrorResponse):
-    def __init__(self, data: Any, *, headers: Headers = {}):
-        super().__init__(WebsocketStatus.server_error, data, headers=headers)
+# >>>
+# I/O
+# >>>
 
+# >>> Send Lifecycle
 
-# >>> I/O
 
+@trace_function(tracer)
+async def accept_connection(
+    send: WebsocketSendCallable,
+    /,
+    *,
+    subprotocol: str | None = None,
+    headers: Headers | None = None,
+) -> None:
+    if headers is None:
+        headers = {}
+
+    headers_to_send: list[tuple[bytes, bytes]] = []
+    for k, v in headers.items():
+        if isinstance(k, str):
+            k = k.encode("latin-1")
+        if isinstance(v, str):
+            v = v.encode("latin-1")
+        headers_to_send.append((k, v))
 
-async def receive_websocket_data(receive: WebsocketReceiveCallable):
-    res = b""
-    with tracer.start_as_current_span("read websocket message") as s:
-        msg = await receive()
+    await send(
+        WebsocketAcceptEvent(
+            type="websocket.accept", subprotocol=subprotocol, headers=headers_to_send
+        )
+    )
 
-        s.set_attribute("type", msg.type)
 
-        if msg.type == "websocket.connect":
-            # todo(ayush): allow upgrades here as well?
-            raise WebsocketBadMessage("connection has already been established")
+@trace_function_with_span(tracer)
+async def close_connection(
+    s: Span, send: WebsocketSendCallable, /, *, status: WebsocketStatus, data: str
+) -> None:
+    s.set_attributes({"status": status.name, "reason": data})
+
+    await send(
+        WebsocketCloseEvent(type="websocket.close", code=status.value, reason=data)
+    )
+
+    current_websocket_session_span().set_attribute("websocket.close_reason", data)
 
-        if msg.type == "websocket.disconnect":
-            raise WebsocketConnectionClosedError()
 
-        if msg.bytes is not None:
-            res = msg.bytes
-        elif msg.text is not None:
-            res = msg.text.encode("utf-8")
-        else:
-            raise WebsocketBadMessage("empty message")
+# >>> Receive
 
-        s.set_attributes({"size": len(res)})
-        return res
+# todo(maximsmol): add max message length limit by default
 
 
 @trace_function(tracer)
-async def receive_websocket_json(receive: WebsocketReceiveCallable) -> Any:
-    res = await receive_websocket_data(receive)
+async def receive_data(receive: WebsocketReceiveCallable) -> bytes | str:
+    msg = await receive()
 
-    p = simdjson.Parser()
-    try:
-        return p.parse(res, True)
-    except ValueError as e:
-        raise WebsocketBadMessage("Failed to parse JSON") from e
+    if msg["type"] == "websocket.connect":
+        # todo(ayush): allow upgrades here as well?
+        raise ValueError("ASGI protocol violation: duplicate websocket.connect event")
+
+    if msg["type"] == "websocket.disconnect":
+        raise WebsocketConnectionClosedError(WebsocketStatus(msg["code"]))
+
+    if msg["bytes"] is not None:
+        res = msg["bytes"]
+    elif msg["text"] is not None:
+        res = msg["text"]
+    else:
+        raise WebsocketBadMessage("empty message")
+
+    return res
 
 
 @trace_function(tracer)
-async def receive_websocket_class_ext(
+async def receive_json(receive: WebsocketReceiveCallable) -> Any:
+    return orjson.loads(await receive_data(receive))
+
+
+@trace_function(tracer)
+async def receive_class_ext(
     receive: WebsocketReceiveCallable, cls: type[T]
 ) -> tuple[Any, T]:
-    data = await receive_websocket_json(receive)
+    data = await receive_json(receive)
 
     try:
         return data, validate(data, cls)
@@ -232,68 +244,28 @@ async def receive_websocket_class_ext(
 
 @trace_function(tracer)
 async def receive_websocket_class(receive: WebsocketReceiveCallable, cls: type[T]) -> T:
-    return (await receive_websocket_class_ext(receive, cls))[1]
-
+    return (await receive_class_ext(receive, cls))[1]
 
-@trace_function_with_span(tracer)
-async def send_websocket_data(
-    s: Span,
-    send: WebsocketSendCallable,
-    data: str | bytes,
-):
-    if isinstance(data, str):
-        data = data.encode("utf-8")
 
-    s.set_attribute("body.size", len(data))
+# >>> Send
 
-    await send(WebsocketSendEvent(type="websocket.send", bytes=data, text=None))
-
-    current_websocket_request_span().set_attribute(
-        "websocket.sent_message_content_length", len(data)
-    )
-
-
-@trace_function_with_span(tracer)
-async def accept_websocket_connection(
-    s: Span,
-    send: WebsocketSendCallable,
-    receive: WebsocketReceiveCallable,
-    /,
-    *,
-    subprotocol: str | None = None,
-    headers: Headers = {},
-):
-    msg = await receive()
-    if msg.type != "websocket.connect":
-        raise WebsocketBadMessage("cannot accept connection without connection request")
 
-    headers_to_send: list[tuple[bytes, bytes]] = []
-
-    for k, v in headers.items():
-        if isinstance(k, str):
-            k = k.encode("latin-1")
-        if isinstance(v, str):
-            v = v.encode("latin-1")
-        headers_to_send.append((k, v))
+@trace_function(tracer)
+async def send_data(send: WebsocketSendCallable, data: str | bytes, /) -> None:
+    if isinstance(data, bytes):
+        await send(WebsocketSendEvent(type="websocket.send", bytes=data, text=None))
+    else:
+        await send(WebsocketSendEvent(type="websocket.send", bytes=None, text=data))
 
-    await send(
-        WebsocketAcceptEvent(
-            type="websocket.accept", subprotocol=subprotocol, headers=headers_to_send
-        )
-    )
 
+@trace_function(tracer)
+async def send_json(send: WebsocketSendCallable, data: Any, /) -> None:
+    return await send_data(send, orjson.dumps(data))
 
-@trace_function_with_span(tracer)
-async def close_websocket_connection(
-    s: Span,
-    send: WebsocketSendCallable,
-    /,
-    *,
-    status: WebsocketStatus,
-    data: str,
-):
-    s.set_attribute("reason", data)
 
-    await send(WebsocketCloseEvent("websocket.close", status.value, data))
+@trace_function(tracer)
+async def send_websocket_auto(send: WebsocketSendCallable, data: Any, /) -> None:
+    if isinstance(data, str | bytes):
+        return await send_data(send, data)
 
-    current_websocket_request_span().set_attribute("websocket.http.close_reason", data)
+    return await send_json(send, data)