Merge pull request #209 from Maxteabag/mssql-ad-default-preflight-hint

Surface actionable hint when Azure AD auth fails for mssql

Author: Maxteabag

sqlit/domains/connections/providers/mssql/adapter.py

@@ -22,6 +22,52 @@
     from sqlit.domains.connections.domain.config import AuthType, ConnectionConfig
 
 
+class AzureAdAuthError(Exception):
+    """Raised when the Azure AD credential chain cannot produce a SQL token.
+
+    Surfaces the actionable hint (e.g. "run 'az login'") that would otherwise
+    be buried in the ODBC driver's generic "Login failed for user ''" error.
+    """
+
+
+def _format_azure_ad_hint(exc: Exception) -> str:
+    """Build a short, actionable error message from a credential-chain failure.
+
+    Picks out the most useful sub-error (typically the AzureCliCredential
+    line saying "Please run 'az login'") and prepends a one-line hint.
+    Falls back to the full message if no specific line stands out.
+    """
+    text = str(exc)
+    primary = _first_actionable_line(text)
+    hint = "Run 'az login' (or set AZURE_CLIENT_ID/SECRET/TENANT environment variables)."
+    if primary:
+        return f"Azure AD authentication failed.\n  {primary}\n{hint}"
+    return f"Azure AD authentication failed.\n{hint}\n\nDetails:\n{text}"
+
+
+def _first_actionable_line(text: str) -> str:
+    """Return the most actionable line from the credential-chain dump.
+
+    Walks the needles in priority order — "Please run 'az login'" beats
+    everything else because it tells the user exactly what to do. Lower-signal
+    lines like "Environment variables are not fully configured" are passive
+    and don't make this list; if no high-signal needle matches, the caller
+    falls back to dumping the full chain.
+    """
+    needles = (
+        "Please run 'az login'",
+        "Please run `az login`",
+        "azd auth login",
+        "Az.Account module",
+    )
+    lines = text.splitlines()
+    for needle in needles:
+        for line in lines:
+            if needle in line:
+                return line.strip()
+    return ""
+
+
 class SQLServerAdapter(DatabaseAdapter):
     """Adapter for Microsoft SQL Server using the mssql-python driver."""
 
@@ -179,6 +225,8 @@ def connect(self, config: ConnectionConfig) -> Any:
             package_name=self.install_package,
         )
 
+        self._preflight_azure_credentials(config)
+
         conn_str = self._build_connection_string(config)
         # Append extra_options to connection string
         for key, value in config.extra_options.items():
@@ -188,6 +236,42 @@ def connect(self, config: ConnectionConfig) -> Any:
         conn.autocommit = True
         return conn
 
+    def _preflight_azure_credentials(self, config: ConnectionConfig) -> None:
+        """Try to acquire a SQL Entra token before opening the SQL connection.
+
+        Without this, an expired/missing `az login` session surfaces as the
+        ODBC driver's generic "Login failed for user ''" — the real cause
+        ("Please run 'az login'") is buried in stderr. Acquiring the token
+        ourselves lets us raise an actionable AzureAdAuthError.
+
+        Silently no-ops if azure-identity isn't installed (the driver still
+        handles auth — we just lose the nicer error message).
+        """
+        import logging
+
+        from sqlit.domains.connections.domain.config import AuthType
+
+        if self.get_auth_type(config) != AuthType.AD_DEFAULT:
+            return
+        try:
+            from azure.core.exceptions import ClientAuthenticationError
+            from azure.identity import DefaultAzureCredential
+        except ImportError:
+            return
+
+        # azure-identity logs the full credential-chain dump to stderr at
+        # WARNING level on failure. Our own error already names the actionable
+        # cause, so silence the library's noise for the duration of this call.
+        azure_logger = logging.getLogger("azure.identity")
+        prior_level = azure_logger.level
+        azure_logger.setLevel(logging.ERROR)
+        try:
+            DefaultAzureCredential().get_token("https://database.windows.net/.default")
+        except ClientAuthenticationError as exc:
+            raise AzureAdAuthError(_format_azure_ad_hint(exc)) from exc
+        finally:
+            azure_logger.setLevel(prior_level)
+
     def get_databases(self, conn: Any) -> list[str]:
         """Get list of databases from SQL Server."""
         cursor = conn.cursor()

tests/unit/test_mssql_adapter.py

@@ -235,3 +235,115 @@ def test_get_columns_returns_primary_keys(self, adapter):
         assert result[0].is_primary_key is True
         assert result[1].name == "name"
         assert result[1].is_primary_key is False
