From 3b7ff79f055468b04d030f0b4fa4dc66426df6ce Mon Sep 17 00:00:00 2001
From: turtleDev <saravjeetamansingh@gmail.com>
Date: Thu, 9 Jul 2026 18:04:31 +0300
Subject: [PATCH] bruin: support Entra ID auth for Fabric connections

Fabric connections were routed straight to the MSSQL factory, which
indexes raw["username"] and raw["password"] unconditionally. Bruin's
Fabric connection also accepts a service principal (client_id,
client_secret, tenant_id) or the DefaultAzureCredential chain, and
those payloads carry no username, so building the client raised
KeyError: 'username'.

pymssql cannot carry a Microsoft Entra ID access token, so give Fabric
its own factory: the two Entra ID modes acquire a token via
azure-identity and hand it to pyodbc through the msodbcsql
SQL_COPT_SS_ACCESS_TOKEN pre-connect attribute, while SQL
authentication keeps using the shared pymssql path.

The ODBC driver defaults to the newest installed
"ODBC Driver NN for SQL Server" and can be pinned with a driver field
on the connection. A payload with none of the three credential sets now
raises ConnectionParseError naming the fields it expects.
---
 README.md                |  17 ++++++
 pyproject.toml           |   2 +-
 src/bruin/_connection.py | 106 +++++++++++++++++++++++++++++++-
 tests/conftest.py        |  24 ++++++++
 tests/test_connection.py | 126 +++++++++++++++++++++++++++++++++++++++
 5 files changed, 272 insertions(+), 3 deletions(-)

diff --git a/README.md b/README.md
index 97a8a68..cb06fff 100644
--- a/README.md
+++ b/README.md
@@ -35,6 +35,7 @@ bruin-sdk[snowflake]    # Snowflake
 bruin-sdk[postgres]     # PostgreSQL / Redshift
 bruin-sdk[redshift]     # Redshift (alias for postgres extra)
 bruin-sdk[mssql]        # Microsoft SQL Server
+bruin-sdk[fabric]       # Microsoft Fabric Warehouse
 bruin-sdk[mysql]        # MySQL
 bruin-sdk[duckdb]       # DuckDB
 bruin-sdk[sheets]       # Google Sheets (for GCP connections)
@@ -209,12 +210,28 @@ conn.client  # Lazy-initialized database client
 | `postgres` | `psycopg2.connection` | `bruin-sdk[postgres]` |
 | `redshift` | `psycopg2.connection` | `bruin-sdk[redshift]` |
 | `mssql` | `pymssql.Connection` | `bruin-sdk[mssql]` |
+| `fabric` | `pyodbc.Connection` or `pymssql.Connection` | `bruin-sdk[fabric]` |
 | `mysql` | `mysql.connector.Connection` | `bruin-sdk[mysql]` |
 | `duckdb` | `duckdb.DuckDBPyConnection` | `bruin-sdk[duckdb]` |
 | `generic` | N/A (raises error) | — |
 
 Client creation is **lazy** — the actual database connection is only established when `.client` is first accessed.
 
