latch-asgi 0.3.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,10 +1,10 @@
 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 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 ..asgi_iface import (
@@ -17,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
 
@@ -42,34 +34,68 @@ def current_http_request_span() -> Span:
 
 
 class HTTPErrorResponse(RuntimeError):
-    def __init__(self, status: HTTPStatus, data: Any, *, headers: Headers = {}):
+    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]
@@ -87,47 +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:
-    return orjson.loads(await receive_data(receive))
-
-
-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"))
     ]
@@ -152,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(
@@ -163,8 +165,8 @@ async def send_json(
     /,
     *,
     content_type: str = "application/json",
-    headers: Headers = {},
-):
+    headers: Headers | None = None,
+) -> None:
     return await send_http_data(
         send, status, orjson.dumps(data), content_type=content_type, headers=headers
     )
@@ -177,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 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
-import orjson
+from opentelemetry.util.types import AttributeValue
 
 from ..asgi_iface import (
     WebsocketAcceptEvent,
@@ -18,20 +19,18 @@ 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
-    """
+    """https://www.rfc-editor.org/rfc/rfc6455.html#section-7.4.1"""
 
     normal = 1000
     """
@@ -134,63 +133,108 @@ class WebsocketStatus(int, Enum):
 
 
 class WebsocketErrorResponse(RuntimeError):
-    def __init__(self, status: WebsocketStatus, data: Any, *, headers: Headers = {}):
-        super().__init__()
+    def __init__(self: Self, status: WebsocketStatus, data: Any) -> None:
         self.status = status
         self.data = data
-        self.headers = headers
 
 
 class WebsocketBadMessage(WebsocketErrorResponse):
-    def __init__(self, data: Any, *, headers: Headers = {}):
-        super().__init__(WebsocketStatus.policy_violation, data, headers=headers)
+    def __init__(self: Self, data: Any) -> None:
+        super().__init__(WebsocketStatus.policy_violation, data)
 
 
 class WebsocketInternalServerError(WebsocketErrorResponse):
-    def __init__(self, data: Any, *, headers: Headers = {}):
-        super().__init__(WebsocketStatus.server_error, data, headers=headers)
+    def __init__(self: Self, data: Any) -> None:
+        super().__init__(WebsocketStatus.server_error, data)
+
+
+class WebsocketConnectionClosedError(RuntimeError):
+    def __init__(self: Self, code: WebsocketStatus) -> None:
+        self.code = code
+
+
+# >>>
+# I/O
+# >>>
+
+# >>> Send Lifecycle
+
+
+@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))
 
+    await send(
+        WebsocketAcceptEvent(
+            type="websocket.accept", subprotocol=subprotocol, headers=headers_to_send
+        )
+    )
+
+
+@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)
+    )
 
-class WebsocketConnectionClosedError(RuntimeError): ...
+    current_websocket_session_span().set_attribute("websocket.close_reason", data)
 
 
-# >>> I/O
+# >>> Receive
 
+# todo(maximsmol): add max message length limit by default
 
-async def receive_websocket_data(receive: WebsocketReceiveCallable):
-    with tracer.start_as_current_span("read websocket message") as s:
-        msg = await receive()
 
-        s.set_attribute("type", msg.type)
+@trace_function(tracer)
+async def receive_data(receive: WebsocketReceiveCallable) -> bytes | str:
+    msg = await receive()
 
-        if msg.type == "websocket.connect":
-            # todo(ayush): allow upgrades here as well?
-            raise WebsocketBadMessage("connection has already been established")
+    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()
+    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.encode("utf-8")
-        else:
-            raise WebsocketBadMessage("empty message")
+    if msg["bytes"] is not None:
+        res = msg["bytes"]
+    elif msg["text"] is not None:
+        res = msg["text"]
+    else:
+        raise WebsocketBadMessage("empty message")
 
-        s.set_attribute("size", len(res))
-        return res
+    return res
 
 
 @trace_function(tracer)
-async def receive_websocket_json(receive: WebsocketReceiveCallable) -> Any:
-    orjson.loads(await receive_websocket_data(receive))
+async def receive_json(receive: WebsocketReceiveCallable) -> Any:
+    return orjson.loads(await receive_data(receive))
 
 
 @trace_function(tracer)
-async def receive_websocket_class_ext(
+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)
@@ -200,67 +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))
-
-    await send(WebsocketSendEvent(type="websocket.send", bytes=data, text=None))
-
-    current_websocket_request_span().set_attribute(
-        "websocket.sent_message_content_length", len(data)
-    )
+# >>> Send
 
 
 @trace_function(tracer)
-async def accept_websocket_connection(
-    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))
+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)