Second attempt at fixing the diagnostics issues (#1755)

* Second attempt at fixing the diagnostics issues

* Second attempt at fixing the diagnostics issues

Author: sfc-gh-cconner

DESCRIPTION.md

@@ -7,6 +7,8 @@ https://docs.snowflake.com/
 Source code is also available at: https://github.com/snowflakedb/snowflake-connector-python
 
 # Release Notes
+- v3.2.2(TBD)
+  - Fixed issue with connection diagnostics failing to complete certificate checks.
 
 - v3.3.0(October 10,2023)
 

src/snowflake/connector/connection.py

@@ -10,6 +10,7 @@
 import pathlib
 import re
 import sys
+import traceback
 import uuid
 import warnings
 import weakref
@@ -626,6 +627,7 @@ def connect(self, **kwargs) -> None:
             TelemetryService.get_instance().update_context(kwargs)
 
         if self.enable_connection_diag:
+            exceptions_dict = {}
             connection_diag = ConnectionDiagnostic(
                 account=self.account,
                 host=self.host,
@@ -640,15 +642,22 @@ def connect(self, **kwargs) -> None:
                 connection_diag.run_test()
                 self.__open_connection()
                 connection_diag.cursor = self.cursor()
-            except Exception as e:
-                logger.warning(f"Exception during connection test: {e}")
-
+            except Exception:
+                exceptions_dict["connection_test"] = traceback.format_exc()
+                logger.warning(
+                    f"""Exception during connection test:\n{exceptions_dict["connection_test"]} """
+                )
             try:
                 connection_diag.run_post_test()
-            except Exception as e:
-                logger.warning(f"Exception during post connection test: {e}")
+            except Exception:
+                exceptions_dict["post_test"] = traceback.format_exc()
+                logger.warning(
+                    f"""Exception during post connection test:\n{exceptions_dict["post_test"]} """
+                )
             finally:
                 connection_diag.generate_report()
+                if exceptions_dict:
+                    raise Exception(str(exceptions_dict))
         else:
             self.__open_connection()
 

src/snowflake/connector/connection_diagnostic.py

@@ -30,6 +30,19 @@
     import winreg
 
 
+def _decode_dict(d):
+    result = {}
+    for key, value in d.items():
+        if isinstance(key, bytes):
+            key = key.decode()
+        if isinstance(value, bytes):
+            value = value.decode()
+        elif isinstance(value, dict):
+            value = _decode_dict(value)
+        result.update({key: value})
+    return result
+
+
 class ConnectionDiagnostic:
     """Implementation of a connection test utility for Snowflake connector
 
@@ -360,10 +373,12 @@ def __https_host_report(
                             f"{host}:{port}: URL Check: Failed: Certificate mismatch: Host not in subject or alt names",
                         )
                 self.__append_message(host_type, f"{host}: Cert info:")
-                self.__append_message(
-                    host_type, f"{host}: subject: {result['subject']}"
-                )
-                self.__append_message(host_type, f"{host}: issuer: {result['issuer']}")
+
+                subject_str = _decode_dict(result["subject"])
+                self.__append_message(host_type, f"{host}: subject: {subject_str}")
+
+                issuer_str = _decode_dict(result["issuer"])
+                self.__append_message(host_type, f"{host}: issuer: {issuer_str}")
                 self.__append_message(
                     host_type, f"{host}: serialNumber: {result['serialNumber']}"
                 )
@@ -414,7 +429,9 @@ def __https_host_report(
 
     def __get_issuer_string(self, issuer: dict[bytes, bytes]) -> str:
         issuer_str: str = (
-            re.sub('[{}"]', "", json.dumps(issuer)).replace(": ", "=").replace(",", ";")
+            re.sub('[{}"]', "", json.dumps(_decode_dict(issuer)))
+            .replace(": ", "=")
+            .replace(",", ";")
         )
         return issuer_str
 

test/unit/test_connection_diagnostic.py

@@ -0,0 +1,59 @@
+#!/usr/bin/env python
+#
+# Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
+#
+
+from __future__ import annotations
+
+import snowflake.connector
+from snowflake.connector.connection_diagnostic import ConnectionDiagnostic, _decode_dict
+
+
+def test_https_host_report(caplog):
+    connection_diag = ConnectionDiagnostic(
+        account="test", host="test.snowflakecomputing.com"
+    )
+    https_host_report = connection_diag._ConnectionDiagnostic__https_host_report
+    https_host_report(
+        host="client-telemetry.snowflakecomputing.com",
+        port=443,
+        host_type="OUT_OF_BAND_TELEMETRY",
+    )
+
+    assert "DNS:client-telemetry.snowflakecomputing.com" in ".".join(
+        connection_diag.test_results["OUT_OF_BAND_TELEMETRY"]
+    )
+
+
+def test_decode_dict():
+    test_dict = {b"CN": b"client-telemetry.snowflakecomputing.com"}
+    result = _decode_dict(test_dict)
+
+    assert result == {
+        "CN": "client-telemetry.snowflakecomputing.com"
+    }, "_decode_dict method failed, binary dict not converted."
+
+
+def test_exception_raise_during_diag_fail(monkeypatch, caplog):
+    def mock_run_post_test(self):
+        raise ValueError("Diagnostic Test Failure")
+
+    monkeypatch.setattr(
+        snowflake.connector.connection_diagnostic.ConnectionDiagnostic,
+        "run_post_test",
+        mock_run_post_test,
+    )
+
+    try:
+        snowflake.connector.connect(
+            account="testaccount",
+            user="testuser",
+            password="testpassword",
+            database="TESTDB",
+            warehouse="TESTWH",
+            enable_connection_diag=True,
+        )
+    except BaseException:
+        pass
+
+    assert "Diagnostic Test Failure" in caplog.text