avoid double serializing json by capturing request size in rest http client and adjusting surrounding calls to be able to pass the info back up. additional improvements based on pr feedback + tests

Author: chrishagglund-ship-it

src/conductor/client/http/api_client.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import base64
 import datetime
 import logging
@@ -6,7 +8,7 @@
 import re
 import tempfile
 import time
-from typing import Dict
+from typing import Dict, Optional, TYPE_CHECKING
 import uuid
 
 import six
@@ -20,6 +22,9 @@
 from conductor.client.http.rest import AuthorizationException
 from conductor.client.http.thread import AwaitableThread
 
+if TYPE_CHECKING:
+    from conductor.client.telemetry.metrics_collector_base import MetricsCollectorBase
+
 logger = logging.getLogger(
     Configuration.get_logging_formatted_name(
         __name__
@@ -45,7 +50,7 @@ def __init__(
             header_name=None,
             header_value=None,
             cookie=None,
-            metrics_collector=None
+            metrics_collector: Optional[MetricsCollectorBase] = None
     ):
         if configuration is None:
             configuration = Configuration()
@@ -188,7 +193,8 @@ def __call_api_no_retry(
             return (return_data)
         else:
             return (return_data, response_data.status,
-                    response_data.getheaders())
+                    response_data.getheaders(),
+                    getattr(response_data, 'request_body_size', 0))
 
     def sanitize_for_serialization(self, obj):
         """Builds a JSON POST object.

src/conductor/client/http/async_api_client.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import base64
 import datetime
 import logging
@@ -6,7 +8,7 @@
 import re
 import tempfile
 import time
-from typing import Dict
+from typing import Dict, Optional, TYPE_CHECKING
 import uuid
 
 import six
@@ -19,6 +21,9 @@
 from conductor.client.http import async_rest
 from conductor.client.http.async_rest import AuthorizationException
 
+if TYPE_CHECKING:
+    from conductor.client.telemetry.metrics_collector_base import MetricsCollectorBase
+
 logger = logging.getLogger(
     Configuration.get_logging_formatted_name(
         __name__
@@ -47,7 +52,7 @@ def __init__(
             header_name=None,
             header_value=None,
             cookie=None,
-            metrics_collector=None
+            metrics_collector: Optional[MetricsCollectorBase] = None
     ):
         if configuration is None:
             configuration = Configuration()

src/conductor/client/http/rest.py

@@ -29,6 +29,7 @@ def __init__(self, resp):
         self.reason = getattr(resp, 'reason_phrase', '') or self._get_reason_phrase(resp.status_code)
         self.data = resp.text                # eagerly read body
         self.headers = resp.headers
+        self.request_body_size = 0
         # Break httpx Response <-> BoundSyncStream reference cycle (issue #395)
         resp.stream = None
         try:
@@ -219,36 +220,42 @@ def request(self, method, url, query_params=None, headers=None,
                     "httpx client was closed before request; re-established a fresh client"
                 )
 
+        # Serialize the request body once, before the retry loop.
+        # The body never changes between attempts; only the httpx client does.
+        request_body = None
+        request_body_size = 0
+        request_url = url
+        has_body = method in ['POST', 'PUT', 'PATCH', 'OPTIONS', 'DELETE']
+
+        if has_body:
+            if query_params:
+                request_url = url + '?' + urlencode(query_params)
+            if re.search('json', headers['Content-Type'], re.IGNORECASE) or isinstance(body, str):
+                request_body = '{}'
+                if body is not None:
+                    request_body = json.dumps(body)
+                    if isinstance(body, str):
+                        request_body = request_body.strip('"')
+                request_body_size = len(request_body.encode('utf-8'))
+            else:
+                msg = """Cannot prepare a request message for provided
+                         arguments. Please check that your arguments match
+                         declared content type."""
+                raise ApiException(status=0, reason=msg)
+
         for attempt in range(2):
             # Snapshot the client we're about to use so that if the request
             # fails we can ask for a reset only if the client is still this
             # one (i.e. nobody else healed it in the meantime).
             client_at_send = self.connection
             try:
-                # For `POST`, `PUT`, `PATCH`, `OPTIONS`, `DELETE`
-                if method in ['POST', 'PUT', 'PATCH', 'OPTIONS', 'DELETE']:
-                    request_url = url
-                    if query_params:
-                        request_url += '?' + urlencode(query_params)
-                    if re.search('json', headers['Content-Type'], re.IGNORECASE) or isinstance(body, str):
-                        request_body = '{}'
-                        if body is not None:
-                            request_body = json.dumps(body)
-                            if isinstance(body, str):
-                                request_body = request_body.strip('"')
-                        r = client_at_send.request(
-                            method, request_url,
-                            content=request_body,
-                            timeout=timeout,
-                            headers=headers
-                        )
-                    else:
-                        # Cannot generate the request from given parameters
-                        msg = """Cannot prepare a request message for provided
-                                 arguments. Please check that your arguments match
-                                 declared content type."""
-                        raise ApiException(status=0, reason=msg)
-                # For `GET`, `HEAD`
+                if has_body:
+                    r = client_at_send.request(
+                        method, request_url,
+                        content=request_body,
+                        timeout=timeout,
+                        headers=headers
+                    )
                 else:
                     r = client_at_send.request(
                         method, url,
@@ -309,6 +316,7 @@ def request(self, method, url, query_params=None, headers=None,
 
         if _preload_content:
             r = RESTResponse(r)
+            r.request_body_size = request_body_size
 
         if r.status == 401 or r.status == 403:
             raise AuthorizationException(http_resp=r)

src/conductor/client/orkes/orkes_base_client.py

@@ -1,6 +1,12 @@
+from __future__ import annotations
+
 import logging
+from typing import Optional, TYPE_CHECKING
 
 from conductor.client.configuration.configuration import Configuration
+
+if TYPE_CHECKING:
+    from conductor.client.telemetry.metrics_collector_base import MetricsCollectorBase
 from conductor.client.http.api.application_resource_api import ApplicationResourceApi
 from conductor.client.http.api.authorization_resource_api import AuthorizationResourceApi
 from conductor.client.http.api.gateway_auth_resource_api import GatewayAuthResourceApi
@@ -22,7 +28,7 @@
 
 
 class OrkesBaseClient(object):
-    def __init__(self, configuration: Configuration, metrics_collector=None):
+    def __init__(self, configuration: Configuration, metrics_collector: Optional[MetricsCollectorBase] = None):
         self.metrics_collector = metrics_collector
         self.api_client = ApiClient(configuration, metrics_collector=metrics_collector)
         self.logger = logging.getLogger(

src/conductor/client/orkes/orkes_workflow_client.py

@@ -1,5 +1,5 @@
 from __future__ import annotations
-from typing import Optional, List, Dict
+from typing import Optional, List, Dict, TYPE_CHECKING
 
 from conductor.client.configuration.configuration import Configuration
 from conductor.client.http.models import SkipTaskRequest, WorkflowStatus, \
@@ -14,12 +14,15 @@
 from conductor.client.orkes.orkes_base_client import OrkesBaseClient
 from conductor.client.workflow_client import WorkflowClient
 
+if TYPE_CHECKING:
+    from conductor.client.telemetry.metrics_collector_base import MetricsCollectorBase
+
 
 class OrkesWorkflowClient(OrkesBaseClient, WorkflowClient):
     def __init__(
             self,
             configuration: Configuration,
-            metrics_collector=None,
+            metrics_collector: Optional[MetricsCollectorBase] = None,
     ):
         super(OrkesWorkflowClient, self).__init__(configuration, metrics_collector=metrics_collector)
 
@@ -39,13 +42,10 @@ def start_workflow_by_name(
         if priority:
             kwargs.update({"priority": priority})
 
-        if self.metrics_collector is not None:
-            try:
-                self.metrics_collector.measure_workflow_input_payload_size(name, version, input)
-            except Exception:
-                pass
         try:
-            return self.workflowResourceApi.start_workflow1(input, name, **kwargs)
+            data, _status, _headers, request_body_size = (
+                self.workflowResourceApi.start_workflow1_with_http_info(input, name, **kwargs)
+            )
         except Exception as e:
             if self.metrics_collector is not None:
                 try:
@@ -54,18 +54,20 @@ def start_workflow_by_name(
                     pass
             raise
 
-    def start_workflow(self, start_workflow_request: StartWorkflowRequest) -> str:
         if self.metrics_collector is not None:
             try:
-                self.metrics_collector.measure_workflow_input_payload_size(
-                    start_workflow_request.name,
-                    start_workflow_request.version,
-                    getattr(start_workflow_request, "input", None),
-                )
+                version_str = str(version) if version is not None else ""
+                self.metrics_collector.record_workflow_input_payload_size(name, version_str, request_body_size)
             except Exception:
                 pass
+
+        return data
+
+    def start_workflow(self, start_workflow_request: StartWorkflowRequest) -> str:
         try:
-            return self.workflowResourceApi.start_workflow(start_workflow_request)
+            data, _status, _headers, request_body_size = (
+                self.workflowResourceApi.start_workflow_with_http_info(start_workflow_request)
+            )
         except Exception as e:
             if self.metrics_collector is not None:
                 try:
@@ -74,6 +76,17 @@ def start_workflow(self, start_workflow_request: StartWorkflowRequest) -> str:
                     pass
             raise
 
+        if self.metrics_collector is not None:
+            try:
+                version_str = str(start_workflow_request.version) if start_workflow_request.version is not None else ""
+                self.metrics_collector.record_workflow_input_payload_size(
+                    start_workflow_request.name, version_str, request_body_size,
+                )
+            except Exception:
+                pass
+
+        return data
+
     def execute_workflow(
             self,
             start_workflow_request: StartWorkflowRequest,

src/conductor/client/orkes_clients.py

@@ -1,5 +1,12 @@
+from __future__ import annotations
+
+from typing import Optional, TYPE_CHECKING
+
 from conductor.client.authorization_client import AuthorizationClient
 from conductor.client.configuration.configuration import Configuration
+
+if TYPE_CHECKING:
+    from conductor.client.telemetry.metrics_collector_base import MetricsCollectorBase
 from conductor.client.integration_client import IntegrationClient
 from conductor.client.metadata_client import MetadataClient
 from conductor.client.orkes.orkes_integration_client import OrkesIntegrationClient
@@ -21,7 +28,7 @@
 
 
 class OrkesClients:
-    def __init__(self, configuration: Configuration = None, metrics_collector=None):
+    def __init__(self, configuration: Configuration = None, metrics_collector: Optional[MetricsCollectorBase] = None):
         if configuration is None:
             configuration = Configuration()
         self.configuration = configuration

src/conductor/client/telemetry/canonical_metrics_collector.py

@@ -194,12 +194,13 @@ def record_task_update_time(self, task_type: str, time_spent: float, status: str
 
     def record_api_request_time(self, method: str, uri: str, status: str, time_spent: float,
                                 metric_uri: str = None) -> None:
+        resolved_uri = (metric_uri or uri).replace("{tasktype}", "{taskType}")
         self._observe_histogram(
             name=MetricName.API_REQUEST_TIME_CANONICAL,
             documentation=MetricDocumentation.API_REQUEST_TIME_CANONICAL,
             labels={
                 MetricLabel.METHOD: method,
-                MetricLabel.URI: metric_uri or uri,
+                MetricLabel.URI: resolved_uri,
                 MetricLabel.STATUS: status,
             },
             value=time_spent,

src/conductor/client/workflow/executor/workflow_executor.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 import uuid
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, TYPE_CHECKING
 
 from typing_extensions import Self
 
@@ -25,9 +25,12 @@
 )
 from conductor.client.orkes.orkes_workflow_client import OrkesWorkflowClient
 
+if TYPE_CHECKING:
+    from conductor.client.telemetry.metrics_collector_base import MetricsCollectorBase
+
 
 class WorkflowExecutor:
-    def __init__(self, configuration: Configuration, metrics_collector=None) -> Self:
+    def __init__(self, configuration: Configuration, metrics_collector: Optional[MetricsCollectorBase] = None) -> Self:
         api_client = ApiClient(configuration, metrics_collector=metrics_collector)
         self.metadata_client = MetadataResourceApi(api_client)
         self.task_client = TaskResourceApi(api_client)

tests/unit/api_client/test_api_client_coverage.py

@@ -683,6 +683,40 @@ def test_request_with_metrics_collector(self):
                 self.assertEqual(call_args[1]['method'], 'GET')
                 self.assertEqual(call_args[1]['status'], '200')
 
+    def test_request_passes_metric_uri_to_collector(self):
+        """metric_uri kwarg is forwarded to record_api_request_time on success."""
+        metrics_collector = Mock()
+        with patch.object(ApiClient, '_ApiClient__refresh_auth_token'):
+            client = ApiClient(configuration=self.config, metrics_collector=metrics_collector)
+
+            with patch.object(client.rest_client, 'GET', return_value=Mock(status=200)):
+                client.request(
+                    'GET', 'http://localhost:8080/api/workflow/test-id',
+                    metric_uri='/api/workflow/{workflowId}',
+                )
+
+                call_args = metrics_collector.record_api_request_time.call_args
+                self.assertEqual(call_args[1]['metric_uri'], '/api/workflow/{workflowId}')
+
+    def test_request_passes_metric_uri_to_collector_on_error(self):
+        """metric_uri kwarg is forwarded to record_api_request_time on error."""
+        metrics_collector = Mock()
+        with patch.object(ApiClient, '_ApiClient__refresh_auth_token'):
+            client = ApiClient(configuration=self.config, metrics_collector=metrics_collector)
+
+            error = Exception('Test error')
+            error.status = 500
+
+            with patch.object(client.rest_client, 'GET', side_effect=error):
+                with self.assertRaises(Exception):
+                    client.request(
+                        'GET', 'http://localhost:8080/api/workflow/test-id',
+                        metric_uri='/api/workflow/{workflowId}',
+                    )
+
+                call_args = metrics_collector.record_api_request_time.call_args
+                self.assertEqual(call_args[1]['metric_uri'], '/api/workflow/{workflowId}')
+
     def test_request_with_metrics_collector_on_error(self):
         """Test request method with metrics collector on error"""
         metrics_collector = Mock()