onehouse-python-sdk 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,25 @@
1
+ __pycache__/
2
+ *.py[cod]
3
+ *.pyo
4
+ *.pyd
5
+ *.egg
6
+ *.egg-info/
7
+ dist/
8
+ build/
9
+ .eggs/
10
+ .venv/
11
+ venv/
12
+ env/
13
+ .env
14
+ *.env
15
+ .tox/
16
+ .pytest_cache/
17
+ .coverage
18
+ .mypy_cache/
19
+ .ruff_cache/
20
+ *.so
21
+
22
+ # Test & coverage reports
23
+ coverage.xml
24
+ test-results.xml
25
+ htmlcov/
@@ -0,0 +1,134 @@
1
+ Metadata-Version: 2.4
2
+ Name: onehouse-python-sdk
3
+ Version: 0.1.0
4
+ Summary: Onehouse Python SDK — modular data-plane clients for LakeBase and beyond
5
+ Project-URL: Homepage, https://github.com/onehouseinc/onehouse-python-sdk
6
+ Project-URL: Repository, https://github.com/onehouseinc/onehouse-python-sdk
7
+ Project-URL: Issues, https://github.com/onehouseinc/onehouse-python-sdk/issues
8
+ Author-email: Onehouse Team <eng@onehouse.ai>
9
+ License: Proprietary
10
+ Keywords: auth,azure,lakebase,oauth2,oidc,onehouse,postgresql,saml
11
+ Classifier: Development Status :: 4 - Beta
12
+ Classifier: Intended Audience :: Developers
13
+ Classifier: License :: Other/Proprietary License
14
+ Classifier: Programming Language :: Python :: 3
15
+ Classifier: Programming Language :: Python :: 3.9
16
+ Classifier: Programming Language :: Python :: 3.10
17
+ Classifier: Programming Language :: Python :: 3.11
18
+ Classifier: Programming Language :: Python :: 3.12
19
+ Classifier: Topic :: Database
20
+ Classifier: Topic :: Software Development :: Libraries :: Python Modules
21
+ Requires-Python: >=3.9
22
+ Provides-Extra: dev
23
+ Requires-Dist: pytest-cov>=4.0; extra == 'dev'
24
+ Requires-Dist: pytest-mock>=3.0; extra == 'dev'
25
+ Requires-Dist: pytest>=7.0; extra == 'dev'
26
+ Provides-Extra: lakebase
27
+ Requires-Dist: certifi>=2024.1; extra == 'lakebase'
28
+ Requires-Dist: psycopg2-binary>=2.9; extra == 'lakebase'
29
+ Description-Content-Type: text/markdown
30
+
31
+ # onehouse-python-sdk
32
+
33
+ Python SDK for connecting to [Onehouse](https://www.onehouse.ai) data-plane services.
34
+
35
+ The base package has **zero required dependencies**. Drivers ship as optional extras so you only install what you need.
36
+
37
+ ```bash
38
+ pip install onehouse-python-sdk[lakebase]
39
+ ```
40
+
41
+ ## LakeBase
42
+
43
+ LakeBase is Onehouse's PostgreSQL-compatible managed lakehouse engine. This SDK provides a `psycopg2`-backed client that handles the browser-based and federated authentication flows LakeBase clusters require, on top of standard username/password auth.
44
+
45
+ ### Installation
46
+
47
+ ```bash
48
+ pip install onehouse-python-sdk[lakebase]
49
+ ```
50
+
51
+ Requires Python 3.9+.
52
+
53
+ ### Quickstart
54
+
55
+ ```python
56
+ from onehouse_python_sdk import LakebaseClient
57
+
58
+ # Username / password
59
+ with LakebaseClient().connect(
60
+ host="<cluster-host>",
61
+ port=5432,
62
+ dbname="mydb",
63
+ user="admin",
64
+ password="secret",
65
+ ) as client:
66
+ rows = client.fetchall("SELECT * FROM mytable WHERE id = %s", (42,))
67
+ ```
68
+
69
+ ### Authentication flows
70
+
71
+ | Flow | Parameters |
72
+ |------|------------|
73
+ | Username / password | `user`, `password` |
74
+ | OIDC Device Flow (Okta / Auth0) | `browser_auth=true`, `oidc_client_id`, `oidc_issuer_url`, `oidc_iam_role` |
75
+ | Azure AD OAuth2 | `browser_auth=true`, `azure_oauth_tenant_id`, `azure_oauth_client_id`, `azure_oauth_client_secret` |
76
+ | Azure Entra ID SAML | `browser_auth=true`, `azure_tenant_id`, `azure_entity_id` |
77
+ | Built-in login form | `browser_auth=true` _(default when no IdP params are set)_ |
78
+ | External redirect | `browser_auth=true`, `auth_redirect_url` |
79
+
80
+ ```python
81
+ # OIDC Device Flow
82
+ client = LakebaseClient().connect(
83
+ host="<cluster-host>",
84
+ port=5432,
85
+ dbname="mydb",
86
+ browser_auth="true",
87
+ oidc_client_id="0oaXXXXX",
88
+ oidc_issuer_url="https://myorg.okta.com",
89
+ oidc_iam_role="arn:aws:iam::123456789012:role/LakeBaseRole",
90
+ )
91
+
92
+ # Azure AD OAuth2
93
+ client = LakebaseClient().connect(
94
+ host="<cluster-host>",
95
+ port=5432,
96
+ dbname="mydb",
97
+ browser_auth="true",
98
+ azure_oauth_tenant_id="your-tenant-id",
99
+ azure_oauth_client_id="your-client-id",
100
+ azure_oauth_client_secret="your-client-secret",
101
+ )
102
+ ```
103
+
104
+ DSN string form is also supported:
105
+
106
+ ```python
107
+ client = LakebaseClient().connect(
108
+ "postgresql://<cluster-host>:5432/mydb"
109
+ "?browser_auth=true"
110
+ "&oidc_client_id=0oaXXXX"
111
+ "&oidc_issuer_url=https://myorg.okta.com"
112
+ "&oidc_iam_role=arn:aws:iam::123456789012:role/LakeBaseRole"
113
+ )
114
+ ```
115
+
116
+ ### Client API
117
+
118
+ All clients implement the `SqlClient` interface:
119
+
120
+ | Method | Description |
121
+ |--------|-------------|
122
+ | `connect(dsn=None, **kwargs) → self` | Establish connection, returns self for chaining |
123
+ | `execute(sql, params)` | Run a statement, return rowcount |
124
+ | `fetchall(sql, params)` | Run a query, return `list[tuple]` |
125
+ | `fetchone(sql, params)` | Run a query, return first row or `None` |
126
+ | `cursor()` | Raw cursor for advanced use |
127
+ | `raw_connection` | Underlying `psycopg2` connection |
128
+ | `close()` | Close the connection |
129
+ | `__enter__` / `__exit__` | Context manager — closes on exit |
130
+
131
+ ### Notes
132
+
133
+ - **Credential caching** — auth tokens are cached for 4 minutes per connection parameters to avoid repeated browser prompts within the same process.
134
+ - **Callback port** — the local auth callback server defaults to port `8888` at path `/lakebase`. Override with `auth_callback_port` and `auth_callback_path`.
@@ -0,0 +1,104 @@
1
+ # onehouse-python-sdk
2
+
3
+ Python SDK for connecting to [Onehouse](https://www.onehouse.ai) data-plane services.
4
+
5
+ The base package has **zero required dependencies**. Drivers ship as optional extras so you only install what you need.
6
+
7
+ ```bash
8
+ pip install onehouse-python-sdk[lakebase]
9
+ ```
10
+
11
+ ## LakeBase
12
+
13
+ LakeBase is Onehouse's PostgreSQL-compatible managed lakehouse engine. This SDK provides a `psycopg2`-backed client that handles the browser-based and federated authentication flows LakeBase clusters require, on top of standard username/password auth.
14
+
15
+ ### Installation
16
+
17
+ ```bash
18
+ pip install onehouse-python-sdk[lakebase]
19
+ ```
20
+
21
+ Requires Python 3.9+.
22
+
23
+ ### Quickstart
24
+
25
+ ```python
26
+ from onehouse_python_sdk import LakebaseClient
27
+
28
+ # Username / password
29
+ with LakebaseClient().connect(
30
+ host="<cluster-host>",
31
+ port=5432,
32
+ dbname="mydb",
33
+ user="admin",
34
+ password="secret",
35
+ ) as client:
36
+ rows = client.fetchall("SELECT * FROM mytable WHERE id = %s", (42,))
37
+ ```
38
+
39
+ ### Authentication flows
40
+
41
+ | Flow | Parameters |
42
+ |------|------------|
43
+ | Username / password | `user`, `password` |
44
+ | OIDC Device Flow (Okta / Auth0) | `browser_auth=true`, `oidc_client_id`, `oidc_issuer_url`, `oidc_iam_role` |
45
+ | Azure AD OAuth2 | `browser_auth=true`, `azure_oauth_tenant_id`, `azure_oauth_client_id`, `azure_oauth_client_secret` |
46
+ | Azure Entra ID SAML | `browser_auth=true`, `azure_tenant_id`, `azure_entity_id` |
47
+ | Built-in login form | `browser_auth=true` _(default when no IdP params are set)_ |
48
+ | External redirect | `browser_auth=true`, `auth_redirect_url` |
49
+
50
+ ```python
51
+ # OIDC Device Flow
52
+ client = LakebaseClient().connect(
53
+ host="<cluster-host>",
54
+ port=5432,
55
+ dbname="mydb",
56
+ browser_auth="true",
57
+ oidc_client_id="0oaXXXXX",
58
+ oidc_issuer_url="https://myorg.okta.com",
59
+ oidc_iam_role="arn:aws:iam::123456789012:role/LakeBaseRole",
60
+ )
61
+
62
+ # Azure AD OAuth2
63
+ client = LakebaseClient().connect(
64
+ host="<cluster-host>",
65
+ port=5432,
66
+ dbname="mydb",
67
+ browser_auth="true",
68
+ azure_oauth_tenant_id="your-tenant-id",
69
+ azure_oauth_client_id="your-client-id",
70
+ azure_oauth_client_secret="your-client-secret",
71
+ )
72
+ ```
73
+
74
+ DSN string form is also supported:
75
+
76
+ ```python
77
+ client = LakebaseClient().connect(
78
+ "postgresql://<cluster-host>:5432/mydb"
79
+ "?browser_auth=true"
80
+ "&oidc_client_id=0oaXXXX"
81
+ "&oidc_issuer_url=https://myorg.okta.com"
82
+ "&oidc_iam_role=arn:aws:iam::123456789012:role/LakeBaseRole"
83
+ )
84
+ ```
85
+
86
+ ### Client API
87
+
88
+ All clients implement the `SqlClient` interface:
89
+
90
+ | Method | Description |
91
+ |--------|-------------|
92
+ | `connect(dsn=None, **kwargs) → self` | Establish connection, returns self for chaining |
93
+ | `execute(sql, params)` | Run a statement, return rowcount |
94
+ | `fetchall(sql, params)` | Run a query, return `list[tuple]` |
95
+ | `fetchone(sql, params)` | Run a query, return first row or `None` |
96
+ | `cursor()` | Raw cursor for advanced use |
97
+ | `raw_connection` | Underlying `psycopg2` connection |
98
+ | `close()` | Close the connection |
99
+ | `__enter__` / `__exit__` | Context manager — closes on exit |
100
+
101
+ ### Notes
102
+
103
+ - **Credential caching** — auth tokens are cached for 4 minutes per connection parameters to avoid repeated browser prompts within the same process.
104
+ - **Callback port** — the local auth callback server defaults to port `8888` at path `/lakebase`. Override with `auth_callback_port` and `auth_callback_path`.
@@ -0,0 +1,64 @@
1
+ [build-system]
2
+ requires = ["hatchling"]
3
+ build-backend = "hatchling.build"
4
+
5
+ [project]
6
+ name = "onehouse-python-sdk"
7
+ version = "0.1.0"
8
+ description = "Onehouse Python SDK — modular data-plane clients for LakeBase and beyond"
9
+ readme = "README.md"
10
+ license = { text = "Proprietary" }
11
+ requires-python = ">=3.9"
12
+ authors = [{ name = "Onehouse Team", email = "eng@onehouse.ai" }]
13
+ keywords = ["onehouse", "lakebase", "postgresql", "oidc", "saml", "azure", "oauth2", "auth"]
14
+ classifiers = [
15
+ "Development Status :: 4 - Beta",
16
+ "Intended Audience :: Developers",
17
+ "License :: Other/Proprietary License",
18
+ "Programming Language :: Python :: 3",
19
+ "Programming Language :: Python :: 3.9",
20
+ "Programming Language :: Python :: 3.10",
21
+ "Programming Language :: Python :: 3.11",
22
+ "Programming Language :: Python :: 3.12",
23
+ "Topic :: Database",
24
+ "Topic :: Software Development :: Libraries :: Python Modules",
25
+ ]
26
+ dependencies = []
27
+
28
+ [project.optional-dependencies]
29
+ lakebase = [
30
+ "psycopg2-binary>=2.9",
31
+ "certifi>=2024.1",
32
+ ]
33
+ dev = [
34
+ "pytest>=7.0",
35
+ "pytest-cov>=4.0",
36
+ "pytest-mock>=3.0",
37
+ ]
38
+
39
+ [project.urls]
40
+ Homepage = "https://github.com/onehouseinc/onehouse-python-sdk"
41
+ Repository = "https://github.com/onehouseinc/onehouse-python-sdk"
42
+ Issues = "https://github.com/onehouseinc/onehouse-python-sdk/issues"
43
+
44
+ [tool.hatch.build.targets.wheel]
45
+ packages = ["src/onehouse_python_sdk"]
46
+
47
+ [tool.hatch.build.targets.sdist]
48
+ include = ["src/"]
49
+
50
+ [tool.pytest.ini_options]
51
+ testpaths = ["tests"]
52
+ addopts = "--cov=onehouse_python_sdk --cov-report=term-missing --cov-report=html"
53
+ markers = [
54
+ "live: requires a running LakeBase cluster (set LAKEBASE_TEST_DSN)",
55
+ ]
56
+
57
+ [tool.coverage.run]
58
+ source = ["src/onehouse_python_sdk"]
59
+
60
+ [tool.coverage.report]
61
+ exclude_lines = [
62
+ "pragma: no cover",
63
+ "if __name__ == .__main__.:",
64
+ ]
@@ -0,0 +1,8 @@
1
+ """Onehouse Python SDK — modular data-plane and control-plane clients."""
2
+
3
+ from . import exceptions
4
+ from .base import SqlClient
5
+ from .lakebase.client import LakebaseClient
6
+
7
+ __all__ = ["SqlClient", "LakebaseClient", "exceptions"]
8
+ __version__ = "0.1.0"
@@ -0,0 +1,50 @@
1
+ from abc import ABC, abstractmethod
2
+ from typing import Any, Optional
3
+
4
+
5
+ class SqlClient(ABC):
6
+ """Abstract base for all Onehouse SQL clients.
7
+
8
+ Subclasses implement connect() and raw_connection. The convenience methods
9
+ (execute, fetchall, fetchone) delegate to a raw DB-API 2.0 cursor so that
10
+ simple one-liners work without boilerplate, while raw_connection provides
11
+ an escape hatch for advanced use cases (streaming, metadata, etc.).
12
+ """
13
+
14
+ @abstractmethod
15
+ def connect(self, **kwargs) -> "SqlClient":
16
+ """Establish the underlying connection and return self."""
17
+ ...
18
+
19
+ @property
20
+ @abstractmethod
21
+ def raw_connection(self) -> Any:
22
+ """The underlying DB-API 2.0 connection object."""
23
+ ...
24
+
25
+ def cursor(self):
26
+ return self.raw_connection.cursor()
27
+
28
+ def execute(self, sql: str, params: Optional[tuple] = None) -> int:
29
+ cur = self.cursor()
30
+ cur.execute(sql, params)
31
+ return cur.rowcount
32
+
33
+ def fetchall(self, sql: str, params: Optional[tuple] = None) -> list:
34
+ cur = self.cursor()
35
+ cur.execute(sql, params)
36
+ return cur.fetchall()
37
+
38
+ def fetchone(self, sql: str, params: Optional[tuple] = None) -> Optional[tuple]:
39
+ cur = self.cursor()
40
+ cur.execute(sql, params)
41
+ return cur.fetchone()
42
+
43
+ def close(self) -> None:
44
+ self.raw_connection.close()
45
+
46
+ def __enter__(self) -> "SqlClient":
47
+ return self
48
+
49
+ def __exit__(self, *_) -> None:
50
+ self.close()
@@ -0,0 +1,10 @@
1
+ class OnehouseSdkError(Exception):
2
+ pass
3
+
4
+
5
+ class ConnectionError(OnehouseSdkError):
6
+ pass
7
+
8
+
9
+ class QueryError(OnehouseSdkError):
10
+ pass
@@ -0,0 +1,6 @@
1
+ """LakeBase subpackage — psycopg2 backend with browser-based auth."""
2
+
3
+ from .client import LakebaseClient
4
+ from .connection import connect
5
+
6
+ __all__ = ["LakebaseClient", "connect"]
@@ -0,0 +1,24 @@
1
+ from .browser_auth import (
2
+ AuthenticationError,
3
+ BrowserAuthHandler,
4
+ DEFAULT_CALLBACK_PATH,
5
+ DEFAULT_CALLBACK_PORT,
6
+ )
7
+ from .oidc_device_flow import OidcAuthenticationError, OidcDeviceFlowHandler
8
+ from .azure_oauth import AzureOAuthError, AzureOAuthHandler
9
+ from .saml import build_azure_entra_login_url, generate_saml_authn_request
10
+ from .local_auth_server import LocalAuthServer
11
+
12
+ __all__ = [
13
+ "AuthenticationError",
14
+ "AzureOAuthError",
15
+ "AzureOAuthHandler",
16
+ "BrowserAuthHandler",
17
+ "DEFAULT_CALLBACK_PATH",
18
+ "DEFAULT_CALLBACK_PORT",
19
+ "LocalAuthServer",
20
+ "OidcAuthenticationError",
21
+ "OidcDeviceFlowHandler",
22
+ "build_azure_entra_login_url",
23
+ "generate_saml_authn_request",
24
+ ]
@@ -0,0 +1,23 @@
1
+ """Shared SSL context that trusts the OS native cert store when available."""
2
+
3
+ import ssl
4
+
5
+
6
+ def default_ssl_context() -> ssl.SSLContext:
7
+ """Return an SSLContext for outbound HTTPS (token endpoints, IdP).
8
+
9
+ Prefers the OS native trust store via the optional ``truststore`` package
10
+ — required on corporate environments that MITM HTTPS with an internal CA
11
+ (Zscaler, Netskope, etc.) deployed to the Windows/macOS root store but
12
+ not to certifi's bundle. Falls back to certifi, then Python defaults.
13
+ """
14
+ try:
15
+ import truststore
16
+ return truststore.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
17
+ except ImportError:
18
+ pass
19
+ try:
20
+ import certifi
21
+ return ssl.create_default_context(cafile=certifi.where())
22
+ except ImportError:
23
+ return ssl.create_default_context()
@@ -0,0 +1,194 @@
1
+ """Azure AD OAuth2 Authorization Code flow (form_post response mode)."""
2
+
3
+ import base64
4
+ import gzip
5
+ import json
6
+ import logging
7
+ import os
8
+ import secrets
9
+ import urllib.parse
10
+ import urllib.request
11
+ import webbrowser
12
+ import zlib
13
+ from typing import Optional
14
+
15
+ from ._ssl import default_ssl_context
16
+ from .local_auth_server import LocalAuthServer
17
+
18
+ logger = logging.getLogger(__name__)
19
+
20
+ _DEFAULT_TIMEOUT_SECONDS = 120
21
+
22
+
23
+ class AzureOAuthError(Exception):
24
+ pass
25
+
26
+
27
+ class AzureOAuthHandler:
28
+ """Authenticates via Azure AD OAuth2 Authorization Code flow.
29
+
30
+ Azure POSTs the authorization code (form_post response mode), so a local
31
+ callback server is required. The token exchange requests a SAML2 token
32
+ (``requested_token_type=urn:ietf:params:oauth:token-type:saml2``);
33
+ the result is GZIP-compressed and base64-encoded for the PostgreSQL password.
34
+ """
35
+
36
+ def __init__(
37
+ self,
38
+ tenant_id: str,
39
+ client_id: str,
40
+ client_secret: str,
41
+ scope: str = "openid",
42
+ iam_role_arn: Optional[str] = None,
43
+ callback_path: str = "/lakebase",
44
+ ) -> None:
45
+ self._tenant_id = tenant_id
46
+ self._client_id = client_id
47
+ self._client_secret = client_secret
48
+ self._scope = scope or "openid"
49
+ self._iam_role_arn = iam_role_arn
50
+ self._callback_path = callback_path
51
+
52
+ def authenticate(self, callback_port: int, auth_server: LocalAuthServer) -> dict:
53
+ """Opens browser, waits for callback, exchanges code for token.
54
+
55
+ Returns dict with keys ``user``, ``password``, and optionally ``options``.
56
+ """
57
+ state = secrets.token_urlsafe(32)
58
+ redirect_uri = f"http://localhost:{callback_port}{self._callback_path}"
59
+ auth_url = self._build_authorization_url(redirect_uri, state)
60
+
61
+ webbrowser.open(auth_url)
62
+ logger.info("=" * 80)
63
+ logger.info("AZURE AD OAUTH2 AUTHENTICATION REQUIRED")
64
+ logger.info("=" * 80)
65
+ logger.info("Please complete authentication in the opened browser window.")
66
+ logger.info("\nIf the browser didn't open automatically, visit:")
67
+ logger.info(" %s", auth_url)
68
+ logger.info("=" * 80)
69
+
70
+ event = auth_server.start()
71
+ if not event.wait(timeout=_DEFAULT_TIMEOUT_SECONDS):
72
+ raise AzureOAuthError(
73
+ f"Azure OAuth2 timed out after {_DEFAULT_TIMEOUT_SECONDS}s"
74
+ )
75
+
76
+ callback_data = auth_server.get_result()
77
+ returned_state = callback_data.get("state")
78
+ if returned_state != state:
79
+ raise AzureOAuthError("State parameter mismatch — possible CSRF attack")
80
+
81
+ # SAML response posted directly (some Azure Entra configurations)
82
+ saml_assertion = callback_data.get("SAMLResponse")
83
+ if saml_assertion:
84
+ password = base64.b64encode(
85
+ saml_assertion.encode("utf-8")
86
+ ).decode("ascii")
87
+ else:
88
+ code = callback_data.get("code")
89
+ if not code:
90
+ error = callback_data.get("error", "unknown_error")
91
+ desc = callback_data.get("error_description", "No description")
92
+ raise AzureOAuthError(f"OAuth2 error: {error} — {desc}")
93
+ password = self._exchange_code_for_token(code, redirect_uri)
94
+
95
+ result = {"user": "user", "password": password}
96
+ if self._iam_role_arn:
97
+ result["options"] = f"-c onehouse.lakebase_iam_role={self._iam_role_arn}"
98
+ return result
99
+
100
+ def _build_authorization_url(self, redirect_uri: str, state: str) -> str:
101
+ base = (
102
+ f"https://login.microsoftonline.com/{self._tenant_id}/oauth2/v2.0/authorize"
103
+ )
104
+ params = {
105
+ "client_id": self._client_id,
106
+ "response_type": "code",
107
+ "response_mode": "form_post",
108
+ "redirect_uri": redirect_uri,
109
+ "scope": self._scope,
110
+ "state": state,
111
+ }
112
+ return f"{base}?{urllib.parse.urlencode(params)}"
113
+
114
+ def _exchange_code_for_token(self, code: str, redirect_uri: str) -> str:
115
+ token_url = (
116
+ f"https://login.microsoftonline.com/{self._tenant_id}/oauth2/token"
117
+ )
118
+ params = {
119
+ "grant_type": "authorization_code",
120
+ "code": code,
121
+ "client_id": self._client_id,
122
+ "client_secret": self._client_secret,
123
+ "redirect_uri": redirect_uri,
124
+ "resource": self._client_id,
125
+ "requested_token_type": "urn:ietf:params:oauth:token-type:saml2",
126
+ }
127
+ body = urllib.parse.urlencode(params).encode()
128
+ req = urllib.request.Request(
129
+ token_url,
130
+ data=body,
131
+ headers={"Content-Type": "application/x-www-form-urlencoded"},
132
+ method="POST",
133
+ )
134
+ try:
135
+ with urllib.request.urlopen(req, timeout=30, context=default_ssl_context()) as resp:
136
+ data = json.loads(resp.read().decode())
137
+ except urllib.error.HTTPError as exc:
138
+ raise AzureOAuthError(
139
+ f"Token exchange failed ({exc.code}): {exc.read().decode()}"
140
+ ) from exc
141
+
142
+ if "access_token" in data:
143
+ try:
144
+ # Azure may use URL-safe base64 (-/_ in place of +/) and strip padding.
145
+ tok = data["access_token"]
146
+ tok_padded = tok + "=" * (-len(tok) % 4)
147
+ saml_bytes = base64.urlsafe_b64decode(tok_padded)
148
+ try:
149
+ # Some tenants return DEFLATE-compressed SAML; others plain XML bytes.
150
+ saml_content = zlib.decompress(saml_bytes, -zlib.MAX_WBITS).decode("utf-8")
151
+ except zlib.error:
152
+ saml_content = saml_bytes.decode("utf-8")
153
+ # Wrap in a proper SAML Response envelope with the
154
+ # required <samlp:Status> Success element. STS
155
+ # AssumeRoleWithSAML requires a full Response, not a
156
+ # bare Assertion.
157
+ if "<samlp:Response" not in saml_content and "<Response" not in saml_content:
158
+ saml_content = (
159
+ '<samlp:Response xmlns:samlp="urn:oasis:names:tc:SAML:2.0:protocol">'
160
+ '<samlp:Status><samlp:StatusCode '
161
+ 'Value="urn:oasis:names:tc:SAML:2.0:status:Success"/>'
162
+ '</samlp:Status>'
163
+ f'{saml_content}'
164
+ '</samlp:Response>'
165
+ )
166
+ # Wire encoding: base64(gzip(base64(xml))) — the XML is
167
+ # base64'd first, then gzipped, then base64'd again.
168
+ # Confirmed by decoding a password emitted by v7 JDBC.
169
+ inner_b64 = base64.b64encode(saml_content.encode("utf-8")).decode("ascii")
170
+ compressed = gzip.compress(inner_b64.encode("utf-8"))
171
+ return "gzip:" + base64.b64encode(compressed).decode("ascii")
172
+ except Exception:
173
+ logger.warning("Could not parse access_token as SAML; treating as opaque token")
174
+
175
+ if "id_token" in data:
176
+ if self._iam_role_arn:
177
+ self._log_id_token_claims(data["id_token"])
178
+ return data["id_token"]
179
+
180
+ raise AzureOAuthError("Token response missing id_token and access_token")
181
+
182
+ def _log_id_token_claims(self, id_token: str) -> None:
183
+ """Decode and log JWT claims for diagnostics (no signature validation)."""
184
+ try:
185
+ parts = id_token.split(".")
186
+ if len(parts) != 3:
187
+ return
188
+ # Add padding
189
+ payload_b64 = parts[1] + "=" * (-len(parts[1]) % 4)
190
+ payload = json.loads(base64.urlsafe_b64decode(payload_b64).decode("utf-8"))
191
+ logger.info("Azure ID token claims: iss=%s aud=%s sub=%s",
192
+ payload.get("iss"), payload.get("aud"), payload.get("sub"))
193
+ except Exception:
194
+ pass