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

latch_asgi/asgi_iface.py

@@ -1,196 +1,47 @@
-import dataclasses
-from dataclasses import dataclass
-from typing import (
-    Awaitable,
-    Callable,
-    Iterable,
-    Literal,
-    TypeAlias,
-    get_args,
-    get_origin,
+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,
 )
 
-import hypercorn.typing as htyping
-from latch_data_validation.data_validation import validate
-
-
-def type_str(x: type) -> str:
-    for f in dataclasses.fields(x):
-        if f.name != "type":
-            continue
-
-        o = get_origin(f.type)
-        if o is not Literal:
-            raise ValueError("'type' field type is not a Literal")
-
-        res = get_args(f.type)[0]
-        if not isinstance(res, str):
-            raise ValueError("'type' field Literal is not a string")
-
-        return res
-
-    raise ValueError("'type' field not found")
-
-
-@dataclass(frozen=True)
-class ASGIVersions:
-    spec_version: str
-    version: Literal["2.0"] | Literal["3.0"]
-
-    def as_dict(self):
-        return validate(dataclasses.asdict(self), htyping.ASGIVersions)
-
-
 # >>> Lifespan
-@dataclass(frozen=True)
-class LifespanScope:
-    type: Literal["lifespan"]
-    asgi: ASGIVersions
-
-    def as_dict(self):
-        return validate(dataclasses.asdict(self), htyping.LifespanScope)
-
-
-@dataclass(frozen=True)
-class LifespanStartupEvent:
-    type: Literal["lifespan.startup"]
-
-    def as_dict(self):
-        return validate(dataclasses.asdict(self), htyping.LifespanStartupEvent)
-
-
-@dataclass(frozen=True)
-class LifespanShutdownEvent:
-    type: Literal["lifespan.shutdown"]
-
-    def as_dict(self):
-        return validate(dataclasses.asdict(self), htyping.LifespanShutdownEvent)
-
-
 LifespanReceiveEvent: TypeAlias = LifespanStartupEvent | LifespanShutdownEvent
 LifespanReceiveCallable: TypeAlias = Callable[[], Awaitable[LifespanReceiveEvent]]
 
-
-@dataclass(frozen=True)
-class LifespanStartupCompleteEvent:
-    type: Literal["lifespan.startup.complete"]
-
-    def as_dict(self):
-        return validate(dataclasses.asdict(self), htyping.LifespanStartupCompleteEvent)
-
-
-@dataclass(frozen=True)
-class LifespanStartupFailedEvent:
-    type: Literal["lifespan.startup.failed"]
-    message: str
-
-    def as_dict(self):
-        return validate(dataclasses.asdict(self), htyping.LifespanStartupFailedEvent)
-
-
-LifespanStartupSendEvent: TypeAlias = (
-    LifespanStartupCompleteEvent | LifespanStartupFailedEvent
-)
-
-
-@dataclass(frozen=True)
-class LifespanShutdownCompleteEvent:
-    type: Literal["lifespan.shutdown.complete"]
-
-    def as_dict(self):
-        return validate(dataclasses.asdict(self), htyping.LifespanShutdownCompleteEvent)
-
-
-@dataclass(frozen=True)
-class LifespanShutdownFailedEvent:
-    type: Literal["lifespan.shutdown.failed"]
-    message: str
-
-    def as_dict(self):
-        return validate(dataclasses.asdict(self), htyping.LifespanShutdownFailedEvent)
-
-
 LifespanShutdownSendEvent: TypeAlias = (
     LifespanShutdownCompleteEvent | LifespanShutdownFailedEvent
 )
+LifespanStartupSendEvent: TypeAlias = (
+    LifespanStartupCompleteEvent | LifespanStartupFailedEvent
+)
 
 LifespanSendEvent: TypeAlias = LifespanStartupSendEvent | LifespanShutdownSendEvent
 LifespanSendCallable: TypeAlias = Callable[[LifespanSendEvent], Awaitable[None]]
 
-
 # >>> HTTP
-@dataclass(frozen=True)
-class HTTPScope:
-    type: Literal["http"]
-    asgi: ASGIVersions
-    http_version: str
-    method: str
-    scheme: str
-    path: str
-    raw_path: bytes
-    query_string: bytes
-    root_path: str
-    headers: Iterable[tuple[bytes, bytes]]
-    client: tuple[str, int] | None
-    server: tuple[str, int | None] | None
-    extensions: dict[str, dict]
-
-    def as_dict(self):
-        return validate(dataclasses.asdict(self), htyping.HTTPScope)
-
-
-@dataclass(frozen=True)
-class HTTPRequestEvent:
-    type: Literal["http.request"]
-    body: bytes
-    more_body: bool
-
-    def as_dict(self):
-        return validate(dataclasses.asdict(self), htyping.HTTPRequestEvent)
-
-
-@dataclass(frozen=True)
-class HTTPDisconnectEvent:
-    type: Literal["http.disconnect"]
-
-    def as_dict(self):
-        return validate(dataclasses.asdict(self), htyping.HTTPDisconnectEvent)
-
 
 HTTPReceiveEvent: TypeAlias = HTTPRequestEvent | HTTPDisconnectEvent
 HTTPReceiveCallable: TypeAlias = Callable[[], Awaitable[HTTPReceiveEvent]]
 
