Skip to content

Commit d4f5eae

Browse files
authored
Merge pull request #6 from qaspen-python/feature/add_dsn_support
Added dsn support
2 parents 31e5415 + d9228e7 commit d4f5eae

File tree

4 files changed

+65
-15
lines changed

4 files changed

+65
-15
lines changed

README.md

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,32 @@ async def main() -> None:
5454
```
5555
Please take into account that each new execute gets new connection from connection pool.
5656

57+
### DSN support
58+
You can separate specify `host`, `port`, `username`, etc or specify everything in one `DSN`.
59+
**Please note that if you specify DSN any other argument doesn't take into account.**
60+
```py
61+
from typing import Any
62+
63+
from psqlpy import PSQLPool
64+
65+
66+
db_pool = PSQLPool(
67+
dsn="postgres://postgres:postgres@localhost:5432/postgres",
68+
max_db_pool_size=2,
69+
)
70+
71+
async def main() -> None:
72+
await db_pool.startup()
73+
74+
res: list[dict[str, Any]] = await db_pool.execute(
75+
"SELECT * FROM users",
76+
)
77+
78+
print(res)
79+
# You don't need to close Database Pool by yourself,
80+
# rust does it instead.
81+
```
82+
5783
## Query parameters
5884
You can pass parameters into queries.
5985
Parameters can be passed in any `execute` method as the second parameter, it must be a list.

python/psqlpy/_internal/__init__.pyi

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -439,6 +439,7 @@ class PSQLPool:
439439

440440
def __init__(
441441
self: Self,
442+
dsn: Optional[str] = None,
442443
username: Optional[str] = None,
443444
password: Optional[str] = None,
444445
host: Optional[str] = None,
@@ -459,6 +460,8 @@ class PSQLPool:
459460
- Create new instance of `Transaction`
460461
461462
### Parameters:
463+
- `dsn`: full dsn connection string.
464+
`postgres://postgres:postgres@localhost:5432/postgres?target_session_attrs=read-write`
462465
- `username`: username of the user in postgres
463466
- `password`: password of the user in postgres
464467
- `host`: host of postgres

python/tests/test_connection_pool.py

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,17 @@
33
from psqlpy import Connection, PSQLPool, QueryResult
44

55

6+
@pytest.mark.anyio
7+
async def test_pool_dsn_startup() -> None:
8+
"""Test that connection pool can startup with dsn."""
9+
pg_pool = PSQLPool(
10+
dsn="postgres://postgres:postgres@localhost:5432/psqlpy_test",
11+
)
12+
await pg_pool.startup()
13+
14+
await pg_pool.execute("SELECT 1")
15+
16+
617
@pytest.mark.anyio
718
async def test_pool_execute(
819
psql_pool: PSQLPool,

src/driver/connection_pool.rs

Lines changed: 25 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
use deadpool_postgres::{Manager, ManagerConfig, Pool, RecyclingMethod};
22
use pyo3::{pyclass, pymethods, PyAny, Python};
3-
use std::{sync::Arc, vec};
3+
use std::{str::FromStr, sync::Arc, vec};
44
use tokio_postgres::{types::ToSql, NoTls};
55

66
use crate::{
@@ -16,6 +16,7 @@ use super::connection::Connection;
1616
///
1717
/// It is not exposed to python.
1818
pub struct RustPSQLPool {
19+
dsn: Option<String>,
1920
username: Option<String>,
2021
password: Option<String>,
2122
host: Option<String>,
@@ -29,6 +30,7 @@ impl RustPSQLPool {
2930
/// Create new `RustPSQLPool`.
3031
#[must_use]
3132
pub fn new(
33+
dsn: Option<String>,
3234
username: Option<String>,
3335
password: Option<String>,
3436
host: Option<String>,
@@ -37,6 +39,7 @@ impl RustPSQLPool {
3739
max_db_pool_size: Option<usize>,
3840
) -> Self {
3941
RustPSQLPool {
42+
dsn,
4043
username,
4144
password,
4245
host,
@@ -115,6 +118,7 @@ impl RustPSQLPool {
115118
/// `max_db_pool_size` is less than 2 or it's impossible to build db pool.
116119
pub async fn inner_startup(&self) -> RustPSQLDriverPyResult<()> {
117120
let db_pool_arc = self.db_pool.clone();
121+
let dsn = self.dsn.clone();
118122
let password = self.password.clone();
119123
let username = self.username.clone();
120124
let db_host = self.host.clone();
@@ -137,22 +141,26 @@ impl RustPSQLPool {
137141
}
138142
}
139143

140-
let mut pg_config = tokio_postgres::Config::new();
141-
142-
if let (Some(password), Some(username)) = (password, username) {
143-
pg_config.password(&password);
144-
pg_config.user(&username);
145-
}
146-
if let Some(db_host) = db_host {
147-
pg_config.host(&db_host);
148-
}
144+
let mut pg_config: tokio_postgres::Config;
145+
if let Some(dsn_string) = dsn {
146+
pg_config = tokio_postgres::Config::from_str(&dsn_string)?;
147+
} else {
148+
pg_config = tokio_postgres::Config::new();
149+
if let (Some(password), Some(username)) = (password, username) {
150+
pg_config.password(&password);
151+
pg_config.user(&username);
152+
}
153+
if let Some(db_host) = db_host {
154+
pg_config.host(&db_host);
155+
}
149156

150-
if let Some(db_port) = db_port {
151-
pg_config.port(db_port);
152-
}
157+
if let Some(db_port) = db_port {
158+
pg_config.port(db_port);
159+
}
153160

154-
if let Some(db_name) = db_name {
155-
pg_config.dbname(&db_name);
161+
if let Some(db_name) = db_name {
162+
pg_config.dbname(&db_name);
163+
}
156164
}
157165

158166
let mgr_config = ManagerConfig {
@@ -180,6 +188,7 @@ impl PSQLPool {
180188
#[new]
181189
#[must_use]
182190
pub fn new(
191+
dsn: Option<String>,
183192
username: Option<String>,
184193
password: Option<String>,
185194
host: Option<String>,
@@ -189,6 +198,7 @@ impl PSQLPool {
189198
) -> Self {
190199
PSQLPool {
191200
rust_psql_pool: Arc::new(tokio::sync::RwLock::new(RustPSQLPool {
201+
dsn,
192202
username,
193203
password,
194204
host,

0 commit comments

Comments
 (0)