class HTTPRequestHandler(BaseHTTPRequestHandler):
"""Custom HTTPRequestHandler class."""
@wtt.spanned(
span_namer=wtt.SpanNamer(use_this_arg="self.command"),
kind=wtt.SpanKind.SERVER,
carrier="self.headers",
)
def do_GET(self) -> None: # pylint: disable=invalid-name
"""Handle GET command."""
logging.info("Example HTTP Server Handler with 'kind=SpanKind.SERVER'")
all_responses = {
"ham": "ham sandwich!",
"cheese": "cheese sandwich!",
"bacon": "BLT!",
}
if self.path in all_responses:
self.send_response(200)
# send header first
self.send_header("Content-type", "text-html")
self.end_headers()
# send file content to client
self.wfile.write(all_responses[self.path].encode())
else:
self.send_error(404, "Ingredient not found!")
logging.critical("Server shutting down (to end testing)")
sys.exit(0)
class ExternalClass:
"""Handle methods using a Span instance as a instance variable."""
def __init__(self, span: wtt.Span) -> None:
self.span = span
self.name = "TheExternalClass!"
@wtt.spanned(
span_namer=wtt.SpanNamer(
literal_name="DISJOINT_SPANNED_METHOD",
use_this_arg="self.name",
use_function_name=False,
)
) # sets current_span
def disjoint_spanned_method(self) -> None:
"""Do some things with a new disjoint span."""
print("disjoint_spanned_method")
assert wtt.get_current_span().is_recording()
assert self.span.is_recording()
assert self.span != wtt.get_current_span()
@wtt.evented(all_args=True)
def inner_event_1(name: str, height: int) -> None:
print(name)
print(height)
@wtt.evented(span="span", all_args=True)
def inner_event_2(name: str, height: int, span: wtt.Span) -> None:
assert span.is_recording()
print(name)
print(height)
inner_event_1("Bar", 185)
inner_event_2("Foo", 177, self.span)
import coloredlogs # type: ignore[import]
if "examples" not in os.listdir():
raise RuntimeError("Script needs to be ran from root of repository.")
sys.path.append(".")
import wipac_telemetry.tracing_tools as wtt # noqa: E402 # pylint: disable=C0413,E0401
ADDRESS = "127.0.0.1"
PORT = 2000
########################################################################################
@wtt.spanned(span_namer=wtt.SpanNamer(literal_name="TheBestClient"),
kind=wtt.SpanKind.CLIENT)
def client() -> None:
"""Run HTTP client."""
logging.info("Example HTTP Client with 'kind=SpanKind.CLIENT'")
logging.debug("http client is starting...")
time.sleep(0.5)
conn = http.client.HTTPConnection(ADDRESS, PORT) # create a connection
logging.debug("http client is running...")
time.sleep(0.5)
for msg in ["ham", "cheese", "bacon", "nails"]:
wtt.add_event(
"Outgoing Server Request",
class RestHandler(tornado.web.RequestHandler):
"""Default REST handler."""
def __init__(self, *args, **kwargs) -> None:
self.server_header = ''
try:
super().__init__(*args, **kwargs)
except Exception:
logging.error('error', exc_info=True)
def initialize(self,
debug=False,
auth=None,
auth_url=None,
module_auth_key='',
server_header='',
route_stats=None,
**kwargs):
super().initialize(**kwargs)
self.debug = debug
self.auth = auth
self.auth_url = auth_url
self.auth_data = {}
self.auth_key = None
self.module_auth_key = module_auth_key
self.server_header = server_header
self.route_stats = route_stats
@wtt.spanned(
span_namer=wtt.SpanNamer(use_this_arg='self.request.method'),
kind=wtt.SpanKind.SERVER,
these=["self.request.method", "self.request.path"],
carrier="self.request.headers",
)
async def _execute(self, *args: Any, **kwargs: Any) -> None:
"""Call implemented methods.
NOTE: This is the closest common call-stack ancestor of
- `prepare()`,
- "method handlers" (`get()`, `post()`, ...),
- `on_finish()`,
- etc.
"""
return await super()._execute(*args, **kwargs)
def set_default_headers(self):
self._headers['Server'] = self.server_header
def get_template_namespace(self):
namespace = super().get_template_namespace()
namespace['version'] = rest_tools.__version__
return namespace
def get_current_user(self):
"""Get the current user, and set auth-related attributes."""
try:
type, token = self.request.headers['Authorization'].split(' ', 1)
if type.lower() != 'bearer':
raise Exception('bad header type')
logger.debug('token: %r', token)
data = self.auth.validate(token)
self.auth_data = data
self.auth_key = token
if "role" in self.auth_data:
wtt.get_current_span().set_attribute('self.auth_data.role',
self.auth_data['role'])
return data['sub']
# Auth Failed
except Exception:
if self.debug and 'Authorization' in self.request.headers:
logger.info('Authorization: %r',
self.request.headers['Authorization'])
logger.info('failed auth', exc_info=True)
return None
@wtt.evented()
def prepare(self):
"""Prepare before http-method request handlers."""
if self.route_stats is not None:
stat = self.route_stats[self.request.path]
if stat.is_overloaded():
backoff = stat.get_backoff_time()
logger.warn('Server is overloaded, backoff %r', backoff)
self.set_header('Retry-After', backoff)
raise tornado.web.HTTPError(503, reason="server overloaded")
self.start_time = time.time()
@wtt.evented()
def on_finish(self):
"""Cleanup after http-method request handlers."""
if self.route_stats is not None and self.get_status() < 500:
stat = self.route_stats[self.request.path]
stat.append(time.time() - self.start_time)
@wtt.evented(all_args=True)
def write_error(self, status_code=500, **kwargs):
"""Write out custom error json."""
data = {
'code': status_code,
'error': self._reason,
}
self.write(data)
self.finish()
def get_json_body_argument(
self,
name: str,
default: Any = arghandler.NO_DEFAULT,
type: Optional[type] = None,
choices: Optional[List[Any]] = None,
forbiddens: Optional[List[Any]] = None,
) -> Any:
"""Get argument from the JSON-decoded request-body.
If no `default` is provided, and the argument is not present, raise `400`.
Arguments:
name -- the argument's name
Keyword Arguments:
default -- a default value to use if the argument is not present
type -- optionally, type-check the argument's value (raise `400` for invalid value)
choices -- a list of valid argument values (raise `400`, if arg's value is not in list)
forbiddens -- a list of disallowed argument values (raise `400`, if arg's value is in list)
Returns:
Any -- the argument's value, unaltered
"""
return arghandler.ArgumentHandler.get_json_body_argument(
self.request.body, name, default, type, choices, forbiddens)
def get_argument(
self,
name: str,
default: Any = arghandler.NO_DEFAULT,
strip: bool = True,
type: Optional[type] = None,
choices: Optional[List[Any]] = None,
forbiddens: Optional[List[Any]] = None,
) -> Any:
"""Get argument from query base-arguments / JSON-decoded request-body.
If no `default` is provided, and the argument is not present, raise `400`.
Arguments:
name -- the argument's name
Keyword Arguments:
default -- a default value to use if the argument is not present
strip {`bool`} -- whether to `str.strip()` the arg's value (default: {`True`})
type -- optionally, type-cast/check the argument's value (raise `400` for invalid value)
choices -- a list of valid argument values (raise `400`, if arg's value is not in list)
forbiddens -- a list of disallowed argument values (raise `400`, if arg's value is in list)
Returns:
Any -- the argument's value, possibly stripped/type-casted
"""
return arghandler.ArgumentHandler.get_argument(self.request.body,
super().get_argument,
name, default, strip,
type, choices,
forbiddens)
class RestClient:
"""A REST client with token handling.
Args:
address (str): base address of REST API
token (str): (optional) access token, or a function generating an access token
timeout (int): (optional) request timeout (default: 60s)
retries (int): (optional) number of retries to attempt (default: 10)
username (str): (optional) auth-basic username
password (str): (optional) auth-basic password
"""
def __init__(self,
address: str,
token: Optional[Union[str, bytes,
Callable[[], Union[str,
bytes]]]] = None,
timeout: float = 60.0,
retries: int = 10,
**kwargs: Any) -> None:
self.address = address
self.timeout = timeout
self.retries = retries
self.kwargs = kwargs
self.logger = logging.getLogger('RestClient')
# token handling
self._token_expire_delay_offset = 5
self.access_token: Optional[Union[str, bytes]] = None
self.token_func: Optional[Callable[[], Union[str, bytes]]] = None
if token:
if isinstance(token, (str, bytes)):
self.access_token = token
elif callable(token):
self.token_func = token
self.session = self.open() # start session
def open(self, sync: bool = False) -> requests.Session:
"""Open the http session."""
self.logger.debug('establish REST http session')
if sync:
self.session = Session(self.retries)
else:
self.session = AsyncSession(self.retries)
self.session.headers = { # type: ignore[assignment]
'Content-Type': 'application/json',
}
if 'username' in self.kwargs and 'password' in self.kwargs:
self.session.auth = (self.kwargs['username'],
self.kwargs['password'])
if 'sslcert' in self.kwargs:
if 'sslkey' in self.kwargs:
self.session.cert = (self.kwargs['sslcert'],
self.kwargs['sslkey'])
else:
self.session.cert = self.kwargs['sslcert']
if 'cacert' in self.kwargs:
self.session.verify = self.kwargs['cacert']
return self.session
def close(self) -> None:
"""Close the http session."""
self.logger.info('close REST http session')
if self.session:
self.session.close()
def _get_token(self) -> None:
if self.access_token:
# check if expired
try:
# NOTE: PyJWT mis-type-hinted arg #1 as a str, but byte is also fine
# https://github.com/jpadilla/pyjwt/pull/605#issuecomment-772082918
data = jwt.decode(
self.access_token, # type: ignore[arg-type]
algorithms=['RS256', 'RS512'],
options={"verify_signature": False})
# account for an X second delay over the wire, so expire sooner
if data['exp'] < time.time() + self._token_expire_delay_offset:
raise Exception()
return
except Exception:
self.access_token = None
self.logger.debug('token expired')
try:
self.access_token = self.token_func() # type: ignore[misc]
except Exception:
self.logger.warning('acquiring access token failed')
raise
def _prepare(
self,
method: str,
path: str,
args: Optional[Dict[str, Any]] = None,
headers: Optional[Dict[str, str]] = None,
) -> Tuple[str, Dict[str, Any]]:
"""Internal method for preparing requests."""
if not args:
args = {}
# auto-inject the current span's info into the HTTP headers
if wtt.get_current_span().is_recording():
wtt.propagations.inject_span_carrier(
self.session.headers) # type: ignore[arg-type]
if path.startswith('/'):
path = path[1:]
url = os.path.join(self.address, path)
kwargs: Dict[str, Any] = {'timeout': self.timeout}
if method in ('GET', 'HEAD'):
# args should be urlencoded
kwargs['params'] = args
else:
kwargs['json'] = args
if self.token_func:
self._get_token()
if not headers:
headers = {}
if self.access_token:
headers['Authorization'] = 'Bearer ' + _to_str(self.access_token)
if headers:
kwargs['headers'] = headers
return (url, kwargs)
def _decode(self, content: Union[str, bytes, bytearray]) -> JSONType:
"""Internal method for translating response from json."""
if not content:
self.logger.info('request returned empty string')
return None
try:
return json_decode(content)
except Exception:
self.logger.info('json data: %r', content)
raise
@wtt.spanned(span_namer=wtt.SpanNamer(use_this_arg='method'),
these=['method', 'path', 'self.address'],
kind=wtt.SpanKind.CLIENT)
async def request(
self,
method: str,
path: str,
args: Optional[Dict[str, Any]] = None,
headers: Optional[Dict[str, str]] = None,
) -> JSONType:
"""Send request to REST Server.
Async request - use with coroutines.
Args:
method (str): the http method
path (str): the url path on the server
args (dict): any arguments to pass
headers (dict): any headers to pass to the request
Returns:
dict: json dict or raw string
"""
url, kwargs = self._prepare(method, path, args, headers)
try:
# session: AsyncSession; So, self.session.request() -> Future
r: requests.Response = await asyncio.wrap_future(
self.session.request(method, url,
**kwargs)) # type: ignore[arg-type]
r.raise_for_status()
return self._decode(r.content)
except requests.exceptions.HTTPError as e:
if method == 'DELETE' and e.response.status_code == 404:
raise # skip the logging for an expected error
self.logger.info('bad request: %s %s %r',
method,
path,
args,
exc_info=True)
raise
@wtt.spanned(span_namer=wtt.SpanNamer(use_this_arg='method'),
these=['method', 'path', 'self.address'],
kind=wtt.SpanKind.CLIENT)
def request_seq(
self,
method: str,
path: str,
args: Optional[Dict[str, Any]] = None,
headers: Optional[Dict[str, str]] = None,
) -> JSONType:
"""Send request to REST Server.
Sequential version of `request`.
Args:
method (str): the http method
path (str): the url path on the server
args (dict): any arguments to pass
headers (dict): any headers to pass to the request
Returns:
dict: json dict or raw string
"""
s = self.session
try:
self.open(sync=True)
url, kwargs = self._prepare(method, path, args, headers)
r = self.session.request(method, url, **kwargs)
r.raise_for_status()
return self._decode(r.content)
finally:
self.session = s
@wtt.spanned(span_namer=wtt.SpanNamer(use_this_arg='method'),
these=['method', 'path', 'self.address'],
kind=wtt.SpanKind.CLIENT)
def request_stream(
self,
method: str,
path: str,
args: Optional[Dict[str, Any]] = None,
headers: Optional[Dict[str, str]] = None,
chunk_size: Optional[int] = 8096,
) -> Generator[JSONType, None, None]:
"""Send request to REST Server, and stream results back.
`chunk_size=None` will read data as it arrives
in whatever size the chunks are received. `chunk_size`<`1`
will be treated as `chunk_size=None`
Args:
method (str): the http method
path (str): the url path on the server
args (dict): any arguments to pass
headers (dict): any headers to pass to the request
chunk_size (int): chunk size (see above)
Returns:
dict: json dict or raw string
"""
if chunk_size is not None and chunk_size < 1:
chunk_size = None
s = self.session
try:
self.open(sync=True)
url, kwargs = self._prepare(method, path, args, headers)
resp = self.session.request(method, url, stream=True, **kwargs)
resp.raise_for_status()
for line in resp.iter_lines(chunk_size=chunk_size,
delimiter=b'\n'):
decoded = self._decode(line.strip())
if decoded: # skip `None`
yield decoded
finally:
self.session = s
class DemoClass:
"""Handle methods using a Span instance as a instance variable."""
def __init__(self) -> None:
self.span: Optional[wtt.Span] = None
self.name = "DEMO!"
@wtt.spanned(
span_namer=wtt.SpanNamer(
literal_name="TheBestPrepare", use_this_arg="self.name"
),
behavior=wtt.SpanBehavior.ONLY_END_ON_EXCEPTION,
)
def prepare(self) -> None:
"""Do some things and start an independent span."""
self.span = wtt.get_current_span()
assert self.span.is_recording()
self.span.add_event("(method) started span from instance method")
time.sleep(3)
@wtt.respanned(None, wtt.SpanBehavior.END_ON_EXIT, all_args=True)
def illegal(num: int) -> None:
# this would end the span before the caller does
pass
@wtt.respanned(None, wtt.SpanBehavior.ONLY_END_ON_EXCEPTION, all_args=True)
def legal_and_rare(num: int) -> None:
# this could end the span, which may or may not be wanted
pass
@wtt.respanned(None, wtt.SpanBehavior.DONT_END, all_args=True)
def legal_and_fine(num: int) -> None:
# this is okay, and a quick way to add to the span
pass
legal_and_rare(11)
legal_and_fine(22)
try:
illegal(33)
except wtt.spans.InvalidSpanBehavior:
assert 1
else:
assert 0
@wtt.respanned(
"self.span", wtt.SpanBehavior.ONLY_END_ON_EXCEPTION, attributes={"a": 2}
)
def process(self) -> None:
"""Do some things and reuse a span."""
time.sleep(2)
assert self.span == wtt.get_current_span()
assert wtt.get_current_span().is_recording()
# this won't end the span
try:
raise KeyError()
except: # noqa: E722 # pylint: disable=bare-except
pass
@wtt.respanned("self.span", wtt.SpanBehavior.DONT_END)
def process_with_exception(self) -> None:
"""Do some things and reuse a span."""
time.sleep(1)
assert self.span == wtt.get_current_span()
assert wtt.get_current_span().is_recording()
# this won't end the span so traces won't be sent,
# unless the exception is excepted by the caller
raise Exception("An exception!")
@wtt.respanned(
"self.span", wtt.SpanBehavior.END_ON_EXIT, attributes={"b": 3}
) # auto-ends
def finish(self) -> None:
"""Do some things, reuse a span, then close that span."""
time.sleep(1)
assert self.span == wtt.get_current_span()
assert wtt.get_current_span().is_recording()
##################
# NOT OKAY - just avoid __del__ + Span.end(), interaction seems undefined
# def __del__(self) -> None:
# """ERROR: ending an instance-dependent span in __del__ is not supported, will break."""
# if self.span:
# self.span.add_event("__del__")
# self.span.end()
# USE THIS INSTEAD
def end(self) -> None:
"""Clean up."""
if self.span:
self.span.add_event("manually ending span")
self.span.end()
print(name)
print(height)
inner_event_1("Bar", 185)
inner_event_2("Foo", 177, self.span)
# NOT OKAY - just avoid __del__ + Span.end(), interaction seems undefined
# def __del__(self) -> None:
# """Clean up."""
# if self.span:
# self.span.add_event("__del__")
# self.span.end()
@wtt.spanned(
span_namer=wtt.SpanNamer(use_function_name=False),
attributes={"a": 1},
) # auto end-on-exit
def injected_span_pass_to_instance() -> ExternalClass:
"""Inject a span then pass onto an instance."""
assert wtt.get_current_span().is_recording()
wtt.get_current_span().set_attribute("Random Int", random.randint(0, 9))
instance = ExternalClass(wtt.get_current_span())
instance.disjoint_spanned_method()
return instance
########################################################################################
msg = "Hello World!"
print(msg)
logging.info(msg)
class Example2:
"""An example with an class instance method."""
@wtt.spanned()
def example_2_instance_method(self) -> None:
"""Print and log simple message."""
msg = "Hello World!"
print(msg)
logging.info(msg)
@wtt.spanned(wtt.SpanNamer("my-span"))
def example_3_with_name() -> None:
"""Print and log simple message."""
msg = "Hello World!"
print(msg)
logging.info(msg)
@wtt.spanned()
def example_4_with_an_uncaught_error() -> None:
"""Print and log simple message."""
msg = "Hello World! I'm about to raise a FileNotFoundError"
print(msg)
logging.info(msg)
raise FileNotFoundError("My FileNotFoundError message")