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.
- onehouse_python_sdk-0.1.0/.gitignore +25 -0
- onehouse_python_sdk-0.1.0/PKG-INFO +134 -0
- onehouse_python_sdk-0.1.0/README.md +104 -0
- onehouse_python_sdk-0.1.0/pyproject.toml +64 -0
- onehouse_python_sdk-0.1.0/src/onehouse_python_sdk/__init__.py +8 -0
- onehouse_python_sdk-0.1.0/src/onehouse_python_sdk/base.py +50 -0
- onehouse_python_sdk-0.1.0/src/onehouse_python_sdk/exceptions.py +10 -0
- onehouse_python_sdk-0.1.0/src/onehouse_python_sdk/lakebase/__init__.py +6 -0
- onehouse_python_sdk-0.1.0/src/onehouse_python_sdk/lakebase/auth/__init__.py +24 -0
- onehouse_python_sdk-0.1.0/src/onehouse_python_sdk/lakebase/auth/_ssl.py +23 -0
- onehouse_python_sdk-0.1.0/src/onehouse_python_sdk/lakebase/auth/azure_oauth.py +194 -0
- onehouse_python_sdk-0.1.0/src/onehouse_python_sdk/lakebase/auth/browser_auth.py +249 -0
- onehouse_python_sdk-0.1.0/src/onehouse_python_sdk/lakebase/auth/local_auth_server.py +237 -0
- onehouse_python_sdk-0.1.0/src/onehouse_python_sdk/lakebase/auth/oidc_device_flow.py +167 -0
- onehouse_python_sdk-0.1.0/src/onehouse_python_sdk/lakebase/auth/saml.py +43 -0
- onehouse_python_sdk-0.1.0/src/onehouse_python_sdk/lakebase/client.py +29 -0
- onehouse_python_sdk-0.1.0/src/onehouse_python_sdk/lakebase/connection.py +145 -0
|
@@ -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,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,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
|