-
-@dataclass(frozen=True)
-class HTTPResponseStartEvent:
-    type: Literal["http.response.start"]
-    status: int
-    headers: Iterable[tuple[bytes, bytes]]
-
-    def as_dict(self):
-        return validate(dataclasses.asdict(self), htyping.HTTPResponseStartEvent)
-
-
-@dataclass(frozen=True)
-class HTTPResponseBodyEvent:
-    type: Literal["http.response.body"]
-    body: bytes
-    more_body: bool
-
-    def as_dict(self):
-        return validate(dataclasses.asdict(self), htyping.HTTPResponseBodyEvent)
-
-
-@dataclass(frozen=True)
-class HTTPServerPushEvent:
-    type: Literal["http.response.push"]
-    path: str
-    headers: Iterable[tuple[bytes, bytes]]
-
-    def as_dict(self):
-        return validate(dataclasses.asdict(self), htyping.HTTPServerPushEvent)
-
-
 HTTPSendEvent: TypeAlias = (
     HTTPResponseStartEvent
     | HTTPResponseBodyEvent
@@ -199,104 +50,12 @@ HTTPSendEvent: TypeAlias = (
 )
 HTTPSendCallable: TypeAlias = Callable[[HTTPSendEvent], Awaitable[None]]
 
-
 # >>> Websocket
-@dataclass(frozen=True)
-class WebsocketScope:
-    type: Literal["websocket"]
-    asgi: ASGIVersions
-    http_version: str
-    scheme: str
-    path: str
-    raw_path: bytes
-    query_string: bytes
-    root_path: str
-    headers: Iterable[tuple[bytes, bytes]]
-    client: tuple[str, int] | None
-    server: tuple[str, int | None] | None
-    subprotocols: Iterable[str]
-    extensions: dict[str, object]
-
-    def as_dict(self):
-        return validate(dataclasses.asdict(self), htyping.WebsocketScope)
-
-
-@dataclass(frozen=True)
-class WebsocketConnectEvent:
-    type: Literal["websocket.connect"]
-
-    def as_dict(self):
-        return validate(dataclasses.asdict(self), htyping.WebsocketConnectEvent)
-
-
-@dataclass(frozen=True)
-class WebsocketAcceptEvent:
-    type: Literal["websocket.accept"]
-    subprotocol: str | None
-    headers: Iterable[tuple[bytes, bytes]]
-
-    def as_dict(self):
-        return validate(dataclasses.asdict(self), htyping.WebsocketAcceptEvent)
-
-
-@dataclass(frozen=True)
-class WebsocketReceiveEvent:
-    type: Literal["websocket.receive"]
-    bytes: bytes | None
-    text: str | None
-
-    def as_dict(self):
-        return validate(dataclasses.asdict(self), htyping.WebsocketReceiveEvent)
-
-
-@dataclass(frozen=True)
-class WebsocketSendEvent:
-    type: Literal["websocket.send"]
-    bytes: bytes | None
-    text: str | None
-
-    def as_dict(self):
-        return validate(dataclasses.asdict(self), htyping.WebsocketSendEvent)
-
-
-@dataclass(frozen=True)
-class WebsocketResponseStartEvent:
-    type: Literal["websocket.http.response.start"]
-    status: int
-    headers: Iterable[tuple[bytes, bytes]]
-
-    def as_dict(self):
-        return validate(dataclasses.asdict(self), htyping.WebsocketResponseStartEvent)
-
-
-@dataclass(frozen=True)
-class WebsocketResponseBodyEvent:
-    type: Literal["websocket.http.response.body"]
-    body: bytes
-    more_body: bool
-
-    def as_dict(self):
-        return validate(dataclasses.asdict(self), htyping.WebsocketResponseBodyEvent)
-
-
-@dataclass(frozen=True)
-class WebsocketDisconnectEvent:
-    type: Literal["websocket.disconnect"]
-    code: int
-
-    def as_dict(self):
-        return validate(dataclasses.asdict(self), htyping.WebsocketDisconnectEvent)
-
-
-@dataclass(frozen=True)
-class WebsocketCloseEvent:
-    type: Literal["websocket.close"]
-    code: int
-    reason: str | None
-
-    def as_dict(self):
-        return validate(dataclasses.asdict(self), htyping.WebsocketCloseEvent)
 
+WebsocketReceiveEventT: TypeAlias = (
+    WebsocketConnectEvent | WebsocketReceiveEvent | WebsocketDisconnectEvent
+)
+WebsocketReceiveCallable: TypeAlias = Callable[[], Awaitable[WebsocketReceiveEventT]]
 
 WebsocketSendEventT: TypeAlias = (
     WebsocketAcceptEvent
@@ -307,22 +66,7 @@ WebsocketSendEventT: TypeAlias = (
 )
 WebsocketSendCallable: TypeAlias = Callable[[WebsocketSendEventT], Awaitable[None]]
 
+# >>> WWW
 
-WebsocketReceiveEventT: TypeAlias = (
-    WebsocketConnectEvent | WebsocketReceiveEvent | WebsocketDisconnectEvent
-)
-WebsocketReceiveCallable: TypeAlias = Callable[[], Awaitable[WebsocketReceiveEventT]]
-
-
-WWWScope: TypeAlias = HTTPScope | WebsocketScope
-Scope: TypeAlias = HTTPScope | WebsocketScope | LifespanScope
-
-WWWSendCallable: TypeAlias = HTTPSendCallable | WebsocketSendCallable
-SendCallable: TypeAlias = (
-    HTTPSendCallable | WebsocketSendCallable | LifespanSendCallable
-)
-
-WWWReceiveCallable: TypeAlias = HTTPReceiveCallable | WebsocketReceiveCallable
-ReceiveCallable: TypeAlias = (
-    HTTPReceiveCallable | WebsocketReceiveCallable | LifespanReceiveCallable
-)
+WWWReceiveCallable = HTTPReceiveCallable | WebsocketReceiveCallable
+WWWSendCallable = HTTPSendCallable | WebsocketSendCallable

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/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/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/context/http.py

@@ -1,10 +1,12 @@
+from collections.abc import Awaitable, Callable
 from dataclasses import dataclass
-from typing import Any, Awaitable, Callable, TypeAlias, TypeVar
+from typing import Any, Self, TypeAlias, TypeVar
 
-from latch_o11y.o11y import AttributesDict, dict_to_attrs, trace_app_function
+from hypercorn.typing import HTTPScope
+from latch_o11y.o11y import trace_app_function
 
-from ..asgi_iface import HTTPReceiveCallable, HTTPScope, HTTPSendCallable
-from ..framework.http import HTTPMethod, current_http_request_span, receive_class_ext
+from ..asgi_iface import HTTPReceiveCallable, HTTPSendCallable
+from ..framework.http import HTTPMethod, receive_class_ext
 from . import common
 
 T = TypeVar("T")
@@ -12,31 +14,15 @@ T = TypeVar("T")
 
 @dataclass
 class Context(common.Context[HTTPScope, HTTPReceiveCallable, HTTPSendCallable]):
-    def __post_init__(self):
-        super().__post_init__()
-
-        if self.auth.oauth_sub is not None:
-            current_http_request_span().set_attribute("enduser.id", self.auth.oauth_sub)
-
-    def add_request_span_attrs(self, data: AttributesDict, prefix: str):
-        current_http_request_span().set_attributes(dict_to_attrs(data, prefix))
-
     @trace_app_function
-    async def receive_request_payload(self, cls: type[T]) -> T:
+    async def receive_request_payload(self: Self, cls: type[T]) -> T:
         json, res = await receive_class_ext(self.receive, cls)
 
-        # todo(maximsmol): datadog has shit support for events
-        # current_http_request_span().add_event(
-        #     "request payload", dict_to_attrs(json, "data")
-        # )
-        self.add_request_span_attrs(json, "http.request_payload")
+        self.add_request_span_attrs(json, "http.request.body.data")
 
         return res
 
 
 HandlerResult: TypeAlias = Any | None
-Handler: TypeAlias = Callable[
-    [Context],
-    Awaitable[HandlerResult],
-]
+Handler: TypeAlias = Callable[[Context], Awaitable[HandlerResult]]
 Route: TypeAlias = Handler | tuple[list[HTTPMethod], Handler]

latch_asgi/context/websocket.py

@@ -1,10 +1,19 @@
+from collections.abc import Awaitable, Callable
 from dataclasses import dataclass
-from typing import Awaitable, Callable, TypeAlias, TypeVar
-
-from latch_o11y.o11y import AttributesDict, dict_to_attrs
-
-from ..asgi_iface import WebsocketReceiveCallable, WebsocketScope, WebsocketSendCallable
-from ..framework.websocket import WebsocketStatus, current_websocket_request_span
+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")
@@ -14,21 +23,28 @@ T = TypeVar("T")
 class Context(
     common.Context[WebsocketScope, WebsocketReceiveCallable, WebsocketSendCallable]
 ):
-    def __post_init__(self):
-        super().__post_init__()
+    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"))
 
-        if self.auth.oauth_sub is not None:
-            current_websocket_request_span().set_attribute(
-                "enduser.id", self.auth.oauth_sub
-            )
+        return res
 
-    def add_request_span_attrs(self, data: AttributesDict, prefix: str):
-        current_websocket_request_span().set_attributes(dict_to_attrs(data, prefix))
+    @trace_app_function
+    async def send_message(self: Self, data: Any) -> None:
+        await send_websocket_auto(self.send, data)
 
 
-HandlerResult = str | tuple[WebsocketStatus, str]
-Handler: TypeAlias = Callable[
-    [Context],
-    Awaitable[HandlerResult],
-]
+HandlerResult = str
+Handler: TypeAlias = Callable[[Context], Awaitable[HandlerResult]]
 Route: TypeAlias = Handler