Skip to content

Commit 1cb4c45

Browse files
committed
postgres: implement query/query_raw using rust-postgres query_typed
1 parent 67e8e11 commit 1cb4c45

File tree

1 file changed

+84
-61
lines changed
  • quaint/src/connector/postgres/native

1 file changed

+84
-61
lines changed

quaint/src/connector/postgres/native/mod.rs

+84-61
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@ use futures::{future::FutureExt, lock::Mutex};
2727
use lru_cache::LruCache;
2828
use native_tls::{Certificate, Identity, TlsConnector};
2929
use postgres_native_tls::MakeTlsConnector;
30-
use postgres_types::{Kind as PostgresKind, Type as PostgresType};
30+
use postgres_types::{Kind as PostgresKind, Type as PostgresType, ToSql};
3131
use std::hash::{DefaultHasher, Hash, Hasher};
3232
use std::{
3333
fmt::{Debug, Display},
@@ -540,29 +540,37 @@ impl Queryable for PostgreSql {
540540
sql,
541541
params,
542542
move || async move {
543-
let stmt = self.fetch_cached(sql, &[]).await?;
544-
545-
if stmt.params().len() != params.len() {
546-
let kind = ErrorKind::IncorrectNumberOfParameters {
547-
expected: stmt.params().len(),
548-
actual: params.len(),
549-
};
550-
551-
return Err(Error::builder(kind).build());
552-
}
543+
let converted_params = conversion::conv_params(params);
544+
let param_types = conversion::params_to_types(params);
545+
let params_with_types: Vec<(&(dyn ToSql + Sync), PostgresType)> = converted_params
546+
.iter()
547+
.zip(param_types)
548+
.map(|(value, ty)| (*value as &(dyn ToSql + Sync), ty))
549+
.collect();
553550

551+
// Execute the query using `query_typed`
554552
let rows = self
555-
.perform_io(self.client.0.query(&stmt, conversion::conv_params(params).as_slice()))
553+
.perform_io(self.client.0.query_typed(sql, params_with_types.as_slice()))
556554
.await?;
557555

558-
let col_types = stmt
559-
.columns()
560-
.iter()
561-
.map(|c| PGColumnType::from_pg_type(c.type_()))
562-
.map(ColumnType::from)
563-
.collect::<Vec<_>>();
564-
let mut result = ResultSet::new(stmt.to_column_names(), col_types, Vec::new());
556+
// Extract column information from the first row, if available
557+
let (col_types, column_names) = if let Some(row) = rows.first() {
558+
let columns = row.columns();
559+
let col_types = columns
560+
.iter()
561+
.map(|c| PGColumnType::from_pg_type(c.type_()))
562+
.map(ColumnType::from)
563+
.collect::<Vec<_>>();
564+
let column_names = columns.iter().map(|c| c.name().to_string()).collect();
565+
566+
(col_types, column_names)
567+
} else {
568+
(Vec::new(), Vec::new())
569+
};
565570

571+
let mut result = ResultSet::new(column_names, col_types, Vec::new());
572+
573+
// Process each row in the result set
566574
for row in rows {
567575
result.rows.push(row.get_result_row()?);
568576
}
@@ -582,28 +590,35 @@ impl Queryable for PostgreSql {
582590
sql,
583591
params,
584592
move || async move {
585-
let stmt = self.fetch_cached(sql, params).await?;
586-
587-
if stmt.params().len() != params.len() {
588-
let kind = ErrorKind::IncorrectNumberOfParameters {
589-
expected: stmt.params().len(),
590-
actual: params.len(),
591-
};
592-
593-
return Err(Error::builder(kind).build());
594-
}
595-
596-
let col_types = stmt
597-
.columns()
593+
let converted_params = conversion::conv_params(params);
594+
let param_types = conversion::params_to_types(params);
595+
let params_with_types: Vec<(&(dyn ToSql + Sync), PostgresType)> = converted_params
598596
.iter()
599-
.map(|c| PGColumnType::from_pg_type(c.type_()))
600-
.map(ColumnType::from)
601-
.collect::<Vec<_>>();
597+
.zip(param_types)
598+
.map(|(value, ty)| (*value as &(dyn ToSql + Sync), ty))
599+
.collect();
600+
601+
// Execute the query using `query_typed`
602602
let rows = self
603-
.perform_io(self.client.0.query(&stmt, conversion::conv_params(params).as_slice()))
603+
.perform_io(self.client.0.query_typed(sql, params_with_types.as_slice()))
604604
.await?;
605605

606-
let mut result = ResultSet::new(stmt.to_column_names(), col_types, Vec::new());
606+
// Extract column information from the first row, if available
607+
let (col_types, column_names) = if let Some(row) = rows.first() {
608+
let columns = row.columns();
609+
let col_types = columns
610+
.iter()
611+
.map(|c| PGColumnType::from_pg_type(c.type_()))
612+
.map(ColumnType::from)
613+
.collect::<Vec<_>>();
614+
let column_names = columns.iter().map(|c| c.name().to_string()).collect();
615+
616+
(col_types, column_names)
617+
} else {
618+
(Vec::new(), Vec::new())
619+
};
620+
621+
let mut result = ResultSet::new(column_names, col_types, Vec::new());
607622

608623
for row in rows {
609624
result.rows.push(row.get_result_row()?);
@@ -705,20 +720,24 @@ impl Queryable for PostgreSql {
705720
sql,
706721
params,
707722
move || async move {
708-
let stmt = self.fetch_cached(sql, &[]).await?;
709-
710-
if stmt.params().len() != params.len() {
711-
let kind = ErrorKind::IncorrectNumberOfParameters {
712-
expected: stmt.params().len(),
713-
actual: params.len(),
714-
};
715-
716-
return Err(Error::builder(kind).build());
717-
}
723+
let converted_params = conversion::conv_params(params);
724+
let param_types = conversion::params_to_types(params);
725+
let params_with_types: Vec<(&(dyn ToSql + Sync), PostgresType)> = converted_params
726+
.iter()
727+
.zip(param_types)
728+
.map(|(value, ty)| (*value as &(dyn ToSql + Sync), ty))
729+
.collect();
718730

719731
let changes = self
720-
.perform_io(self.client.0.execute(&stmt, conversion::conv_params(params).as_slice()))
721-
.await?;
732+
.perform_io(self.client.0.query_typed_raw::<&(dyn ToSql + Sync), _>(
733+
sql,
734+
params_with_types.as_slice().iter()
735+
.map(|(v, t)| (*v, t.clone()))
736+
.collect::<Vec<_>>()
737+
))
738+
.await?
739+
.rows_affected()
740+
.unwrap_or(0);
722741

723742
Ok(changes)
724743
},
@@ -735,20 +754,24 @@ impl Queryable for PostgreSql {
735754
sql,
736755
params,
737756
move || async move {
738-
let stmt = self.fetch_cached(sql, params).await?;
739-
740-
if stmt.params().len() != params.len() {
741-
let kind = ErrorKind::IncorrectNumberOfParameters {
742-
expected: stmt.params().len(),
743-
actual: params.len(),
744-
};
745-
746-
return Err(Error::builder(kind).build());
747-
}
757+
let converted_params = conversion::conv_params(params);
758+
let param_types = conversion::params_to_types(params);
759+
let params_with_types: Vec<(&(dyn ToSql + Sync), PostgresType)> = converted_params
760+
.iter()
761+
.zip(param_types)
762+
.map(|(value, ty)| (*value as &(dyn ToSql + Sync), ty))
763+
.collect();
748764

749765
let changes = self
750-
.perform_io(self.client.0.execute(&stmt, conversion::conv_params(params).as_slice()))
751-
.await?;
766+
.perform_io(self.client.0.query_typed_raw::<&(dyn ToSql + Sync), _>(
767+
sql,
768+
params_with_types.as_slice().iter()
769+
.map(|(v, t)| (*v, t.clone()))
770+
.collect::<Vec<_>>()
771+
))
772+
.await?
773+
.rows_affected()
774+
.unwrap_or(0);
752775

753776
Ok(changes)
754777
},

0 commit comments

Comments
 (0)