Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 17 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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:
Expand Down
2 changes: 1 addition & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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"]
Expand Down
106 changes: 104 additions & 2 deletions src/bruin/_connection.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
import json
import logging
import os
import struct

from bruin.exceptions import (
ConnectionNotFoundError,
Expand All @@ -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."""

Expand Down Expand Up @@ -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."""
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand Down
24 changes: 24 additions & 0 deletions tests/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
126 changes: 126 additions & 0 deletions tests/test_connection.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"})
Expand Down
Loading