+#### Fabric connections
+
+Fabric supports three authentication modes, selected by the fields present on the connection:
+
+| Fields | Mode | Driver |
+|--------|------|--------|
+| `use_azure_default_credential: true` | `DefaultAzureCredential` (e.g. `az login`, managed identity) | `pyodbc` |
+| `client_id` + `client_secret` + `tenant_id` | Microsoft Entra ID service principal | `pyodbc` |
+| `username` + `password` | SQL authentication | `pymssql` |
+
+The Microsoft Entra ID modes acquire an access token and hand it to the driver, which requires an
+[ODBC Driver for SQL Server](https://learn.microsoft.com/sql/connect/odbc/download-odbc-driver-for-sql-server)
+(msodbcsql18 or newer) on the machine running the asset. The newest installed driver is used unless
+the connection sets `driver` explicitly.
+
 #### GCP connections
 
 GCP connections have extra methods since one connection can access multiple Google services:
diff --git a/pyproject.toml b/pyproject.toml
index de5c354..8f322c8 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -41,7 +41,7 @@ athena = ["pyathena"]
 trino = ["trino"]
 motherduck = ["duckdb"]
 synapse = ["pymssql"]
-fabric = ["pymssql"]
+fabric = ["pymssql", "pyodbc", "azure-identity"]
 oracle = ["oracledb"]
 db2 = ["ibm-db"]
 hana = ["hdbcli"]
diff --git a/src/bruin/_connection.py b/src/bruin/_connection.py
index 3cfc022..b07e7fd 100644
--- a/src/bruin/_connection.py
+++ b/src/bruin/_connection.py
@@ -1,6 +1,7 @@
 import json
 import logging
 import os
+import struct
 
 from bruin.exceptions import (
     ConnectionNotFoundError,
@@ -11,6 +12,13 @@
 logger = logging.getLogger("bruin")
 
 
+def _as_bool(value) -> bool:
+    """Interpret a connection field as a boolean, tolerating string payloads."""
+    if isinstance(value, str):
+        return value.lower() == "true"
+    return bool(value)
+
+
 class Connection:
     """A Bruin-managed database connection with lazy client initialization."""
 
@@ -73,7 +81,7 @@ def __init__(self, name: str, raw: dict):
 
     def _uses_adc(self) -> bool:
         """Return True when the connection is configured for Application Default Credentials."""
-        return self.raw.get("use_application_default_credentials") in ("true", "True", True)
+        return _as_bool(self.raw.get("use_application_default_credentials"))
 
     def _parse_sa_info(self):
         """Parse the service account JSON from the connection payload."""
@@ -194,7 +202,7 @@ def _create_client(conn_type: str, raw):
         "redshift": _create_redshift,
         "mssql": _create_mssql,
         "synapse": _create_mssql,
-        "fabric": _create_mssql,
+        "fabric": _create_fabric,
         "mysql": _create_mysql,
         "duckdb": _create_duckdb,
         "databricks": _create_databricks,
@@ -308,6 +316,100 @@ def _create_mssql(raw: dict):
     )
 
 
+_FABRIC_TOKEN_SCOPE = "https://database.windows.net/.default"
+
+# SQL_COPT_SS_ACCESS_TOKEN — the msodbcsql pre-connect attribute that carries an Entra ID token.
+_SQL_COPT_SS_ACCESS_TOKEN = 1256
+
+
+def _fabric_driver(raw: dict) -> str:
+    """Pick the ODBC driver to talk to Fabric with."""
+    import pyodbc
+
+    installed = pyodbc.drivers()
+    configured = raw.get("driver")
+    if configured:
+        if configured not in installed:
+            raise ConnectionTypeError(
+                f"ODBC driver '{configured}' is not installed. Available drivers: "
+                f"{', '.join(installed) or '(none)'}."
+            )
+        return configured
+
+    candidates = sorted(
+        (d for d in installed if d.startswith("ODBC Driver ") and d.endswith(" for SQL Server")),
+        reverse=True,
+    )
+    if not candidates:
+        raise ConnectionTypeError(
+            "No 'ODBC Driver NN for SQL Server' is installed, which is required for Microsoft "
+            "Entra ID authentication against Fabric. Install msodbcsql18: "
+            "https://learn.microsoft.com/sql/connect/odbc/download-odbc-driver-for-sql-server"
+        )
+    return candidates[0]
+
+
+def _fabric_token_struct(raw: dict) -> bytes:
+    """Fetch an Entra ID access token and pack it the way msodbcsql expects."""
+    try:
+        from azure.identity import ClientSecretCredential, DefaultAzureCredential
+    except ImportError:
+        raise ImportError(
+            "Install bruin-sdk[fabric] to use Microsoft Entra ID authentication: "
+            "pip install 'bruin-sdk[fabric]'"
+        )
+
+    if _as_bool(raw.get("use_azure_default_credential")):
+        logger.debug("Using DefaultAzureCredential for Fabric connection")
+        credential = DefaultAzureCredential()
+    else:
+        logger.debug("Using service principal %s for Fabric connection", raw["client_id"])
+        credential = ClientSecretCredential(
+            tenant_id=raw["tenant_id"],
+            client_id=raw["client_id"],
+            client_secret=raw["client_secret"],
+        )
+
+    token = credential.get_token(_FABRIC_TOKEN_SCOPE).token.encode("utf-16-le")
+    return struct.pack(f"<I{len(token)}s", len(token), token)
+
+
+def _create_fabric(raw: dict):
+    """Connect to a Microsoft Fabric warehouse.
+
+    Fabric accepts three auth modes.  Entra ID (service principal or
+    ``DefaultAzureCredential``) requires an access token, which only the
+    msodbcsql driver can carry — so those go through ``pyodbc``.  SQL
+    authentication falls back to the shared MSSQL path.
+    """
+    has_service_principal = all(raw.get(k) for k in ("client_id", "client_secret", "tenant_id"))
+    if not has_service_principal and not _as_bool(raw.get("use_azure_default_credential")):
+        if raw.get("username"):
+            return _create_mssql(raw)
+        raise ConnectionParseError(
+            "Fabric connection is missing credentials. Set 'use_azure_default_credential', "
+            "or 'client_id' + 'client_secret' + 'tenant_id', or 'username' + 'password'."
+        )
+
+    try:
+        import pyodbc
+    except ImportError:
+        raise ImportError(
+            "Install bruin-sdk[fabric] to use Fabric connections: pip install 'bruin-sdk[fabric]'"
+        )
+
+    conn_str = (
+        f"DRIVER={{{_fabric_driver(raw)}}};"
+        f"SERVER={raw['host']},{raw.get('port', 1433)};"
+        f"DATABASE={raw.get('database', '')};"
+        f"Encrypt=yes;TrustServerCertificate=no"
+    )
+    return pyodbc.connect(
+        conn_str,
+        attrs_before={_SQL_COPT_SS_ACCESS_TOKEN: _fabric_token_struct(raw)},
+    )
+
+
 def _create_mysql(raw: dict):
     try:
         import mysql.connector
diff --git a/tests/conftest.py b/tests/conftest.py
index 4047e0a..0b84693 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -174,6 +174,30 @@ def fabric_connection_json():
     }
 
 
