Add key authentication option

snowflake/assets/configuration/spec.yaml

@@ -114,6 +114,20 @@ files:
         value:
           type: string
           example: /path/to/token
+      - name: private_key_path
+        description: |
+          The path to the file that contains the private key used to connect to Snowflake.
+          The key is re-read at every check run.
+        value:
+          type: string
+          display_default: null
+          example: /path/to/private_key
+      - name: private_key_password
+        secret: true
+        description: |
+          The password protecting the private key file in the `private_key_path` option.
+        value:
+          type: string
       - name: client_session_keep_alive
         description: |
           If set to true, Snowflake keeps the session active indefinitely as long as the connection is active,

snowflake/datadog_checks/snowflake/check.py

@@ -4,6 +4,8 @@
 from contextlib import closing
 
 import snowflake.connector as sf
+from cryptography.hazmat.backends import default_backend
+from cryptography.hazmat.primitives import serialization
 
 from datadog_checks.base import AgentCheck, ConfigurationError, to_native_string
 from datadog_checks.base.utils.db import QueryManager
@@ -77,10 +79,30 @@ def __init__(self, *args, **kwargs):
         self._query_manager = QueryManager(self, self.execute_query_raw, queries=self.metric_queries, tags=self._tags)
         self.check_initializations.append(self._query_manager.compile_queries)
 
-    def renew_token(self):
-        self.log.debug("Renewing Snowflake client token")
-        with open(self._config.token_path, 'rb', encoding="UTF-8") as f:
-            self._config.token = f.read()
+    def read_token(self):
+        if self._config.token_path:
+            self.log.debug("Renewing Snowflake client token")
+            with open(self._config.token_path, 'rb', encoding="UTF-8") as f:
+                self._config.token = f.read()
+
+    def read_key(self):
+        if self._config.private_key_path:
+            self.log.debug("Reading Snowflake client key for key pair authentication")
+            # https://docs.snowflake.com/en/user-guide/python-connector-example.html#using-key-pair-authentication-key-pair-rotation
+            with open(self._config.private_key_path, "rb") as key:
+                p_key = serialization.load_pem_private_key(
+                    key.read(), password=self._config.private_key_password, backend=default_backend()
+                )
+
+                pkb = p_key.private_bytes(
+                    encoding=serialization.Encoding.DER,
+                    format=serialization.PrivateFormat.PKCS8,
+                    encryption_algorithm=serialization.NoEncryption(),
+                )
+
+                return pkb
+
+        return None
 
     def check(self, _):
         if self.instance.get('user'):
@@ -126,8 +148,7 @@ def connect(self):
             self.proxy_port,
         )
 
-        if self._config.token_path:
-            self.renew_token()
+        self.read_token()
 
         try:
             conn = sf.connect(
@@ -145,6 +166,7 @@ def connect(self):
                 ocsp_response_cache_filename=self._config.ocsp_response_cache_filename,
                 authenticator=self._config.authenticator,
                 token=self._config.token,
+                private_key=self.read_key(),
                 client_session_keep_alive=self._config.client_keep_alive,
                 proxy_host=self.proxy_host,
                 proxy_port=self.proxy_port,

snowflake/datadog_checks/snowflake/config.py

@@ -3,7 +3,7 @@
 # Licensed under a 3-clause BSD style license (see LICENSE)
 from typing import List, Optional
 
-from datadog_checks.base import ConfigurationError, is_affirmative
+from datadog_checks.base import ConfigurationError, ensure_bytes, is_affirmative
 
 
 class Config(object):
@@ -19,7 +19,7 @@ class Config(object):
         'snowflake.logins',
     ]
 
-    AUTHENTICATION_MODES = ['snowflake', 'oauth']
+    AUTHENTICATION_MODES = ['snowflake', 'oauth', 'snowflake_jwt']
 
     def __init__(self, instance=None):
         if instance is None:
@@ -41,6 +41,8 @@ def __init__(self, instance=None):
         authenticator = instance.get('authenticator', 'snowflake')
         token = instance.get('token', None)
         token_path = instance.get('token_path', None)
