latch-asgi 0.3.0__tar.gz → 1.0.0__tar.gz

{latch_asgi-0.3.0 → latch_asgi-1.0.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: latch-asgi
-Version: 0.3.0
+Version: 1.0.0
 Summary: ASGI python server
 Author-Email: Max Smolin <max@latch.bio>
 License: CC0-1.0
@@ -13,7 +13,6 @@ Requires-Dist: PyJWT[crypto]<3.0.0,>=2.6.0
 Requires-Dist: orjson<4.0.0,>=3.8.5
 Requires-Dist: opentelemetry-sdk<2.0.0,>=1.15.0
 Requires-Dist: opentelemetry-api<2.0.0,>=1.15.0
-Requires-Dist: opentelemetry-instrumentation-asgi<1.0,>=0.36b0
 Description-Content-Type: text/markdown
 
 # python-asgi

latch_asgi-1.0.0/latch_asgi/asgi_iface.py

@@ -0,0 +1,72 @@
+from collections.abc import Awaitable, Callable
+from typing import TypeAlias
+
+from hypercorn.typing import (
+    HTTPDisconnectEvent,
+    HTTPRequestEvent,
+    HTTPResponseBodyEvent,
+    HTTPResponseStartEvent,
+    HTTPServerPushEvent,
+    LifespanShutdownCompleteEvent,
+    LifespanShutdownEvent,
+    LifespanShutdownFailedEvent,
+    LifespanStartupCompleteEvent,
+    LifespanStartupEvent,
+    LifespanStartupFailedEvent,
+    WebsocketAcceptEvent,
+    WebsocketCloseEvent,
+    WebsocketConnectEvent,
+    WebsocketDisconnectEvent,
+    WebsocketReceiveEvent,
+    WebsocketResponseBodyEvent,
+    WebsocketResponseStartEvent,
+    WebsocketSendEvent,
+)
+
+# >>> Lifespan
+LifespanReceiveEvent: TypeAlias = LifespanStartupEvent | LifespanShutdownEvent
+LifespanReceiveCallable: TypeAlias = Callable[[], Awaitable[LifespanReceiveEvent]]
+
+LifespanShutdownSendEvent: TypeAlias = (
+    LifespanShutdownCompleteEvent | LifespanShutdownFailedEvent
+)
+LifespanStartupSendEvent: TypeAlias = (
+    LifespanStartupCompleteEvent | LifespanStartupFailedEvent
+)
+
+LifespanSendEvent: TypeAlias = LifespanStartupSendEvent | LifespanShutdownSendEvent
+LifespanSendCallable: TypeAlias = Callable[[LifespanSendEvent], Awaitable[None]]
+
+# >>> HTTP
+
+HTTPReceiveEvent: TypeAlias = HTTPRequestEvent | HTTPDisconnectEvent
+HTTPReceiveCallable: TypeAlias = Callable[[], Awaitable[HTTPReceiveEvent]]
+
+HTTPSendEvent: TypeAlias = (
+    HTTPResponseStartEvent
+    | HTTPResponseBodyEvent
+    | HTTPServerPushEvent
+    | HTTPDisconnectEvent
+)
+HTTPSendCallable: TypeAlias = Callable[[HTTPSendEvent], Awaitable[None]]
+
+# >>> Websocket
+
+WebsocketReceiveEventT: TypeAlias = (
+    WebsocketConnectEvent | WebsocketReceiveEvent | WebsocketDisconnectEvent
+)
+WebsocketReceiveCallable: TypeAlias = Callable[[], Awaitable[WebsocketReceiveEventT]]
+
+WebsocketSendEventT: TypeAlias = (
+    WebsocketAcceptEvent
+    | WebsocketSendEvent
+    | WebsocketResponseBodyEvent
+    | WebsocketResponseStartEvent
+    | WebsocketCloseEvent
+)
+WebsocketSendCallable: TypeAlias = Callable[[WebsocketSendEventT], Awaitable[None]]
+
+# >>> WWW
+
+WWWReceiveCallable = HTTPReceiveCallable | WebsocketReceiveCallable
+WWWSendCallable = HTTPSendCallable | WebsocketSendCallable

{latch_asgi-0.3.0 → latch_asgi-1.0.0}/latch_asgi/auth.py

@@ -3,7 +3,7 @@
 import re
 from dataclasses import dataclass
 from http import HTTPStatus
-from typing import Literal
+from typing import Literal, Self
 
 import jwt
 from jwt import PyJWKClient
@@ -37,14 +37,10 @@ class _HTTPUnauthorized(HTTPErrorResponse):
     """WARNING: HTTPForbidden is the correct error to use in virtually all cases"""
 
     def __init__(
-        self,
+        self: Self,
         error_description: str,
-        error: (
-            Literal["invalid_request"]
-            | Literal["invalid_token"]
-            | Literal["insufficient_scope"]
-        ),
-    ):
+        error: (Literal["invalid_request", "invalid_token", "insufficient_scope"]),
+    ) -> None:
         escaped_description = error_description.replace('"', '\\"')
         super().__init__(
             HTTPStatus.UNAUTHORIZED,
@@ -63,7 +59,7 @@ class Authorization:
     execution_token: str | None = None
     sdk_token: str | None = None
 
-    def unauthorized_if_none(self):
+    def unauthorized_if_none(self: Self) -> None:
         if self.oauth_sub is not None:
             return
         if self.execution_token is not None:
@@ -71,10 +67,7 @@ class Authorization:
         if self.sdk_token is not None:
             return
 
-        raise _HTTPUnauthorized(
-            "Authenticaton required",
-            error="invalid_request",
-        )
+        raise _HTTPUnauthorized("Authenticaton required", error="invalid_request")
 
 
 @trace_app_function
@@ -121,8 +114,7 @@ def get_signer_sub(auth_header: str) -> Authorization:
             jwt_key = jwk_client.get_signing_key_from_jwt(oauth_token).key
         except jwt.exceptions.InvalidTokenError as e:
             raise _HTTPUnauthorized(
-                error_description="JWT decoding failed",
-                error="invalid_token",
+                error_description="JWT decoding failed", error="invalid_token"
             ) from e
         except jwt.exceptions.PyJWKClientError:
             # fixme(maximsmol): gut this abomination

{latch_asgi-0.3.0 → latch_asgi-1.0.0}/latch_asgi/config.py

@@ -1,4 +1,5 @@
 from dataclasses import dataclass
+
 from latch_config.config import read_config
 
 
@@ -8,4 +9,5 @@ class AuthConfig:
     self_signed_jwk: str
     allow_spoofing: bool = False
 
+
 config = read_config(AuthConfig, "auth_")

{latch_asgi-0.3.0 → latch_asgi-1.0.0}/latch_asgi/context/common.py

@@ -1,61 +1,72 @@
 from dataclasses import dataclass, field
-from typing import Generic, TypeVar
-
-from latch_o11y.o11y import AttributesDict, app_tracer, trace_app_function
-
-from ..asgi_iface import WWWReceiveCallable, WWWScope, WWWSendCallable
+from typing import Generic, Self, TypeVar
+
+from hypercorn.typing import WWWScope
+from latch_o11y.o11y import (
+    AttributesDict,
+    app_tracer,
+    dict_to_attrs,
+    trace_app_function,
+)
+from opentelemetry.util.types import AttributeValue
+
+from ..asgi_iface import WWWReceiveCallable, WWWSendCallable
 from ..auth import Authorization, get_signer_sub
+from ..framework.common import otel_header_whitelist
+from ..framework.http import current_http_request_span
 
-# todo(ayush): this sucks
 Scope = TypeVar("Scope", bound=WWWScope)
-Send = TypeVar("Send", bound=WWWSendCallable)
-Receive = TypeVar("Receive", bound=WWWReceiveCallable)
+SendCallable = TypeVar("SendCallable", bound=WWWSendCallable)
+ReceiveCallable = TypeVar("ReceiveCallable", bound=WWWReceiveCallable)
 
 
 @dataclass
-class Context(Generic[Scope, Receive, Send]):
+class Context(Generic[Scope, ReceiveCallable, SendCallable]):
     scope: Scope
-    receive: Receive
-    send: Send
+    receive: ReceiveCallable
+    send: SendCallable
 
     auth: Authorization = field(default_factory=Authorization, init=False)
 
     _header_cache: dict[bytes, bytes] = field(default_factory=dict, init=False)
     _db_response_idx: int = field(default=0, init=False)
 
-    def __post_init__(self):
+    def __post_init__(self: Self) -> None:
         with app_tracer.start_as_current_span("find Authentication header"):
             auth_header = self.header_str("authorization")
 
         if auth_header is not None:
             self.auth = get_signer_sub(auth_header)
 
-    def header(self, x: str | bytes):
+        if self.auth.oauth_sub is not None:
+            current_http_request_span().set_attribute("enduser.id", self.auth.oauth_sub)
+
+    def header(self: Self, x: str | bytes) -> bytes | None:
         if isinstance(x, str):
-            x = x.encode("utf-8")
+            x = x.encode("latin-1")
 
         if x in self._header_cache:
             return self._header_cache[x]
 
-        for k, v in self.scope.headers:
+        for k, v in self.scope["headers"]:
             self._header_cache[k] = v
             if k == x:
                 return v
 
         return None
 
-    def header_str(self, x: str | bytes):
+    def header_str(self: Self, x: str | bytes) -> str | None:
         res = self.header(x)
         if res is None:
             return None
 
         return res.decode("latin-1")
 
-    def add_request_span_attrs(self, data: AttributesDict, prefix: str):
-        raise NotImplementedError()
+    def add_request_span_attrs(self: Self, data: AttributesDict, prefix: str) -> None:
+        current_http_request_span().set_attributes(dict_to_attrs(data, prefix))
 
     @trace_app_function
-    def add_db_response(self, data: AttributesDict):
+    def add_db_response(self: Self, data: AttributesDict) -> None:
         # todo(maximsmol): datadog has shit support for events
         # current_http_request_span().add_event(
         #     f"database response {self._db_response_idx}", dict_to_attrs(data, "data")

latch_asgi-1.0.0/latch_asgi/context/http.py

@@ -0,0 +1,28 @@
+from collections.abc import Awaitable, Callable
+from dataclasses import dataclass
+from typing import Any, Self, TypeAlias, TypeVar
+
+from hypercorn.typing import HTTPScope
+from latch_o11y.o11y import trace_app_function
+
+from ..asgi_iface import HTTPReceiveCallable, HTTPSendCallable
+from ..framework.http import HTTPMethod, receive_class_ext
+from . import common
+
+T = TypeVar("T")
+
+
+@dataclass
+class Context(common.Context[HTTPScope, HTTPReceiveCallable, HTTPSendCallable]):
+    @trace_app_function
+    async def receive_request_payload(self: Self, cls: type[T]) -> T:
+        json, res = await receive_class_ext(self.receive, cls)
+
+        self.add_request_span_attrs(json, "http.request.body.data")
+
+        return res
+
+
+HandlerResult: TypeAlias = Any | None
+Handler: TypeAlias = Callable[[Context], Awaitable[HandlerResult]]
+Route: TypeAlias = Handler | tuple[list[HTTPMethod], Handler]

latch_asgi-1.0.0/latch_asgi/context/websocket.py

@@ -0,0 +1,50 @@
+from collections.abc import Awaitable, Callable
+from dataclasses import dataclass
+from typing import Any, Self, TypeAlias, TypeVar
+
+from hypercorn.typing import WebsocketScope
+from latch_o11y.o11y import AttributesDict, dict_to_attrs, trace_app_function
+from opentelemetry.trace import get_current_span
+
+from ..asgi_iface import WebsocketReceiveCallable, WebsocketSendCallable
+from ..framework.common import Headers
+from ..framework.websocket import (
+    accept_connection,
+    current_websocket_session_span,
+    receive_class_ext,
+    send_websocket_auto,
+)
+from . import common
+
+T = TypeVar("T")
+
+
+@dataclass
+class Context(
+    common.Context[WebsocketScope, WebsocketReceiveCallable, WebsocketSendCallable]
+):
+    def add_session_span_attrs(self: Self, data: AttributesDict, prefix: str) -> None:
+        current_websocket_session_span().set_attributes(dict_to_attrs(data, prefix))
+
+    @trace_app_function
+    async def accept_connection(
+        self: Self, *, subprotocol: str | None = None, headers: Headers | None = None
+    ) -> None:
+        await accept_connection(self.send, subprotocol=subprotocol, headers=headers)
+
+    @trace_app_function
+    async def receive_message(self: Self, cls: type[T]) -> T:
+        json, res = await receive_class_ext(self.receive, cls)
+
+        get_current_span().set_attributes(dict_to_attrs(json, "payload"))
+
+        return res
+
+    @trace_app_function
+    async def send_message(self: Self, data: Any) -> None:
+        await send_websocket_auto(self.send, data)
+
+
+HandlerResult = str
+Handler: TypeAlias = Callable[[Context], Awaitable[HandlerResult]]
+Route: TypeAlias = Handler

{latch_asgi-0.3.0 → latch_asgi-1.0.0}/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-1.0.0/latch_asgi/framework/common.py

@@ -0,0 +1,40 @@
+from typing import TypeAlias
+
+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-0.3.0 → latch_asgi-1.0.0}/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)