Repository: litestar-org/litestar

Issue:
Bug: Prometheus client causes immediate websocket disconnect 
### Description

When using the Prometheus integration, websockets always disconnect immediately after opening the connection.

I've done a bit of digging, and the error arises in `PrometheusMiddleware._get_default_labels` [here](https://github.com/litestar-org/litestar/blob/main/litestar/contrib/prometheus/middleware.py#L107-L122), where `request.method` raises an attribute error for websockets. One possible fix that I've tested locally is to change it to:
```python
"method": (
    request.method
    if request.scope["type"] == "http"
    else request.scope["type"]
),
```

I'd be happy to open a PR for this change or some other fix if someone has a better idea.

### URL to code causing the issue

_No response_

### MCVE

```python
from litestar import Litestar, websocket_listener
from litestar.testing import TestClient
from litestar.contrib.prometheus import PrometheusConfig, PrometheusController

config = PrometheusConfig()

@websocket_listener("/")
def websocket_handler(data: str):
    return f"Hello, {data}!"

app_with_prometheus = Litestar(
    route_handlers=[websocket_handler, PrometheusController],
    middleware=[config.middleware]
)

app_without_prometheus = Litestar(route_handlers=[websocket_handler])

def test_app(app_to_test):
    with TestClient(app=app_to_test) as test_client:
        with test_client.websocket_connect("/") as socket:
            socket.send("World")
            response = socket.receive()
            print(response)

test_app(app_without_prometheus)
test_app(app_with_prometheus)
```


### Steps to reproduce

```bash
1. Run the MCVE
2. Observe the traceback
```


### Screenshots

_No response_

### Logs

```bash
Traceback (most recent call last):
  File "/Users/erik/test.py", line 27, in <module>
    test_app(app_with_prometheus)
  File "/Users/erik/test.py", line 20, in test_app
    with test_client.websocket_connect("/") as socket:
  File "/Users/erik/Library/Caches/pypoetry/virtualenvs/*****-ocilal_E-py3.11/lib/python3.11/site-packages/litestar/testing/websocket_test_session.py", line 52, in __enter__
    message = self.receive(timeout=self.client.timeout.read)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/erik/Library/Caches/pypoetry/virtualenvs/*****-ocilal_E-py3.11/lib/python3.11/site-packages/litestar/testing/websocket_test_session.py", line 195, in receive
    raise WebSocketDisconnect(
litestar.exceptions.websocket_exceptions.WebSocketDisconnect: KeyError('method')
```


### Litestar Version

2.0.0b2

### Platform

- [ ] Linux
- [X] Mac
- [ ] Windows
- [ ] Other (Please specify in the description above)

