SNOW-2306184: config refactor - clean tmp files

Author: sfc-gh-mraba

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

@@ -26,7 +26,7 @@
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, field
 from datetime import datetime
-from typing import Any, Callable, Dict, List, Optional
+from typing import Any, Callable, Dict, List, Literal, Optional
 
 
 @dataclass(frozen=True)
@@ -84,9 +84,20 @@ class ValueSource(ABC):
     Precedence is determined by the order sources are provided to the resolver.
     """
 
+    # Allowed source names for config resolution
+    SourceName = Literal[
+        "snowsql_config",
+        "cli_config_toml",
+        "connections_toml",
+        "snowsql_env",
+        "connection_specific_env",
+        "cli_env",
+        "cli_arguments",
+    ]
+
     @property
     @abstractmethod
-    def source_name(self) -> str:
+    def source_name(self) -> SourceName:
         """
         Unique identifier for this source.
         Examples: "cli_arguments", "snowsql_config", "cli_env"

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

@@ -404,13 +404,11 @@ def resolve(self, key: Optional[str] = None, default: Any = None) -> Dict[str, A
         # Identify sources that connections.toml replaces
         # connections.toml only replaces cli_config_toml, not SnowSQL config
         cli_config_source = "cli_config_toml"
-        connections_file_source = None
         connections_to_replace: set[str] = set()
 
         # First pass: find connections.toml and identify connections to replace
         for source in self._sources:
             if hasattr(source, "is_connections_file") and source.is_connections_file:
-                connections_file_source = source
                 connections_to_replace = source.get_defined_connections()
                 break
 

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

@@ -32,7 +32,7 @@
 import logging
 import os
 from pathlib import Path
-from typing import Any, Dict, Optional
+from typing import Any, Dict, Final, Optional
 
 from snowflake.cli.api.config_ng.core import ConfigValue, ValueSource
 
@@ -117,7 +117,7 @@ def __init__(self):
             ]
 
     @property
-    def source_name(self) -> str:
+    def source_name(self) -> "ValueSource.SourceName":
         return "snowsql_config"
 
     def discover(self, key: Optional[str] = None) -> Dict[str, ConfigValue]:
@@ -225,7 +225,7 @@ def __init__(self):
             ]
 
     @property
-    def source_name(self) -> str:
+    def source_name(self) -> "ValueSource.SourceName":
         return "cli_config_toml"
 
     def discover(self, key: Optional[str] = None) -> Dict[str, ConfigValue]:
@@ -309,7 +309,7 @@ def __init__(self):
             self._file_path = Path.home() / ".snowflake" / "connections.toml"
 
     @property
-    def source_name(self) -> str:
+    def source_name(self) -> "ValueSource.SourceName":
         return "connections_toml"
 
     @property
@@ -448,7 +448,7 @@ class SnowSQLEnvironment(ValueSource):
     }
 
     @property
-    def source_name(self) -> str:
+    def source_name(self) -> "ValueSource.SourceName":
         return "snowsql_env"
 
     def discover(self, key: Optional[str] = None) -> Dict[str, ConfigValue]:
@@ -485,7 +485,7 @@ def supports_key(self, key: str) -> bool:
 
 
 # Base configuration keys that can be set via environment
