mirror of
https://github.com/qdrant/qdrant-client.git
synced 2026-07-23 11:11:01 -05:00
07f19f7dee
* implement thread-local header overwrite for setting headers for individual requests * mypy * fix: remove redundant import in init, remove redundant import in qdrant remote --------- Co-authored-by: George Panchuk <george.panchuk@qdrant.tech>
48 lines
1.4 KiB
Python
48 lines
1.4 KiB
Python
from contextlib import asynccontextmanager, contextmanager
|
|
from contextvars import ContextVar
|
|
from typing import AsyncIterator, Awaitable, Callable, Iterator
|
|
|
|
from httpx import Request, Response
|
|
|
|
_context_headers: ContextVar[dict[str, str]] = ContextVar("_context_headers", default={})
|
|
|
|
|
|
def get_context_headers() -> dict[str, str]:
|
|
return _context_headers.get()
|
|
|
|
|
|
@contextmanager
|
|
def headers(extra_headers: dict[str, str]) -> Iterator[None]:
|
|
current = _context_headers.get()
|
|
merged = {**current, **extra_headers}
|
|
token = _context_headers.set(merged)
|
|
try:
|
|
yield
|
|
finally:
|
|
_context_headers.reset(token)
|
|
|
|
|
|
@asynccontextmanager
|
|
async def async_headers(extra_headers: dict[str, str]) -> AsyncIterator[None]:
|
|
current = _context_headers.get()
|
|
merged = {**current, **extra_headers}
|
|
token = _context_headers.set(merged)
|
|
try:
|
|
yield
|
|
finally:
|
|
_context_headers.reset(token)
|
|
|
|
|
|
def rest_headers_middleware(request: Request, call_next: Callable[[Request], Response]) -> Response:
|
|
for key, value in get_context_headers().items():
|
|
request.headers[key] = value
|
|
return call_next(request)
|
|
|
|
|
|
async def async_rest_headers_middleware(
|
|
request: Request, call_next: Callable[[Request], Awaitable[Response]]
|
|
) -> Response:
|
|
for key, value in get_context_headers().items():
|
|
request.headers[key] = value
|
|
return await call_next(request)
|