Merge pull request #7 from taranisag/converting-to-async

Supporting int status code
taranisag · May 28, 2024 · 363d3bb · 363d3bb
2 parents a975e9b + 524084b
commit 363d3bb
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/py_async_grpc_prometheus/prometheus_async_server_interceptor.py b/py_async_grpc_prometheus/prometheus_async_server_interceptor.py
@@ -31,6 +31,7 @@ def __init__(self,
     self._metrics = server_metrics.init_metrics(registry)
     self._skip_exceptions = skip_exceptions
     self._log_exceptions = log_exceptions
+    self._code_to_status_mapping = {x.value[0]: x for x in grpc.StatusCode}
 
   async def intercept_service(self, continuation, handler_call_details):
     """
@@ -164,7 +165,8 @@ async def new_response_stream_behavior(request_or_iterator, servicer_context):
   def _compute_status_code(self, servicer_context):
     if servicer_context.code() is None:
       return StatusCode.OK
-
+    elif isinstance(servicer_context.code(), int):
+      return self._code_to_status_mapping[servicer_context.code()]
     return servicer_context.code()
 
   def _compute_error_code(self, grpc_exception):

diff --git a/setup.py b/setup.py
@@ -6,7 +6,7 @@
   long_description = fh.read()
 
 setup(name="py_async_grpc_prometheus",
-      version="0.1.1",
+      version="0.1.2",
       description="Python async gRPC Prometheus Interceptors",
       long_description=long_description,
       long_description_content_type="text/markdown",