+        private_key_path = instance.get('private_key_path', None)
+        private_key_password = ensure_bytes(instance.get('private_key_password', None))
         client_keep_alive = instance.get('client_session_keep_alive', False)
         aggregate_last_24_hours = instance.get('aggregate_last_24_hours', False)
 
@@ -56,8 +58,8 @@ def __init__(self, instance=None):
             if is_affirmative(passcode_in_password) and passcode is None:
                 raise ConfigurationError('MFA enabled, please specify a passcode')
 
-            if password is None:
-                raise ConfigurationError('Must specify a password if using snowflake authentication')
+            if password is None and private_key_path is None:
+                raise ConfigurationError('Must specify a password or a private key if using snowflake authentication')
 
         elif authenticator == 'oauth':
             if token is None and token_path is None:
@@ -92,5 +94,7 @@ def __init__(self, instance=None):
         self.authenticator = authenticator
         self.token = token
         self.token_path = token_path
+        self.private_key_path = private_key_path
+        self.private_key_password = private_key_password
         self.client_keep_alive = client_keep_alive
         self.aggregate_last_24_hours = aggregate_last_24_hours

snowflake/datadog_checks/snowflake/config_models/defaults.py

@@ -90,6 +90,14 @@ def instance_password(field, value):
     return get_default_field_value(field, value)
 
 
+def instance_private_key_password(field, value):
+    return get_default_field_value(field, value)
+
+
+def instance_private_key_path(field, value):
+    return get_default_field_value(field, value)
+
+
 def instance_schema_(field, value):
     return 'ACCOUNT_USAGE'
 

snowflake/datadog_checks/snowflake/config_models/instance.py

@@ -47,6 +47,8 @@ class Config:
     ocsp_response_cache_filename: Optional[str]
     only_custom_queries: Optional[bool]
     password: Optional[str]
+    private_key_password: Optional[str]
+    private_key_path: Optional[str]
     role: str
     schema_: Optional[str] = Field(None, alias='schema')
     service: Optional[str]

snowflake/datadog_checks/snowflake/config_models/validators.py

@@ -8,4 +8,10 @@ def initialize_instance(values, **kwargs):
     if 'username' not in values and 'user' in values:
         values['username'] = values['user']
 
+    if 'private_key_password' in values and 'private_key_path' not in values:
+        raise Exception(
+            'Option `private_key_password` is set but not option `private_key_path`. '
+            'Set `private_key_path` or remove `private_key_password` entry.'
+        )
+
     return values

snowflake/datadog_checks/snowflake/data/conf.yaml.example

@@ -122,6 +122,17 @@ instances:
     #
     # token_path: /path/to/token
 
+    ## @param private_key_path - string - optional
+    ## The path to the file that contains the private key used to connect to Snowflake.
+    ## The key is re-read at every check run.
+    #
+    # private_key_path: /path/to/private_key
+
+    ## @param private_key_password - string - optional
+    ## The password protecting the private key file in the `private_key_path` option.
+    #
+    # private_key_password: <PRIVATE_KEY_PASSWORD>
+
     ## @param client_session_keep_alive - boolean - optional - default: false
     ## If set to true, Snowflake keeps the session active indefinitely as long as the connection is active,
     ## even if there is no activity from the user.

snowflake/tests/common.py

@@ -42,5 +42,8 @@
         'pip install --upgrade setuptools',
         'pip install /home/snowflake_connector_patch',
     ],
-    'docker_volumes': ['{}:/home/snowflake_connector_patch'.format(os.path.join(HERE, 'snowflake_connector_patch'))],
+    'docker_volumes': [
+        '{}:/home/snowflake_connector_patch'.format(os.path.join(HERE, 'snowflake_connector_patch')),
+        '{}:/home/keys'.format(os.path.join(HERE, 'keys')),
+    ],
 }

snowflake/tests/keys/rsa_key_example.p8

