Dependency injection of a reference to JWTAuth

[aranvir]

Hi, I want to make a reference to my JWTAuth object available via dependency injection. My motivation is that I want to create it only once and thus not use an import statement to fetch it from the module. After some digging in this I realized that it's probably not an issue, since JWTAuth does not manage any state so as long as the inputs don't change it would have been okay to instantiate it multiple times.

However, I was very confused why I could not (easily) provide a reference via dependency injection. I attached a full example at the end of the post. In short:

• When I type my dependency injection as JWTAuth, I get an NameError: name 'ASGIConnection' is not defined exception. The stack trace does not make any sense to me.
• When I type the dependency injection as Any, it works fine.
• If I add Dependency(skip_validation=True) it also works.

So, I guess it has something to do with the inner workings of litestar but I did not expect that behavior at all and it took me too long to figure out the source of the issue :D so I was wondering: is it intended or an oversight?

@post("/login", exclude_from_auth=True)
async def login(data: UserLogin, jwt_auth_service: JWTAuth) -> Response[User]:
    ... # THIS WILL NOT WORK -> NameError: name 'ASGIConnection' is not defined

@post("/login", exclude_from_auth=True)
async def login(data: UserLogin, jwt_auth_service: Any) -> Response[User]:
    ... # THIS WILL WORK

@post("/login", exclude_from_auth=True)
async def login(data: UserLogin, jwt_auth_service: Annotated[JWTAuth, Dependency(skip_validation=True)]) -> Response[User]:
    ... # THIS WILL ALSO WORK

import os
from typing import Any

from dataclasses import dataclass
from litestar import Litestar, get, post, exceptions, Response, Request
from litestar.di import Provide
from litestar.security.jwt import JWTAuth, Token
from litestar.connection import ASGIConnection
from litestar.dto import DataclassDTO, DTOConfig


### MODEL
@dataclass
class UserLogin:
    user_name: str
    user_password: str


@dataclass
class User(UserLogin):
    user_id: int


class UserDTO(DataclassDTO[User]):
    config = DTOConfig(
        exclude={
            "user_password"
        }
    )


### DATABASE
DATABASE: dict[int, User] = {
    1: User(user_id=1, user_name="test", user_password="test"),
}

### JWT AUTH
async def retrieve_user_handler(token: Token, connection: ASGIConnection) -> User:
    try:
        user_id = int(token.sub)
    except ValueError:
        raise exceptions.NotAuthorizedException("Bad token id.")
    user = DATABASE.get(user_id, None)
    if user is None:
        raise exceptions.NotAuthorizedException("Unknown user.")

    return user


def setup_jwt() -> tuple[JWTAuth, Provide]:
    jwt_auth = JWTAuth[User](
        retrieve_user_handler=retrieve_user_handler,
        token_secret=os.getenv("JWT_SECRET", "my-super-secret"),
        # we are specifying which endpoints should be excluded from authentication. In this case the login endpoint
        # and our openAPI docs.
        exclude=[
            "/schema",
        ],
    )
    def _fetch() -> JWTAuth:
        return jwt_auth

    service = Provide(_fetch, sync_to_thread=False)
    return jwt_auth, service


jwt_service, jwt_fetch = setup_jwt()


### ROUTES
@post("/login", exclude_from_auth=True)
#async def login(data: UserLogin, jwt_auth_service: Any) -> Response[User]:
# async def login(data: UserLogin, jwt_auth_service: JWTAuth) -> Response[User]:
async def login(data: UserLogin, jwt_auth_service: Annotated[JWTAuth, Dependency(skip_validation=True)]) -> Response[User]:
    for user in DATABASE.values():
        if user.user_name == data.user_name:
            if user.user_password == data.user_password:
                return jwt_auth_service.login(str(user.user_id), token_extras={"user_name": user.user_name})
            raise exceptions.NotAuthorizedException("Bad password")
    raise exceptions.NotAuthorizedException("User not found")


@get("/current", return_dto=UserDTO)
async def get_user(request: Request) -> User:
    return request.user

### APP
app = Litestar(
    route_handlers=[
        login,
        get_user,
    ],
    on_app_init=[
        jwt_service.on_app_init,
    ],
    dependencies={
        "jwt_auth_service": jwt_fetch,
    },
)

[provinzkraut]

I can reproduce this with

from litestar import post
from litestar.connection import ASGIConnection
from litestar.di import Provide
from litestar.security.jwt import JWTAuth, Token
from litestar.testing import create_test_client


async def retrieve_user_handler(token: Token, connection: ASGIConnection) -> None:
    pass


jwt_auth = JWTAuth(retrieve_user_handler=retrieve_user_handler, token_secret="")


@post("/", exclude_from_auth=True)
async def login(jwt: JWTAuth) -> None:
    return None


with create_test_client(
    [login],
    dependencies={
        "jwt": Provide(lambda: jwt_auth, sync_to_thread=False, use_cache=True),
    },
    raise_server_exceptions=True,
) as client:
    print(client.post("/"))

@aranvir just a tip: When providing bug reports, it's always best to keep the example as minimal as possible :)

---

What's interesting about this is that even if you do signature_types=[ASGIConnection]), there's still a name error. This suggests to me there's something weird going on somewhere.

[provinzkraut]

I have a more simplified version:

from __future__ import annotations

import dataclasses
from typing import TYPE_CHECKING
from litestar import get
from litestar.testing import create_test_client

if TYPE_CHECKING:
    from litestar.connection import ASGIConnection

@dataclasses.dataclass
class Foo:
    something: ASGIConnection

@get("/")
async def handler(foo: Foo | None = None) -> None:
    pass

with create_test_client([handler], raise_server_exceptions=True, debug=True) as client:
    client.get("/")

the interesting thing is that it only works if the injected class is a dataclass.

[provinzkraut]

Okay so this isn't really a bug but working as expected, although it's certainly not working how it would be desired.

What's happening is simply that when msgspec is trying to validate Foo, it can't resolve ASGIConnection because we have put it in a if TYPE_CHECKING block. This isn't a huge problem for the simplified example, because we can simply move it out there, however, when this is a builtin-type from Litestar, like JWTAuth, this isn't really possible. The reason why the signature_types don't help us here either is because Litestar doesn't handle this part of the validation.

The only reasonable fix I can see here is for msgspec to add a way to extend the namespace it uses to inspect the annotations.

[aranvir]

Thanks a lot for this deep dive and the explanation! Is that then something you'd raise to msgspec to add, since I guess that would be a useful thing to have in general.

Also: I'll take this as inspiration for more concise examples in the future ;)

[provinzkraut]

Is that then something you'd raise to msgspec to add, since I guess that would be a useful thing to have in general.

Yeah, that would be the course of action here. Also trying to see if we can't easily make a PR for that ourselves :)

[euri10]

facing the same on trying to use the jwt_auth instance in a handler in the context of a app factory (see https://discord.com/channels/919193495116337154/1331963919605956608 but the discussion here is better described than my ramblings)
is there an issue I can follow here or on msgspec to add to my '#TODO: comments ?