Don't use connection to store task handler credentials (#9381)

mik-laj · web-flow · commit 5b680e27e811 · 2020-06-19T16:27:25.000+02:00
diff --git a/airflow/config_templates/airflow_local_settings.py b/airflow/config_templates/airflow_local_settings.py
@@ -218,15 +218,15 @@
 
         DEFAULT_LOGGING_CONFIG['handlers'].update(WASB_REMOTE_HANDLERS)
     elif REMOTE_BASE_LOG_FOLDER.startswith('stackdriver://'):
-        gcp_conn_id = conf.get('core', 'REMOTE_LOG_CONN_ID', fallback=None)
+        key_path = conf.get('logging', 'STACKDRIVER_KEY_PATH', fallback=None)
         # stackdriver://github.com/airflow-tasks => airflow-tasks
         log_name = urlparse(REMOTE_BASE_LOG_FOLDER).path[1:]
         STACKDRIVER_REMOTE_HANDLERS = {
             'task': {
                 'class': 'airflow.utils.log.stackdriver_task_handler.StackdriverTaskHandler',
                 'formatter': 'airflow',
                 'name': log_name,
-                'gcp_conn_id': gcp_conn_id
+                'gcp_key_path': key_path
             }
         }
 
diff --git a/airflow/config_templates/config.yml b/airflow/config_templates/config.yml
@@ -394,6 +394,16 @@
       type: string
       example: ~
       default: ""
+    - name: stackdriver_key_path
+      description: |
+        Path to GCP Credential JSON file. If ommited, authorization based on `the Application Default
+        Credentials
+        <https://cloud.google.com/docs/authentication/production#finding_credentials_automatically>`__ will
+        be used.
+      version_added: ~
+      type: string
+      example: ~
+      default: ""
     - name: remote_base_log_folder
       description: |
         Storage bucket URL for remote logging
diff --git a/airflow/config_templates/default_airflow.cfg b/airflow/config_templates/default_airflow.cfg
@@ -225,6 +225,12 @@ remote_logging = False
 # location.
 remote_log_conn_id =
 
+# Path to GCP Credential JSON file. If ommited, authorization based on `the Application Default
+# Credentials
+# <https://cloud.google.com/docs/authentication/production#finding_credentials_automatically>`__ will
+# be used.
+stackdriver_key_path =
+
 # Storage bucket URL for remote logging
 # S3 buckets should start with "s3://"
 # Cloudwatch log groups should start with "cloudwatch://"
diff --git a/airflow/providers/google/cloud/utils/credentials_provider.py b/airflow/providers/google/cloud/utils/credentials_provider.py
@@ -23,7 +23,7 @@
 import logging
 import tempfile
 from contextlib import ExitStack, contextmanager
-from typing import Dict, Optional, Sequence, Tuple
+from typing import Collection, Dict, Optional, Sequence, Tuple
 from urllib.parse import urlencode
 
 import google.auth
