Added enable_diag, diag_log_path and diag_allowlist_path options to t… (

#893)

RELEASE-NOTES.md

@@ -6,6 +6,7 @@
 * Added support for fully qualified name (`database.schema.name`) in `name` parameter in streamlit project definition
 * Added support for fully qualified image repository names in `spcs image-repository` commands.
 * Added `--if-not-exists` option to `create` commands for `service`, and `compute-pool`. Added `--replace` and `--if-not-exists` options for `image-repository create`.
+* Added support for python connector diagnostic report.
 
 ## Fixes and improvements
 * Adding `--image-name` option for image name argument in `spcs image-repository list-tags` for consistency with other commands.

src/snowflake/cli/api/cli_global_context.py

@@ -24,6 +24,9 @@ def __init__(self):
         self._private_key_path: Optional[str] = None
         self._warehouse: Optional[str] = None
         self._mfa_passcode: Optional[str] = None
+        self._enable_diag: Optional[bool] = False
+        self._diag_log_path: Optional[Path] = None
+        self._diag_allowlist_path: Optional[Path] = None
         self._temporary_connection: bool = False
 
     def __setattr__(self, key, value):
@@ -118,6 +121,27 @@ def mfa_passcode(self) -> Optional[str]:
     def set_mfa_passcode(self, value: Optional[str]):
         self._mfa_passcode = value
 
+    @property
+    def enable_diag(self) -> Optional[bool]:
+        return self._enable_diag
+
+    def set_enable_diag(self, value: Optional[bool]):
+        self._enable_diag = value
+
+    @property
+    def diag_log_path(self) -> Optional[Path]:
+        return self._diag_log_path
+
+    def set_diag_log_path(self, value: Optional[Path]):
+        self._diag_log_path = value
+
+    @property
+    def diag_allowlist_path(self) -> Optional[Path]:
+        return self._diag_allowlist_path
+
+    def set_diag_allowlist_path(self, value: Optional[Path]):
+        self._diag_allowlist_path = value
+
     @property
     def temporary_connection(self) -> bool:
         return self._temporary_connection
@@ -150,6 +174,9 @@ def _build_connection(self):
         return connect_to_snowflake(
             temporary_connection=self.temporary_connection,
             mfa_passcode=self._mfa_passcode,
+            enable_diag=self._enable_diag,
+            diag_log_path=self._diag_log_path,
+            diag_allowlist_path=self._diag_allowlist_path,
             connection_name=self.connection_name,
             **self._collect_not_empty_connection_attributes(),
         )

src/snowflake/cli/api/commands/decorators.py

@@ -12,6 +12,9 @@
     ConnectionOption,
     DatabaseOption,
     DebugOption,
+    DiagAllowlistPathOption,
+    DiagLogPathOption,
+    EnableDiagOption,
     MfaPasscodeOption,
     OutputFormatOption,
     PasswordOption,
@@ -242,6 +245,24 @@ def _evaluate_param_type(
         annotation=Optional[str],
         default=MfaPasscodeOption,
     ),
+    inspect.Parameter(
+        "enable_diag",
+        inspect.Parameter.KEYWORD_ONLY,
+        annotation=Optional[bool],
+        default=EnableDiagOption,
+    ),
+    inspect.Parameter(
+        "diag_log_path",
+        inspect.Parameter.KEYWORD_ONLY,
+        annotation=Optional[str],
+        default=DiagLogPathOption,
+    ),
+    inspect.Parameter(
+        "diag_allowlist_path",
+        inspect.Parameter.KEYWORD_ONLY,
+        annotation=Optional[str],
+        default=DiagAllowlistPathOption,
+    ),
 ]
 
 GLOBAL_OPTIONS = [

src/snowflake/cli/api/commands/flags.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 
+import tempfile
 from inspect import signature
+from pathlib import Path
 from typing import Any, Callable, List, Optional, Tuple
 
 import click
@@ -265,6 +267,42 @@ def _password_callback(value: str):
     rich_help_panel=_CONNECTION_SECTION,
 )
 
+EnableDiagOption = typer.Option(
+    False,
+    "--enable-diag",
+    help="Run python connector diagnostic test",
+    callback=_callback(lambda: cli_context_manager.connection_context.set_enable_diag),
+    show_default=False,
+    is_flag=True,
+    rich_help_panel=_CONNECTION_SECTION,
+)
+
+DiagLogPathOption: Path = typer.Option(
+    tempfile.gettempdir(),
+    "--diag-log-path",
+    help="Diagnostic report path",
+    callback=_callback(
+        lambda: cli_context_manager.connection_context.set_diag_log_path
+    ),
+    show_default=False,
+    rich_help_panel=_CONNECTION_SECTION,
+    exists=True,
+    writable=True,
+)
+
+DiagAllowlistPathOption: Path = typer.Option(
+    None,
+    "--diag-allowlist-path",
+    help="Diagnostic report path to optional allowlist",
+    callback=_callback(
+        lambda: cli_context_manager.connection_context.set_diag_allowlist_path
+    ),
+    show_default=False,
+    rich_help_panel=_CONNECTION_SECTION,
+    exists=True,
+    file_okay=True,
+)
+
 OutputFormatOption = typer.Option(
     OutputFormat.TABLE.value,
     "--format",

src/snowflake/cli/app/snow_connector.py

@@ -27,6 +27,9 @@
 def connect_to_snowflake(
     temporary_connection: bool = False,
     mfa_passcode: Optional[str] = None,
+    enable_diag: Optional[bool] = False,
+    diag_log_path: Optional[str] = None,
+    diag_allowlist_path: Optional[str] = None,
     connection_name: Optional[str] = None,
     **overrides,
 ) -> SnowflakeConnection:
@@ -65,6 +68,15 @@ def connect_to_snowflake(
     if mfa_passcode:
         connection_parameters["passcode"] = mfa_passcode
 
+    if enable_diag:
+        connection_parameters["enable_connection_diag"] = enable_diag
+        if diag_log_path:
+            connection_parameters["connection_diag_log_path"] = diag_log_path
+        if diag_allowlist_path:
+            connection_parameters[
+                "connection_diag_allowlist_path"
+            ] = diag_allowlist_path
+
     try:
         # Whatever output is generated when creating connection,
         # we don't want it in our output. This is particularly important

src/snowflake/cli/plugins/connection/commands.py

@@ -10,6 +10,9 @@
 from snowflake.cli.api.commands.flags import (
     PLAIN_PASSWORD_MSG,
     ConnectionOption,
+    DiagAllowlistPathOption,
+    DiagLogPathOption,
+    EnableDiagOption,
     MfaPasscodeOption,
 )
 from snowflake.cli.api.commands.snow_typer import SnowTyper
@@ -227,7 +230,12 @@ def add(
 
 @app.command(requires_connection=False)
 def test(
-    connection: str = ConnectionOption, mfa_passcode: str = MfaPasscodeOption, **options
+    connection: str = ConnectionOption,
+    mfa_passcode: str = MfaPasscodeOption,
+    enable_diag: bool = EnableDiagOption,
+    diag_log_path: str = DiagLogPathOption,
+    diag_allowlist_path: str = DiagAllowlistPathOption,
+    **options,
 ) -> CommandResult:
     """
     Tests the connection to Snowflake.
@@ -265,6 +273,11 @@ def test(
         "Warehouse": f'{conn.warehouse or "not set"}',
     }
 
+    if enable_diag:
+        result[
+            "Diag Report Location"
+        ] = f"{diag_log_path}/SnowflakeConnectionTestReport.txt"
+
     return ObjectResult(result)