response.py 6.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163
  1. # Copyright (c) 2012-2013 Mitch Garnaat http://garnaat.org/
  2. # Copyright 2012-2014 Amazon.com, Inc. or its affiliates. All Rights Reserved.
  3. #
  4. # Licensed under the Apache License, Version 2.0 (the "License"). You
  5. # may not use this file except in compliance with the License. A copy of
  6. # the License is located at
  7. #
  8. # http://aws.amazon.com/apache2.0/
  9. #
  10. # or in the "license" file accompanying this file. This file is
  11. # distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
  12. # ANY KIND, either express or implied. See the License for the specific
  13. # language governing permissions and limitations under the License.
  14. import sys
  15. import logging
  16. from botocore import ScalarTypes
  17. from botocore.hooks import first_non_none_response
  18. from botocore.compat import json, set_socket_timeout, XMLParseError
  19. from botocore.exceptions import IncompleteReadError, ReadTimeoutError
  20. from urllib3.exceptions import ReadTimeoutError as URLLib3ReadTimeoutError
  21. from botocore import parsers
  22. logger = logging.getLogger(__name__)
  23. class StreamingBody(object):
  24. """Wrapper class for an http response body.
  25. This provides a few additional conveniences that do not exist
  26. in the urllib3 model:
  27. * Set the timeout on the socket (i.e read() timeouts)
  28. * Auto validation of content length, if the amount of bytes
  29. we read does not match the content length, an exception
  30. is raised.
  31. """
  32. _DEFAULT_CHUNK_SIZE = 1024
  33. def __init__(self, raw_stream, content_length):
  34. self._raw_stream = raw_stream
  35. self._content_length = content_length
  36. self._amount_read = 0
  37. def set_socket_timeout(self, timeout):
  38. """Set the timeout seconds on the socket."""
  39. # The problem we're trying to solve is to prevent .read() calls from
  40. # hanging. This can happen in rare cases. What we'd like to ideally
  41. # do is set a timeout on the .read() call so that callers can retry
  42. # the request.
  43. # Unfortunately, this isn't currently possible in requests.
  44. # See: https://github.com/kennethreitz/requests/issues/1803
  45. # So what we're going to do is reach into the guts of the stream and
  46. # grab the socket object, which we can set the timeout on. We're
  47. # putting in a check here so in case this interface goes away, we'll
  48. # know.
  49. try:
  50. # To further complicate things, the way to grab the
  51. # underlying socket object from an HTTPResponse is different
  52. # in py2 and py3. So this code has been pushed to botocore.compat.
  53. set_socket_timeout(self._raw_stream, timeout)
  54. except AttributeError:
  55. logger.error("Cannot access the socket object of "
  56. "a streaming response. It's possible "
  57. "the interface has changed.", exc_info=True)
  58. raise
  59. def read(self, amt=None):
  60. """Read at most amt bytes from the stream.
  61. If the amt argument is omitted, read all data.
  62. """
  63. try:
  64. chunk = self._raw_stream.read(amt)
  65. except URLLib3ReadTimeoutError as e:
  66. # TODO: the url will be None as urllib3 isn't setting it yet
  67. raise ReadTimeoutError(endpoint_url=e.url, error=e)
  68. self._amount_read += len(chunk)
  69. if amt is None or (not chunk and amt > 0):
  70. # If the server sends empty contents or
  71. # we ask to read all of the contents, then we know
  72. # we need to verify the content length.
  73. self._verify_content_length()
  74. return chunk
  75. def __iter__(self):
  76. """Return an iterator to yield 1k chunks from the raw stream.
  77. """
  78. return self.iter_chunks(self._DEFAULT_CHUNK_SIZE)
  79. def __next__(self):
  80. """Return the next 1k chunk from the raw stream.
  81. """
  82. current_chunk = self.read(self._DEFAULT_CHUNK_SIZE)
  83. if current_chunk:
  84. return current_chunk
  85. raise StopIteration()
  86. next = __next__
  87. def iter_lines(self, chunk_size=1024, keepends=False):
  88. """Return an iterator to yield lines from the raw stream.
  89. This is achieved by reading chunk of bytes (of size chunk_size) at a
  90. time from the raw stream, and then yielding lines from there.
  91. """
  92. pending = b''
  93. for chunk in self.iter_chunks(chunk_size):
  94. lines = (pending + chunk).splitlines(True)
  95. for line in lines[:-1]:
  96. yield line.splitlines(keepends)[0]
  97. pending = lines[-1]
  98. if pending:
  99. yield pending.splitlines(keepends)[0]
  100. def iter_chunks(self, chunk_size=_DEFAULT_CHUNK_SIZE):
  101. """Return an iterator to yield chunks of chunk_size bytes from the raw
  102. stream.
  103. """
  104. while True:
  105. current_chunk = self.read(chunk_size)
  106. if current_chunk == b"":
  107. break
  108. yield current_chunk
  109. def _verify_content_length(self):
  110. # See: https://github.com/kennethreitz/requests/issues/1855
  111. # Basically, our http library doesn't do this for us, so we have
  112. # to do this ourself.
  113. if self._content_length is not None and \
  114. self._amount_read != int(self._content_length):
  115. raise IncompleteReadError(
  116. actual_bytes=self._amount_read,
  117. expected_bytes=int(self._content_length))
  118. def close(self):
  119. """Close the underlying http response stream."""
  120. self._raw_stream.close()
  121. def get_response(operation_model, http_response):
  122. protocol = operation_model.metadata['protocol']
  123. response_dict = {
  124. 'headers': http_response.headers,
  125. 'status_code': http_response.status_code,
  126. }
  127. # TODO: Unfortunately, we have to have error logic here.
  128. # If it looks like an error, in the streaming response case we
  129. # need to actually grab the contents.
  130. if response_dict['status_code'] >= 300:
  131. response_dict['body'] = http_response.content
  132. elif operation_model.has_streaming_output:
  133. response_dict['body'] = StreamingBody(
  134. http_response.raw, response_dict['headers'].get('content-length'))
  135. else:
  136. response_dict['body'] = http_response.content
  137. parser = parsers.create_parser(protocol)
  138. return http_response, parser.parse(response_dict,
  139. operation_model.output_shape)