+
+
+class TestMSSQLAdapterAzureAdPreflight:
+    """Pre-flight Entra-token check for ad_default auth converts the
+    DefaultAzureCredential chain failure into an actionable AzureAdAuthError
+    *before* the ODBC driver gets to emit its generic "Login failed for user ''".
+    """
+
+    @pytest.fixture
+    def ad_default_config(self):
+        """Minimal ConnectionConfig stub with auth_type=ad_default."""
+        from sqlit.domains.connections.domain.config import (
+            ConnectionConfig,
+            TcpEndpoint,
+        )
+
+        endpoint = TcpEndpoint(host="example.database.windows.net", port="1433", database="mydb")
+        return ConnectionConfig(
+            name="t",
+            db_type="mssql",
+            endpoint=endpoint,
+            options={"auth_type": "ad_default"},
+        )
+
+    def test_preflight_surfaces_az_login_hint_on_credential_failure(self, ad_default_config):
+        """When DefaultAzureCredential.get_token raises, we raise AzureAdAuthError
+        with the actionable 'Please run az login' line and a one-line hint."""
+        from sqlit.domains.connections.providers.mssql.adapter import (
+            AzureAdAuthError,
+            SQLServerAdapter,
+        )
+
+        chain_message = (
+            "DefaultAzureCredential failed to retrieve a token from the included credentials.\n"
+            "Attempted credentials:\n"
+            "\tEnvironmentCredential: EnvironmentCredential authentication unavailable.\n"
+            "\tAzureCliCredential: Please run 'az login' to set up an account\n"
+            "\tAzurePowerShellCredential: Az.Account module >= 2.2.0 is not installed\n"
+        )
+
+        fake_azure_core = MagicMock()
+        fake_azure_core_exceptions = MagicMock()
+
+        class _ClientAuthError(Exception):
+            pass
+
+        fake_azure_core_exceptions.ClientAuthenticationError = _ClientAuthError
+
+        fake_azure_identity = MagicMock()
+        fake_credential = MagicMock()
+        fake_credential.get_token.side_effect = _ClientAuthError(chain_message)
+        fake_azure_identity.DefaultAzureCredential.return_value = fake_credential
+
+        with patch.dict(
+            "sys.modules",
+            {
+                "azure": MagicMock(),
+                "azure.core": fake_azure_core,
+                "azure.core.exceptions": fake_azure_core_exceptions,
+                "azure.identity": fake_azure_identity,
+            },
+        ):
+            adapter = SQLServerAdapter()
+            with pytest.raises(AzureAdAuthError) as exc_info:
+                adapter._preflight_azure_credentials(ad_default_config)
+
+        message = str(exc_info.value)
+        assert "Please run 'az login'" in message
+        assert "Azure AD authentication failed" in message
+        # The verbose chain dump should NOT be included when we extracted a
+        # specific actionable line — keep the error tight, like sqlcmd does.
+        assert "DefaultAzureCredential failed to retrieve" not in message
+
+    def test_preflight_noop_when_azure_identity_missing(self, ad_default_config):
+        """If azure-identity is not installed, we silently skip and let the
+        ODBC driver handle auth (preserving the existing behavior)."""
+        from sqlit.domains.connections.providers.mssql.adapter import SQLServerAdapter
+
+        import builtins
+
+        real_import = builtins.__import__
+
+        def _block_azure(name, *args, **kwargs):
+            if name.startswith("azure."):
+                raise ImportError(f"No module named {name!r}")
+            return real_import(name, *args, **kwargs)
+
+        adapter = SQLServerAdapter()
+        with patch("builtins.__import__", side_effect=_block_azure):
+            # Should not raise
+            adapter._preflight_azure_credentials(ad_default_config)
+
+    def test_preflight_skipped_for_non_ad_default_auth(self):
+        """Pre-flight only runs for ad_default; sql/ad_password etc. must be
+        untouched even if azure-identity would fail."""
+        from sqlit.domains.connections.domain.config import (
+            ConnectionConfig,
+            TcpEndpoint,
+        )
+        from sqlit.domains.connections.providers.mssql.adapter import SQLServerAdapter
+
+        endpoint = TcpEndpoint(host="h", port="1433", database="d", username="u", password="p")
+        config = ConnectionConfig(
+            name="t",
+            db_type="mssql",
+            endpoint=endpoint,
+            options={"auth_type": "sql"},
+        )
+
+        adapter = SQLServerAdapter()
+        # Even if azure-identity would explode, this never touches it.
+        adapter._preflight_azure_credentials(config)