@@ -179,7 +179,8 @@ def provide_gcp_conn_and_credentials(
 def get_credentials_and_project_id(
     key_path: Optional[str] = None,
     keyfile_dict: Optional[Dict[str, str]] = None,
-    scopes: Optional[Sequence[str]] = None,
+    # See: https://github.com/PyCQA/pylint/issues/2377
+    scopes: Optional[Collection[str]] = None,  # pylint: disable=unsubscriptable-object
     delegate_to: Optional[str] = None
 ) -> Tuple[google.auth.credentials.Credentials, str]:
     """
diff --git a/airflow/utils/log/stackdriver_task_handler.py b/airflow/utils/log/stackdriver_task_handler.py
@@ -18,7 +18,7 @@
 Handler that integrates with Stackdriver
 """
 import logging
-from typing import Dict, List, Optional, Tuple, Type
+from typing import Collection, Dict, List, Optional, Tuple, Type
 
 from cached_property import cached_property
 from google.api_core.gapic_v1.client_info import ClientInfo
@@ -28,10 +28,16 @@
 
 from airflow import version
 from airflow.models import TaskInstance
+from airflow.providers.google.cloud.utils.credentials_provider import get_credentials_and_project_id
 
 DEFAULT_LOGGER_NAME = "airflow"
 _GLOBAL_RESOURCE = Resource(type="global", labels={})
 
+_DEFAULT_SCOPESS = frozenset([
+    "https://www.googleapis.com/auth/logging.read",
+    "https://www.googleapis.com/auth/logging.write"
+])
+
 
 class StackdriverTaskHandler(logging.Handler):
     """Handler that directly makes Stackdriver logging API calls.
@@ -45,11 +51,14 @@ class StackdriverTaskHandler(logging.Handler):
 
     This handler supports both an asynchronous and synchronous transport.
 
-    :param gcp_conn_id: Connection ID that will be used for authorization to the Google Cloud Platform.
-        If omitted, authorization based on `the Application Default Credentials
+
+    :param gcp_key_path: Path to GCP Credential JSON file.
+        If ommited, authorization based on `the Application Default Credentials
         <https://cloud.google.com/docs/authentication/production#finding_credentials_automatically>`__ will
         be used.
-    :type gcp_conn_id: str
+    :type gcp_key_path: str
+    :param scopes: OAuth scopes for the credentials,
+    :type scopes: Sequence[str]
     :param name: the name of the custom log in Stackdriver Logging. Defaults
         to 'airflow'. The name of the Python logger will be represented
          in the ``python_logger`` field.
@@ -74,14 +83,18 @@ class StackdriverTaskHandler(logging.Handler):
 
     def __init__(
         self,
-        gcp_conn_id: Optional[str] = None,
+        gcp_key_path: Optional[str] = None,
+        # See: https://github.com/PyCQA/pylint/issues/2377
+        scopes: Optional[Collection[str]] = _DEFAULT_SCOPESS,  # pylint: disable=unsubscriptable-object
         name: str = DEFAULT_LOGGER_NAME,
         transport: Type[Transport] = BackgroundThreadTransport,
         resource: Resource = _GLOBAL_RESOURCE,
         labels: Optional[Dict[str, str]] = None,
     ):
         super().__init__()
-        self.gcp_conn_id = gcp_conn_id
+        self.gcp_key_path: Optional[str] = gcp_key_path
+        # See: https://github.com/PyCQA/pylint/issues/2377
+        self.scopes: Optional[Collection[str]] = scopes  # pylint: disable=unsubscriptable-object
         self.name: str = name
         self.transport_type: Type[Transport] = transport
         self.resource: Resource = resource
@@ -91,14 +104,10 @@ def __init__(
     @cached_property
     def _client(self) -> gcp_logging.Client:
         """Google Cloud Library API client"""
-        if self.gcp_conn_id:
-            from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
-
-            hook = GoogleBaseHook(gcp_conn_id=self.gcp_conn_id)
-            credentials = hook._get_credentials()  # pylint: disable=protected-access
-        else:
-            # Use Application Default Credentials
-            credentials = None
+        credentials = get_credentials_and_project_id(
+            key_path=self.gcp_key_path,
+            scopes=self.scopes,
+        )
         client = gcp_logging.Client(
             credentials=credentials,
             client_info=ClientInfo(client_library_version='airflow_v' + version.version)
diff --git a/docs/howto/write-logs.rst b/docs/howto/write-logs.rst
@@ -309,6 +309,12 @@ For integration with Stackdriver, this option should start with ``stackdriver://
 The path section of the URL specifies the name of the log e.g. ``stackdriver://airflow-tasks`` writes
 logs under the name ``airflow-tasks``.
 
+You can set ``stackdriver_key_path`` option in the ``[logging]`` section to specify the path to `the service
+account key file <https://cloud.google.com/iam/docs/service-accounts>`__.
+If ommited, authorization based on `the Application Default Credentials
+<https://cloud.google.com/docs/authentication/production#finding_credentials_automatically>`__ will
+be used.
+
 By using the ``logging_config_class`` option you can get :ref:`advanced features <write-logs-advanced>` of
 this handler. Details are available in the handler's documentation -
 :class:`~airflow.utils.log.stackdriver_task_handler.StackdriverTaskHandler`.
diff --git a/tests/utils/log/test_stackdriver_task_handler.py b/tests/utils/log/test_stackdriver_task_handler.py
@@ -36,11 +36,11 @@ def _create_list_response(messages, token):
 
 class TestStackdriverLoggingHandlerStandalone(unittest.TestCase):
 
+    @mock.patch('airflow.utils.log.stackdriver_task_handler.get_credentials_and_project_id')
     @mock.patch('airflow.utils.log.stackdriver_task_handler.gcp_logging.Client')
-    def test_should_pass_message_to_client(self, mock_client):
+    def test_should_pass_message_to_client(self, mock_client, mock_get_creds_and_project_id):
         transport_type = mock.MagicMock()
         stackdriver_task_handler = StackdriverTaskHandler(
-            gcp_conn_id=None,
             transport=transport_type,
             labels={"key": 'value'}
         )
@@ -54,7 +54,10 @@ def test_should_pass_message_to_client(self, mock_client):
         transport_type.return_value.send.assert_called_once_with(
             mock.ANY, 'test-message', labels={"key": 'value'}, resource=Resource(type='global', labels={})
         )
-        mock_client.assert_called_once()
+        mock_client.assert_called_once_with(
+            credentials=mock_get_creds_and_project_id.return_value,
+            client_info=mock.ANY
+        )
 
 
 class TestStackdriverLoggingHandlerTask(unittest.TestCase):
@@ -73,8 +76,9 @@ def setUp(self) -> None:
         self.ti.state = State.RUNNING
         self.addCleanup(self.dag.clear)
 
+    @mock.patch('airflow.utils.log.stackdriver_task_handler.get_credentials_and_project_id')
     @mock.patch('airflow.utils.log.stackdriver_task_handler.gcp_logging.Client')
-    def test_should_set_labels(self, mock_client):
+    def test_should_set_labels(self, mock_client, mock_get_creds_and_project_id):
         self.stackdriver_task_handler.set_context(self.ti)
         self.logger.addHandler(self.stackdriver_task_handler)
 
@@ -92,8 +96,9 @@ def test_should_set_labels(self, mock_client):
             mock.ANY, 'test-message', labels=labels, resource=resource
         )
 
+    @mock.patch('airflow.utils.log.stackdriver_task_handler.get_credentials_and_project_id')
     @mock.patch('airflow.utils.log.stackdriver_task_handler.gcp_logging.Client')
-    def test_should_append_labels(self, mock_client):
+    def test_should_append_labels(self, mock_client, mock_get_creds_and_project_id):
         self.stackdriver_task_handler = StackdriverTaskHandler(
             transport=self.transport_mock,
             labels={"product.googleapis.com/task_id": "test-value"}
@@ -116,11 +121,12 @@ def test_should_append_labels(self, mock_client):
             mock.ANY, 'test-message', labels=labels, resource=resource
         )
 
+    @mock.patch('airflow.utils.log.stackdriver_task_handler.get_credentials_and_project_id')
     @mock.patch(
         'airflow.utils.log.stackdriver_task_handler.gcp_logging.Client',
         **{'return_value.project': 'asf-project'}  # type: ignore
     )
-    def test_should_read_logs_for_all_try(self, mock_client):
+    def test_should_read_logs_for_all_try(self, mock_client, mock_get_creds_and_project_id):
         mock_client.return_value.list_entries.return_value = _create_list_response(["MSG1", "MSG2"], None)
 
         logs, metadata = self.stackdriver_task_handler.read(self.ti)
@@ -135,11 +141,12 @@ def test_should_read_logs_for_all_try(self, mock_client):
         self.assertEqual(['MSG1\nMSG2'], logs)
         self.assertEqual([{'end_of_log': True}], metadata)
 
+    @mock.patch('airflow.utils.log.stackdriver_task_handler.get_credentials_and_project_id')
     @mock.patch(  # type: ignore
         'airflow.utils.log.stackdriver_task_handler.gcp_logging.Client',
         **{'return_value.project': 'asf-project'}  # type: ignore
     )
-    def test_should_read_logs_for_task_with_quote(self, mock_client):
+    def test_should_read_logs_for_task_with_quote(self, mock_client, mock_get_creds_and_project_id):
         mock_client.return_value.list_entries.return_value = _create_list_response(["MSG1", "MSG2"], None)
         self.ti.task_id = "K\"OT"
         logs, metadata = self.stackdriver_task_handler.read(self.ti)
@@ -154,11 +161,12 @@ def test_should_read_logs_for_task_with_quote(self, mock_client):
         self.assertEqual(['MSG1\nMSG2'], logs)
         self.assertEqual([{'end_of_log': True}], metadata)
 
+    @mock.patch('airflow.utils.log.stackdriver_task_handler.get_credentials_and_project_id')
     @mock.patch(
         'airflow.utils.log.stackdriver_task_handler.gcp_logging.Client',
         **{'return_value.project': 'asf-project'}  # type: ignore
     )
-    def test_should_read_logs_for_single_try(self, mock_client):
+    def test_should_read_logs_for_single_try(self, mock_client, mock_get_creds_and_project_id):
         mock_client.return_value.list_entries.return_value = _create_list_response(["MSG1", "MSG2"], None)
 
         logs, metadata = self.stackdriver_task_handler.read(self.ti, 3)
@@ -174,8 +182,9 @@ def test_should_read_logs_for_single_try(self, mock_client):
         self.assertEqual(['MSG1\nMSG2'], logs)
         self.assertEqual([{'end_of_log': True}], metadata)
 
+    @mock.patch('airflow.utils.log.stackdriver_task_handler.get_credentials_and_project_id')
     @mock.patch('airflow.utils.log.stackdriver_task_handler.gcp_logging.Client')
-    def test_should_read_logs_with_pagination(self, mock_client):
+    def test_should_read_logs_with_pagination(self, mock_client, mock_get_creds_and_project_id):
         mock_client.return_value.list_entries.side_effect = [
             _create_list_response(["MSG1", "MSG2"], "TOKEN1"),
             _create_list_response(["MSG3", "MSG4"], None),
@@ -195,8 +204,9 @@ def test_should_read_logs_with_pagination(self, mock_client):
         self.assertEqual(['MSG3\nMSG4'], logs)
         self.assertEqual([{'end_of_log': True}], metadata2)
 
+    @mock.patch('airflow.utils.log.stackdriver_task_handler.get_credentials_and_project_id')
     @mock.patch('airflow.utils.log.stackdriver_task_handler.gcp_logging.Client')
-    def test_should_read_logs_with_download(self, mock_client):
+    def test_should_read_logs_with_download(self, mock_client, mock_get_creds_and_project_id):
         mock_client.return_value.list_entries.side_effect = [
             _create_list_response(["MSG1", "MSG2"], "TOKEN1"),
             _create_list_response(["MSG3", "MSG4"], None),
@@ -207,11 +217,12 @@ def test_should_read_logs_with_download(self, mock_client):
         self.assertEqual(['MSG1\nMSG2\nMSG3\nMSG4'], logs)
         self.assertEqual([{'end_of_log': True}], metadata1)
 
+    @mock.patch('airflow.utils.log.stackdriver_task_handler.get_credentials_and_project_id')
     @mock.patch(
         'airflow.utils.log.stackdriver_task_handler.gcp_logging.Client',
         **{'return_value.project': 'asf-project'}  # type: ignore
     )
-    def test_should_read_logs_with_custom_resources(self, mock_client):
+    def test_should_read_logs_with_custom_resources(self, mock_client, mock_get_creds_and_project_id):
         resource = Resource(
             type="cloud_composer_environment",
             labels={
@@ -245,31 +256,26 @@ def test_should_read_logs_with_custom_resources(self, mock_client):
         self.assertEqual(['TEXT\nTEXT'], logs)
         self.assertEqual([{'end_of_log': True}], metadata)
 
-
-class TestStackdriverTaskHandlerAuthorization(unittest.TestCase):
-
+    @mock.patch('airflow.utils.log.stackdriver_task_handler.get_credentials_and_project_id')
     @mock.patch('airflow.utils.log.stackdriver_task_handler.gcp_logging.Client')
-    def test_should_fallback_to_adc(self, mock_client):
+    def test_should_use_credentials(self, mock_client, mock_get_creds_and_project_id):
         stackdriver_task_handler = StackdriverTaskHandler(
-            gcp_conn_id=None
+            gcp_key_path="KEY_PATH",
         )
 
         client = stackdriver_task_handler._client
 
-        mock_client.assert_called_once_with(credentials=None, client_info=mock.ANY)
-        self.assertEqual(mock_client.return_value, client)
-
-    @mock.patch("airflow.providers.google.common.hooks.base_google.GoogleBaseHook")
-    @mock.patch('airflow.utils.log.stackdriver_task_handler.gcp_logging.Client')
-    def test_should_support_gcp_conn_id(self, mock_client, mock_hook):
-        stackdriver_task_handler = StackdriverTaskHandler(
-            gcp_conn_id="test-gcp-conn"
+        mock_get_creds_and_project_id.assert_called_once_with(
+            key_path='KEY_PATH',
+            scopes=frozenset(
+                {
+                    'https://www.googleapis.com/auth/logging.write',
+                    'https://www.googleapis.com/auth/logging.read'
+                }
+            )
         )
-
-        client = stackdriver_task_handler._client
-
         mock_client.assert_called_once_with(
-            credentials=mock_hook.return_value._get_credentials.return_value,
+            credentials=mock_get_creds_and_project_id.return_value,
             client_info=mock.ANY
         )
         self.assertEqual(mock_client.return_value, client)
