Resolving comments

Author: jahnvi480

mssql_python/auth.py

@@ -7,7 +7,7 @@
 import platform
 import struct
 from typing import Tuple, Dict, Optional, Union
-from mssql_python.logging_config import get_logger
+from mssql_python.logging_config import get_logger, ENABLE_LOGGING
 from mssql_python.constants import AuthType
 
 logger = get_logger()
@@ -22,42 +22,38 @@ def get_token_struct(token: str) -> bytes:
         return struct.pack(f"<I{len(token_bytes)}s", len(token_bytes), token_bytes)
 
     @staticmethod
-    def get_default_token() -> bytes:
-        """Get token using DefaultAzureCredential"""
-        from azure.identity import DefaultAzureCredential
-
-        try:
-            # DefaultAzureCredential will automatically use the best available method
-            # based on the environment (e.g., managed identity, environment variables)
-            credential = DefaultAzureCredential()
-            token = credential.get_token("https://database.windows.net/.default").token
-            return AADAuth.get_token_struct(token)
-        except Exception as e:
-            raise RuntimeError(f"Failed to create DefaultAzureCredential: {e}")  
-
-    @staticmethod
-    def get_device_code_token() -> bytes:
-        """Get token using DeviceCodeCredential"""
-        from azure.identity import DeviceCodeCredential
-            
-        try:
-            credential = DeviceCodeCredential()
-            token = credential.get_token("https://database.windows.net/.default").token
-            return AADAuth.get_token_struct(token)
-        except Exception as e:
-            raise RuntimeError(f"Failed to create DeviceCodeCredential: {e}")
-
-    @staticmethod
-    def get_interactive_token() -> bytes:
-        """Get token using InteractiveBrowserCredential"""
-        from azure.identity import InteractiveBrowserCredential
+    def get_token(auth_type: str) -> bytes:
+        """Get token using the specified authentication type"""
+        from azure.identity import (
+            DefaultAzureCredential, 
+            DeviceCodeCredential, 
+            InteractiveBrowserCredential
+        )
+        from azure.core.exceptions import ClientAuthenticationError
+        
+        # Mapping of auth types to credential classes
+        credential_map = {
+            "default": DefaultAzureCredential,
+            "devicecode": DeviceCodeCredential,
+            "interactive": InteractiveBrowserCredential,
+        }
+        
+        credential_class = credential_map[auth_type]
 
         try:
-            credential = InteractiveBrowserCredential()
+            credential = credential_class()
             token = credential.get_token("https://database.windows.net/.default").token
             return AADAuth.get_token_struct(token)
+        except ClientAuthenticationError as e:
+            # Re-raise with more specific context about Azure AD authentication failure
+            raise RuntimeError(
+                f"Azure AD authentication failed for {credential_class.__name__}: {e}. "
+                f"This could be due to invalid credentials, missing environment variables, "
+                f"user cancellation, network issues, or unsupported configuration."
+            ) from e
         except Exception as e:
-            raise RuntimeError(f"Failed to create InteractiveBrowserCredential: {e}")
+            # Catch any other unexpected exceptions
+            raise RuntimeError(f"Failed to create {credential_class.__name__}: {e}") from e
 
 def process_auth_parameters(parameters: list) -> Tuple[list, Optional[str]]:
     """
@@ -120,15 +116,14 @@ def get_auth_token(auth_type: str) -> Optional[bytes]:
     if not auth_type:
         return None
 
-    if auth_type == "default":
-        return AADAuth.get_default_token()
-    elif auth_type == "devicecode":
-        return AADAuth.get_device_code_token()
-    # If interactive authentication is requested, use InteractiveBrowserCredential
-    # but only if not on Windows, since in Windows: AADInteractive is supported.
-    elif auth_type == "interactive" and platform.system().lower() != "windows":
-        return AADAuth.get_interactive_token()
-    return None
+    # Handle platform-specific logic for interactive auth
+    if auth_type == "interactive" and platform.system().lower() == "windows":
+        return None  # Let Windows handle AADInteractive natively
+        
+    try:
+        return AADAuth.get_token(auth_type)
+    except (ValueError, RuntimeError):
+        return None
 
 def process_connection_string(connection_string: str) -> Tuple[str, Optional[Dict]]:
     """

