Skip to content

Commit 602f00e

Browse files
chandr-andrchandr-andr
committed
Full value converter refactor
1 parent af51f43 commit 602f00e

18 files changed

+301
-144
lines changed

‎python/tests/test_value_converter.py

Lines changed: 29 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -139,13 +139,18 @@ async def test_as_class(
139139
),
140140
("BOOL", True, True),
141141
("INT2", SmallInt(12), 12),
142+
("INT2", 12, 12),
142143
("INT4", Integer(121231231), 121231231),
144+
("INT4", 121231231, 121231231),
143145
("INT8", BigInt(99999999999999999), 99999999999999999),
146+
("INT8", 99999999999999999, 99999999999999999),
144147
("MONEY", Money(99999999999999999), 99999999999999999),
148+
("MONEY", 99999999999999999, 99999999999999999),
145149
("NUMERIC(5, 2)", Decimal("120.12"), Decimal("120.12")),
146-
("FLOAT8", 32.12329864501953, 32.12329864501953),
147150
("FLOAT4", Float32(32.12329864501953), 32.12329864501953),
151+
("FLOAT4", 32.12329864501953, 32.12329864501953),
148152
("FLOAT8", Float64(32.12329864501953), 32.12329864501953),
153+
("FLOAT8", 32.12329864501953, 32.12329864501953),
149154
("DATE", now_datetime.date(), now_datetime.date()),
150155
("TIME", now_datetime.time(), now_datetime.time()),
151156
("TIMESTAMP", now_datetime, now_datetime),
@@ -426,6 +431,29 @@ async def test_as_class(
426431
[[{"array": "json"}], [{"one more": "test"}]],
427432
],
428433
),
434+
(
435+
"JSON ARRAY",
436+
[
437+
{
438+
"test": ["something", 123, "here"],
439+
"nested": ["JSON"],
440+
},
441+
{
442+
"test": ["something", 123, "here"],
443+
"nested": ["JSON"],
444+
},
445+
],
446+
[
447+
{
448+
"test": ["something", 123, "here"],
449+
"nested": ["JSON"],
450+
},
451+
{
452+
"test": ["something", 123, "here"],
453+
"nested": ["JSON"],
454+
},
455+
],
456+
),
429457
(
430458
"JSON ARRAY",
431459
[

‎src/driver/connection.rs

Lines changed: 18 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ pub struct Connection {
2525
db_client: Option<Arc<PsqlpyConnection>>,
2626
db_pool: Option<Pool>,
2727
pg_config: Arc<Config>,
28+
prepare: bool,
2829
}
2930

3031
impl Connection {
@@ -33,11 +34,13 @@ impl Connection {
3334
db_client: Option<Arc<PsqlpyConnection>>,
3435
db_pool: Option<Pool>,
3536
pg_config: Arc<Config>,
37+
prepare: bool,
3638
) -> Self {
3739
Connection {
3840
db_client,
3941
db_pool,
4042
pg_config,
43+
prepare,
4144
}
4245
}
4346

@@ -54,7 +57,7 @@ impl Connection {
5457

5558
impl Default for Connection {
5659
fn default() -> Self {
57-
Connection::new(None, None, Arc::new(Config::default()))
60+
Connection::new(None, None, Arc::new(Config::default()), true)
5861
}
5962
}
6063

@@ -138,11 +141,16 @@ impl Connection {
138141
}
139142

140143
async fn __aenter__<'a>(self_: Py<Self>) -> PSQLPyResult<Py<Self>> {
141-
let (db_client, db_pool) = pyo3::Python::with_gil(|gil| {
144+
let (db_client, db_pool, prepare) = pyo3::Python::with_gil(|gil| {
142145
let self_ = self_.borrow(gil);
143-
(self_.db_client.clone(), self_.db_pool.clone())
146+
(
147+
self_.db_client.clone(),
148+
self_.db_pool.clone(),
149+
self_.prepare,
150+
)
144151
});
145152

153+
let db_pool_2 = db_pool.clone();
146154
if db_client.is_some() {
147155
return Ok(self_);
148156
}
@@ -155,7 +163,11 @@ impl Connection {
155163
.await??;
156164
pyo3::Python::with_gil(|gil| {
157165
let mut self_ = self_.borrow_mut(gil);
158-
self_.db_client = Some(Arc::new(PsqlpyConnection::PoolConn(db_connection)));
166+
self_.db_client = Some(Arc::new(PsqlpyConnection::PoolConn(
167+
db_connection,
168+
db_pool_2.unwrap(),
169+
prepare,
170+
)));
159171
});
160172
return Ok(self_);
161173
}
@@ -209,7 +221,8 @@ impl Connection {
209221
let db_client = pyo3::Python::with_gil(|gil| self_.borrow(gil).db_client.clone());
210222

211223
if let Some(db_client) = db_client {
212-
return db_client.execute(querystring, parameters, prepared).await;
224+
let res = db_client.execute(querystring, parameters, prepared).await;
225+
return res;
213226
}
214227

215228
Err(RustPSQLDriverError::ConnectionClosedError)

‎src/driver/connection_pool.rs

Lines changed: 32 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
use crate::runtime::tokio_runtime;
22
use deadpool_postgres::{Manager, ManagerConfig, Pool, RecyclingMethod};
3+
use postgres_types::Type;
34
use pyo3::{pyclass, pyfunction, pymethods, Py, PyAny};
45
use std::sync::Arc;
56
use tokio_postgres::Config;
@@ -46,6 +47,7 @@ use super::{
4647
ca_file=None,
4748
max_db_pool_size=None,
4849
conn_recycling_method=None,
50+
prepare=None,
4951
))]
5052
#[allow(clippy::too_many_arguments)]
5153
pub fn connect(
@@ -75,6 +77,7 @@ pub fn connect(
7577
ca_file: Option<String>,
7678
max_db_pool_size: Option<usize>,
7779
conn_recycling_method: Option<ConnRecyclingMethod>,
80+
prepare: Option<bool>,
7881
) -> PSQLPyResult<ConnectionPool> {
7982
if let Some(max_db_pool_size) = max_db_pool_size {
8083
if max_db_pool_size < 2 {
@@ -139,6 +142,7 @@ pub fn connect(
139142
pg_config: Arc::new(pg_config),
140143
ca_file: ca_file,
141144
ssl_mode: ssl_mode,
145+
prepare: prepare.unwrap_or(true),
142146
})
143147
}
144148

@@ -207,6 +211,7 @@ pub struct ConnectionPool {
207211
pg_config: Arc<Config>,
208212
ca_file: Option<String>,
209213
ssl_mode: Option<SslMode>,
214+
prepare: bool,
210215
}
211216

212217
impl ConnectionPool {
@@ -216,14 +221,20 @@ impl ConnectionPool {
216221
pg_config: Config,
217222
ca_file: Option<String>,
218223
ssl_mode: Option<SslMode>,
224+
prepare: Option<bool>,
219225
) -> Self {
220226
ConnectionPool {
221227
pool: pool,
222228
pg_config: Arc::new(pg_config),
223229
ca_file: ca_file,
224230
ssl_mode: ssl_mode,
231+
prepare: prepare.unwrap_or(true),
225232
}
226233
}
234+
235+
pub fn remove_prepared_stmt(&mut self, query: &str, types: &[Type]) {
236+
self.pool.manager().statement_caches.remove(query, types);
237+
}
227238
}
228239

229240
#[pymethods]
@@ -260,6 +271,7 @@ impl ConnectionPool {
260271
conn_recycling_method=None,
261272
ssl_mode=None,
262273
ca_file=None,
274+
prepare=None,
263275
))]
264276
#[allow(clippy::too_many_arguments)]
265277
pub fn new(
@@ -289,6 +301,7 @@ impl ConnectionPool {
289301
conn_recycling_method: Option<ConnRecyclingMethod>,
290302
ssl_mode: Option<SslMode>,
291303
ca_file: Option<String>,
304+
prepare: Option<bool>,
292305
) -> PSQLPyResult<Self> {
293306
connect(
294307
dsn,
@@ -317,6 +330,7 @@ impl ConnectionPool {
317330
ca_file,
318331
max_db_pool_size,
319332
conn_recycling_method,
333+
prepare,
320334
)
321335
}
322336

@@ -360,43 +374,55 @@ impl ConnectionPool {
360374

361375
#[must_use]
362376
pub fn acquire(&self) -> Connection {
363-
Connection::new(None, Some(self.pool.clone()), self.pg_config.clone())
377+
Connection::new(
378+
None,
379+
Some(self.pool.clone()),
380+
self.pg_config.clone(),
381+
self.prepare,
382+
)
364383
}
365384

366385
#[must_use]
367386
#[allow(clippy::needless_pass_by_value)]
368387
pub fn listener(self_: pyo3::Py<Self>) -> Listener {
369-
let (pg_config, ca_file, ssl_mode) = pyo3::Python::with_gil(|gil| {
388+
let (pg_config, ca_file, ssl_mode, prepare) = pyo3::Python::with_gil(|gil| {
370389
let b_gil = self_.borrow(gil);
371390
(
372391
b_gil.pg_config.clone(),
373392
b_gil.ca_file.clone(),
374393
b_gil.ssl_mode,
394+
b_gil.prepare,
375395
)
376396
});
377397

378-
Listener::new(pg_config, ca_file, ssl_mode)
398+
Listener::new(pg_config, ca_file, ssl_mode, prepare)
379399
}
380400

381401
/// Return new single connection.
382402
///
383403
/// # Errors
384404
/// May return Err Result if cannot get new connection from the pool.
385405
pub async fn connection(self_: pyo3::Py<Self>) -> PSQLPyResult<Connection> {
386-
let (db_pool, pg_config) = pyo3::Python::with_gil(|gil| {
406+
let (db_pool, pg_config, prepare) = pyo3::Python::with_gil(|gil| {
387407
let slf = self_.borrow(gil);
388-
(slf.pool.clone(), slf.pg_config.clone())
408+
(slf.pool.clone(), slf.pg_config.clone(), slf.prepare)
389409
});
410+
let db_pool_2 = db_pool.clone();
390411
let db_connection = tokio_runtime()
391412
.spawn(async move {
392413
Ok::<deadpool_postgres::Object, RustPSQLDriverError>(db_pool.get().await?)
393414
})
394415
.await??;
395416

396417
Ok(Connection::new(
397-
Some(Arc::new(PsqlpyConnection::PoolConn(db_connection))),
418+
Some(Arc::new(PsqlpyConnection::PoolConn(
419+
db_connection,
420+
db_pool_2.clone(),
421+
prepare,
422+
))),
398423
None,
399424
pg_config,
425+
prepare,
400426
))
401427
}
402428

‎src/driver/connection_pool_builder.rs

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ pub struct ConnectionPoolBuilder {
1818
conn_recycling_method: Option<RecyclingMethod>,
1919
ca_file: Option<String>,
2020
ssl_mode: Option<common_options::SslMode>,
21+
prepare: Option<bool>,
2122
}
2223

2324
#[pymethods]
@@ -31,6 +32,7 @@ impl ConnectionPoolBuilder {
3132
conn_recycling_method: None,
3233
ca_file: None,
3334
ssl_mode: None,
35+
prepare: None,
3436
}
3537
}
3638

@@ -68,6 +70,7 @@ impl ConnectionPoolBuilder {
6870
self.config.clone(),
6971
self.ca_file.clone(),
7072
self.ssl_mode,
73+
self.prepare,
7174
))
7275
}
7376

@@ -80,6 +83,15 @@ impl ConnectionPoolBuilder {
8083
self_
8184
}
8285

86+
/// Set ca_file for ssl_mode in PostgreSQL.
87+
fn prepare(self_: Py<Self>, prepare: bool) -> Py<Self> {
88+
Python::with_gil(|gil| {
89+
let mut self_ = self_.borrow_mut(gil);
90+
self_.prepare = Some(prepare);
91+
});
92+
self_
93+
}
94+
8395
/// Set size to the connection pool.
8496
///
8597
/// # Error

0 commit comments

Comments
 (0)