+@pytest.fixture
+def fabric_service_principal_connection_json():
+    """Fabric connection using a Microsoft Entra ID service principal (no username)."""
+    return {
+        "host": "sql-endpoint-guid.datawarehouse.fabric.microsoft.com",
+        "port": 1433,
+        "database": "warehouse",
+        "client_id": "app-id",
+        "client_secret": "app-secret",
+        "tenant_id": "tenant-id",
+    }
+
+
+@pytest.fixture
+def fabric_azure_default_connection_json():
+    """Fabric connection using the DefaultAzureCredential chain (no username)."""
+    return {
+        "host": "sql-endpoint-guid.datawarehouse.fabric.microsoft.com",
+        "port": 1433,
+        "database": "warehouse",
+        "use_azure_default_credential": True,
+    }
+
+
 @pytest.fixture
 def oracle_connection_json():
     return {
diff --git a/tests/test_connection.py b/tests/test_connection.py
index 862991e..9cd5145 100644
--- a/tests/test_connection.py
+++ b/tests/test_connection.py
@@ -883,6 +883,132 @@ def test_redshift_uses_port_5439(self):
         assert kwargs["sslmode"] == "allow"
 
 
+class TestFabricAuth:
+    """Fabric picks its driver and credential from the fields on the connection."""
+
+    @staticmethod
+    def _mock_azure(token="access-token"):
+        azure = MagicMock()
+        identity = MagicMock()
+        azure.identity = identity
+        for factory in (identity.ClientSecretCredential, identity.DefaultAzureCredential):
+            factory.return_value.get_token.return_value.token = token
+        return azure, identity
+
+    @staticmethod
+    def _mock_pyodbc(drivers=("ODBC Driver 17 for SQL Server", "ODBC Driver 18 for SQL Server")):
+        pyodbc = MagicMock()
+        pyodbc.drivers.return_value = list(drivers)
+        return pyodbc
+
+    def _connect(self, raw, pyodbc=None):
+        pyodbc = pyodbc or self._mock_pyodbc()
+        azure, identity = self._mock_azure()
+        modules = {"pyodbc": pyodbc, "azure": azure, "azure.identity": identity}
+        with patch.dict("sys.modules", modules):
+            from bruin._connection import _create_fabric
+
+            _create_fabric(raw)
+        return pyodbc, identity
+
+    def test_service_principal_uses_pyodbc_with_access_token(
+        self, fabric_service_principal_connection_json
+    ):
+        pyodbc, identity = self._connect(fabric_service_principal_connection_json)
+
+        identity.ClientSecretCredential.assert_called_once_with(
+            tenant_id="tenant-id",
+            client_id="app-id",
+            client_secret="app-secret",
+        )
+        identity.DefaultAzureCredential.assert_not_called()
+
+        conn_str = pyodbc.connect.call_args[0][0]
+        assert "DRIVER={ODBC Driver 18 for SQL Server};" in conn_str
+        assert "SERVER=sql-endpoint-guid.datawarehouse.fabric.microsoft.com,1433;" in conn_str
+        assert "DATABASE=warehouse;" in conn_str
+
+        # The token is length-prefixed UTF-16-LE, as msodbcsql expects.
+        token_struct = pyodbc.connect.call_args[1]["attrs_before"][1256]
+        expected = "access-token".encode("utf-16-le")
+        assert token_struct == len(expected).to_bytes(4, "little") + expected
+
+    def test_azure_default_credential(self, fabric_azure_default_connection_json):
+        _, identity = self._connect(fabric_azure_default_connection_json)
+        identity.DefaultAzureCredential.assert_called_once_with()
+        identity.ClientSecretCredential.assert_not_called()
+
+    def test_azure_default_credential_as_string(self, fabric_azure_default_connection_json):
+        """Bruin may serialize the flag as a string rather than a JSON bool."""
+        fabric_azure_default_connection_json["use_azure_default_credential"] = "true"
+        _, identity = self._connect(fabric_azure_default_connection_json)
+        identity.DefaultAzureCredential.assert_called_once_with()
+
+    def test_sql_auth_falls_back_to_mssql(self, fabric_connection_json):
+        mock_pymssql = MagicMock()
+        with patch.dict("sys.modules", {"pymssql": mock_pymssql}):
+            from bruin._connection import _create_fabric
+
+            _create_fabric(fabric_connection_json)
+        mock_pymssql.connect.assert_called_once_with(
+            server="fabric.example.com",
+            port=1433,
+            user="admin",
+            password="s3cret",
+            database="warehouse",
+        )
+
+    def test_service_principal_preferred_over_username(
+        self, fabric_service_principal_connection_json
+    ):
+        fabric_service_principal_connection_json["username"] = "admin"
+        _, identity = self._connect(fabric_service_principal_connection_json)
+        identity.ClientSecretCredential.assert_called_once()
+
+    def test_partial_service_principal_without_username_raises(
+        self, fabric_service_principal_connection_json
+    ):
+        del fabric_service_principal_connection_json["client_secret"]
+        from bruin._connection import _create_fabric
+
+        with pytest.raises(ConnectionParseError, match="missing credentials"):
+            _create_fabric(fabric_service_principal_connection_json)
+
+    def test_explicit_driver_is_honored(self, fabric_service_principal_connection_json):
+        fabric_service_principal_connection_json["driver"] = "ODBC Driver 17 for SQL Server"
+        pyodbc, _ = self._connect(fabric_service_principal_connection_json)
+        assert "DRIVER={ODBC Driver 17 for SQL Server};" in pyodbc.connect.call_args[0][0]
+
+    def test_missing_explicit_driver_raises(self, fabric_service_principal_connection_json):
+        fabric_service_principal_connection_json["driver"] = "Nonexistent Driver"
+        with pytest.raises(ConnectionTypeError, match="is not installed"):
+            self._connect(fabric_service_principal_connection_json)
+
+    def test_no_odbc_driver_installed_raises(self, fabric_service_principal_connection_json):
+        with pytest.raises(ConnectionTypeError, match="msodbcsql18"):
+            self._connect(
+                fabric_service_principal_connection_json,
+                pyodbc=self._mock_pyodbc(drivers=["PostgreSQL Unicode"]),
+            )
+
+    def test_missing_pyodbc_raises_import_error(self, fabric_service_principal_connection_json):
+        with patch.dict("sys.modules", {"pyodbc": None}):
+            from bruin._connection import _create_fabric
+
+            with pytest.raises(ImportError, match=r"bruin-sdk\[fabric\]"):
+                _create_fabric(fabric_service_principal_connection_json)
+
+    def test_missing_azure_identity_raises_import_error(
+        self, fabric_service_principal_connection_json
+    ):
+        modules = {"pyodbc": self._mock_pyodbc(), "azure": None, "azure.identity": None}
+        with patch.dict("sys.modules", modules):
+            from bruin._connection import _create_fabric
+
+            with pytest.raises(ImportError, match=r"bruin-sdk\[fabric\]"):
+                _create_fabric(fabric_service_principal_connection_json)
+
+
 class TestRepr:
     def test_connection_repr(self):
         conn = Connection("my_pg", "postgres", {"host": "localhost"})