Skip to content
Open
Show file tree
Hide file tree
Changes from 6 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 1 addition & 2 deletions CHANGES.rst
Original file line number Diff line number Diff line change
Expand Up @@ -472,8 +472,7 @@ Features



- Added :attr:`~aiohttp.ClientResponse.output_size` and
:attr:`~aiohttp.ClientResponse.upload_complete` -- by :user:`Dreamsorcerer`.
- Added ``ClientResponse.output_size`` and ``ClientResponse.upload_complete`` -- by :user:`Dreamsorcerer`.


*Related issues and pull requests on GitHub:*
Expand Down
2 changes: 2 additions & 0 deletions CHANGES/13427.deprecation.rst
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
Deprecated ``ClientResponse.output_size`` and ``ClientResponse.upload_complete``;
use ``aiohttp.UploadTracker`` instead -- by :user:`Dreamsorcerer`.
1 change: 1 addition & 0 deletions CHANGES/13427.feature.rst
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Added :class:`aiohttp.UploadTracker` for observing a client request's upload progress -- by :user:`Dreamsorcerer`.
2 changes: 2 additions & 0 deletions aiohttp/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,7 @@
TCPConnector,
TooManyRedirects,
UnixConnector,
UploadTracker,
WSMessageTypeError,
WSServerHandshakeError,
request,
Expand Down Expand Up @@ -156,6 +157,7 @@
"TCPConnector",
"TooManyRedirects",
"UnixConnector",
"UploadTracker",
"NamedPipeConnector",
"WSServerHandshakeError",
"request",
Expand Down
6 changes: 6 additions & 0 deletions aiohttp/abc.py
Original file line number Diff line number Diff line change
Expand Up @@ -204,6 +204,12 @@ class AbstractStreamWriter(ABC):
buffer_size: int = 0
output_size: int = 0
length: int | None = 0
# Called with each accepted body chunk's byte count (before any
# transport-level transformation such as compression or chunked
# framing). Assigned by the client request machinery for upload
# progress tracking; write()/write_eof() implementations should
# invoke it for every body chunk they accept.
on_body_write: Callable[[int], None] | None = None

