chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1 @@
|
||||
from .client import Client, GzipCompressor, ConnectException, Code # noqa: F401
|
||||
@@ -0,0 +1,534 @@
|
||||
import gzip
|
||||
import inspect
|
||||
import json
|
||||
import logging
|
||||
import struct
|
||||
import typing
|
||||
|
||||
from httpcore import (
|
||||
ConnectionPool,
|
||||
AsyncConnectionPool,
|
||||
RemoteProtocolError,
|
||||
Response,
|
||||
)
|
||||
from enum import Flag, Enum
|
||||
from typing import Callable, Optional, Dict, Any, Generator, Tuple
|
||||
from google.protobuf import json_format
|
||||
|
||||
|
||||
class EnvelopeFlags(Flag):
|
||||
compressed = 0b00000001
|
||||
end_stream = 0b00000010
|
||||
|
||||
|
||||
class Code(Enum):
|
||||
canceled = "canceled"
|
||||
unknown = "unknown"
|
||||
invalid_argument = "invalid_argument"
|
||||
deadline_exceeded = "deadline_exceeded"
|
||||
not_found = "not_found"
|
||||
already_exists = "already_exists"
|
||||
permission_denied = "permission_denied"
|
||||
resource_exhausted = "resource_exhausted"
|
||||
failed_precondition = "failed_precondition"
|
||||
aborted = "aborted"
|
||||
out_of_range = "out_of_range"
|
||||
unimplemented = "unimplemented"
|
||||
internal = "internal"
|
||||
unavailable = "unavailable"
|
||||
data_loss = "data_loss"
|
||||
unauthenticated = "unauthenticated"
|
||||
|
||||
|
||||
def make_error_from_http_code(http_code: int):
|
||||
error_code_map = {
|
||||
400: Code.invalid_argument,
|
||||
401: Code.unauthenticated,
|
||||
403: Code.permission_denied,
|
||||
404: Code.not_found,
|
||||
409: Code.already_exists,
|
||||
413: Code.resource_exhausted,
|
||||
429: Code.resource_exhausted,
|
||||
499: Code.canceled,
|
||||
500: Code.internal,
|
||||
501: Code.unimplemented,
|
||||
502: Code.unavailable,
|
||||
503: Code.unavailable,
|
||||
504: Code.deadline_exceeded,
|
||||
505: Code.unimplemented,
|
||||
}
|
||||
|
||||
return error_code_map.get(http_code, Code.unknown)
|
||||
|
||||
|
||||
class ConnectException(Exception):
|
||||
def __init__(self, status: Code, message: str):
|
||||
self.status = status
|
||||
self.message = message
|
||||
|
||||
|
||||
envelope_header_length = 5
|
||||
envelope_header_pack = ">BI"
|
||||
|
||||
|
||||
def encode_envelope(*, flags: EnvelopeFlags, data):
|
||||
return encode_envelope_header(flags=flags.value, data=data) + data
|
||||
|
||||
|
||||
def encode_envelope_header(*, flags, data):
|
||||
return struct.pack(envelope_header_pack, flags, len(data))
|
||||
|
||||
|
||||
def decode_envelope_header(header):
|
||||
flags, data_len = struct.unpack(envelope_header_pack, header)
|
||||
return EnvelopeFlags(flags), data_len
|
||||
|
||||
|
||||
def error_for_response(http_resp: Response):
|
||||
try:
|
||||
error = json.loads(http_resp.content)
|
||||
return make_error(error)
|
||||
except (json.decoder.JSONDecodeError, KeyError):
|
||||
error = {"code": http_resp.status, "message": http_resp.content.decode("utf-8")}
|
||||
return make_error(error)
|
||||
|
||||
|
||||
def make_error(error):
|
||||
status = None
|
||||
try:
|
||||
code_value = error.get("code")
|
||||
# return error code from http status code
|
||||
if isinstance(code_value, int):
|
||||
status = make_error_from_http_code(code_value)
|
||||
else:
|
||||
status = Code(code_value)
|
||||
except (KeyError, ValueError):
|
||||
status = Code.unknown
|
||||
|
||||
return ConnectException(status, error.get("message", ""))
|
||||
|
||||
|
||||
def _sync_retry(func, exc, retries):
|
||||
def retry(*args, **kwargs):
|
||||
for _ in range(retries):
|
||||
try:
|
||||
return func(*args, **kwargs)
|
||||
except exc:
|
||||
continue
|
||||
|
||||
return func(*args, **kwargs)
|
||||
|
||||
return retry
|
||||
|
||||
|
||||
def _async_retry(func, exc, retries):
|
||||
async def retry(*args, **kwargs):
|
||||
for _ in range(retries):
|
||||
try:
|
||||
return await func(*args, **kwargs)
|
||||
except exc:
|
||||
continue
|
||||
|
||||
return await func(*args, **kwargs)
|
||||
|
||||
return retry
|
||||
|
||||
|
||||
def _retry(exc: typing.Type[Exception], retries: int):
|
||||
def decorator(func):
|
||||
if inspect.iscoroutinefunction(func):
|
||||
return _async_retry(func, exc, retries)
|
||||
|
||||
return _sync_retry(func, exc, retries)
|
||||
|
||||
return decorator
|
||||
|
||||
|
||||
class GzipCompressor:
|
||||
name = "gzip"
|
||||
decompress = gzip.decompress
|
||||
compress = gzip.compress
|
||||
|
||||
|
||||
class JSONCodec:
|
||||
content_type = "json"
|
||||
|
||||
@staticmethod
|
||||
def encode(msg):
|
||||
return json_format.MessageToJson(msg).encode("utf8")
|
||||
|
||||
@staticmethod
|
||||
def decode(data, *, msg_type):
|
||||
msg = msg_type()
|
||||
json_format.Parse(data.decode("utf8"), msg, ignore_unknown_fields=True)
|
||||
return msg
|
||||
|
||||
|
||||
class ProtobufCodec:
|
||||
content_type = "proto"
|
||||
|
||||
@staticmethod
|
||||
def encode(msg):
|
||||
return msg.SerializeToString()
|
||||
|
||||
@staticmethod
|
||||
def decode(data, *, msg_type):
|
||||
msg = msg_type()
|
||||
msg.ParseFromString(data)
|
||||
return msg
|
||||
|
||||
|
||||
class Client:
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
pool: Optional[ConnectionPool] = None,
|
||||
async_pool: Optional[AsyncConnectionPool] = None,
|
||||
url: str,
|
||||
response_type,
|
||||
compressor=None,
|
||||
json: Optional[bool] = False,
|
||||
headers: Optional[Dict[str, str]] = None,
|
||||
logger: Optional[logging.Logger] = None,
|
||||
):
|
||||
if headers is None:
|
||||
headers = {}
|
||||
|
||||
self.pool = pool
|
||||
self.async_pool = async_pool
|
||||
self.url = url
|
||||
self._codec = JSONCodec if json else ProtobufCodec
|
||||
self._response_type = response_type
|
||||
self._compressor = compressor
|
||||
self._headers = headers
|
||||
self._connection_retries = 3
|
||||
self._logger = logger
|
||||
|
||||
def _log_request(self) -> None:
|
||||
if self._logger is not None:
|
||||
self._logger.info(f"Request: POST {self.url}")
|
||||
|
||||
def _log_response(self, status: int) -> None:
|
||||
if self._logger is None:
|
||||
return
|
||||
if status >= 400:
|
||||
self._logger.error(f"Response: {status} {self.url}")
|
||||
else:
|
||||
self._logger.info(f"Response: {status} {self.url}")
|
||||
|
||||
def _log_stream_message(self) -> None:
|
||||
if self._logger is not None:
|
||||
self._logger.debug(f"Response stream: {self.url}")
|
||||
|
||||
def _prepare_unary_request(
|
||||
self,
|
||||
req,
|
||||
request_timeout=None,
|
||||
headers: Optional[dict] = None,
|
||||
**opts,
|
||||
) -> dict:
|
||||
data = self._codec.encode(req)
|
||||
|
||||
if self._compressor is not None:
|
||||
data = self._compressor.compress(data)
|
||||
|
||||
if headers is None:
|
||||
headers = {}
|
||||
|
||||
extensions = (
|
||||
None
|
||||
if request_timeout is None
|
||||
else {
|
||||
"timeout": {
|
||||
"connect": request_timeout,
|
||||
"pool": request_timeout,
|
||||
"read": request_timeout,
|
||||
"write": request_timeout,
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
return {
|
||||
"method": "POST",
|
||||
"url": self.url,
|
||||
"content": data,
|
||||
"extensions": extensions,
|
||||
"headers": {
|
||||
**self._headers,
|
||||
**headers,
|
||||
**opts.get("headers", {}),
|
||||
"connect-protocol-version": "1",
|
||||
"content-encoding": (
|
||||
"identity" if self._compressor is None else self._compressor.name
|
||||
),
|
||||
"content-type": f"application/{self._codec.content_type}",
|
||||
},
|
||||
}
|
||||
|
||||
def _process_unary_response(
|
||||
self,
|
||||
http_resp: Response,
|
||||
):
|
||||
self._log_response(http_resp.status)
|
||||
|
||||
if http_resp.status != 200:
|
||||
raise error_for_response(http_resp)
|
||||
|
||||
content = http_resp.content
|
||||
|
||||
if self._compressor is not None:
|
||||
content = self._compressor.decompress(content)
|
||||
|
||||
return self._codec.decode(
|
||||
content,
|
||||
msg_type=self._response_type,
|
||||
)
|
||||
|
||||
@_retry(RemoteProtocolError, 3)
|
||||
async def acall_unary(
|
||||
self,
|
||||
req,
|
||||
request_timeout=None,
|
||||
headers: Optional[dict] = None,
|
||||
**opts,
|
||||
):
|
||||
if self.async_pool is None:
|
||||
raise ValueError("async_pool is required")
|
||||
|
||||
req_data = self._prepare_unary_request(
|
||||
req,
|
||||
request_timeout,
|
||||
headers,
|
||||
**opts,
|
||||
)
|
||||
|
||||
self._log_request()
|
||||
res = await self.async_pool.request(**req_data)
|
||||
return self._process_unary_response(res)
|
||||
|
||||
@_retry(RemoteProtocolError, 3)
|
||||
def call_unary(
|
||||
self,
|
||||
req,
|
||||
request_timeout=None,
|
||||
headers: Optional[dict] = None,
|
||||
**opts,
|
||||
):
|
||||
if self.pool is None:
|
||||
raise ValueError("pool is required")
|
||||
|
||||
req_data = self._prepare_unary_request(
|
||||
req,
|
||||
request_timeout,
|
||||
headers,
|
||||
**opts,
|
||||
)
|
||||
|
||||
self._log_request()
|
||||
res = self.pool.request(**req_data)
|
||||
return self._process_unary_response(res)
|
||||
|
||||
def _create_stream_timeout(self, timeout: Optional[float]):
|
||||
if timeout:
|
||||
return {"connect-timeout-ms": str(int(timeout * 1000))}
|
||||
return {}
|
||||
|
||||
def _prepare_server_stream_request(
|
||||
self,
|
||||
req,
|
||||
request_timeout=None,
|
||||
timeout=None,
|
||||
headers: Optional[dict] = None,
|
||||
**opts,
|
||||
) -> dict:
|
||||
headers = headers or {}
|
||||
data = self._codec.encode(req)
|
||||
flags = EnvelopeFlags(0)
|
||||
|
||||
# `request_timeout` bounds connection setup and request sending, but NOT the
|
||||
# stream read: a stream can stay open for the whole command `timeout` (minutes
|
||||
# or, when disabled, indefinitely), so we deliberately leave `read` unset.
|
||||
# The command `timeout` is enforced server-side via the `connect-timeout-ms`
|
||||
# header (see `_create_stream_timeout`), which returns a clean `deadline_exceeded`.
|
||||
# This mirrors the JS SDK, which has no per-chunk read timeout either — setting
|
||||
# `read` to the command `timeout` would race that server response and surface a
|
||||
# raw transport `ReadTimeout` instead.
|
||||
timeout_ext = {}
|
||||
if request_timeout is not None:
|
||||
timeout_ext["connect"] = request_timeout
|
||||
timeout_ext["pool"] = request_timeout
|
||||
timeout_ext["write"] = request_timeout
|
||||
extensions = {"timeout": timeout_ext} if timeout_ext else None
|
||||
|
||||
if self._compressor is not None:
|
||||
data = self._compressor.compress(data)
|
||||
flags |= EnvelopeFlags.compressed
|
||||
|
||||
stream_timeout = self._create_stream_timeout(timeout)
|
||||
|
||||
return {
|
||||
"method": "POST",
|
||||
"url": self.url,
|
||||
"content": encode_envelope(
|
||||
flags=flags,
|
||||
data=data,
|
||||
),
|
||||
"extensions": extensions,
|
||||
"headers": {
|
||||
**self._headers,
|
||||
**headers,
|
||||
**opts.get("headers", {}),
|
||||
**stream_timeout,
|
||||
"connect-protocol-version": "1",
|
||||
"connect-content-encoding": (
|
||||
"identity" if self._compressor is None else self._compressor.name
|
||||
),
|
||||
"content-type": f"application/connect+{self._codec.content_type}",
|
||||
},
|
||||
}
|
||||
|
||||
# Note: no retry here — generator functions don't execute until iterated, so a
|
||||
# call-level retry never fires, and retrying mid-stream would replay delivered events.
|
||||
async def acall_server_stream(
|
||||
self,
|
||||
req,
|
||||
request_timeout=None,
|
||||
timeout=None,
|
||||
headers: Optional[dict] = None,
|
||||
**opts,
|
||||
):
|
||||
if self.async_pool is None:
|
||||
raise ValueError("async_pool is required")
|
||||
|
||||
req_data = self._prepare_server_stream_request(
|
||||
req,
|
||||
request_timeout,
|
||||
timeout,
|
||||
headers,
|
||||
**opts,
|
||||
)
|
||||
|
||||
parser = ServerStreamParser(
|
||||
decode=self._codec.decode,
|
||||
response_type=self._response_type,
|
||||
)
|
||||
|
||||
self._log_request()
|
||||
async with self.async_pool.stream(**req_data) as http_resp:
|
||||
if http_resp.status != 200:
|
||||
self._log_response(http_resp.status)
|
||||
await http_resp.aread()
|
||||
raise error_for_response(http_resp)
|
||||
|
||||
async for chunk in http_resp.aiter_stream():
|
||||
for parsed in parser.parse(chunk):
|
||||
self._log_stream_message()
|
||||
yield parsed
|
||||
|
||||
def call_server_stream(
|
||||
self,
|
||||
req,
|
||||
request_timeout=None,
|
||||
timeout=None,
|
||||
headers: Optional[dict] = None,
|
||||
**opts,
|
||||
):
|
||||
if self.pool is None:
|
||||
raise ValueError("pool is required")
|
||||
|
||||
req_data = self._prepare_server_stream_request(
|
||||
req,
|
||||
request_timeout,
|
||||
timeout,
|
||||
headers,
|
||||
**opts,
|
||||
)
|
||||
|
||||
parser = ServerStreamParser(
|
||||
decode=self._codec.decode,
|
||||
response_type=self._response_type,
|
||||
)
|
||||
|
||||
self._log_request()
|
||||
with self.pool.stream(**req_data) as http_resp:
|
||||
if http_resp.status != 200:
|
||||
self._log_response(http_resp.status)
|
||||
http_resp.read()
|
||||
raise error_for_response(http_resp)
|
||||
|
||||
for chunk in http_resp.iter_stream():
|
||||
for parsed in parser.parse(chunk):
|
||||
self._log_stream_message()
|
||||
yield parsed
|
||||
|
||||
def call_client_stream(self, req, **opts):
|
||||
raise NotImplementedError("client stream not supported")
|
||||
|
||||
def acall_client_stream(self, req, **opts):
|
||||
raise NotImplementedError("client stream not supported")
|
||||
|
||||
def call_bidi_stream(self, req, **opts):
|
||||
raise NotImplementedError("bidi stream not supported")
|
||||
|
||||
def acall_bidi_stream(self, req, **opts):
|
||||
raise NotImplementedError("bidi stream not supported")
|
||||
|
||||
|
||||
DataLen = int
|
||||
|
||||
|
||||
class ServerStreamParser:
|
||||
def __init__(
|
||||
self,
|
||||
decode: Callable,
|
||||
response_type: Any,
|
||||
):
|
||||
self.decode = decode
|
||||
self.response_type = response_type
|
||||
|
||||
self.buffer: bytes = b""
|
||||
self._header: Optional[tuple[EnvelopeFlags, DataLen]] = None
|
||||
|
||||
def shift_buffer(self, size: int):
|
||||
buffer = self.buffer[:size]
|
||||
self.buffer = self.buffer[size:]
|
||||
return buffer
|
||||
|
||||
@property
|
||||
def header(self) -> Tuple[EnvelopeFlags, DataLen]:
|
||||
if self._header:
|
||||
return self._header
|
||||
|
||||
header_data = self.shift_buffer(envelope_header_length)
|
||||
self._header = decode_envelope_header(header_data)
|
||||
|
||||
return self._header
|
||||
|
||||
@header.deleter
|
||||
def header(self):
|
||||
self._header = None
|
||||
|
||||
def parse(self, chunk: bytes) -> Generator[Any, None, None]:
|
||||
self.buffer += chunk
|
||||
|
||||
# Once the header is consumed, the remaining payload can be shorter
|
||||
# than the header length, so only require a full header when we still
|
||||
# need to read one.
|
||||
while self._header is not None or len(self.buffer) >= envelope_header_length:
|
||||
flags, data_len = self.header
|
||||
|
||||
if data_len > len(self.buffer):
|
||||
break
|
||||
|
||||
data = self.shift_buffer(data_len)
|
||||
|
||||
if EnvelopeFlags.end_stream in flags:
|
||||
data = json.loads(data)
|
||||
|
||||
if "error" in data:
|
||||
raise make_error(data["error"])
|
||||
|
||||
return
|
||||
|
||||
yield self.decode(data, msg_type=self.response_type)
|
||||
del self.header
|
||||
Reference in New Issue
Block a user