<!-- POLAR PLEDGE BADGE START -->
## Funding
* You can sponsor this specific effort via a [Polar.sh](https://polar.sh) pledge below
* We receive the pledge once the issue is completed & verified

<a href="https://polar.sh/litestar-org/litestar/issues/1934">
<picture>
  <source media="(prefers-color-scheme: dark)" srcset="https://polar.sh/api/github/litestar-org/litestar/issues/1934/pledge.svg?darkmode=1">
  <img alt="Fund with Polar" src="https://polar.sh/api/github/litestar-org/litestar/issues/1934/pledge.svg">
</picture>
</a>
<!-- POLAR PLEDGE BADGE END -->

Bug: Prometheus client causes immediate websocket disconnect 
### Description

When using the Prometheus integration, websockets always disconnect immediately after opening the connection.

I've done a bit of digging, and the error arises in `PrometheusMiddleware._get_default_labels` [here](https://github.com/litestar-org/litestar/blob/main/litestar/contrib/prometheus/middleware.py#L107-L122), where `request.method` raises an attribute error for websockets. One possible fix that I've tested locally is to change it to:
```python
"method": (
    request.method
    if request.scope["type"] == "http"
    else request.scope["type"]
),
```

I'd be happy to open a PR for this change or some other fix if someone has a better idea.

### URL to code causing the issue

_No response_

### MCVE

```python
from litestar import Litestar, websocket_listener
from litestar.testing import TestClient
from litestar.contrib.prometheus import PrometheusConfig, PrometheusController

config = PrometheusConfig()

@websocket_listener("/")
def websocket_handler(data: str):
    return f"Hello, {data}!"

app_with_prometheus = Litestar(
    route_handlers=[websocket_handler, PrometheusController],
    middleware=[config.middleware]
)

app_without_prometheus = Litestar(route_handlers=[websocket_handler])

def test_app(app_to_test):
    with TestClient(app=app_to_test) as test_client:
        with test_client.websocket_connect("/") as socket:
            socket.send("World")
            response = socket.receive()
            print(response)

test_app(app_without_prometheus)
test_app(app_with_prometheus)
```


### Steps to reproduce

```bash
1. Run the MCVE
2. Observe the traceback
```


### Screenshots

_No response_

### Logs

```bash
Traceback (most recent call last):
  File "/Users/erik/test.py", line 27, in <module>
    test_app(app_with_prometheus)
  File "/Users/erik/test.py", line 20, in test_app
    with test_client.websocket_connect("/") as socket:
  File "/Users/erik/Library/Caches/pypoetry/virtualenvs/*****-ocilal_E-py3.11/lib/python3.11/site-packages/litestar/testing/websocket_test_session.py", line 52, in __enter__
    message = self.receive(timeout=self.client.timeout.read)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/erik/Library/Caches/pypoetry/virtualenvs/*****-ocilal_E-py3.11/lib/python3.11/site-packages/litestar/testing/websocket_test_session.py", line 195, in receive
    raise WebSocketDisconnect(
litestar.exceptions.websocket_exceptions.WebSocketDisconnect: KeyError('method')
```


### Litestar Version

2.0.0b2

### Platform

- [ ] Linux
- [X] Mac
- [ ] Windows
- [ ] Other (Please specify in the description above)

<!-- POLAR PLEDGE BADGE START -->
## Funding
* You can sponsor this specific effort via a [Polar.sh](https://polar.sh) pledge below
* We receive the pledge once the issue is completed & verified

<a href="https://polar.sh/litestar-org/litestar/issues/1934">
<picture>
  <source media="(prefers-color-scheme: dark)" srcset="https://polar.sh/api/github/litestar-org/litestar/issues/1934/pledge.svg?darkmode=1">
  <img alt="Fund with Polar" src="https://polar.sh/api/github/litestar-org/litestar/issues/1934/pledge.svg">
</picture>
</a>
<!-- POLAR PLEDGE BADGE END -->



Discussion:
@JacobCoffee can you try to reproduce this? 
`"method": request.method if request.scope["type"] == "http" else request.scope["type"],` does indeed solve this but I am not certain it is the right place to fix. Investigating...
@erik-hasse would you like to open a PR with your fix?  This will be okay.
@JacobCoffee can you try to reproduce this? 
`"method": request.method if request.scope["type"] == "http" else request.scope["type"],` does indeed solve this but I am not certain it is the right place to fix. Investigating...
@erik-hasse would you like to open a PR with your fix?  This will be okay.

Resolving patch:
diff --git a/litestar/contrib/prometheus/middleware.py b/litestar/contrib/prometheus/middleware.py
--- a/litestar/contrib/prometheus/middleware.py
+++ b/litestar/contrib/prometheus/middleware.py
@@ -5,6 +5,7 @@
 from typing import TYPE_CHECKING, Any, Callable, ClassVar, cast
 
 from litestar.connection.request import Request
+from litestar.enums import ScopeType
 from litestar.exceptions import MissingDependencyException
 from litestar.middleware.base import AbstractMiddleware
 
@@ -115,7 +116,7 @@ def _get_default_labels(self, request: Request[Any, Any, Any]) -> dict[str, str
         """
 
         return {
-            "method": request.method,
+            "method": request.method if request.scope["type"] == ScopeType.HTTP else request.scope["type"],
             "path": request.url.path,
             "status_code": 200,
             "app_name": self._config.app_name,


Test patch:
diff --git a/tests/unit/test_contrib/prometheus/test_prometheus_exporter_middleware.py b/tests/unit/test_contrib/prometheus/test_prometheus_exporter_middleware.py
--- a/tests/unit/test_contrib/prometheus/test_prometheus_exporter_middleware.py
+++ b/tests/unit/test_contrib/prometheus/test_prometheus_exporter_middleware.py
@@ -5,7 +5,7 @@
 
 from prometheus_client import REGISTRY
 
-from litestar import get, post
+from litestar import get, post, websocket_listener
 from litestar.contrib.prometheus import PrometheusConfig, PrometheusController, PrometheusMiddleware
 from litestar.status_codes import HTTP_200_OK
 from litestar.testing import create_test_client
@@ -172,3 +172,26 @@ def test() -> dict:
             """litestar_requests_total{app_name="litestar",method="GET",path="/test",status_code="200"} 1.0 # {trace_id="1234"} 1.0"""
             in metrics
         )
+
+
+def test_prometheus_with_websocket() -> None:
+    config = create_config()
+
+    @websocket_listener("/test")
+    def test(data: str) -> dict:
+        return {"hello": data}
+
+    with create_test_client([test, PrometheusController], middleware=[config.middleware]) as client:
+        with client.websocket_connect("/test") as websocket:
+            websocket.send_text("litestar")
+            websocket.receive_json()
+
+        metrix_exporter_response = client.get("/metrics")
+
+        assert metrix_exporter_response.status_code == HTTP_200_OK
+        metrics = metrix_exporter_response.content.decode()
+
+        assert (
+            """litestar_requests_total{app_name="litestar",method="websocket",path="/test",status_code="200"} 1.0"""
+            in metrics
+        )