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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: latch-asgi
-Version: 1.0.0
+Version: 1.0.1
 Summary: ASGI python server
 Author-Email: Max Smolin <max@latch.bio>
 License: CC0-1.0

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

@@ -1,5 +1,5 @@
 from dataclasses import dataclass, field
-from typing import Generic, Self, TypeVar
+from typing import ClassVar, Generic, Self, TypeVar, cast
 
 from hypercorn.typing import WWWScope
 from latch_o11y.o11y import (
@@ -8,12 +8,11 @@ from latch_o11y.o11y import (
     dict_to_attrs,
     trace_app_function,
 )
-from opentelemetry.util.types import AttributeValue
+from opentelemetry import context
+from opentelemetry.trace.span import Span
 
 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
 
 Scope = TypeVar("Scope", bound=WWWScope)
 SendCallable = TypeVar("SendCallable", bound=WWWSendCallable)
@@ -31,6 +30,8 @@ class Context(Generic[Scope, ReceiveCallable, SendCallable]):
     _header_cache: dict[bytes, bytes] = field(default_factory=dict, init=False)
     _db_response_idx: int = field(default=0, init=False)
 
+    _request_span_key: ClassVar[str]
+
     def __post_init__(self: Self) -> None:
         with app_tracer.start_as_current_span("find Authentication header"):
             auth_header = self.header_str("authorization")
@@ -39,7 +40,7 @@ class Context(Generic[Scope, ReceiveCallable, SendCallable]):
             self.auth = get_signer_sub(auth_header)
 
         if self.auth.oauth_sub is not None:
-            current_http_request_span().set_attribute("enduser.id", self.auth.oauth_sub)
+            self.current_request_span().set_attribute("enduser.id", self.auth.oauth_sub)
 
     def header(self: Self, x: str | bytes) -> bytes | None:
         if isinstance(x, str):
@@ -62,8 +63,11 @@ class Context(Generic[Scope, ReceiveCallable, SendCallable]):
 
         return res.decode("latin-1")
 
+    def current_request_span(self: Self) -> Span:
+        return cast(Span, context.get_value(self._request_span_key))
+
     def add_request_span_attrs(self: Self, data: AttributesDict, prefix: str) -> None:
-        current_http_request_span().set_attributes(dict_to_attrs(data, prefix))
+        self.current_request_span().set_attributes(dict_to_attrs(data, prefix))
 
     @trace_app_function
     def add_db_response(self: Self, data: AttributesDict) -> None:

{latch_asgi-1.0.0 → latch_asgi-1.0.1}/latch_asgi/context/http.py

@@ -6,7 +6,7 @@ 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 ..framework.http import HTTPMethod, http_request_span_key, receive_class_ext
 from . import common
 
 T = TypeVar("T")
@@ -14,6 +14,8 @@ T = TypeVar("T")
 
 @dataclass
 class Context(common.Context[HTTPScope, HTTPReceiveCallable, HTTPSendCallable]):
+    _request_span_key = http_request_span_key
+
     @trace_app_function
     async def receive_request_payload(self: Self, cls: type[T]) -> T:
         json, res = await receive_class_ext(self.receive, cls)

{latch_asgi-1.0.0 → latch_asgi-1.0.1}/latch_asgi/context/websocket.py

@@ -3,16 +3,16 @@ 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 latch_o11y.o11y import 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,
+    websocket_session_span_key,
 )
 from . import common
 
@@ -23,8 +23,7 @@ T = TypeVar("T")
 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))
