Support stamp specific redirects with Azure exporters (#1078)

Author: lzchen

contrib/opencensus-ext-azure/CHANGELOG.md

@@ -12,6 +12,8 @@
 ([#1062](https://github.com/census-instrumentation/opencensus-python/pull/1062))
 - Implement feature and instrumentation metrics via Statsbeat
 ([#1076](https://github.com/census-instrumentation/opencensus-python/pull/1076))
+- Support stamp specific redirect in exporters
+([#1078](https://github.com/census-instrumentation/opencensus-python/pull/1078))
 
 ## 1.0.8
 Released 2021-05-13

contrib/opencensus-ext-azure/opencensus/ext/azure/common/transport.py

@@ -22,13 +22,22 @@
 from azure.core.exceptions import ClientAuthenticationError
 from azure.identity._exceptions import CredentialUnavailableError
 
+try:
+    from urllib.parse import urlparse
+except ImportError:
+    from urlparse import urlparse
+
+
 logger = logging.getLogger(__name__)
+
+_MAX_CONSECUTIVE_REDIRECTS = 10
 _MONITOR_OAUTH_SCOPE = "https://monitor.azure.com//.default"
 _requests_lock = threading.Lock()
 _requests_map = {}
 
 
 class TransportMixin(object):
+
     def _check_stats_collection(self):
         return not os.environ.get("APPLICATIONINSIGHTS_STATSBEAT_DISABLED_ALL") and (not hasattr(self, '_is_stats') or not self._is_stats)  # noqa: E501
 
@@ -66,10 +75,7 @@ def _transmit(self, envelopes):
             if self.options.credential:
                 token = self.options.credential.get_token(_MONITOR_OAUTH_SCOPE)
                 headers["Authorization"] = "Bearer {}".format(token.token)
-                # Use new api for aad scenario
-                endpoint += '/v2.1/track'
-            else:
-                endpoint += '/v2/track'
+            endpoint += '/v2.1/track'
             if self._check_stats_collection():
                 with _requests_lock:
                     _requests_map['count'] = _requests_map.get('count', 0) + 1  # noqa: E501
@@ -79,6 +85,7 @@ def _transmit(self, envelopes):
                 headers=headers,
                 timeout=self.options.timeout,
                 proxies=json.loads(self.options.proxies),
+                allow_redirects=False,
             )
         except requests.Timeout:
             logger.warning(
@@ -127,6 +134,7 @@ def _transmit(self, envelopes):
             except Exception:
                 pass
         if response.status_code == 200:
+            self._consecutive_redirects = 0
             if self._check_stats_collection():
                 with _requests_lock:
                     _requests_map['success'] = _requests_map.get('success', 0) + 1  # noqa: E501
@@ -209,6 +217,26 @@ def _transmit(self, envelopes):
                 with _requests_lock:
                     _requests_map['retry'] = _requests_map.get('retry', 0) + 1  # noqa: E501
             return self.options.minimum_retry_interval
+        # Redirect
+        if response.status_code in (307, 308):
+            self._consecutive_redirects += 1
+            if self._consecutive_redirects < _MAX_CONSECUTIVE_REDIRECTS:
+                if response.headers:
+                    location = response.headers.get("location")
+                    if location:
+                        url = urlparse(location)
+                        if url.scheme and url.netloc:
+                            # Change the host to the new redirected host
+                            self.options.endpoint = "{}://{}".format(url.scheme, url.netloc)  # noqa: E501
+                            # Attempt to export again
+                            return self._transmit(envelopes)
+                logger.error(
+                    "Error parsing redirect information."
+                )
+            logger.error(
+                "Error sending telemetry because of circular redirects."
+                " Please check the integrity of your connection string."
+            )
         logger.error(
             'Non-retryable server side error %s: %s.',
             response.status_code,

contrib/opencensus-ext-azure/opencensus/ext/azure/log_exporter/__init__.py

@@ -66,6 +66,8 @@ def __init__(self, **options):
         # start statsbeat on exporter instantiation
         if not os.environ.get("APPLICATIONINSIGHTS_STATSBEAT_DISABLED_ALL"):
             statsbeat_metrics.collect_statsbeat_metrics(self.options)
+        # For redirects
+        self._consecutive_redirects = 0  # To prevent circular redirects
 
     def _export(self, batch, event=None):  # pragma: NO COVER
         try:

contrib/opencensus-ext-azure/opencensus/ext/azure/metrics_exporter/__init__.py

@@ -60,6 +60,8 @@ def __init__(self, is_stats=False, **options):
             )
         self._atexit_handler = atexit.register(self.shutdown)
         self.exporter_thread = None
+        # For redirects
+        self._consecutive_redirects = 0  # To prevent circular redirects
         super(MetricsExporter, self).__init__()
 
     def export_metrics(self, metrics):

contrib/opencensus-ext-azure/opencensus/ext/azure/trace_exporter/__init__.py

@@ -77,6 +77,8 @@ def __init__(self, **options):
         # start statsbeat on exporter instantiation
         if not os.environ.get("APPLICATIONINSIGHTS_STATSBEAT_DISABLED_ALL"):
             statsbeat_metrics.collect_statsbeat_metrics(self.options)
+        # For redirects
+        self._consecutive_redirects = 0  # To prevent circular redirects
 
     def span_data_to_envelope(self, sd):
         envelope = Envelope(

contrib/opencensus-ext-azure/tests/test_transport_mixin.py

@@ -25,6 +25,7 @@
 from opencensus.ext.azure.common import Options
 from opencensus.ext.azure.common.storage import LocalFileStorage
 from opencensus.ext.azure.common.transport import (
+    _MAX_CONSECUTIVE_REDIRECTS,
     _MONITOR_OAUTH_SCOPE,
     TransportMixin,
 )
@@ -47,9 +48,10 @@ def func(*_args, **_kwargs):
 
 
 class MockResponse(object):
-    def __init__(self, status_code, text):
+    def __init__(self, status_code, text, headers=None):
         self.status_code = status_code
         self.text = text
+        self.headers = headers
 
 
 # pylint: disable=W0212
@@ -197,7 +199,8 @@ def test_transmission_auth(self):
                     data=data,
                     headers=headers,
                     timeout=10.0,
-                    proxies={}
+                    proxies={},
+                    allow_redirects=False,
                 )
             credential.get_token.assert_called_with(_MONITOR_OAUTH_SCOPE)
             self.assertIsNone(mixin.storage.get())
@@ -318,3 +321,30 @@ def test_transmission_400(self):
                 post.return_value = MockResponse(400, '{}')
                 mixin._transmit_from_storage()
             self.assertEqual(len(os.listdir(mixin.storage.path)), 0)
+
+    def test_transmission_307(self):
+        mixin = TransportMixin()
+        mixin.options = Options()
+        mixin._consecutive_redirects = 0
+        mixin.options.endpoint = "test.endpoint"
+        with LocalFileStorage(os.path.join(TEST_FOLDER, self.id())) as stor:
+            mixin.storage = stor
+            mixin.storage.put([1, 2, 3])
+            with mock.patch('requests.post') as post:
+                post.return_value = MockResponse(307, '{}', {"location": "https://example.com"})  # noqa: E501
+                mixin._transmit_from_storage()
+            self.assertEqual(post.call_count, _MAX_CONSECUTIVE_REDIRECTS)
+            self.assertEqual(len(os.listdir(mixin.storage.path)), 0)
+            self.assertEqual(mixin.options.endpoint, "https://example.com")
+
+    def test_transmission_307_circular_reference(self):
+        mixin = TransportMixin()
+        mixin.options = Options()
+        mixin._consecutive_redirects = 0
+        mixin.options.endpoint = "https://example.com"
+        with mock.patch('requests.post') as post:
+            post.return_value = MockResponse(307, '{}', {"location": "https://example.com"})  # noqa: E501
+            result = mixin._transmit([1, 2, 3])
+            self.assertEqual(result, -307)
+        self.assertEqual(post.call_count, _MAX_CONSECUTIVE_REDIRECTS)
+        self.assertEqual(mixin.options.endpoint, "https://example.com")