dbx-tools-postgres 0.6.78__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.
- dbx_tools_postgres-0.6.78/PKG-INFO +188 -0
- dbx_tools_postgres-0.6.78/README.md +174 -0
- dbx_tools_postgres-0.6.78/pyproject.toml +26 -0
- dbx_tools_postgres-0.6.78/pyproject.toml.orig +28 -0
- dbx_tools_postgres-0.6.78/src/dbx_tools/postgres/__init__.py +109 -0
- dbx_tools_postgres-0.6.78/src/dbx_tools/postgres/address.py +114 -0
- dbx_tools_postgres-0.6.78/src/dbx_tools/postgres/advisory_lock.py +223 -0
- dbx_tools_postgres-0.6.78/src/dbx_tools/postgres/engine.py +425 -0
- dbx_tools_postgres-0.6.78/src/dbx_tools/postgres/topic_bus.py +380 -0
|
@@ -0,0 +1,188 @@
|
|
|
1
|
+
Metadata-Version: 2.3
|
|
2
|
+
Name: dbx-tools-postgres
|
|
3
|
+
Version: 0.6.78
|
|
4
|
+
Summary: WorkspaceClient-backed Lakebase Postgres resolution, SQLAlchemy engines, advisory locks, and LISTEN/NOTIFY topic bus
|
|
5
|
+
Requires-Dist: asyncpg>=0.30,<1
|
|
6
|
+
Requires-Dist: databricks-sdk>=0.63.0,<1
|
|
7
|
+
Requires-Dist: greenlet>=3.2,<4
|
|
8
|
+
Requires-Dist: dbx-tools-core==0.6.78
|
|
9
|
+
Requires-Dist: psycopg[binary]>=3.2.9,<4
|
|
10
|
+
Requires-Dist: sqlalchemy>=2.0.41,<3
|
|
11
|
+
Requires-Python: >=3.10
|
|
12
|
+
Project-URL: Source, https://github.com/reggie-db/dbx-tools/tree/main/packages/py/postgres
|
|
13
|
+
Description-Content-Type: text/markdown
|
|
14
|
+
|
|
15
|
+
# `dbx-tools-postgres`
|
|
16
|
+
|
|
17
|
+
Python Lakebase/Postgres connection setup, advisory locks, and topic fan-out for
|
|
18
|
+
services that already hold a Databricks `WorkspaceClient`. This package is
|
|
19
|
+
currently an unpublished workspace package, and is the Python counterpart to
|
|
20
|
+
`@dbx-tools/postgres` plus `@dbx-tools/appkit`'s address parsing.
|
|
21
|
+
|
|
22
|
+
Install directly from this monorepo:
|
|
23
|
+
|
|
24
|
+
```bash
|
|
25
|
+
pip install "dbx-tools-postgres @ git+https://github.com/reggie-db/dbx-tools.git@main#subdirectory=packages/py/postgres"
|
|
26
|
+
```
|
|
27
|
+
|
|
28
|
+
Key features:
|
|
29
|
+
|
|
30
|
+
- accepts the same Postgres URI, Lakebase resource path, hostname, and project-id
|
|
31
|
+
address shapes as `@dbx-tools/appkit`;
|
|
32
|
+
- resolves missing autoscaling endpoint fields through
|
|
33
|
+
`WorkspaceClient.api_client`;
|
|
34
|
+
- resolves provisioned Lakebase instance DNS through
|
|
35
|
+
`WorkspaceClient.database`;
|
|
36
|
+
- injects a fresh database credential on SQLAlchemy's `do_connect` event rather
|
|
37
|
+
than storing an expiring password in the engine URL, using the SDK's
|
|
38
|
+
provisioned-instance API or the Autoscaling `/postgres/credentials` endpoint;
|
|
39
|
+
- supports sync psycopg and asyncpg SQLAlchemy engines;
|
|
40
|
+
- derives advisory-lock ids from the same stable structured keys as the Node
|
|
41
|
+
package and holds one checked-out connection for the full critical section;
|
|
42
|
+
- provides blocking and try-lock context managers for session and transaction
|
|
43
|
+
locks, with sync and async SQLAlchemy variants;
|
|
44
|
+
- fans messages out to every process on a channel with `PostgresTopicBus`, using
|
|
45
|
+
the same lifecycle and wire envelope as the Node package.
|
|
46
|
+
|
|
47
|
+
```python
|
|
48
|
+
from databricks.sdk import WorkspaceClient
|
|
49
|
+
from dbx_tools.postgres import PostgresEngineConfig, create_async_engine
|
|
50
|
+
|
|
51
|
+
engine = create_async_engine(
|
|
52
|
+
WorkspaceClient(),
|
|
53
|
+
PostgresEngineConfig(instance_name="my-lakebase", database="databricks_postgres"),
|
|
54
|
+
pool_pre_ping=True,
|
|
55
|
+
pool_recycle=1800,
|
|
56
|
+
)
|
|
57
|
+
```
|
|
58
|
+
|
|
59
|
+
Pass `credential_provider=` to either engine factory to inject credentials from
|
|
60
|
+
another source while retaining the same connect-time rotation behavior.
|
|
61
|
+
|
|
62
|
+
```python
|
|
63
|
+
from dbx_tools.postgres import advisory_transaction_lock
|
|
64
|
+
|
|
65
|
+
with advisory_transaction_lock(engine, ["schema-install", "v2"]) as connection:
|
|
66
|
+
connection.exec_driver_sql("CREATE TABLE IF NOT EXISTS ...")
|
|
67
|
+
```
|
|
68
|
+
|
|
69
|
+
## Topic bus
|
|
70
|
+
|
|
71
|
+
`PostgresTopicBus` is async Postgres topic fan-out built on `LISTEN`/`NOTIFY`. Its
|
|
72
|
+
public lifecycle and wire shape match `@dbx-tools/postgres`'s `PostgresTopicBus`,
|
|
73
|
+
so Node and Python services can share a channel:
|
|
74
|
+
|
|
75
|
+
- `PostgresTopicBus(engine, options)`;
|
|
76
|
+
- `channelName`;
|
|
77
|
+
- `await start()`;
|
|
78
|
+
- `await broadcast(topic, TopicPublishInput(...))`;
|
|
79
|
+
- `await listen(topic, listener)` returning an async unsubscribe function;
|
|
80
|
+
- `await close()`;
|
|
81
|
+
- envelope fields `id`, `topic`, `type`, `metadata`, `body`, and `publishedAt`.
|
|
82
|
+
|
|
83
|
+
Channel derivation ports the Node stable-key and FNV rules, so equivalent channel
|
|
84
|
+
parts resolve to the same PostgreSQL identifier in Python and Node.
|
|
85
|
+
|
|
86
|
+
```python
|
|
87
|
+
from dbx_tools.postgres import PostgresTopicBus, TopicPublishInput
|
|
88
|
+
|
|
89
|
+
bus = PostgresTopicBus(engine, channel=["billing", "production"])
|
|
90
|
+
|
|
91
|
+
unsubscribe = await bus.listen("invoice.updated", handle_invoice)
|
|
92
|
+
await bus.broadcast(
|
|
93
|
+
"invoice.updated",
|
|
94
|
+
TopicPublishInput(type="invoice.updated", body={"invoice_id": "inv-7"}),
|
|
95
|
+
)
|
|
96
|
+
```
|
|
97
|
+
|
|
98
|
+
Delivery is live and unstored, like PostgreSQL `LISTEN`/`NOTIFY` itself. Use a
|
|
99
|
+
table or queue when consumers need replay or acknowledgements.
|
|
100
|
+
|
|
101
|
+
## Databricks notebooks and Spark
|
|
102
|
+
|
|
103
|
+
Verified end to end against a Lakebase endpoint on serverless notebook compute.
|
|
104
|
+
[`packages/example/notebooks/bus-lakebase.py`](../../example/notebooks/bus-lakebase.py)
|
|
105
|
+
is the runnable version of everything below.
|
|
106
|
+
|
|
107
|
+
Two things about the Databricks Python runtime change how the bus is called, and
|
|
108
|
+
neither is a limitation of the bus itself:
|
|
109
|
+
|
|
110
|
+
- **A notebook kernel already runs an event loop**, so `asyncio.run` in a cell
|
|
111
|
+
raises `RuntimeError: asyncio.run() cannot be called from a running event
|
|
112
|
+
loop`. Drive the coroutine on a short-lived thread with its own loop rather
|
|
113
|
+
than reaching for `nest_asyncio` — the bus holds a dedicated `LISTEN`
|
|
114
|
+
connection bound to whichever loop started it, so one loop per bus lifetime is
|
|
115
|
+
the invariant to preserve.
|
|
116
|
+
- **Install with `%pip install`, not `pip install --target`.** A `--target`
|
|
117
|
+
install leaves the runtime's preloaded `typing_extensions` ahead of the new
|
|
118
|
+
one on `sys.path`, and importing `dbx_tools.postgres` then fails with
|
|
119
|
+
`ImportError: cannot import name 'TypeAliasType'`. `%pip` restarts the Python
|
|
120
|
+
process, which resolves it.
|
|
121
|
+
|
|
122
|
+
### Publishing from a Spark UDF
|
|
123
|
+
|
|
124
|
+
Publishing from executors works. Listening from them does not, and should not be
|
|
125
|
+
attempted: a UDF invocation is short-lived, while `listen` keeps a connection
|
|
126
|
+
open until `close`.
|
|
127
|
+
|
|
128
|
+
Executors have no Databricks credentials, so they cannot build a
|
|
129
|
+
`WorkspaceClient`. This is where connect-time credential injection pays off — the
|
|
130
|
+
driver mints the Lakebase token once and the UDF closes over it, so the executor
|
|
131
|
+
builds a plain SQLAlchemy engine and installs the token as its provider:
|
|
132
|
+
|
|
133
|
+
```python
|
|
134
|
+
from sqlalchemy import URL
|
|
135
|
+
from sqlalchemy.ext.asyncio import create_async_engine
|
|
136
|
+
from dbx_tools.postgres import (
|
|
137
|
+
PostgresTopicBus,
|
|
138
|
+
TopicPublishInput,
|
|
139
|
+
install_credential_injection,
|
|
140
|
+
)
|
|
141
|
+
|
|
142
|
+
# driver: resolve once, capture in the closure
|
|
143
|
+
resolved = resolve_postgres_connection(workspace_client, config)
|
|
144
|
+
token = workspace_client.api_client.do(
|
|
145
|
+
"POST",
|
|
146
|
+
"/api/2.0/postgres/credentials",
|
|
147
|
+
body={"endpoint": resolved.endpoint},
|
|
148
|
+
)["token"]
|
|
149
|
+
|
|
150
|
+
|
|
151
|
+
@udf(returnType=StringType())
|
|
152
|
+
def publish(key: str) -> str:
|
|
153
|
+
async def run() -> str:
|
|
154
|
+
engine = create_async_engine(
|
|
155
|
+
URL.create(
|
|
156
|
+
"postgresql+asyncpg",
|
|
157
|
+
username=resolved.user,
|
|
158
|
+
host=resolved.host,
|
|
159
|
+
port=resolved.port,
|
|
160
|
+
database=resolved.database,
|
|
161
|
+
query={"ssl": resolved.ssl_mode},
|
|
162
|
+
)
|
|
163
|
+
)
|
|
164
|
+
install_credential_injection(engine.sync_engine, lambda: token)
|
|
165
|
+
bus = PostgresTopicBus(engine, channel="app-events")
|
|
166
|
+
try:
|
|
167
|
+
message = await bus.broadcast(
|
|
168
|
+
"row.processed", TopicPublishInput(type="row.processed", body={"key": key})
|
|
169
|
+
)
|
|
170
|
+
return message.id
|
|
171
|
+
finally:
|
|
172
|
+
await bus.close()
|
|
173
|
+
await engine.dispose()
|
|
174
|
+
|
|
175
|
+
return asyncio.run(run())
|
|
176
|
+
```
|
|
177
|
+
|
|
178
|
+
Constraints worth knowing before this reaches production:
|
|
179
|
+
|
|
180
|
+
- A captured token EXPIRES (about an hour). Re-mint per job run; a long-running
|
|
181
|
+
streaming query needs a provider that refreshes instead of a captured string.
|
|
182
|
+
- Each UDF call opens and closes its own connection, so batch the publish at
|
|
183
|
+
partition scope (`mapInPandas`, `foreachPartition`) rather than per row.
|
|
184
|
+
- `spark.sparkContext.broadcast` is unavailable on serverless (Spark Connect).
|
|
185
|
+
A plain closure over driver-side values serializes with the UDF and is enough.
|
|
186
|
+
- Delivery stays live and unstored: if no listener is connected when the UDF
|
|
187
|
+
publishes, the message is gone. Write to a table when executors produce
|
|
188
|
+
results a consumer must not miss.
|
|
@@ -0,0 +1,174 @@
|
|
|
1
|
+
# `dbx-tools-postgres`
|
|
2
|
+
|
|
3
|
+
Python Lakebase/Postgres connection setup, advisory locks, and topic fan-out for
|
|
4
|
+
services that already hold a Databricks `WorkspaceClient`. This package is
|
|
5
|
+
currently an unpublished workspace package, and is the Python counterpart to
|
|
6
|
+
`@dbx-tools/postgres` plus `@dbx-tools/appkit`'s address parsing.
|
|
7
|
+
|
|
8
|
+
Install directly from this monorepo:
|
|
9
|
+
|
|
10
|
+
```bash
|
|
11
|
+
pip install "dbx-tools-postgres @ git+https://github.com/reggie-db/dbx-tools.git@main#subdirectory=packages/py/postgres"
|
|
12
|
+
```
|
|
13
|
+
|
|
14
|
+
Key features:
|
|
15
|
+
|
|
16
|
+
- accepts the same Postgres URI, Lakebase resource path, hostname, and project-id
|
|
17
|
+
address shapes as `@dbx-tools/appkit`;
|
|
18
|
+
- resolves missing autoscaling endpoint fields through
|
|
19
|
+
`WorkspaceClient.api_client`;
|
|
20
|
+
- resolves provisioned Lakebase instance DNS through
|
|
21
|
+
`WorkspaceClient.database`;
|
|
22
|
+
- injects a fresh database credential on SQLAlchemy's `do_connect` event rather
|
|
23
|
+
than storing an expiring password in the engine URL, using the SDK's
|
|
24
|
+
provisioned-instance API or the Autoscaling `/postgres/credentials` endpoint;
|
|
25
|
+
- supports sync psycopg and asyncpg SQLAlchemy engines;
|
|
26
|
+
- derives advisory-lock ids from the same stable structured keys as the Node
|
|
27
|
+
package and holds one checked-out connection for the full critical section;
|
|
28
|
+
- provides blocking and try-lock context managers for session and transaction
|
|
29
|
+
locks, with sync and async SQLAlchemy variants;
|
|
30
|
+
- fans messages out to every process on a channel with `PostgresTopicBus`, using
|
|
31
|
+
the same lifecycle and wire envelope as the Node package.
|
|
32
|
+
|
|
33
|
+
```python
|
|
34
|
+
from databricks.sdk import WorkspaceClient
|
|
35
|
+
from dbx_tools.postgres import PostgresEngineConfig, create_async_engine
|
|
36
|
+
|
|
37
|
+
engine = create_async_engine(
|
|
38
|
+
WorkspaceClient(),
|
|
39
|
+
PostgresEngineConfig(instance_name="my-lakebase", database="databricks_postgres"),
|
|
40
|
+
pool_pre_ping=True,
|
|
41
|
+
pool_recycle=1800,
|
|
42
|
+
)
|
|
43
|
+
```
|
|
44
|
+
|
|
45
|
+
Pass `credential_provider=` to either engine factory to inject credentials from
|
|
46
|
+
another source while retaining the same connect-time rotation behavior.
|
|
47
|
+
|
|
48
|
+
```python
|
|
49
|
+
from dbx_tools.postgres import advisory_transaction_lock
|
|
50
|
+
|
|
51
|
+
with advisory_transaction_lock(engine, ["schema-install", "v2"]) as connection:
|
|
52
|
+
connection.exec_driver_sql("CREATE TABLE IF NOT EXISTS ...")
|
|
53
|
+
```
|
|
54
|
+
|
|
55
|
+
## Topic bus
|
|
56
|
+
|
|
57
|
+
`PostgresTopicBus` is async Postgres topic fan-out built on `LISTEN`/`NOTIFY`. Its
|
|
58
|
+
public lifecycle and wire shape match `@dbx-tools/postgres`'s `PostgresTopicBus`,
|
|
59
|
+
so Node and Python services can share a channel:
|
|
60
|
+
|
|
61
|
+
- `PostgresTopicBus(engine, options)`;
|
|
62
|
+
- `channelName`;
|
|
63
|
+
- `await start()`;
|
|
64
|
+
- `await broadcast(topic, TopicPublishInput(...))`;
|
|
65
|
+
- `await listen(topic, listener)` returning an async unsubscribe function;
|
|
66
|
+
- `await close()`;
|
|
67
|
+
- envelope fields `id`, `topic`, `type`, `metadata`, `body`, and `publishedAt`.
|
|
68
|
+
|
|
69
|
+
Channel derivation ports the Node stable-key and FNV rules, so equivalent channel
|
|
70
|
+
parts resolve to the same PostgreSQL identifier in Python and Node.
|
|
71
|
+
|
|
72
|
+
```python
|
|
73
|
+
from dbx_tools.postgres import PostgresTopicBus, TopicPublishInput
|
|
74
|
+
|
|
75
|
+
bus = PostgresTopicBus(engine, channel=["billing", "production"])
|
|
76
|
+
|
|
77
|
+
unsubscribe = await bus.listen("invoice.updated", handle_invoice)
|
|
78
|
+
await bus.broadcast(
|
|
79
|
+
"invoice.updated",
|
|
80
|
+
TopicPublishInput(type="invoice.updated", body={"invoice_id": "inv-7"}),
|
|
81
|
+
)
|
|
82
|
+
```
|
|
83
|
+
|
|
84
|
+
Delivery is live and unstored, like PostgreSQL `LISTEN`/`NOTIFY` itself. Use a
|
|
85
|
+
table or queue when consumers need replay or acknowledgements.
|
|
86
|
+
|
|
87
|
+
## Databricks notebooks and Spark
|
|
88
|
+
|
|
89
|
+
Verified end to end against a Lakebase endpoint on serverless notebook compute.
|
|
90
|
+
[`packages/example/notebooks/bus-lakebase.py`](../../example/notebooks/bus-lakebase.py)
|
|
91
|
+
is the runnable version of everything below.
|
|
92
|
+
|
|
93
|
+
Two things about the Databricks Python runtime change how the bus is called, and
|
|
94
|
+
neither is a limitation of the bus itself:
|
|
95
|
+
|
|
96
|
+
- **A notebook kernel already runs an event loop**, so `asyncio.run` in a cell
|
|
97
|
+
raises `RuntimeError: asyncio.run() cannot be called from a running event
|
|
98
|
+
loop`. Drive the coroutine on a short-lived thread with its own loop rather
|
|
99
|
+
than reaching for `nest_asyncio` — the bus holds a dedicated `LISTEN`
|
|
100
|
+
connection bound to whichever loop started it, so one loop per bus lifetime is
|
|
101
|
+
the invariant to preserve.
|
|
102
|
+
- **Install with `%pip install`, not `pip install --target`.** A `--target`
|
|
103
|
+
install leaves the runtime's preloaded `typing_extensions` ahead of the new
|
|
104
|
+
one on `sys.path`, and importing `dbx_tools.postgres` then fails with
|
|
105
|
+
`ImportError: cannot import name 'TypeAliasType'`. `%pip` restarts the Python
|
|
106
|
+
process, which resolves it.
|
|
107
|
+
|
|
108
|
+
### Publishing from a Spark UDF
|
|
109
|
+
|
|
110
|
+
Publishing from executors works. Listening from them does not, and should not be
|
|
111
|
+
attempted: a UDF invocation is short-lived, while `listen` keeps a connection
|
|
112
|
+
open until `close`.
|
|
113
|
+
|
|
114
|
+
Executors have no Databricks credentials, so they cannot build a
|
|
115
|
+
`WorkspaceClient`. This is where connect-time credential injection pays off — the
|
|
116
|
+
driver mints the Lakebase token once and the UDF closes over it, so the executor
|
|
117
|
+
builds a plain SQLAlchemy engine and installs the token as its provider:
|
|
118
|
+
|
|
119
|
+
```python
|
|
120
|
+
from sqlalchemy import URL
|
|
121
|
+
from sqlalchemy.ext.asyncio import create_async_engine
|
|
122
|
+
from dbx_tools.postgres import (
|
|
123
|
+
PostgresTopicBus,
|
|
124
|
+
TopicPublishInput,
|
|
125
|
+
install_credential_injection,
|
|
126
|
+
)
|
|
127
|
+
|
|
128
|
+
# driver: resolve once, capture in the closure
|
|
129
|
+
resolved = resolve_postgres_connection(workspace_client, config)
|
|
130
|
+
token = workspace_client.api_client.do(
|
|
131
|
+
"POST",
|
|
132
|
+
"/api/2.0/postgres/credentials",
|
|
133
|
+
body={"endpoint": resolved.endpoint},
|
|
134
|
+
)["token"]
|
|
135
|
+
|
|
136
|
+
|
|
137
|
+
@udf(returnType=StringType())
|
|
138
|
+
def publish(key: str) -> str:
|
|
139
|
+
async def run() -> str:
|
|
140
|
+
engine = create_async_engine(
|
|
141
|
+
URL.create(
|
|
142
|
+
"postgresql+asyncpg",
|
|
143
|
+
username=resolved.user,
|
|
144
|
+
host=resolved.host,
|
|
145
|
+
port=resolved.port,
|
|
146
|
+
database=resolved.database,
|
|
147
|
+
query={"ssl": resolved.ssl_mode},
|
|
148
|
+
)
|
|
149
|
+
)
|
|
150
|
+
install_credential_injection(engine.sync_engine, lambda: token)
|
|
151
|
+
bus = PostgresTopicBus(engine, channel="app-events")
|
|
152
|
+
try:
|
|
153
|
+
message = await bus.broadcast(
|
|
154
|
+
"row.processed", TopicPublishInput(type="row.processed", body={"key": key})
|
|
155
|
+
)
|
|
156
|
+
return message.id
|
|
157
|
+
finally:
|
|
158
|
+
await bus.close()
|
|
159
|
+
await engine.dispose()
|
|
160
|
+
|
|
161
|
+
return asyncio.run(run())
|
|
162
|
+
```
|
|
163
|
+
|
|
164
|
+
Constraints worth knowing before this reaches production:
|
|
165
|
+
|
|
166
|
+
- A captured token EXPIRES (about an hour). Re-mint per job run; a long-running
|
|
167
|
+
streaming query needs a provider that refreshes instead of a captured string.
|
|
168
|
+
- Each UDF call opens and closes its own connection, so batch the publish at
|
|
169
|
+
partition scope (`mapInPandas`, `foreachPartition`) rather than per row.
|
|
170
|
+
- `spark.sparkContext.broadcast` is unavailable on serverless (Spark Connect).
|
|
171
|
+
A plain closure over driver-side values serializes with the UDF and is enough.
|
|
172
|
+
- Delivery stays live and unstored: if no listener is connected when the UDF
|
|
173
|
+
publishes, the message is gone. Write to a table when executors produce
|
|
174
|
+
results a consumer must not miss.
|
|
@@ -0,0 +1,26 @@
|
|
|
1
|
+
[project]
|
|
2
|
+
name = "dbx-tools-postgres"
|
|
3
|
+
version = "0.6.78"
|
|
4
|
+
description = "WorkspaceClient-backed Lakebase Postgres resolution, SQLAlchemy engines, advisory locks, and LISTEN/NOTIFY topic bus"
|
|
5
|
+
readme = "README.md"
|
|
6
|
+
requires-python = ">=3.10"
|
|
7
|
+
dependencies = [
|
|
8
|
+
"asyncpg>=0.30,<1",
|
|
9
|
+
"databricks-sdk>=0.63.0,<1",
|
|
10
|
+
"greenlet>=3.2,<4",
|
|
11
|
+
"dbx-tools-core==0.6.78",
|
|
12
|
+
"psycopg[binary]>=3.2.9,<4",
|
|
13
|
+
"sqlalchemy>=2.0.41,<3",
|
|
14
|
+
]
|
|
15
|
+
|
|
16
|
+
[project.urls]
|
|
17
|
+
Source = "https://github.com/reggie-db/dbx-tools/tree/main/packages/py/postgres"
|
|
18
|
+
|
|
19
|
+
[build-system]
|
|
20
|
+
requires = ["uv_build>=0.11.28,<0.12.0"]
|
|
21
|
+
build-backend = "uv_build"
|
|
22
|
+
|
|
23
|
+
[tool.uv.build-backend]
|
|
24
|
+
module-name = "dbx_tools.postgres"
|
|
25
|
+
module-root = "src"
|
|
26
|
+
namespace = true
|
|
@@ -0,0 +1,28 @@
|
|
|
1
|
+
# ~~ Generated by projen. To modify, edit .projenrc.js and run "bunx projen".
|
|
2
|
+
|
|
3
|
+
[project]
|
|
4
|
+
name = "dbx-tools-postgres"
|
|
5
|
+
version = "0.6.78"
|
|
6
|
+
description = "WorkspaceClient-backed Lakebase Postgres resolution, SQLAlchemy engines, advisory locks, and LISTEN/NOTIFY topic bus"
|
|
7
|
+
readme = "README.md"
|
|
8
|
+
requires-python = ">=3.10"
|
|
9
|
+
dependencies = [
|
|
10
|
+
"asyncpg>=0.30,<1",
|
|
11
|
+
"databricks-sdk>=0.63.0,<1",
|
|
12
|
+
"greenlet>=3.2,<4",
|
|
13
|
+
"dbx-tools-core==0.6.78",
|
|
14
|
+
"psycopg[binary]>=3.2.9,<4",
|
|
15
|
+
"sqlalchemy>=2.0.41,<3"
|
|
16
|
+
]
|
|
17
|
+
|
|
18
|
+
[project.urls]
|
|
19
|
+
Source = "https://github.com/reggie-db/dbx-tools/tree/main/packages/py/postgres"
|
|
20
|
+
|
|
21
|
+
[build-system]
|
|
22
|
+
requires = [ "uv_build>=0.11.28,<0.12.0" ]
|
|
23
|
+
build-backend = "uv_build"
|
|
24
|
+
|
|
25
|
+
[tool.uv.build-backend]
|
|
26
|
+
module-name = "dbx_tools.postgres"
|
|
27
|
+
module-root = "src"
|
|
28
|
+
namespace = true
|
|
@@ -0,0 +1,109 @@
|
|
|
1
|
+
from .address import (
|
|
2
|
+
SSL_MODES,
|
|
3
|
+
LakebaseConnectionInputs,
|
|
4
|
+
ParsedAddress,
|
|
5
|
+
SslMode,
|
|
6
|
+
parse_address,
|
|
7
|
+
parse_resource_path,
|
|
8
|
+
parseAddress,
|
|
9
|
+
parseResourcePath,
|
|
10
|
+
)
|
|
11
|
+
from .advisory_lock import (
|
|
12
|
+
ExplicitAdvisoryLockId,
|
|
13
|
+
acquire_advisory_lock,
|
|
14
|
+
acquire_advisory_lock_async,
|
|
15
|
+
advisory_lock,
|
|
16
|
+
advisory_lock_async,
|
|
17
|
+
advisory_lock_id,
|
|
18
|
+
advisory_transaction_lock,
|
|
19
|
+
advisory_transaction_lock_async,
|
|
20
|
+
explicit_advisory_lock_id,
|
|
21
|
+
release_advisory_lock,
|
|
22
|
+
release_advisory_lock_async,
|
|
23
|
+
try_advisory_lock,
|
|
24
|
+
try_advisory_lock_async,
|
|
25
|
+
try_advisory_transaction_lock,
|
|
26
|
+
try_advisory_transaction_lock_async,
|
|
27
|
+
)
|
|
28
|
+
from .engine import (
|
|
29
|
+
CredentialProvider,
|
|
30
|
+
PostgresEngineConfig,
|
|
31
|
+
ResolvedPostgresConnection,
|
|
32
|
+
WorkspaceClientLike,
|
|
33
|
+
autoscaling_credential_provider,
|
|
34
|
+
autoscalingCredentialProvider,
|
|
35
|
+
create_async_engine,
|
|
36
|
+
create_engine,
|
|
37
|
+
createAsyncEngine,
|
|
38
|
+
createEngine,
|
|
39
|
+
install_credential_injection,
|
|
40
|
+
installCredentialInjection,
|
|
41
|
+
resolve_postgres_connection,
|
|
42
|
+
resolvePostgresConnection,
|
|
43
|
+
workspace_credential_provider,
|
|
44
|
+
workspaceCredentialProvider,
|
|
45
|
+
)
|
|
46
|
+
from .topic_bus import (
|
|
47
|
+
PostgresTopicBus,
|
|
48
|
+
PostgresTopicBusOptions,
|
|
49
|
+
SerializableValue,
|
|
50
|
+
TopicListener,
|
|
51
|
+
TopicMessage,
|
|
52
|
+
TopicMetadata,
|
|
53
|
+
TopicMetadataProvider,
|
|
54
|
+
TopicPublishInput,
|
|
55
|
+
channel_name,
|
|
56
|
+
channelName,
|
|
57
|
+
)
|
|
58
|
+
|
|
59
|
+
__all__ = [
|
|
60
|
+
"SSL_MODES",
|
|
61
|
+
"CredentialProvider",
|
|
62
|
+
"ExplicitAdvisoryLockId",
|
|
63
|
+
"LakebaseConnectionInputs",
|
|
64
|
+
"ParsedAddress",
|
|
65
|
+
"PostgresEngineConfig",
|
|
66
|
+
"PostgresTopicBus",
|
|
67
|
+
"PostgresTopicBusOptions",
|
|
68
|
+
"ResolvedPostgresConnection",
|
|
69
|
+
"SerializableValue",
|
|
70
|
+
"SslMode",
|
|
71
|
+
"TopicListener",
|
|
72
|
+
"TopicMessage",
|
|
73
|
+
"TopicMetadata",
|
|
74
|
+
"TopicMetadataProvider",
|
|
75
|
+
"TopicPublishInput",
|
|
76
|
+
"WorkspaceClientLike",
|
|
77
|
+
"acquire_advisory_lock",
|
|
78
|
+
"acquire_advisory_lock_async",
|
|
79
|
+
"advisory_lock",
|
|
80
|
+
"advisory_lock_async",
|
|
81
|
+
"advisory_lock_id",
|
|
82
|
+
"advisory_transaction_lock",
|
|
83
|
+
"advisory_transaction_lock_async",
|
|
84
|
+
"autoscalingCredentialProvider",
|
|
85
|
+
"autoscaling_credential_provider",
|
|
86
|
+
"channelName",
|
|
87
|
+
"channel_name",
|
|
88
|
+
"createAsyncEngine",
|
|
89
|
+
"createEngine",
|
|
90
|
+
"create_async_engine",
|
|
91
|
+
"create_engine",
|
|
92
|
+
"explicit_advisory_lock_id",
|
|
93
|
+
"installCredentialInjection",
|
|
94
|
+
"install_credential_injection",
|
|
95
|
+
"parseAddress",
|
|
96
|
+
"parseResourcePath",
|
|
97
|
+
"parse_address",
|
|
98
|
+
"parse_resource_path",
|
|
99
|
+
"release_advisory_lock",
|
|
100
|
+
"release_advisory_lock_async",
|
|
101
|
+
"resolvePostgresConnection",
|
|
102
|
+
"resolve_postgres_connection",
|
|
103
|
+
"try_advisory_lock",
|
|
104
|
+
"try_advisory_lock_async",
|
|
105
|
+
"try_advisory_transaction_lock",
|
|
106
|
+
"try_advisory_transaction_lock_async",
|
|
107
|
+
"workspaceCredentialProvider",
|
|
108
|
+
"workspace_credential_provider",
|
|
109
|
+
]
|
|
@@ -0,0 +1,114 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import re
|
|
4
|
+
from dataclasses import asdict, dataclass
|
|
5
|
+
from typing import Literal
|
|
6
|
+
from urllib.parse import parse_qs, unquote, urlparse
|
|
7
|
+
|
|
8
|
+
SslMode = Literal["require", "disable", "prefer"]
|
|
9
|
+
SSL_MODES: tuple[SslMode, ...] = ("require", "disable", "prefer")
|
|
10
|
+
|
|
11
|
+
_URL_SCHEME_RE = re.compile(r"^(postgres|postgresql)://", re.IGNORECASE)
|
|
12
|
+
_PROJECT_ID_RE = re.compile(r"^[a-z][a-z0-9-]{0,61}[a-z0-9]$|^[a-z]$")
|
|
13
|
+
_HOSTNAME_HINT_RE = re.compile(r"^[a-z0-9][a-z0-9-]*(\.[a-z0-9][a-z0-9-]*)+$", re.IGNORECASE)
|
|
14
|
+
|
|
15
|
+
|
|
16
|
+
@dataclass(frozen=True, slots=True)
|
|
17
|
+
class LakebaseConnectionInputs:
|
|
18
|
+
project: str | None = None
|
|
19
|
+
branch: str | None = None
|
|
20
|
+
endpoint: str | None = None
|
|
21
|
+
database: str | None = None
|
|
22
|
+
host: str | None = None
|
|
23
|
+
port: int | None = None
|
|
24
|
+
ssl_mode: SslMode | None = None
|
|
25
|
+
|
|
26
|
+
def as_dict(self) -> dict[str, object]:
|
|
27
|
+
return {key: value for key, value in asdict(self).items() if value is not None}
|
|
28
|
+
|
|
29
|
+
|
|
30
|
+
@dataclass(frozen=True, slots=True)
|
|
31
|
+
class ParsedAddress(LakebaseConnectionInputs):
|
|
32
|
+
endpoint_id: str | None = None
|
|
33
|
+
database_resource_id: str | None = None
|
|
34
|
+
user: str | None = None
|
|
35
|
+
|
|
36
|
+
|
|
37
|
+
def parse_address(value: str | None) -> ParsedAddress:
|
|
38
|
+
if not value or not (address := value.strip()):
|
|
39
|
+
return ParsedAddress()
|
|
40
|
+
if _URL_SCHEME_RE.match(address):
|
|
41
|
+
return _parse_uri(address)
|
|
42
|
+
if address.startswith("projects/"):
|
|
43
|
+
return _parse_resource_path_segments(address)
|
|
44
|
+
if "." in address and _HOSTNAME_HINT_RE.fullmatch(address):
|
|
45
|
+
return ParsedAddress(host=address)
|
|
46
|
+
if _PROJECT_ID_RE.fullmatch(address):
|
|
47
|
+
return ParsedAddress(project=address)
|
|
48
|
+
return ParsedAddress()
|
|
49
|
+
|
|
50
|
+
|
|
51
|
+
def parse_resource_path(value: str | None) -> ParsedAddress:
|
|
52
|
+
if not value or not (address := value.strip()).startswith("projects/"):
|
|
53
|
+
return ParsedAddress()
|
|
54
|
+
return _parse_resource_path_segments(address)
|
|
55
|
+
|
|
56
|
+
|
|
57
|
+
def _parse_uri(address: str) -> ParsedAddress:
|
|
58
|
+
try:
|
|
59
|
+
parsed = urlparse(address)
|
|
60
|
+
port = parsed.port
|
|
61
|
+
except ValueError:
|
|
62
|
+
return ParsedAddress()
|
|
63
|
+
if parsed.scheme.lower() not in {"postgres", "postgresql"}:
|
|
64
|
+
return ParsedAddress()
|
|
65
|
+
ssl_mode_value = parse_qs(parsed.query).get("sslmode") or parse_qs(parsed.query).get("sslMode")
|
|
66
|
+
ssl_mode = ssl_mode_value[0].lower() if ssl_mode_value else None
|
|
67
|
+
return ParsedAddress(
|
|
68
|
+
host=parsed.hostname or None,
|
|
69
|
+
port=port,
|
|
70
|
+
user=unquote(parsed.username) if parsed.username else None,
|
|
71
|
+
database=unquote(parsed.path.removeprefix("/")) or None,
|
|
72
|
+
ssl_mode=ssl_mode if ssl_mode in SSL_MODES else None,
|
|
73
|
+
)
|
|
74
|
+
|
|
75
|
+
|
|
76
|
+
def _parse_resource_path_segments(address: str) -> ParsedAddress:
|
|
77
|
+
parts = address.split("/")
|
|
78
|
+
if len(parts) < 2 or parts[0] != "projects" or not parts[1]:
|
|
79
|
+
return ParsedAddress()
|
|
80
|
+
project = parts[1]
|
|
81
|
+
if len(parts) == 2:
|
|
82
|
+
return ParsedAddress(project=project)
|
|
83
|
+
if len(parts) == 4 and parts[2] == "branches" and parts[3]:
|
|
84
|
+
return ParsedAddress(project=project, branch=parts[3])
|
|
85
|
+
if (
|
|
86
|
+
len(parts) == 6
|
|
87
|
+
and parts[2] == "branches"
|
|
88
|
+
and parts[3]
|
|
89
|
+
and parts[4] == "endpoints"
|
|
90
|
+
and parts[5]
|
|
91
|
+
):
|
|
92
|
+
return ParsedAddress(
|
|
93
|
+
project=project,
|
|
94
|
+
branch=parts[3],
|
|
95
|
+
endpoint=address,
|
|
96
|
+
endpoint_id=parts[5],
|
|
97
|
+
)
|
|
98
|
+
if (
|
|
99
|
+
len(parts) == 6
|
|
100
|
+
and parts[2] == "branches"
|
|
101
|
+
and parts[3]
|
|
102
|
+
and parts[4] == "databases"
|
|
103
|
+
and parts[5]
|
|
104
|
+
):
|
|
105
|
+
return ParsedAddress(
|
|
106
|
+
project=project,
|
|
107
|
+
branch=parts[3],
|
|
108
|
+
database_resource_id=parts[5],
|
|
109
|
+
)
|
|
110
|
+
return ParsedAddress()
|
|
111
|
+
|
|
112
|
+
|
|
113
|
+
parseAddress = parse_address
|
|
114
|
+
parseResourcePath = parse_resource_path
|