From 0fa6449d2d7c2d69e4d41c07fcc8836a5985a7fd Mon Sep 17 00:00:00 2001 From: Paul Coignet Date: Tue, 18 Jan 2022 14:08:04 +0100 Subject: [PATCH 01/13] Add token_path config option --- .../snowflake/config_models/defaults.py | 2 +- .../snowflake/config_models/instance.py | 2 +- .../snowflake/data/conf.yaml.example | 14 ++++++++++++++ 3 files changed, 16 insertions(+), 2 deletions(-) diff --git a/snowflake/datadog_checks/snowflake/config_models/defaults.py b/snowflake/datadog_checks/snowflake/config_models/defaults.py index 89ccae578127a..d790eae82c684 100644 --- a/snowflake/datadog_checks/snowflake/config_models/defaults.py +++ b/snowflake/datadog_checks/snowflake/config_models/defaults.py @@ -1,4 +1,4 @@ -# (C) Datadog, Inc. 2021-present +# (C) Datadog, Inc. 2022-present # All rights reserved # Licensed under a 3-clause BSD style license (see LICENSE) diff --git a/snowflake/datadog_checks/snowflake/config_models/instance.py b/snowflake/datadog_checks/snowflake/config_models/instance.py index 6fa5cf0547b7c..cf8d4b2a66aa3 100644 --- a/snowflake/datadog_checks/snowflake/config_models/instance.py +++ b/snowflake/datadog_checks/snowflake/config_models/instance.py @@ -1,4 +1,4 @@ -# (C) Datadog, Inc. 2021-present +# (C) Datadog, Inc. 2022-present # All rights reserved # Licensed under a 3-clause BSD style license (see LICENSE) diff --git a/snowflake/datadog_checks/snowflake/data/conf.yaml.example b/snowflake/datadog_checks/snowflake/data/conf.yaml.example index b75f035990424..3702aa7652088 100644 --- a/snowflake/datadog_checks/snowflake/data/conf.yaml.example +++ b/snowflake/datadog_checks/snowflake/data/conf.yaml.example @@ -106,19 +106,33 @@ instances: ## @param authenticator - string - optional - default: snowflake ## Authenticator for Snowflake. The default `snowflake` uses the internal Snowflake authenticator. +<<<<<<< HEAD ## Use `oauth` to authenticate with OAuth. Also, set the `token` or `token_path` option. +======= + ## Use `oauth` to authenticate with OAuth, be sure to set the `token` or `token_path` option as well. +>>>>>>> 78003da9e (Add token_path config option) # # authenticator: ## @param token - string - optional +<<<<<<< HEAD ## Token used for OAuth connection to Snowflake. You cannot use this alongside `token_path`. +======= + ## Token used for OAuth connection to Snowflake. Can't be used along side with `token_path`. +>>>>>>> 78003da9e (Add token_path config option) # # token: ## @param token_path - string - optional - default: /path/to/token +<<<<<<< HEAD ## Path to the file that contains the token used for OAuth connection to Snowflake. ## You cannot use this alongside `token`. ## The token is re-read at every check run. +======= + ## Path to the file containing the token used for OAuth connection to snowflake. + ## Can't be used along side with `token`. + ## The token will be re-read at each check run. +>>>>>>> 78003da9e (Add token_path config option) # # token_path: /path/to/token From 962c9be5279e2617930e0b95207f56b44d8287e7 Mon Sep 17 00:00:00 2001 From: Paul Coignet Date: Wed, 19 Jan 2022 14:47:25 +0100 Subject: [PATCH 02/13] Sync spec --- .../snowflake/data/conf.yaml.example | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/snowflake/datadog_checks/snowflake/data/conf.yaml.example b/snowflake/datadog_checks/snowflake/data/conf.yaml.example index 3702aa7652088..f6363227a48cf 100644 --- a/snowflake/datadog_checks/snowflake/data/conf.yaml.example +++ b/snowflake/datadog_checks/snowflake/data/conf.yaml.example @@ -106,24 +106,33 @@ instances: ## @param authenticator - string - optional - default: snowflake ## Authenticator for Snowflake. The default `snowflake` uses the internal Snowflake authenticator. +<<<<<<< HEAD <<<<<<< HEAD ## Use `oauth` to authenticate with OAuth. Also, set the `token` or `token_path` option. ======= ## Use `oauth` to authenticate with OAuth, be sure to set the `token` or `token_path` option as well. >>>>>>> 78003da9e (Add token_path config option) +======= + ## Use `oauth` to authenticate with OAuth. Also, set the `token` or `token_path` option. +>>>>>>> 1f30fba28 (Sync spec) # # authenticator: ## @param token - string - optional +<<<<<<< HEAD <<<<<<< HEAD ## Token used for OAuth connection to Snowflake. You cannot use this alongside `token_path`. ======= ## Token used for OAuth connection to Snowflake. Can't be used along side with `token_path`. >>>>>>> 78003da9e (Add token_path config option) +======= + ## Token used for OAuth connection to Snowflake. You cannot use this alongside `token_path`. +>>>>>>> 1f30fba28 (Sync spec) # # token: ## @param token_path - string - optional - default: /path/to/token +<<<<<<< HEAD <<<<<<< HEAD ## Path to the file that contains the token used for OAuth connection to Snowflake. ## You cannot use this alongside `token`. @@ -133,6 +142,11 @@ instances: ## Can't be used along side with `token`. ## The token will be re-read at each check run. >>>>>>> 78003da9e (Add token_path config option) +======= + ## Path to the file that contains the token used for OAuth connection to Snowflake. + ## You cannot use this alongside `token`. + ## The token is re-read at every check run. +>>>>>>> 1f30fba28 (Sync spec) # # token_path: /path/to/token From 29a0b0d9d1ad7d369fffa866e574a3bde88f3a57 Mon Sep 17 00:00:00 2001 From: Paul Coignet Date: Wed, 19 Jan 2022 14:49:22 +0100 Subject: [PATCH 03/13] Add key auth option --- snowflake/assets/configuration/spec.yaml | 13 +++++++++++ snowflake/datadog_checks/snowflake/check.py | 22 +++++++++++++++++++ snowflake/datadog_checks/snowflake/config.py | 10 ++++++--- .../snowflake/config_models/defaults.py | 8 +++++++ .../snowflake/config_models/instance.py | 2 ++ .../snowflake/config_models/validators.py | 6 +++++ .../snowflake/data/conf.yaml.example | 11 ++++++++++ 7 files changed, 69 insertions(+), 3 deletions(-) diff --git a/snowflake/assets/configuration/spec.yaml b/snowflake/assets/configuration/spec.yaml index fbe56db15f604..d0de3ccb37edd 100644 --- a/snowflake/assets/configuration/spec.yaml +++ b/snowflake/assets/configuration/spec.yaml @@ -114,6 +114,19 @@ files: value: type: string example: /path/to/token + - name: private_key_path + description: | + 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 + example: /path/to/private_key + - name: private_key_password + secret: true + description: | + Password protecting the private key file in `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, diff --git a/snowflake/datadog_checks/snowflake/check.py b/snowflake/datadog_checks/snowflake/check.py index 93bc6473fc0fe..d6778e592c81c 100644 --- a/snowflake/datadog_checks/snowflake/check.py +++ b/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 @@ -82,6 +84,25 @@ def renew_token(self): 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'): self._log_deprecation('_config_renamed', 'user', 'username') @@ -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, diff --git a/snowflake/datadog_checks/snowflake/config.py b/snowflake/datadog_checks/snowflake/config.py index e180032921ee2..dc123cd05c650 100644 --- a/snowflake/datadog_checks/snowflake/config.py +++ b/snowflake/datadog_checks/snowflake/config.py @@ -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 = 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 diff --git a/snowflake/datadog_checks/snowflake/config_models/defaults.py b/snowflake/datadog_checks/snowflake/config_models/defaults.py index d790eae82c684..42fe90c1aa6d8 100644 --- a/snowflake/datadog_checks/snowflake/config_models/defaults.py +++ b/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 '/path/to/private_key' + + def instance_schema_(field, value): return 'ACCOUNT_USAGE' diff --git a/snowflake/datadog_checks/snowflake/config_models/instance.py b/snowflake/datadog_checks/snowflake/config_models/instance.py index cf8d4b2a66aa3..a9cc51c69ab27 100644 --- a/snowflake/datadog_checks/snowflake/config_models/instance.py +++ b/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] diff --git a/snowflake/datadog_checks/snowflake/config_models/validators.py b/snowflake/datadog_checks/snowflake/config_models/validators.py index 96996e5bfa2e2..9a6d2b7d5a72b 100644 --- a/snowflake/datadog_checks/snowflake/config_models/validators.py +++ b/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 diff --git a/snowflake/datadog_checks/snowflake/data/conf.yaml.example b/snowflake/datadog_checks/snowflake/data/conf.yaml.example index f6363227a48cf..5e4468e1ea158 100644 --- a/snowflake/datadog_checks/snowflake/data/conf.yaml.example +++ b/snowflake/datadog_checks/snowflake/data/conf.yaml.example @@ -150,6 +150,17 @@ instances: # # token_path: /path/to/token + ## @param private_key_path - string - optional - default: /path/to/private_key + ## 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 + ## Password protecting the private key file in `private_key_path` option. + # + # 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. From 1f9995a9edf527aa00ea10e4e87993a89c18c027 Mon Sep 17 00:00:00 2001 From: Paul Coignet Date: Thu, 20 Jan 2022 14:23:50 +0100 Subject: [PATCH 04/13] Fix tests --- snowflake/tests/test_snowflake.py | 1 + snowflake/tests/test_unit.py | 3 +++ 2 files changed, 4 insertions(+) diff --git a/snowflake/tests/test_snowflake.py b/snowflake/tests/test_snowflake.py index 4890261880f0d..13683e5ccc5cc 100644 --- a/snowflake/tests/test_snowflake.py +++ b/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, diff --git a/snowflake/tests/test_unit.py b/snowflake/tests/test_unit.py index 3787414b40d1e..43167058ba923 100644 --- a/snowflake/tests/test_unit.py +++ b/snowflake/tests/test_unit.py @@ -81,6 +81,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 +114,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, @@ -149,6 +151,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, From 8e22db19d15abb766090fee5450b9a684cdb74db Mon Sep 17 00:00:00 2001 From: Paul Coignet Date: Thu, 20 Jan 2022 15:17:48 +0100 Subject: [PATCH 05/13] Add test --- snowflake/tests/keys/rsa_key_example.p8 | 28 +++++++++++++++++++++++++ snowflake/tests/test_unit.py | 26 ++++++++++++++++++----- 2 files changed, 49 insertions(+), 5 deletions(-) create mode 100644 snowflake/tests/keys/rsa_key_example.p8 diff --git a/snowflake/tests/keys/rsa_key_example.p8 b/snowflake/tests/keys/rsa_key_example.p8 new file mode 100644 index 0000000000000..6d2439379a00a --- /dev/null +++ b/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----- diff --git a/snowflake/tests/test_unit.py b/snowflake/tests/test_unit.py index 43167058ba923..4e383bf33a317 100644 --- a/snowflake/tests/test_unit.py +++ b/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 @@ -79,9 +80,9 @@ def test_default_auth(instance): client_prefetch_threads=4, login_timeout=3, ocsp_response_cache_filename=None, - authenticator='snowflake', + authenticator='oauth', token=None, - private_key=None + private_key=None, client_session_keep_alive=False, proxy_host=None, proxy_port=None, @@ -112,9 +113,9 @@ def test_oauth_auth(oauth_instance): client_prefetch_threads=4, login_timeout=3, ocsp_response_cache_filename=None, - authenticator='oauth', + authenticator='snowflake_jwt', token='testtoken', - private_key=None + private_key=None, client_session_keep_alive=False, proxy_host=None, proxy_port=None, @@ -123,6 +124,21 @@ def test_oauth_auth(oauth_instance): ) +def test_read_key(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 + assert len(check.read_key()) == 1216 + + 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() + + def test_proxy_settings(instance): init_config = { 'proxy_host': 'testhost', @@ -151,7 +167,7 @@ def test_proxy_settings(instance): ocsp_response_cache_filename=None, authenticator='snowflake', token=None, - private_key=None + private_key=None, client_session_keep_alive=False, proxy_host='testhost', proxy_port=8000, From 8b9c993d561bad568ca3c307b96b8fdcf9bccec3 Mon Sep 17 00:00:00 2001 From: Paul Coignet Date: Fri, 21 Jan 2022 15:14:45 +0100 Subject: [PATCH 06/13] Fix tests --- snowflake/tests/test_unit.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/snowflake/tests/test_unit.py b/snowflake/tests/test_unit.py index 4e383bf33a317..a465cb9f7cfe9 100644 --- a/snowflake/tests/test_unit.py +++ b/snowflake/tests/test_unit.py @@ -80,7 +80,7 @@ def test_default_auth(instance): client_prefetch_threads=4, login_timeout=3, ocsp_response_cache_filename=None, - authenticator='oauth', + authenticator='snowflake', token=None, private_key=None, client_session_keep_alive=False, @@ -113,7 +113,7 @@ def test_oauth_auth(oauth_instance): client_prefetch_threads=4, login_timeout=3, ocsp_response_cache_filename=None, - authenticator='snowflake_jwt', + authenticator='oauth', token='testtoken', private_key=None, client_session_keep_alive=False, From 437b2b64f9fe50d3a9a341aafc36861d538de987 Mon Sep 17 00:00:00 2001 From: Paul Coignet Date: Fri, 21 Jan 2022 15:47:22 +0100 Subject: [PATCH 07/13] Sync --- .../snowflake/config_models/defaults.py | 2 +- .../snowflake/config_models/instance.py | 2 +- .../snowflake/data/conf.yaml.example | 28 ------------------- 3 files changed, 2 insertions(+), 30 deletions(-) diff --git a/snowflake/datadog_checks/snowflake/config_models/defaults.py b/snowflake/datadog_checks/snowflake/config_models/defaults.py index 42fe90c1aa6d8..0de8eab6d1f31 100644 --- a/snowflake/datadog_checks/snowflake/config_models/defaults.py +++ b/snowflake/datadog_checks/snowflake/config_models/defaults.py @@ -1,4 +1,4 @@ -# (C) Datadog, Inc. 2022-present +# (C) Datadog, Inc. 2021-present # All rights reserved # Licensed under a 3-clause BSD style license (see LICENSE) diff --git a/snowflake/datadog_checks/snowflake/config_models/instance.py b/snowflake/datadog_checks/snowflake/config_models/instance.py index a9cc51c69ab27..6fb2a0e8c0677 100644 --- a/snowflake/datadog_checks/snowflake/config_models/instance.py +++ b/snowflake/datadog_checks/snowflake/config_models/instance.py @@ -1,4 +1,4 @@ -# (C) Datadog, Inc. 2022-present +# (C) Datadog, Inc. 2021-present # All rights reserved # Licensed under a 3-clause BSD style license (see LICENSE) diff --git a/snowflake/datadog_checks/snowflake/data/conf.yaml.example b/snowflake/datadog_checks/snowflake/data/conf.yaml.example index 5e4468e1ea158..8c456817fa383 100644 --- a/snowflake/datadog_checks/snowflake/data/conf.yaml.example +++ b/snowflake/datadog_checks/snowflake/data/conf.yaml.example @@ -106,47 +106,19 @@ instances: ## @param authenticator - string - optional - default: snowflake ## Authenticator for Snowflake. The default `snowflake` uses the internal Snowflake authenticator. -<<<<<<< HEAD -<<<<<<< HEAD ## Use `oauth` to authenticate with OAuth. Also, set the `token` or `token_path` option. -======= - ## Use `oauth` to authenticate with OAuth, be sure to set the `token` or `token_path` option as well. ->>>>>>> 78003da9e (Add token_path config option) -======= - ## Use `oauth` to authenticate with OAuth. Also, set the `token` or `token_path` option. ->>>>>>> 1f30fba28 (Sync spec) # # authenticator: ## @param token - string - optional -<<<<<<< HEAD -<<<<<<< HEAD - ## Token used for OAuth connection to Snowflake. You cannot use this alongside `token_path`. -======= - ## Token used for OAuth connection to Snowflake. Can't be used along side with `token_path`. ->>>>>>> 78003da9e (Add token_path config option) -======= ## Token used for OAuth connection to Snowflake. You cannot use this alongside `token_path`. ->>>>>>> 1f30fba28 (Sync spec) # # token: ## @param token_path - string - optional - default: /path/to/token -<<<<<<< HEAD -<<<<<<< HEAD - ## Path to the file that contains the token used for OAuth connection to Snowflake. - ## You cannot use this alongside `token`. - ## The token is re-read at every check run. -======= - ## Path to the file containing the token used for OAuth connection to snowflake. - ## Can't be used along side with `token`. - ## The token will be re-read at each check run. ->>>>>>> 78003da9e (Add token_path config option) -======= ## Path to the file that contains the token used for OAuth connection to Snowflake. ## You cannot use this alongside `token`. ## The token is re-read at every check run. ->>>>>>> 1f30fba28 (Sync spec) # # token_path: /path/to/token From f7fc7d3166f17d0e59d887dcd4ac8c7c3ffbdfb0 Mon Sep 17 00:00:00 2001 From: Paul Coignet Date: Mon, 24 Jan 2022 15:06:40 +0100 Subject: [PATCH 08/13] Mock connect --- snowflake/tests/common.py | 5 ++++- snowflake/tests/test_unit.py | 31 +++++++++++++++++++++++++++++-- 2 files changed, 33 insertions(+), 3 deletions(-) diff --git a/snowflake/tests/common.py b/snowflake/tests/common.py index 12b44946e6383..b299e81865d9a 100644 --- a/snowflake/tests/common.py +++ b/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'), + ], } diff --git a/snowflake/tests/test_unit.py b/snowflake/tests/test_unit.py index a465cb9f7cfe9..cbbd56bb10500 100644 --- a/snowflake/tests/test_unit.py +++ b/snowflake/tests/test_unit.py @@ -124,14 +124,41 @@ def test_oauth_auth(oauth_instance): ) -def test_read_key(instance): +def test_key_auth(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 - assert len(check.read_key()) == 1216 + 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]) + check.check(inst) + 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]) From 4526b797d7d9db1922c1090308c050c67afbb38a Mon Sep 17 00:00:00 2001 From: Paul Coignet Date: Mon, 24 Jan 2022 15:12:23 +0100 Subject: [PATCH 09/13] Fix style --- snowflake/tests/common.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/snowflake/tests/common.py b/snowflake/tests/common.py index b299e81865d9a..c1c4e77eb1076 100644 --- a/snowflake/tests/common.py +++ b/snowflake/tests/common.py @@ -44,6 +44,6 @@ ], 'docker_volumes': [ '{}:/home/snowflake_connector_patch'.format(os.path.join(HERE, 'snowflake_connector_patch')), - '{}:/home/keys'.format(os.path.join(HERE), 'keys'), + '{}:/home/keys'.format(os.path.join(HERE, 'keys')), ], } From cb08fed32004f31965a32d360d05006a5cb2c1ae Mon Sep 17 00:00:00 2001 From: Paul Coignet Date: Tue, 25 Jan 2022 10:45:21 +0100 Subject: [PATCH 10/13] Address reviews --- snowflake/assets/configuration/spec.yaml | 5 ++-- snowflake/datadog_checks/snowflake/config.py | 4 +-- .../snowflake/data/conf.yaml.example | 6 ++-- snowflake/tests/keys/rsa_key_pass_example.p8 | 29 +++++++++++++++++++ snowflake/tests/test_unit.py | 10 +++++-- 5 files changed, 45 insertions(+), 9 deletions(-) create mode 100644 snowflake/tests/keys/rsa_key_pass_example.p8 diff --git a/snowflake/assets/configuration/spec.yaml b/snowflake/assets/configuration/spec.yaml index d0de3ccb37edd..1bf7bd2a1183e 100644 --- a/snowflake/assets/configuration/spec.yaml +++ b/snowflake/assets/configuration/spec.yaml @@ -116,15 +116,16 @@ files: example: /path/to/token - name: private_key_path description: | - Path to the file that contains the private key used to connect to Snowflake. + 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: | - Password protecting the private key file in `private_key_path` option. + The password protecting the private key file in the `private_key_path` option. value: type: string - name: client_session_keep_alive diff --git a/snowflake/datadog_checks/snowflake/config.py b/snowflake/datadog_checks/snowflake/config.py index dc123cd05c650..a2b4bc575de85 100644 --- a/snowflake/datadog_checks/snowflake/config.py +++ b/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): @@ -42,7 +42,7 @@ def __init__(self, instance=None): token = instance.get('token', None) token_path = instance.get('token_path', None) private_key_path = instance.get('private_key_path', None) - private_key_password = instance.get('private_key_password', 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) diff --git a/snowflake/datadog_checks/snowflake/data/conf.yaml.example b/snowflake/datadog_checks/snowflake/data/conf.yaml.example index 8c456817fa383..e00b122868114 100644 --- a/snowflake/datadog_checks/snowflake/data/conf.yaml.example +++ b/snowflake/datadog_checks/snowflake/data/conf.yaml.example @@ -122,14 +122,14 @@ instances: # # token_path: /path/to/token - ## @param private_key_path - string - optional - default: /path/to/private_key - ## Path to the file that contains the private key used to connect to Snowflake. + ## @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 - ## Password protecting the private key file in `private_key_path` option. + ## The password protecting the private key file in the `private_key_path` option. # # private_key_password: diff --git a/snowflake/tests/keys/rsa_key_pass_example.p8 b/snowflake/tests/keys/rsa_key_pass_example.p8 new file mode 100644 index 0000000000000..b176b5efc04e2 --- /dev/null +++ b/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----- diff --git a/snowflake/tests/test_unit.py b/snowflake/tests/test_unit.py index cbbd56bb10500..3acf1a0e6b27a 100644 --- a/snowflake/tests/test_unit.py +++ b/snowflake/tests/test_unit.py @@ -124,7 +124,7 @@ def test_oauth_auth(oauth_instance): ) -def test_key_auth(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') @@ -136,7 +136,7 @@ def test_key_auth(instance): with mock.patch('datadog_checks.snowflake.check.sf') as sf: check = SnowflakeCheck(CHECK_NAME, {}, [inst]) - check.check(inst) + dd_run_check(check) sf.connect.assert_called_with( user='testuser', password='pass', @@ -165,6 +165,12 @@ def test_key_auth(instance): 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 = { From 99aa12b973ad1c9b81a011d6c95e1a7e8de305e8 Mon Sep 17 00:00:00 2001 From: Paul Coignet Date: Tue, 25 Jan 2022 10:50:10 +0100 Subject: [PATCH 11/13] Rename renew_token --- snowflake/datadog_checks/snowflake/check.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/snowflake/datadog_checks/snowflake/check.py b/snowflake/datadog_checks/snowflake/check.py index d6778e592c81c..483029c076c23 100644 --- a/snowflake/datadog_checks/snowflake/check.py +++ b/snowflake/datadog_checks/snowflake/check.py @@ -79,10 +79,11 @@ 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: @@ -147,8 +148,7 @@ def connect(self): self.proxy_port, ) - if self._config.token_path: - self.renew_token() + self.read_token() try: conn = sf.connect( From 693234dfdbcfa112a71132ae7392b2c3f8522f3c Mon Sep 17 00:00:00 2001 From: Paul Coignet Date: Tue, 25 Jan 2022 10:52:50 +0100 Subject: [PATCH 12/13] Sync --- snowflake/datadog_checks/snowflake/config_models/defaults.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/snowflake/datadog_checks/snowflake/config_models/defaults.py b/snowflake/datadog_checks/snowflake/config_models/defaults.py index 0de8eab6d1f31..48185318cccf4 100644 --- a/snowflake/datadog_checks/snowflake/config_models/defaults.py +++ b/snowflake/datadog_checks/snowflake/config_models/defaults.py @@ -1,4 +1,4 @@ -# (C) Datadog, Inc. 2021-present +# (C) Datadog, Inc. 2022-present # All rights reserved # Licensed under a 3-clause BSD style license (see LICENSE) @@ -95,7 +95,7 @@ def instance_private_key_password(field, value): def instance_private_key_path(field, value): - return '/path/to/private_key' + return get_default_field_value(field, value) def instance_schema_(field, value): From 8f33951e6f93dba2545a967f4c713e480427f92b Mon Sep 17 00:00:00 2001 From: Paul Coignet Date: Tue, 25 Jan 2022 14:19:22 +0100 Subject: [PATCH 13/13] Fix license year --- snowflake/datadog_checks/snowflake/config_models/defaults.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/snowflake/datadog_checks/snowflake/config_models/defaults.py b/snowflake/datadog_checks/snowflake/config_models/defaults.py index 48185318cccf4..0978291fe8d42 100644 --- a/snowflake/datadog_checks/snowflake/config_models/defaults.py +++ b/snowflake/datadog_checks/snowflake/config_models/defaults.py @@ -1,4 +1,4 @@ -# (C) Datadog, Inc. 2022-present +# (C) Datadog, Inc. 2021-present # All rights reserved # Licensed under a 3-clause BSD style license (see LICENSE)