Support AAD authentication (#1021)

lzchen · web-flow · commit e2cb2af425de · 2021-06-01T14:44:03.000-07:00
diff --git a/.isort.cfg b/.isort.cfg
@@ -13,5 +13,5 @@ line_length=79
 ; docs: https://github.com/timothycrosley/isort#multi-line-output-modes
 multi_line_output=3
 known_future_library = six,six.moves,__future__
-known_third_party=google,mock,pymysql,sqlalchemy,psycopg2,mysql,requests,django,pytest,grpc,flask,bitarray,prometheus_client,psutil,pymongo,wrapt,thrift,retrying,pyramid,werkzeug,gevent
+known_third_party=azure-core,azure-identity,google,mock,pymysql,sqlalchemy,psycopg2,mysql,requests,django,pytest,grpc,flask,bitarray,prometheus_client,psutil,pymongo,wrapt,thrift,retrying,pyramid,werkzeug,gevent
 known_first_party=opencensus
diff --git a/contrib/opencensus-ext-azure/CHANGELOG.md b/contrib/opencensus-ext-azure/CHANGELOG.md
@@ -2,13 +2,16 @@
 
 ## Unreleased
 
+- Enable AAD authorization via TokenCredential
+([#1021](https://github.com/census-instrumentation/opencensus-python/pull/1021))
+
 ## 1.0.8
 Released 2021-05-13
 
 - Fix `logger.exception` with no exception info throwing error
 ([#1006](https://github.com/census-instrumentation/opencensus-python/pull/1006))
 - Add `enable_local_storage` to turn on/off local storage + retry + flushing logic
-([#1006](https://github.com/census-instrumentation/opencensus-python/pull/1006))
+([#1016](https://github.com/census-instrumentation/opencensus-python/pull/1016))
 
 ## 1.0.7
 Released 2021-01-25
diff --git a/contrib/opencensus-ext-azure/examples/traces/credential.py b/contrib/opencensus-ext-azure/examples/traces/credential.py
@@ -0,0 +1,38 @@
+# Copyright 2021, OpenCensus Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from azure.identity import ClientSecretCredential
+
+from opencensus.ext.azure.trace_exporter import AzureExporter
+from opencensus.trace.samplers import ProbabilitySampler
+from opencensus.trace.tracer import Tracer
+
+tenant_id = "<tenant-id>"
+client_id = "<client-id>"
+client_secret = "<client-secret>"
+
+credential = ClientSecretCredential(
+    tenant_id=tenant_id,
+    client_id=client_id,
+    client_secret=client_secret
+)
+
+tracer = Tracer(
+    exporter=AzureExporter(
+        credential=credential, connection_string="<your-connection-string>"),
+    sampler=ProbabilitySampler(1.0)
+)
+
+with tracer.span(name='foo'):
+    print('Hello, World!')
+input(...)
diff --git a/contrib/opencensus-ext-azure/opencensus/ext/azure/common/__init__.py b/contrib/opencensus-ext-azure/opencensus/ext/azure/common/__init__.py
@@ -46,7 +46,14 @@ def process_options(options):
     endpoint = code_cs.get(INGESTION_ENDPOINT) \
         or env_cs.get(INGESTION_ENDPOINT) \
         or 'https://dc.services.visualstudio.com'
-    options.endpoint = endpoint + '/v2/track'
+    options.endpoint = endpoint
+
+    # Authorization
+    # `azure.core.credentials.TokenCredential` class must be valid
+    if options.credential and not hasattr(options.credential, 'get_token'):
+        raise ValueError(
+            'Must pass in valid TokenCredential.'
+        )
 
     # storage path
     if options.storage_path is None:
@@ -101,6 +108,7 @@ def __init__(self, *args, **kwargs):
 
     _default = BaseObject(
         connection_string=None,
+        credential=None,
         enable_local_storage=True,
         enable_standard_metrics=True,
         endpoint='https://dc.services.visualstudio.com/v2/track',
diff --git a/contrib/opencensus-ext-azure/opencensus/ext/azure/common/transport.py b/contrib/opencensus-ext-azure/opencensus/ext/azure/common/transport.py
@@ -16,8 +16,11 @@
 import logging
 
 import requests
+from azure.core.exceptions import ClientAuthenticationError
+from azure.identity._exceptions import CredentialUnavailableError
 
 logger = logging.getLogger(__name__)
+_MONITOR_OAUTH_SCOPE = "https://monitor.azure.com//.default"
 
 
 class TransportMixin(object):
@@ -45,25 +48,45 @@ def _transmit(self, envelopes):
         if not envelopes:
             return 0
         try:
+            headers = {
+                'Accept': 'application/json',
+                'Content-Type': 'application/json; charset=utf-8',
+            }
+            endpoint = self.options.endpoint
+            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'
             response = requests.post(
-                url=self.options.endpoint,
+                url=endpoint,
                 data=json.dumps(envelopes),
-                headers={
-                    'Accept': 'application/json',
-                    'Content-Type': 'application/json; charset=utf-8',
-                },
+                headers=headers,
                 timeout=self.options.timeout,
                 proxies=json.loads(self.options.proxies),
             )
         except requests.Timeout:
             logger.warning(
                 'Request time out. Ingestion may be backed up. Retrying.')
             return self.options.minimum_retry_interval
-        except Exception as ex:  # TODO: consider RequestException
+        except requests.RequestException as ex:
             logger.warning(
                 'Retrying due to transient client side error %s.', ex)
             # client side error (retryable)
             return self.options.minimum_retry_interval
+        except CredentialUnavailableError as ex:
+            logger.warning('Credential error. %s. Dropping telemetry.', ex)
+            return -1
+        except ClientAuthenticationError as ex:
+            logger.warning('Authentication error %s', ex)
+            return self.options.minimum_retry_interval
+        except Exception as ex:
+            logger.warning(
+                'Error when sending request %s. Dropping telemetry.', ex)
+            # Extraneous error (non-retryable)
+            return -1
 
         text = 'N/A'
         data = None
@@ -120,6 +143,24 @@ def _transmit(self, envelopes):
             )
             # server side error (retryable)
             return self.options.minimum_retry_interval
+        # Authentication error
+        if response.status_code == 401:
+            logger.warning(
+                'Authentication error %s: %s.',
+                response.status_code,
+                text,
+            )
+            return self.options.minimum_retry_interval
+        # Forbidden error
+        # Can occur when v2 endpoint is used while AI resource is configured
+        # with disableLocalAuth
+        if response.status_code == 403:
+            logger.warning(
+                'Forbidden error %s: %s.',
+                response.status_code,
+                text,
+            )
+            return self.options.minimum_retry_interval
         logger.error(
             'Non-retryable server side error %s: %s.',
             response.status_code,
diff --git a/contrib/opencensus-ext-azure/setup.py b/contrib/opencensus-ext-azure/setup.py
@@ -39,6 +39,8 @@
     include_package_data=True,
     long_description=open('README.rst').read(),
     install_requires=[
+        'azure-core >= 1.12.0, < 2.0.0',
+        'azure-identity >= 1.5.0, < 2.0.0',
         'opencensus >= 0.8.dev0, < 1.0.0',
         'psutil >= 5.6.3',
         'requests >= 2.19.0',
diff --git a/contrib/opencensus-ext-azure/tests/test_options.py b/contrib/opencensus-ext-azure/tests/test_options.py
@@ -71,7 +71,7 @@ def test_process_options_endpoint_code_cs(self):
             'Authorization=ikey;IngestionEndpoint=456'
         common.process_options(options)
 
-        self.assertEqual(options.endpoint, '123/v2/track')
+        self.assertEqual(options.endpoint, '123')
 
     def test_process_options_endpoint_env_cs(self):
         options = common.Options()
@@ -80,15 +80,15 @@ def test_process_options_endpoint_env_cs(self):
             'Authorization=ikey;IngestionEndpoint=456'
         common.process_options(options)
 
-        self.assertEqual(options.endpoint, '456/v2/track')
+        self.assertEqual(options.endpoint, '456')
 
     def test_process_options_endpoint_default(self):
         options = common.Options()
         options.connection_string = None
         common.process_options(options)
 
         self.assertEqual(options.endpoint,
-                         'https://dc.services.visualstudio.com/v2/track')
+                         'https://dc.services.visualstudio.com')
 
     def test_process_options_proxies_default(self):
         options = common.Options()
diff --git a/contrib/opencensus-ext-azure/tests/test_transport_mixin.py b/contrib/opencensus-ext-azure/tests/test_transport_mixin.py
@@ -19,12 +19,17 @@
 
 import mock
 import requests
+from azure.core.exceptions import ClientAuthenticationError
+from azure.identity._exceptions import CredentialUnavailableError
 
 from opencensus.ext.azure.common import Options
 from opencensus.ext.azure.common.storage import LocalFileStorage
-from opencensus.ext.azure.common.transport import TransportMixin
+from opencensus.ext.azure.common.transport import (
+    _MONITOR_OAUTH_SCOPE,
+    TransportMixin,
+)
 
-TEST_FOLDER = os.path.abspath('.test.storage')
+TEST_FOLDER = os.path.abspath('.test.transport')
 
 
 def setUpModule():
@@ -68,6 +73,39 @@ def test_transmission_pre_timeout(self):
             self.assertIsNone(mixin.storage.get())
             self.assertEqual(len(os.listdir(mixin.storage.path)), 1)
 
+    def test_transmission_pre_req_exception(self):
+        mixin = TransportMixin()
+        mixin.options = Options()
+        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', throw(requests.RequestException)):
+                mixin._transmit_from_storage()
+            self.assertIsNone(mixin.storage.get())
+            self.assertEqual(len(os.listdir(mixin.storage.path)), 1)
+
+    def test_transmission_cred_exception(self):
+        mixin = TransportMixin()
+        mixin.options = Options()
+        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', throw(CredentialUnavailableError)):  # noqa: E501
+                mixin._transmit_from_storage()
+            self.assertIsNone(mixin.storage.get())
+            self.assertEqual(len(os.listdir(mixin.storage.path)), 0)
+
+    def test_transmission_client_exception(self):
+        mixin = TransportMixin()
+        mixin.options = Options()
+        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', throw(ClientAuthenticationError)):
+                mixin._transmit_from_storage()
+            self.assertIsNone(mixin.storage.get())
+            self.assertEqual(len(os.listdir(mixin.storage.path)), 1)
+
     def test_transmission_pre_exception(self):
         mixin = TransportMixin()
         mixin.options = Options()
@@ -77,7 +115,7 @@ def test_transmission_pre_exception(self):
             with mock.patch('requests.post', throw(Exception)):
                 mixin._transmit_from_storage()
             self.assertIsNone(mixin.storage.get())
-            self.assertEqual(len(os.listdir(mixin.storage.path)), 1)
+            self.assertEqual(len(os.listdir(mixin.storage.path)), 0)
 
     @mock.patch('requests.post', return_value=mock.Mock())
     def test_transmission_lease_failure(self, requests_mock):
@@ -120,6 +158,40 @@ def test_transmission_200(self):
             self.assertIsNone(mixin.storage.get())
             self.assertEqual(len(os.listdir(mixin.storage.path)), 0)
 
+    def test_transmission_auth(self):
+        mixin = TransportMixin()
+        mixin.options = Options()
+        url = 'https://dc.services.visualstudio.com'
+        mixin.options.endpoint = url
+        credential = mock.Mock()
+        mixin.options.credential = credential
+        token_mock = mock.Mock()
+        token_mock.token = "test_token"
+        credential.get_token.return_value = token_mock
+        data = '[1, 2, 3]'
+        headers = {
+            'Accept': 'application/json',
+            'Content-Type': 'application/json; charset=utf-8',
+            'Authorization': 'Bearer test_token',
+        }
+        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(200, 'unknown')
+                mixin._transmit_from_storage()
+                post.assert_called_with(
+                    url=url + '/v2.1/track',
+                    data=data,
+                    headers=headers,
+                    timeout=10.0,
+                    proxies={}
+                )
+            credential.get_token.assert_called_with(_MONITOR_OAUTH_SCOPE)
+            self.assertIsNone(mixin.storage.get())
+            self.assertEqual(len(os.listdir(mixin.storage.path)), 0)
+            credential.get_token.assert_called_once()
+
     def test_transmission_206(self):
         mixin = TransportMixin()
         mixin.options = Options()
@@ -201,16 +273,16 @@ def test_transmission_206_bogus(self):
             self.assertIsNone(mixin.storage.get())
             self.assertEqual(len(os.listdir(mixin.storage.path)), 0)
 
-    def test_transmission_400(self):
+    def test_transmission_401(self):
         mixin = TransportMixin()
         mixin.options = Options()
         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(400, '{}')
+                post.return_value = MockResponse(401, '{}')
                 mixin._transmit_from_storage()
-            self.assertEqual(len(os.listdir(mixin.storage.path)), 0)
+            self.assertEqual(len(os.listdir(mixin.storage.path)), 1)
 
     def test_transmission_500(self):
         mixin = TransportMixin()
@@ -223,3 +295,14 @@ def test_transmission_500(self):
                 mixin._transmit_from_storage()
             self.assertIsNone(mixin.storage.get())
             self.assertEqual(len(os.listdir(mixin.storage.path)), 1)
+
+    def test_transmission_400(self):
+        mixin = TransportMixin()
+        mixin.options = Options()
+        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(400, '{}')
+                mixin._transmit_from_storage()
+            self.assertEqual(len(os.listdir(mixin.storage.path)), 0)
diff --git a/tox.ini b/tox.ini
@@ -50,7 +50,7 @@ deps =
   docs: sphinx >= 1.6.3
 
 commands = 
-  unit: py.test --quiet --cov={envdir}/opencensus --cov=context --cov=contrib --cov-report term-missing --cov-config=.coveragerc --cov-fail-under=97 tests/unit/ context/ contrib/
+  unit: py.test --quiet --cov={envdir}/opencensus --cov=context --cov=contrib --cov-report term-missing --cov-config=.coveragerc --cov-fail-under=97 --ignore=contrib/opencensus-ext-datadog tests/unit/ context/ contrib/
   ; TODO system tests
   lint: isort --check-only --diff --recursive .
   lint: flake8 context/ contrib/ opencensus/ tests/ examples/
