SNOW-2306184: config refactor - snowsql env support

Author: sfc-gh-mraba

src/snowflake/cli/api/config.py

@@ -100,14 +100,18 @@ class ConnectionConfig:
     authenticator: Optional[str] = None
     workload_identity_provider: Optional[str] = None
     private_key_file: Optional[str] = None
+    private_key_passphrase: Optional[str] = field(default=None, repr=False)
+    token: Optional[str] = field(default=None, repr=False)
+    session_token: Optional[str] = field(default=None, repr=False)
+    master_token: Optional[str] = field(default=None, repr=False)
     token_file_path: Optional[str] = None
     oauth_client_id: Optional[str] = None
     oauth_client_secret: Optional[str] = None
     oauth_authorization_url: Optional[str] = None
     oauth_token_request_url: Optional[str] = None
     oauth_redirect_uri: Optional[str] = None
     oauth_scope: Optional[str] = None
-    oatuh_enable_pkce: Optional[bool] = None
+    oauth_enable_pkce: Optional[bool] = None
     oauth_enable_refresh_tokens: Optional[bool] = None
     oauth_enable_single_use_refresh_tokens: Optional[bool] = None
     client_store_temporary_credential: Optional[bool] = None

src/snowflake/cli/api/config_ng/sources.py

@@ -388,14 +388,18 @@ class CliEnvironment(ValueSource):
         "private_key_file",
         "private_key_path",  # Used by integration tests
         "private_key_raw",  # Used by integration tests
+        "private_key_passphrase",  # Private key passphrase for encrypted keys
+        "token",  # OAuth token
+        "session_token",  # Session token for session-based authentication
+        "master_token",  # Master token for advanced authentication
         "token_file_path",
         "oauth_client_id",
         "oauth_client_secret",
         "oauth_authorization_url",
         "oauth_token_request_url",
         "oauth_redirect_uri",
         "oauth_scope",
-        "oatuh_enable_pkce",
+        "oauth_enable_pkce",  # Fixed typo: was "oatuh_enable_pkce"
         "oauth_enable_refresh_tokens",
         "oauth_enable_single_use_refresh_tokens",
         "client_store_temporary_credential",

src/snowflake/cli/api/config_provider.py

@@ -70,6 +70,53 @@ def get_all_connections(self) -> dict:
         """Get all connection configurations."""
         ...
 
+    def _transform_private_key_raw(self, connection_dict: dict) -> dict:
+        """
+        Transform private_key_raw to private_key_file for ConnectionContext compatibility.
+
+        The ConnectionContext dataclass doesn't have a private_key_raw field, so it gets
+        filtered out by merge_with_config. To work around this, we write private_key_raw
+        content to a temporary file and return it as private_key_file.
+
+        Args:
+            connection_dict: Connection configuration dictionary
+
+        Returns:
+            Modified connection dictionary with private_key_raw transformed to private_key_file
+        """
+        if "private_key_raw" not in connection_dict:
+            return connection_dict
+
+        # Don't transform if private_key_file is already set
+        if "private_key_file" in connection_dict:
+            return connection_dict
+
+        import os
+        import tempfile
+
+        try:
+            # Create a temporary file with the private key content
+            with tempfile.NamedTemporaryFile(
+                mode="w", suffix=".pem", delete=False
+            ) as f:
+                f.write(connection_dict["private_key_raw"])
+                temp_file_path = f.name
+
+            # Set restrictive permissions on the temporary file
+            os.chmod(temp_file_path, 0o600)
+
+            # Create a copy of the connection dict with the transformation
+            result = connection_dict.copy()
+            result["private_key_file"] = temp_file_path
+            del result["private_key_raw"]
+
+            return result
+
+        except Exception:
+            # If transformation fails, return original dict
+            # The error will be handled downstream
+            return connection_dict
+
 
 class LegacyConfigProvider(ConfigProvider):
     """
@@ -113,7 +160,8 @@ def read_config(self) -> None:
     def get_connection_dict(self, connection_name: str) -> dict:
         from snowflake.cli.api.config import get_connection_dict
 
-        return get_connection_dict(connection_name)
+        result = get_connection_dict(connection_name)
+        return self._transform_private_key_raw(result)
 
     def get_all_connections(self) -> dict:
         from snowflake.cli.api.config import get_all_connections
@@ -377,7 +425,8 @@ def get_connection_dict(self, connection_name: str) -> dict:
         Returns:
             Dictionary of connection parameters
         """
-        return self._get_connection_dict_internal(connection_name)
+        result = self._get_connection_dict_internal(connection_name)
+        return self._transform_private_key_raw(result)
 
     def _get_all_connections_dict(self) -> Dict[str, Dict[str, Any]]:
         """

src/snowflake/cli/api/connections.py

@@ -47,6 +47,7 @@ class ConnectionContext:
     authenticator: Optional[str] = None
     workload_identity_provider: Optional[str] = None
     private_key_file: Optional[str] = None
+    private_key_passphrase: Optional[str] = field(default=None, repr=False)
     warehouse: Optional[str] = None
     mfa_passcode: Optional[str] = None
     token: Optional[str] = None