def install()

in skywalking/plugins/sw_websockets.py [0:0]

• 36 lines of code
• 7 McCabe index (conditional complexity)

def install():
    from websockets.legacy.client import WebSocketClientProtocol
    _protocol_handshake_client = WebSocketClientProtocol.handshake

    async def _sw_protocol_handshake_client(self, wsuri,
                                            origin=None,
                                            available_extensions=None,
                                            available_subprotocols=None,
                                            extra_headers=None):

        span = get_context().new_exit_span(op=wsuri.path or '/', peer=f'{wsuri.host}:{wsuri.port}',
                                           component=Component.Websockets)
        with span:
            carrier = span.inject()
            span.layer = Layer.Http
            if not extra_headers and not self.extra_headers:
                # this method actually uses self.extra_headers, not extra_headers
                self.extra_headers = {}
            for item in carrier:
                self.extra_headers[item.key] = item.val

            span.tag(TagHttpMethod('websocket.connect'))

            scheme = 'wss' if wsuri.secure else 'ws'
            span.tag(TagHttpURL(f'{scheme}://{wsuri.host}:{wsuri.port}{wsuri.path}'))
            status_msg = 'connection open'
            try:
                await _protocol_handshake_client(self,
                                                 wsuri=wsuri,
                                                 origin=origin,
                                                 available_extensions=available_extensions,
                                                 available_subprotocols=available_subprotocols,
                                                 extra_headers=extra_headers)
            except Exception as e:
                span.error_occurred = True
                span.log(e)
                status_msg = 'invalid handshake'
                raise e
            finally:
                span.tag(TagHttpStatusMsg(status_msg))

    WebSocketClientProtocol.handshake = _sw_protocol_handshake_client