Added trace and retry context properties (#96)

* Added trace and retry context properties

* Fixed build errors

* Fixed unit tests

* Fixed line breaks

* Removed empty lines

Author: gavin-aguiar

azure/functions/_abc.py

@@ -24,6 +24,74 @@ def get(self) -> T:
         pass
 
 
+class RpcException:
+    """Rpc Exception object."""
+
+    @property
+    @abc.abstractmethod
+    def source(self) -> str:
+        """Source of the exception."""
+        pass
+
+    @property
+    @abc.abstractmethod
+    def stack_trace(self) -> str:
+        """Stack trace for the exception."""
+        pass
+
+    @property
+    @abc.abstractmethod
+    def message(self) -> str:
+        """Textual message describing the exception."""
+        pass
+
+
+class TraceContext(abc.ABC):
+    """Trace context object."""
+
+    @property
+    @abc.abstractmethod
+    def trace_state(self) -> str:
+        """Gets trace state from trace-context."""
+        pass
+
+    @property
+    @abc.abstractmethod
+    def trace_parent(self) -> str:
+        """Gets trace parent from trace-context."""
+        pass
+
+    @property
+    @abc.abstractmethod
+    def attributes(self) -> typing.Dict[str, str]:
+        """Gets trace-context attributes."""
+        pass
+
+
+class RetryContext(abc.ABC):
+    """Retry Context object.
+       For more information refer: https://aka.ms/azfunc-retries-policies
+    """
+
+    @property
+    @abc.abstractmethod
+    def retry_count(self) -> int:
+        """Gets the current retry count from retry-context."""
+        pass
+
+    @property
+    @abc.abstractmethod
+    def max_retry_count(self) -> int:
+        """Gets the max retry count from retry-context."""
+        pass
+
+    @property
+    @abc.abstractmethod
+    def exception(self) -> RpcException:
+        """Gets the RpcException"""
+        pass
+
+
 class Context(abc.ABC):
     """Function invocation context."""
 
@@ -45,6 +113,18 @@ def function_directory(self) -> str:
         """Function directory."""
         pass
 
+    @property
+    @abc.abstractmethod
+    def trace_context(self) -> TraceContext:
+        """Context for distributed tracing."""
+        pass
+
+    @property
+    @abc.abstractmethod
+    def retry_context(self) -> RetryContext:
+        """Context for retries to the function."""
+        pass
+
 
 class HttpRequest(abc.ABC):
     """HTTP request object."""

azure/functions/_http_asgi.py

@@ -45,7 +45,9 @@ def to_asgi_http_scope(self):
             "client": None,
             "azure_functions.function_directory": self.af_function_directory,
             "azure_functions.function_name": self.af_function_name,
-            "azure_functions.invocation_id": self.af_invocation_id
+            "azure_functions.invocation_id": self.af_invocation_id,
+            "azure_functions.trace_context": self.af_trace_context,
+            "azure_functions.retry_context": self.af_retry_context
         }
         # Notes, missing client name, port
 

azure/functions/_http_wsgi.py

@@ -56,6 +56,8 @@ def __init__(self,
                                              'function_directory', None)
         self.af_function_name = getattr(func_ctx, 'function_name', None)
         self.af_invocation_id = getattr(func_ctx, 'invocation_id', None)
+        self.af_trace_context = getattr(func_ctx, 'trace_context', None)
+        self.af_retry_context = getattr(func_ctx, 'retry_context', None)
 
     def to_environ(self, errors_buffer: StringIO) -> Dict[str, Any]:
         if self._environ_cache is not None:
@@ -80,7 +82,9 @@ def to_environ(self, errors_buffer: StringIO) -> Dict[str, Any]:
             'wsgi.run_once': self.wsgi_run_once,
             'azure_functions.function_directory': self.af_function_directory,
             'azure_functions.function_name': self.af_function_name,
-            'azure_functions.invocation_id': self.af_invocation_id
+            'azure_functions.invocation_id': self.af_invocation_id,
+            'azure_functions.trace_context': self.af_trace_context,
+            'azure_functions.retry_context': self.af_retry_context
         }
         environ.update(self._http_environ)
 

tests/test_http_asgi.py

@@ -4,6 +4,7 @@
 import unittest
 
 import azure.functions as func
+from azure.functions._abc import TraceContext, RetryContext
 from azure.functions._http_asgi import (
     AsgiMiddleware
 )
@@ -107,13 +108,17 @@ def _generate_func_context(
         self,
         invocation_id='123e4567-e89b-12d3-a456-426655440000',
         function_name='httptrigger',
-        function_directory='/home/roger/wwwroot/httptrigger'
+        function_directory='/home/roger/wwwroot/httptrigger',
+        trace_context=TraceContext,
+        retry_context=RetryContext
     ) -> func.Context:
         class MockContext(func.Context):
-            def __init__(self, ii, fn, fd):
+            def __init__(self, ii, fn, fd, tc, rc):
                 self._invocation_id = ii
                 self._function_name = fn
                 self._function_directory = fd
+                self._trace_context = tc
+                self._retry_context = rc
 
             @property
             def invocation_id(self):
@@ -127,7 +132,16 @@ def function_name(self):
             def function_directory(self):
                 return self._function_directory
 
-        return MockContext(invocation_id, function_name, function_directory)
+            @property
+            def trace_context(self):
+                return self._trace_context
+
+            @property
+            def retry_context(self):
+                return self._retry_context
+
+        return MockContext(invocation_id, function_name, function_directory,
+                           trace_context, retry_context)
 
     def test_middleware_calls_app(self):
         app = MockAsgiApplication()

tests/test_http_wsgi.py

@@ -5,6 +5,7 @@
 from io import StringIO, BytesIO
 
 import azure.functions as func
+from azure.functions._abc import TraceContext, RetryContext
 from azure.functions._http_wsgi import (
     WsgiRequest,
     WsgiResponse,
@@ -114,6 +115,10 @@ def test_request_parse_function_context(self):
                          'httptrigger')
         self.assertEqual(environ['azure_functions.function_directory'],
                          '/home/roger/wwwroot/httptrigger')
+        self.assertEqual(environ['azure_functions.trace_context'],
+                         TraceContext)
+        self.assertEqual(environ['azure_functions.retry_context'],
+                         RetryContext)
 
     def test_response_from_wsgi_app(self):
         app = self._generate_wsgi_app()
@@ -215,13 +220,17 @@ def _generate_func_context(
         self,
         invocation_id='123e4567-e89b-12d3-a456-426655440000',
         function_name='httptrigger',
-        function_directory='/home/roger/wwwroot/httptrigger'
+        function_directory='/home/roger/wwwroot/httptrigger',
+        trace_context=TraceContext,
+        retry_context=RetryContext
     ) -> func.Context:
         class MockContext(func.Context):
-            def __init__(self, ii, fn, fd):
+            def __init__(self, ii, fn, fd, tc, rc):
                 self._invocation_id = ii
                 self._function_name = fn
                 self._function_directory = fd
+                self._trace_context = tc
+                self._retry_context = rc
 
             @property
             def invocation_id(self):
@@ -235,7 +244,16 @@ def function_name(self):
             def function_directory(self):
                 return self._function_directory
 
-        return MockContext(invocation_id, function_name, function_directory)
+            @property
+            def trace_context(self):
+                return self._trace_context
+
+            @property
+            def retry_context(self):
+                return self._retry_context
+
+        return MockContext(invocation_id, function_name, function_directory,
+                           trace_context, retry_context)
 
     def _generate_wsgi_app(self,
                            status='200 OK',