@abstractmethod
async def write(
Expand Down
150 changes: 88 additions & 62 deletions aiohttp/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -77,6 +77,7 @@
Fingerprint,
RequestInfo,
ResponseParams,
UploadTracker,
)
from .client_ws import (
DEFAULT_WS_CLIENT_TIMEOUT,
Expand Down Expand Up @@ -143,6 +144,7 @@
"ClientResponse",
"Fingerprint",
"RequestInfo",
"UploadTracker",
# connector
"BaseConnector",
"TCPConnector",
Expand Down Expand Up @@ -194,6 +196,7 @@ class _RequestOptions(TypedDict, total=False):
max_field_size: int | None
max_headers: int | None
middlewares: Sequence[ClientMiddlewareType] | None
upload_tracker: UploadTracker | None


class _WSConnectOptions(TypedDict, total=False):
Expand Down Expand Up @@ -494,6 +497,7 @@ async def _request(
max_field_size: int | None = None,
max_headers: int | None = None,
middlewares: Sequence[ClientMiddlewareType] | None = None,
upload_tracker: UploadTracker | None = None,
Comment on lines 499 to +502

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Threat model omits UploadTracker

This adds a public client request option without updating THREAT_MODEL.md, leaving the security model's public API inventory and required audit trail out of sync with the newly exposed upload-observation surface.

Context Used: AGENTS.md (source)

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

) -> ClientResponse:
# NOTE: timeout clamps existing connect and read timeouts. We cannot
# set the default to None because we need to detect if the user wants
Expand Down Expand Up @@ -522,83 +526,99 @@ async def _request(
else:
data = payload.JsonPayload(json, dumps=self._json_serialize)

redirects = 0
history: list[ClientResponse] = []
version = self._version
params = params or {}
if upload_tracker is not None:
upload_tracker._bind()

# Merge with default headers and transform to CIMultiDict
headers = self._prepare_headers(headers)
real_timeout = (
self._timeout if timeout is sentinel or timeout is None else timeout
)
# timeout is cumulative for all request operations
# (request, redirects, responses, data consuming)
tm = TimeoutHandle(
self._loop,
real_timeout.total,
ceil_threshold=real_timeout.ceil_threshold,
)
handle: asyncio.TimerHandle | None = None

try:
url = self._build_url(str_or_url)
except ValueError as e:
raise InvalidUrlClientError(str_or_url) from e

assert self._connector is not None
if url.scheme not in self._connector.allowed_protocol_schema_set:
raise NonHttpUrlClientError(url)
redirects = 0
history: list[ClientResponse] = []
version = self._version
params = params or {}

skip_headers: Iterable[istr] | None
if skip_auto_headers is not None:
skip_headers = {
istr(i) for i in skip_auto_headers
} | self._skip_auto_headers
elif self._skip_auto_headers:
skip_headers = self._skip_auto_headers
else:
skip_headers = None
# Merge with default headers and transform to CIMultiDict
headers = self._prepare_headers(headers)

if proxy is None:
proxy = self._default_proxy

resolved_proxy_headers: CIMultiDict[str] | None
if proxy is None:
resolved_proxy_headers = None
else:
resolved_proxy_headers = self._prepare_headers(proxy_headers)
try:
proxy = URL(proxy)
url = self._build_url(str_or_url)
except ValueError as e:
raise InvalidURL(proxy) from e
raise InvalidUrlClientError(str_or_url) from e

assert self._connector is not None
if url.scheme not in self._connector.allowed_protocol_schema_set:
raise NonHttpUrlClientError(url)

skip_headers: Iterable[istr] | None
if skip_auto_headers is not None:
skip_headers = {
istr(i) for i in skip_auto_headers
} | self._skip_auto_headers
elif self._skip_auto_headers:
skip_headers = self._skip_auto_headers
else:
skip_headers = None

if timeout is sentinel or timeout is None:
real_timeout: ClientTimeout = self._timeout
else:
real_timeout = timeout
# timeout is cumulative for all request operations
# (request, redirects, responses, data consuming)
tm = TimeoutHandle(
self._loop, real_timeout.total, ceil_threshold=real_timeout.ceil_threshold
)
handle = tm.start()
if proxy is None:
proxy = self._default_proxy

if read_bufsize is None:
read_bufsize = self._read_bufsize
resolved_proxy_headers: CIMultiDict[str] | None
if proxy is None:
resolved_proxy_headers = None
else:
resolved_proxy_headers = self._prepare_headers(proxy_headers)
try:
proxy = URL(proxy)
except ValueError as e:
raise InvalidURL(proxy) from e

if auto_decompress is None:
auto_decompress = self._auto_decompress
handle = tm.start()

if max_line_size is None:
max_line_size = self._max_line_size
if read_bufsize is None:
read_bufsize = self._read_bufsize

if max_field_size is None:
max_field_size = self._max_field_size
if auto_decompress is None:
auto_decompress = self._auto_decompress

if max_headers is None:
max_headers = self._max_headers
if max_line_size is None:
max_line_size = self._max_line_size

traces = [
Trace(
self,
trace_config,
trace_config.trace_config_ctx(trace_request_ctx=trace_request_ctx),
)
for trace_config in self._trace_configs
]
if max_field_size is None:
max_field_size = self._max_field_size

for trace in traces:
await trace.send_request_start(method, url.update_query(params), headers)
if max_headers is None:
max_headers = self._max_headers

traces = [
Trace(
self,
trace_config,
trace_config.trace_config_ctx(trace_request_ctx=trace_request_ctx),
)
for trace_config in self._trace_configs
]

for trace in traces:
await trace.send_request_start(
method, url.update_query(params), headers
)
except BaseException as e:
tm.close()
if handle is not None:
handle.cancel()
if upload_tracker is not None:
upload_tracker._finalize(e)
raise

timer = tm.timer()
req: ClientRequest | None = None
Expand Down Expand Up @@ -709,6 +729,7 @@ async def _request(
traces=traces,
trust_env=self.trust_env,
)
req._upload_tracker = upload_tracker

# Apply middleware (if any) - per-request middleware overrides session middleware
effective_middlewares = (
Expand Down Expand Up @@ -883,6 +904,8 @@ async def _request(
await trace.send_request_end(
method, url.update_query(params), headers, resp
)
if upload_tracker is not None:
upload_tracker._finalize(None)
return resp

except BaseException as e:
Expand All @@ -892,6 +915,9 @@ async def _request(
handle.cancel()
handle = None

if upload_tracker is not None:
upload_tracker._finalize(e)

if req is not None and req._body is not None:
await req._body.close()

Expand Down
Loading
Loading