-_ENV_CONFIG_KEYS = [
+_ENV_CONFIG_KEYS: Final[list[str]] = [
     "account",
     "user",
     "password",
@@ -533,7 +533,7 @@ class ConnectionSpecificEnvironment(ValueSource):
     """
 
     @property
-    def source_name(self) -> str:
+    def source_name(self) -> "ValueSource.SourceName":
         return "connection_specific_env"
 
     def discover(self, key: Optional[str] = None) -> Dict[str, ConfigValue]:
@@ -608,7 +608,7 @@ class CliEnvironment(ValueSource):
     """
 
     @property
-    def source_name(self) -> str:
+    def source_name(self) -> "ValueSource.SourceName":
         return "cli_env"
 
     def discover(self, key: Optional[str] = None) -> Dict[str, ConfigValue]:
@@ -677,7 +677,7 @@ def __init__(self, cli_context: Optional[Dict[str, Any]] = None):
         self._cli_context = cli_context or {}
 
     @property
-    def source_name(self) -> str:
+    def source_name(self) -> "ValueSource.SourceName":
         return "cli_arguments"
 
     def discover(self, key: Optional[str] = None) -> Dict[str, ConfigValue]:

src/snowflake/cli/api/config_provider.py

@@ -14,15 +14,17 @@
 
 from __future__ import annotations
 
+import atexit
 import os
 from abc import ABC, abstractmethod
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Dict, Optional
+from typing import TYPE_CHECKING, Any, Dict, Final, Optional
 
 if TYPE_CHECKING:
+    from snowflake.cli.api.config_ng.core import ValueSource
     from snowflake.cli.api.config_ng.resolver import ConfigurationResolver
 
-ALTERNATIVE_CONFIG_ENV_VAR = "SNOWFLAKE_CLI_CONFIG_V2_ENABLED"
+ALTERNATIVE_CONFIG_ENV_VAR: Final[str] = "SNOWFLAKE_CLI_CONFIG_V2_ENABLED"
 
 
 class ConfigProvider(ABC):
@@ -97,7 +99,6 @@ def _transform_private_key_raw(self, connection_dict: dict) -> dict:
         if "private_key_file" in connection_dict:
             return connection_dict
 
-        import os
         import tempfile
 
         try:
@@ -116,13 +117,35 @@ def _transform_private_key_raw(self, connection_dict: dict) -> dict:
             result["private_key_file"] = temp_file_path
             del result["private_key_raw"]
 
+            # Track created temp file on the provider instance for cleanup
+            temp_files_attr = "_temp_private_key_files"
+            existing = getattr(self, temp_files_attr, None)
+            if existing is None:
+                setattr(self, temp_files_attr, {temp_file_path})
+            else:
+                existing.add(temp_file_path)
+
             return result
 
         except Exception:
             # If transformation fails, return original dict
             # The error will be handled downstream
             return connection_dict
 
+    def cleanup_temp_files(self) -> None:
+        """Delete any temporary files created from private_key_raw transformation."""
+        temp_files = getattr(self, "_temp_private_key_files", None)
+        if not temp_files:
+            return
+        to_remove = list(temp_files)
+        for path in to_remove:
+            try:
+                Path(path).unlink(missing_ok=True)
+            except Exception:
+                # Best-effort cleanup; ignore failures
+                pass
+        temp_files.clear()
+
 
 class LegacyConfigProvider(ConfigProvider):
     """
@@ -409,7 +432,7 @@ def section_exists(self, *path) -> bool:
         )
 
     # Source priority levels (higher number = higher priority)
-    _SOURCE_PRIORITIES = {
+    _SOURCE_PRIORITIES: Final[dict["ValueSource.SourceName", int]] = {
         "snowsql_config": 1,
         "cli_config_toml": 2,
         "connections_toml": 3,
@@ -683,4 +706,23 @@ def reset_config_provider():
     Useful for testing and when config source changes.
     """
     global _config_provider_instance
+    # Cleanup any temp files created by the current provider instance
+    if _config_provider_instance is not None:
+        try:
+            _config_provider_instance.cleanup_temp_files()
+        except Exception:
+            pass
     _config_provider_instance = None
+
+
+def _cleanup_provider_at_exit() -> None:
+    """Process-exit cleanup for provider-managed temporary files."""
+    global _config_provider_instance
+    if _config_provider_instance is not None:
+        try:
+            _config_provider_instance.cleanup_temp_files()
+        except Exception:
+            pass
+
+
+atexit.register(_cleanup_provider_at_exit)

tests/config_ng/test_private_key_cleanup.py

@@ -0,0 +1,39 @@
+"""Tests for temporary private_key_raw file lifecycle and cleanup."""
+
+from pathlib import Path
+
+
+def test_private_key_raw_creates_and_cleans_temp_file(config_ng_setup, tmp_path):
+    priv_key_content = (
+        """-----BEGIN PRIVATE KEY-----\nABC\n-----END PRIVATE KEY-----\n"""
+    )
+
+    cli_config = """
+    [connections.test]
+    user = "cli-user"
+    """
+
+    env_vars = {
+        # Provide private_key_raw via env to trigger transformation
+        "SNOWFLAKE_CONNECTIONS_TEST_PRIVATE_KEY_RAW": priv_key_content,
+    }
+
+    with config_ng_setup(cli_config=cli_config, env_vars=env_vars):
+        from snowflake.cli.api.config import get_connection_dict
+        from snowflake.cli.api.config_provider import (
+            get_config_provider_singleton,
+            reset_config_provider,
+        )
+
+        provider = get_config_provider_singleton()
+
+        conn = get_connection_dict("test")
+        temp_path = Path(conn["private_key_file"])  # should exist now
+        assert temp_path.exists()
+        assert temp_path.read_text() == priv_key_content
+
+        # Reset provider triggers cleanup
+        reset_config_provider()
+
+        # File should be gone after cleanup
+        assert not temp_path.exists()