+    _request_span_key = websocket_session_span_key
 
     @trace_app_function
     async def accept_connection(

{latch_asgi-1.0.0 → latch_asgi-1.0.1}/latch_asgi/framework/http.py

@@ -3,7 +3,7 @@ from typing import Any, Literal, Self, TypeAlias, TypeVar, cast
 
 import orjson
 from latch_data_validation.data_validation import DataValidationError, validate
-from latch_o11y.o11y import trace_function, trace_function_with_span
+from latch_o11y.o11y import trace_function
 from opentelemetry import context
 from opentelemetry.trace.span import Span
 

{latch_asgi-1.0.0 → latch_asgi-1.0.1}/latch_asgi/framework/websocket.py

@@ -6,7 +6,6 @@ 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,
@@ -188,15 +187,15 @@ async def accept_connection(
 
 @trace_function_with_span(tracer)
 async def close_connection(
-    s: Span, send: WebsocketSendCallable, /, *, status: WebsocketStatus, data: str
+    s: Span, send: WebsocketSendCallable, /, *, status: WebsocketStatus, reason: str
 ) -> None:
-    s.set_attributes({"status": status.name, "reason": data})
+    s.set_attributes({"status": status.name, "reason": reason})
 
     await send(
-        WebsocketCloseEvent(type="websocket.close", code=status.value, reason=data)
+        WebsocketCloseEvent(type="websocket.close", code=status.value, reason=reason)
     )
 
-    current_websocket_session_span().set_attribute("websocket.close_reason", data)
+    current_websocket_session_span().set_attribute("websocket.close_reason", reason)
 
 
 # >>> Receive

{latch_asgi-1.0.0 → latch_asgi-1.0.1}/latch_asgi/server.py

@@ -48,6 +48,7 @@ from .framework.http import (
     send_http_data,
 )
 from .framework.websocket import (
+    WebsocketConnectionClosedError,
     WebsocketErrorResponse,
     WebsocketInternalServerError,
     WebsocketStatus,
@@ -214,12 +215,13 @@ class LatchASGIServer:
             # todo(maximsmol): better error message
             await log.info("Not found")
             await close_connection(
-                send, status=WebsocketStatus.policy_violation, data="Not found"
+                send, status=WebsocketStatus.policy_violation, reason="Not found"
             )
             return
 
         s.set_attribute("http.route", scope["path"])
 
+        close_reason: str | None = None
         try:
             try:
                 msg = await receive()
@@ -229,24 +231,30 @@ class LatchASGIServer:
                     )
 
                 ctx = websocket.Context(scope, receive, send)
-                data = await handler(ctx)
+                close_reason = await handler(ctx)
             except WebsocketErrorResponse:
                 raise
+            except WebsocketConnectionClosedError as e:
+                s.set_attribute("websocket.close.code", e.code.name)
             except Exception as e:
                 # todo(maximsmol): better error message
                 raise WebsocketInternalServerError("Internal Error") from e
         except WebsocketErrorResponse as e:
             await close_connection(
-                send, status=e.status, data=orjson.dumps({"error": e.data}).decode()
+                send, status=e.status, reason=orjson.dumps({"error": e.data}).decode()
             )
 
             if e.status != WebsocketStatus.server_error:
                 return
 
-            # await log.exception() # fixme(maximsmol)
             raise
         else:
-            await close_connection(send, status=WebsocketStatus.normal, data=data)
+            if close_reason is None:
+                close_reason = "Session complete"
+
+            await close_connection(
+                send, status=WebsocketStatus.normal, reason=close_reason
+            )
 
     async def scope_http(
         self: Self,
@@ -286,12 +294,6 @@ class LatchASGIServer:
             )
             return
 
-        if handler is None:
-            # todo(maximsmol): better error message
-            await log.info("Not found")
-            await send_http_data(send, HTTPStatus.NOT_FOUND, "Not found")
-            return
-
         try:
             try:
                 ctx = http.Context(scope, receive, send)
@@ -311,7 +313,6 @@ class LatchASGIServer:
             if e.status != HTTPStatus.INTERNAL_SERVER_ERROR:
                 return
 
-            # await log.exception() # fixme(maximsmol)
             raise
 
     async def raw_app(

{latch_asgi-1.0.0 → latch_asgi-1.0.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "latch-asgi"
-version = "1.0.0"
+version = "1.0.1"
 description = "ASGI python server"
 authors = [
     { name = "Max Smolin", email = "max@latch.bio" },