@@ -0,0 +1,28 @@
+-----BEGIN PRIVATE KEY-----
+MIIEvAIBADANBgkqhkiG9w0BAQEFAASCBKYwggSiAgEAAoIBAQCXeqIH2MXr6E7p
+utpe+ZC/31oMMgXi7EPqjeNdbEMkM5c1hCLSyy762mlafkvxqEaL+OIEamZXDtOq
+/cenVogdu5gvP/mMkijxWnMg7CNOtg6ulME2UqOPiMRQRJm9iXFoqx1xLSVSUh08
+0ydbh2k8+yav/eRqaAJyREOE0IvafyeG1ocHTL/D6eHos3epB9T4OLXS6UjOvLoJ
+9eu4GLj7R6GuMB1X15RlmYP+XLWWO2pKjlfBdY8KxQ48Q5GQHpKV10ytuH36VO0q
+UvGRy5svzQEcI7VX/HCLkiFtbL7yQ+gyKACYql2CpBaBR0IDYAyx3KjrFdLSVkZ+
+goWNcTThAgMBAAECggEAYPKDsTHzukA8ASsth4uCMMrp/tQlIE7GSN/2nFwlvI2o
+QQAIqZCQyMvwkZIzWL6yJ/Np9BmE1kMPBWjW1ypyg1UE2XjAZk9FFPMmq/N1xXKP
+gvyBjBrnw31s51KAcuX8R1j8xup18HHNdJhPoCzSFc1HvWtpPRDEQA2giOhQGc40
+xcz34161dj8BAeeESHaIfed7WTwBu1IBTpVmAKPJudlAf3EjiAcO/pclLS7tJwN5
+qsX0K3oHkm4W1aaOLo/G4akUxJ72bS8ny9w8wvJmJKIJ93koHwJ5HnThL5UjGhCD
+cAJPCdUYm1oUtKAsTkIh2SC/sGBH28Q20D5UGHms4QKBgQD4UgvXIQYIcUv0M0qA
+vGxyRDsG474Q8PqnClRAEZc/NA3S0g/B+jq0kyq7J3DLBtLCDEPmjBKfAbm39D2t
+TE4sebo2/CTYmJy5FgiJxdSaEoZOrPAZAJfpkN3lfohN6McIT+0YUTvdgnDk/5xo
+NwPQQEfdN9OtWMqvaFMrr6gf8wKBgQCcKeSH2/t39r79LmZvxW42+dBQvhoDVR4J
+OsTYx058YWdAbWBy2CccCkDdOQnuPBsqIkQy5YQ4E6dt5AqOP9Kcl4xrfyocvQpy
+ULa/sLl5xZTtKf6iyRhOuZS+LSgychoc3Y50xyNKoVAMv2ddvQnEQrEDQizbO6gW
+43idKEKg2wKBgAzukOlKMfs8kz0LcsTTiz5EKWLJd3uAYT1Tv2F6yQqklle1UtbC
+Rk5jH6WRf0FDgLRUWTDneIzJVTesQ44D3EpaqIT2iqCxCfBloloycEj5z/7G6NYU
+ftTOE5BBD64nAj5/kxRiHqEBiwmR+j4/JzawMk3l+2MarauG3lX3FuVbAoGADoVe
+uLtd4MPS8pvz7oS/QOFt23Qx2wl5J4aNc1LlG2+7OCRziXpL+LGDYo7BO6PfKsXQ
+7aKl7sj1EqTXzm5k2SbGaeCDO/TgGc0jkSOPu6EBviPfh6eHWRqsmBp+2GH/x5ta
+ecVipLfnR6gspmzDkbpZ12G55hDgCnDQcFykBW0CgYBeYRRfvZuJVDbA7M3DygfI
+DCNPpgcjaECt5Uz/3N0udBNsoXY9vRakRlhKANLYesMUd76mUjVWu6Kc+EJtyEcC
+q2plTc4gowD/hhxT1Fw95p6yxz8m+eHFc8BWAaHtoGMMgK1/wm1UXYe/hQDu39g9
+ojckZpT64MHda9CvnF814g==
+-----END PRIVATE KEY-----

