Skip to content

Commit f874221

Browse files
committed
use secret detector and add tests
1 parent f7dbdd6 commit f874221

File tree

2 files changed

+42
-9
lines changed

2 files changed

+42
-9
lines changed

src/snowflake/connector/network.py

Lines changed: 4 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020

2121
import OpenSSL.SSL
2222

23+
from snowflake.connector.secret_detector import SecretDetector
2324
from snowflake.connector.vendored.requests.models import PreparedRequest
2425
from snowflake.connector.vendored.urllib3.connectionpool import (
2526
HTTPConnectionPool,
@@ -980,15 +981,9 @@ def handle_invalid_certificate_error(self, conn, full_url, cause) -> None:
980981
def _handle_unknown_error(self, method, full_url, headers, data, conn) -> None:
981982
"""Handles unknown errors."""
982983
if data:
983-
try:
984-
# masking the secret credentials
985-
decoded_data = json.loads(data)
986-
for secret_name in ("PASSWORD", "TOKEN"):
987-
if decoded_data.get("data") and decoded_data["data"].get(secret_name):
988-
decoded_data["data"][secret_name] = "********"
989-
data = json.dumps(decoded_data)
990-
except Exception:
991-
logger.info("data is not JSON")
984+
_, masked_data, err_str = SecretDetector.mask_secrets(data)
985+
if err_str is None:
986+
data = masked_data
992987
logger.error(
993988
f"Failed to get the response. Hanging? "
994989
f"method: {method}, url: {full_url}, headers:{headers}, "

test/unit/test_retry_network.py

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
from __future__ import annotations
77

88
import errno
9+
import logging
910
import os
1011
import time
1112
from unittest.mock import MagicMock, Mock, PropertyMock
@@ -217,3 +218,40 @@ def fake_request_exec(**kwargs):
217218
assert ret == {}
218219
assert cnt.c == 1 # failed on first call - did not retry
219220
assert rest._connection.errorhandler.called # error
221+
222+
223+
def test_secret_masking(caplog):
224+
connection = MagicMock()
225+
connection.errorhandler = Mock(return_value=None)
226+
227+
rest = SnowflakeRestful(
228+
host="testaccount.snowflakecomputing.com", port=443, connection=connection
229+
)
230+
231+
data = (
232+
'{"code": 12345,'
233+
' "data": {"TOKEN": "_Y1ZNETTn5/qfUWj3Jedb", "PASSWORD": "dummy_pass"}'
234+
"}"
235+
)
236+
default_parameters = {
237+
"method": "POST",
238+
"full_url": "https://testaccount.snowflakecomputing.com/",
239+
"headers": {},
240+
"data": data,
241+
}
242+
243+
class NotRetryableException(Exception):
244+
pass
245+
246+
def fake_request_exec(**kwargs):
247+
return None
248+
249+
# inject a fake method
250+
rest._request_exec = fake_request_exec
251+
252+
# first two attempts will fail but third will success
253+
with caplog.at_level(logging.ERROR):
254+
ret = rest.fetch(timeout=10, **default_parameters)
255+
assert '"TOKEN": "****' in caplog.text
256+
assert '"PASSWORD": "****' in caplog.text
257+
assert ret == {}

0 commit comments

Comments
 (0)