tests/test_008_auth.py

@@ -37,23 +37,34 @@ class MockInteractiveBrowserCredential:
         def get_token(self, scope):
             return MockToken()
 
+    # Mock ClientAuthenticationError
+    class MockClientAuthenticationError(Exception):
+        pass
+
     class MockIdentity:
         DefaultAzureCredential = MockDefaultAzureCredential
         DeviceCodeCredential = MockDeviceCodeCredential
         InteractiveBrowserCredential = MockInteractiveBrowserCredential
 
+    class MockCore:
+        class exceptions:
+            ClientAuthenticationError = MockClientAuthenticationError
+
     # Create mock azure module if it doesn't exist
     if 'azure' not in sys.modules:
         sys.modules['azure'] = type('MockAzure', (), {})()
 
-    # Add identity module to azure
+    # Add identity and core modules to azure
     sys.modules['azure.identity'] = MockIdentity()
+    sys.modules['azure.core'] = MockCore()
+    sys.modules['azure.core.exceptions'] = MockCore.exceptions()
 
     yield
 
     # Cleanup
-    if 'azure.identity' in sys.modules:
-        del sys.modules['azure.identity']
+    for module in ['azure.identity', 'azure.core', 'azure.core.exceptions']:
+        if module in sys.modules:
+            del sys.modules[module]
 
 class TestAuthType:
     def test_auth_type_constants(self):
@@ -67,18 +78,55 @@ def test_get_token_struct(self):
         assert isinstance(token_struct, bytes)
         assert len(token_struct) > 4
 
-    def test_get_default_token(self):
-        token_struct = AADAuth.get_default_token()
+    def test_get_token_default(self):
+        token_struct = AADAuth.get_token("default")
         assert isinstance(token_struct, bytes)
 
-    def test_get_device_code_token(self):
-        token_struct = AADAuth.get_device_code_token()
+    def test_get_token_device_code(self):
+        token_struct = AADAuth.get_token("devicecode")
         assert isinstance(token_struct, bytes)
 
-    def test_get_interactive_token(self):
-        token_struct = AADAuth.get_interactive_token()
+    def test_get_token_interactive(self):
+        token_struct = AADAuth.get_token("interactive")
         assert isinstance(token_struct, bytes)
 
+    def test_get_token_credential_mapping(self):
+        # Test that all supported auth types work
+        supported_types = ["default", "devicecode", "interactive"]
+        for auth_type in supported_types:
+            token_struct = AADAuth.get_token(auth_type)
+            assert isinstance(token_struct, bytes)
+            assert len(token_struct) > 4
+
+    def test_get_token_client_authentication_error(self):
+        """Test that ClientAuthenticationError is properly handled"""
+        from azure.core.exceptions import ClientAuthenticationError
+        
+        # Create a mock credential that raises ClientAuthenticationError
+        class MockFailingCredential:
+            def get_token(self, scope):
+                raise ClientAuthenticationError("Mock authentication failed")
+        
+        # Use monkeypatch to mock the credential creation
+        def mock_get_token_failing(auth_type):
+            from azure.core.exceptions import ClientAuthenticationError
+            if auth_type == "default":
+                try:
+                    credential = MockFailingCredential()
+                    token = credential.get_token("https://database.windows.net/.default").token
+                    return AADAuth.get_token_struct(token)
+                except ClientAuthenticationError as e:
+                    raise RuntimeError(
+                        f"Azure AD authentication failed for MockFailingCredential: {e}. "
+                        f"This could be due to invalid credentials, missing environment variables, "
+                        f"user cancellation, network issues, or unsupported configuration."
+                    ) from e
+            else:
+                return AADAuth.get_token(auth_type)
+        
+        with pytest.raises(RuntimeError, match="Azure AD authentication failed"):
+            mock_get_token_failing("default")
+
 class TestProcessAuthParameters:
     def test_empty_parameters(self):
         modified_params, auth_type = process_auth_parameters([])