snowflake/tests/keys/rsa_key_pass_example.p8

@@ -0,0 +1,29 @@
+-----BEGIN ENCRYPTED PRIVATE KEY-----
+MIIE6TAbBgkqhkiG9w0BBQMwDgQIH+eVVXditq8CAggABIIEyLIREkj8xviusabN
+Eb9xkeuJJ2AMUHZjDlrhlUDN8LwZc5REcxvv1Y+aEvAaiVFjVM8NdhLLmreW6/mf
+1P2Ms0p8dUXg4w0ulU0dfoMwDSTz4vy/F66+pRNhctcMCxsbUITtYwe5tFKffJuC
+ovtuj0WbqgDngMiB+LUI9MDxBCKsS0tYoIE9RlNCplWZmNW8dY0rTAwJx0K0D5Hm
+cxZkj4E95wn2Jw57jdnbw1tEZyDAmR+v1aKVCtD39V8rAvBTBlfPAtD5IPLtOnem
+gvPv6EuSjwMmuKREBSkrZr/ibPbPcJwSvHBoQH1YVk5qQjXp+qXjrPA0xGFmYZ5O
+fc4kbMhZfB9QbioVql3wBd1we+3au+lC+Xao0RniOaX3waVOvvlsZFRBCnULuAVb
+YdezqDTpukxhi1Zp9+nXNFf0jQdKbEAcMP9MMe14IGTH0Zqs9g63jzOig5V1wU3a
+/ghalyC9RxnICimo/JObRl3hIgLK2VSoLIxguRdrScKXAFv2RCk6QuoavJwoxfdc
+ElXhHYuXM6+QzCHfaT3G1NPg9aJ40Ycp08Js/5DLzVjhdwBtQuJCpO+c70aIM5qw
+zb/4ZvOU+vFQVZV5cEYJEuHBvQ65wnJPhrt6tWi87xlW2Enbx9jNKThMeT6UtV6M
+Rapy8nQp7Tg48PkuK6m7FrnBOrf6g+QmqpEuO+G7wH4b+CiLzcGiivmg3LkNXBO3
+YRzY1xYLemowwdtLKFVMw4otXQqjRey6TU54MNLMGLg8kmqBmOi0ti+noInWMsT7
+sohpq7KhZtsDICl2sTxSz8UKZJG8vIuhBGoVfl3YKtp+M5yTB4/uGi3urV4jg266
+FRRfvEvRFh6Bvu016L0ClXXnsOyqcd9vAie+bQUY39L/WTetMh8U/9AKDUUvJONa
+NJf4wnnlhfyVx8Tg2DoPf7V0ekK2fPfmvxhlM+RVPKw1GxXzum2q4TKPBZaZJAR6
+aGJwb5FtfM6EmTa5vJBGOxfOf9cVAu67OoMAyAXYQDFM+ruOXzTrIa5No67FRXkt
+LUWgRClSYPg0tEeBiNut4M5ZyxIwNBKG3Y3gFyUbMVpj8EfnkmAYTUJCKeXKnor6
+JmtAAYYGfugBzdyZ74qYdEqai1cMkTE2NNm2xeJBrq5Cg3e1qWnnJk8+OZkFAxNF
+i2eNEzcgwQDhz40M/3gPzEjYtQmaVibtN/kN4/UMGYHrKUIqZ4HeddouZf30wXo2
+ovjhFiXJBE2HgDYDJ42rCDhgMWir9XFgsM8+51Xg0sNBDIAa2PRhnUM46jLMWZo/
+PU/6vK1jKPejeuXAeZwGSFJOdlC7U7X4+zZ8gBtD7YiiqJMmKl/0Go6ACXokB7r7
+i097/NOpH0Q0z4TsC8xEHinDTepmTskAaEHEFWrULQ3kd5I3sym2OHRz/wrxOz6v
+9IWshaEkgp+UmdPeh0qiZwP6D9wuRSBeb5RFTfPebVlDnUbPvcItaI8oS1ux6WJe
+SOA16NvCjeyjjnWKFKKYV2VFbIcZevO+GwsPh0zmtPdAgb2dAATfjb1PLbV2yaCs
+nAjIhDEsG54pmJA+MCdY/auJeDqu8OrdftLR6gCDqNRJOM7c0bM09dstkLXsiadm
+Izh71I5VdjD088Wdeg==
+-----END ENCRYPTED PRIVATE KEY-----

