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.
@@ -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