From cc8b867308a2474eda8f5bf2ce8ba33372ce9bea Mon Sep 17 00:00:00 2001
From: Igor Rogoza <i.rogoza@hh.ru>
Date: Wed, 11 Oct 2023 00:26:08 +0400
Subject: [PATCH] HH-194786 add client_response_hook and get status code from
 context

---
 frontik/integrations/telemetry.py | 18 ++++++++++++++----
 poetry.lock                       | 16 +++++++++++++---
 pyproject.toml                    |  2 +-
 3 files changed, 28 insertions(+), 8 deletions(-)

diff --git a/frontik/integrations/telemetry.py b/frontik/integrations/telemetry.py
index 2dae82920..1e4c4a241 100644
--- a/frontik/integrations/telemetry.py
+++ b/frontik/integrations/telemetry.py
@@ -5,7 +5,7 @@
 from typing import TYPE_CHECKING
 from urllib.parse import urlparse
 
-from http_client import client_request_context
+from http_client import client_request_context, response_status_code_context
 from http_client.options import options as http_client_options
 from opentelemetry import trace
 from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
@@ -71,9 +71,7 @@ def initialize_app(self, app: FrontikApplication) -> Future | None:
         provider.add_span_processor(BatchSpanProcessor(otlp_exporter))
         trace.set_tracer_provider(provider)
 
-        self.aiohttp_instrumentor.instrument(
-            request_hook=_client_request_hook,
-        )
+        self.aiohttp_instrumentor.instrument(request_hook=_client_request_hook, response_hook=_client_response_hook)
 
         self.tornado_instrumentor.instrument(
             server_request_hook=_server_request_hook,
@@ -121,6 +119,18 @@ def _client_request_hook(span: Span, params: aiohttp.TraceRequestStartParams) ->
     return
 
 
+def _client_response_hook(
+    span: Span,
+    params: aiohttp.TraceRequestEndParams | aiohttp.TraceRequestExceptionParams,
+) -> None:
+    if not span or not span.is_recording():
+        return
+    response_status_code: int = response_status_code_context.get(None)
+    if response_status_code is None:
+        return
+    span.set_attribute(SpanAttributes.HTTP_STATUS_CODE, response_status_code)
+
+
 def get_netloc(url: str) -> str:
     parts = urlparse(url)
     if parts.scheme not in ('http', 'https', ''):
diff --git a/poetry.lock b/poetry.lock
index 951c8001a..1c2e10b9c 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -552,8 +552,8 @@ yarl = "1.9.2"
 [package.source]
 type = "git"
 url = "https://github.com/hhru/balancing-http-client.git"
-reference = "2.1.3"
-resolved_reference = "217b4b77aab5286ee1781f90ec569d8bf1211f58"
+reference = "HH-194888"
+resolved_reference = "d806f9d0e49f33d8f905e11c073496dc6bd53565"
 
 [[package]]
 name = "idna"
@@ -759,6 +759,16 @@ files = [
     {file = "MarkupSafe-2.1.3-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:5bbe06f8eeafd38e5d0a4894ffec89378b6c6a625ff57e3028921f8ff59318ac"},
     {file = "MarkupSafe-2.1.3-cp311-cp311-win32.whl", hash = "sha256:dd15ff04ffd7e05ffcb7fe79f1b98041b8ea30ae9234aed2a9168b5797c3effb"},
     {file = "MarkupSafe-2.1.3-cp311-cp311-win_amd64.whl", hash = "sha256:134da1eca9ec0ae528110ccc9e48041e0828d79f24121a1a146161103c76e686"},
+    {file = "MarkupSafe-2.1.3-cp312-cp312-macosx_10_9_universal2.whl", hash = "sha256:f698de3fd0c4e6972b92290a45bd9b1536bffe8c6759c62471efaa8acb4c37bc"},
+    {file = "MarkupSafe-2.1.3-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:aa57bd9cf8ae831a362185ee444e15a93ecb2e344c8e52e4d721ea3ab6ef1823"},
+    {file = "MarkupSafe-2.1.3-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:ffcc3f7c66b5f5b7931a5aa68fc9cecc51e685ef90282f4a82f0f5e9b704ad11"},
+    {file = "MarkupSafe-2.1.3-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:47d4f1c5f80fc62fdd7777d0d40a2e9dda0a05883ab11374334f6c4de38adffd"},
+    {file = "MarkupSafe-2.1.3-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:1f67c7038d560d92149c060157d623c542173016c4babc0c1913cca0564b9939"},
+    {file = "MarkupSafe-2.1.3-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:9aad3c1755095ce347e26488214ef77e0485a3c34a50c5a5e2471dff60b9dd9c"},
+    {file = "MarkupSafe-2.1.3-cp312-cp312-musllinux_1_1_i686.whl", hash = "sha256:14ff806850827afd6b07a5f32bd917fb7f45b046ba40c57abdb636674a8b559c"},
+    {file = "MarkupSafe-2.1.3-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:8f9293864fe09b8149f0cc42ce56e3f0e54de883a9de90cd427f191c346eb2e1"},
+    {file = "MarkupSafe-2.1.3-cp312-cp312-win32.whl", hash = "sha256:715d3562f79d540f251b99ebd6d8baa547118974341db04f5ad06d5ea3eb8007"},
+    {file = "MarkupSafe-2.1.3-cp312-cp312-win_amd64.whl", hash = "sha256:1b8dd8c3fd14349433c79fa8abeb573a55fc0fdd769133baac1f5e07abf54aeb"},
     {file = "MarkupSafe-2.1.3-cp37-cp37m-macosx_10_9_x86_64.whl", hash = "sha256:8e254ae696c88d98da6555f5ace2279cf7cd5b3f52be2b5cf97feafe883b58d2"},
     {file = "MarkupSafe-2.1.3-cp37-cp37m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:cb0932dc158471523c9637e807d9bfb93e06a95cbf010f1a38b98623b929ef2b"},
     {file = "MarkupSafe-2.1.3-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9402b03f1a1b4dc4c19845e5c749e3ab82d5078d16a2a4c2cd2df62d57bb0707"},
@@ -1725,4 +1735,4 @@ testing = ["big-O", "jaraco.functools", "jaraco.itertools", "more-itertools", "p
 [metadata]
 lock-version = "2.0"
 python-versions = "3.11.*"
-content-hash = "59bc322708d74495cc23989e4f3ff571d25574c851e2c21a94b35066b6284bdd"
+content-hash = "e5cf2878b3d6b018062ee4495d1e26bdf776561bab2728de9b4c366ea52782e1"
diff --git a/pyproject.toml b/pyproject.toml
index 5ff99357d..e4cafa905 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -21,7 +21,7 @@ jinja2 = '3.1.2'
 lxml = '4.9.2'
 pydantic = '>=1.10.5,<2'
 tornado = '6.3.2'
-http-client = {git = 'https://github.com/hhru/balancing-http-client.git', tag = '2.1.3'}
+http-client = {git = 'https://github.com/hhru/balancing-http-client.git', branch = 'HH-194888'}
 python-consul2-hh = {git = 'https://github.com/hhru/python-consul2', tag = 'v0.2.9'}
 opentelemetry-sdk = '1.17.0'
 opentelemetry-api = '1.17.0'