snowflake/tests/test_snowflake.py

@@ -183,6 +183,7 @@ def test_token_path(dd_run_check, aggregator):
         'ocsp_response_cache_filename': None,
         'authenticator': 'oauth',
         'client_session_keep_alive': False,
+        'private_key': None,
         'proxy_host': None,
         'proxy_port': None,
         'proxy_user': None,

snowflake/tests/test_unit.py

@@ -2,6 +2,7 @@
 # All rights reserved
 # Licensed under a 3-clause BSD style license (see LICENSE)
 import copy
+import os
 
 import mock
 import pytest
@@ -81,6 +82,7 @@ def test_default_auth(instance):
             ocsp_response_cache_filename=None,
             authenticator='snowflake',
             token=None,
+            private_key=None,
             client_session_keep_alive=False,
             proxy_host=None,
             proxy_port=None,
@@ -113,6 +115,7 @@ def test_oauth_auth(oauth_instance):
             ocsp_response_cache_filename=None,
             authenticator='oauth',
             token='testtoken',
+            private_key=None,
             client_session_keep_alive=False,
             proxy_host=None,
             proxy_port=None,
@@ -121,6 +124,54 @@ def test_oauth_auth(oauth_instance):
         )
 
 
+def test_key_auth(dd_run_check, instance):
+    # Key auth
+    inst = copy.deepcopy(instance)
+    inst['private_key_path'] = os.path.join(os.path.dirname(__file__), 'keys', 'rsa_key_example.p8')
+
+    check = SnowflakeCheck(CHECK_NAME, {}, [inst])
+    # Checking size instead of the read key
+    read_key = check.read_key()
+    assert len(read_key) == 1216
+
+    with mock.patch('datadog_checks.snowflake.check.sf') as sf:
+        check = SnowflakeCheck(CHECK_NAME, {}, [inst])
+        dd_run_check(check)
+        sf.connect.assert_called_with(
+            user='testuser',
+            password='pass',
+            account='test_acct.us-central1.gcp',
+            database='SNOWFLAKE',
+            schema='ACCOUNT_USAGE',
+            warehouse=None,
+            role='ACCOUNTADMIN',
+            passcode_in_password=False,
+            passcode=None,
+            client_prefetch_threads=4,
+            login_timeout=3,
+            ocsp_response_cache_filename=None,
+            authenticator='snowflake',
+            token=None,
+            private_key=read_key,
+            client_session_keep_alive=False,
+            proxy_host=None,
+            proxy_port=None,
+            proxy_user=None,
+            proxy_password=None,
+        )
+
+    inst['private_key_path'] = os.path.join(os.path.dirname(__file__), 'keys', 'wrong_key.p8')
+    check = SnowflakeCheck(CHECK_NAME, {}, [inst])
+    with pytest.raises(FileNotFoundError):
+        check.read_key()
+
+    # Read key protected by a passphrase
+    inst['private_key_path'] = os.path.join(os.path.dirname(__file__), 'keys', 'rsa_key_pass_example.p8')
+    inst['private_key_password'] = 'keypass'
+    check = SnowflakeCheck(CHECK_NAME, {}, [inst])
+    assert len(check.read_key()) == 1218
+
+
 def test_proxy_settings(instance):
     init_config = {
         'proxy_host': 'testhost',
@@ -149,6 +200,7 @@ def test_proxy_settings(instance):
             ocsp_response_cache_filename=None,
             authenticator='snowflake',
             token=None,
+            private_key=None,
             client_session_keep_alive=False,
             proxy_host='testhost',
             proxy_port=8000,