use crate::describe::Describe;
use crate::error::Error;
use crate::executor::{Execute, Executor};
use crate::logger::QueryLogger;
use crate::postgres::message::{
self, Bind, Close, CommandComplete, DataRow, MessageFormat, ParameterDescription, Parse, Query,
RowDescription,
};
use crate::postgres::statement::PgStatementMetadata;
use crate::postgres::type_info::PgType;
use crate::postgres::{
statement::PgStatement, PgArguments, PgConnection, PgQueryResult, PgRow, PgTypeInfo,
PgValueFormat, Postgres,
};
use either::Either;
use futures_core::future::BoxFuture;
use futures_core::stream::BoxStream;
use futures_core::Stream;
use futures_util::{pin_mut, TryStreamExt};
use std::{borrow::Cow, sync::Arc};
async fn prepare(
conn: &mut PgConnection,
sql: &str,
parameters: &[PgTypeInfo],
metadata: Option<Arc<PgStatementMetadata>>,
) -> Result<(u32, Arc<PgStatementMetadata>), Error> {
let id = conn.next_statement_id;
conn.next_statement_id = conn.next_statement_id.wrapping_add(1);
let mut param_types = Vec::with_capacity(parameters.len());
for ty in parameters {
param_types.push(if let PgType::DeclareWithName(name) = &ty.0 {
conn.fetch_type_id_by_name(name).await?
} else {
ty.0.oid()
});
}
conn.wait_until_ready().await?;
conn.stream.write(Parse {
param_types: &*param_types,
query: sql,
statement: id,
});
if metadata.is_none() {
conn.stream.write(message::Describe::Statement(id));
}
conn.write_sync();
conn.stream.flush().await?;
let _ = conn
.stream
.recv_expect(MessageFormat::ParseComplete)
.await?;
let metadata = if let Some(metadata) = metadata {
conn.recv_ready_for_query().await?;
metadata
} else {
let parameters = recv_desc_params(conn).await?;
let rows = recv_desc_rows(conn).await?;
conn.recv_ready_for_query().await?;
let parameters = conn.handle_parameter_description(parameters).await?;
let (columns, column_names) = conn.handle_row_description(rows, true).await?;
conn.wait_until_ready().await?;
Arc::new(PgStatementMetadata {
parameters,
columns,
column_names,
})
};
Ok((id, metadata))
}
async fn recv_desc_params(conn: &mut PgConnection) -> Result<ParameterDescription, Error> {
conn.stream
.recv_expect(MessageFormat::ParameterDescription)
.await
}
async fn recv_desc_rows(conn: &mut PgConnection) -> Result<Option<RowDescription>, Error> {
let rows: Option<RowDescription> = match conn.stream.recv().await? {
message if message.format == MessageFormat::RowDescription => Some(message.decode()?),
message if message.format == MessageFormat::NoData => None,
message => {
return Err(err_protocol!(
"expecting RowDescription or NoData but received {:?}",
message.format
));
}
};
Ok(rows)
}
impl PgConnection {
pub(super) async fn wait_for_close_complete(&mut self, mut count: usize) -> Result<(), Error> {
while count > 0 {
match self.stream.recv().await? {
message if message.format == MessageFormat::PortalSuspended => {
}
message if message.format == MessageFormat::CloseComplete => {
count -= 1;
}
message => {
return Err(err_protocol!(
"expecting PortalSuspended or CloseComplete but received {:?}",
message.format
));
}
}
}
Ok(())
}
pub(crate) fn write_sync(&mut self) {
self.stream.write(message::Sync);
self.pending_ready_for_query_count += 1;
}
async fn get_or_prepare<'a>(
&mut self,
sql: &str,
parameters: &[PgTypeInfo],
store_to_cache: bool,
metadata: Option<Arc<PgStatementMetadata>>,
) -> Result<(u32, Arc<PgStatementMetadata>), Error> {
if let Some(statement) = self.cache_statement.get_mut(sql) {
return Ok((*statement).clone());
}
let statement = prepare(self, sql, parameters, metadata).await?;
if store_to_cache && self.cache_statement.is_enabled() {
if let Some((id, _)) = self.cache_statement.insert(sql, statement.clone()) {
self.stream.write(Close::Statement(id));
self.write_sync();
self.stream.flush().await?;
self.wait_for_close_complete(1).await?;
self.recv_ready_for_query().await?;
}
}
Ok(statement)
}
async fn run<'e, 'c: 'e, 'q: 'e>(
&'c mut self,
query: &'q str,
arguments: Option<PgArguments>,
limit: u8,
persistent: bool,
metadata_opt: Option<Arc<PgStatementMetadata>>,
) -> Result<impl Stream<Item = Result<Either<PgQueryResult, PgRow>, Error>> + 'e, Error> {
let mut logger = QueryLogger::new(query, self.log_settings.clone());
self.wait_until_ready().await?;
let mut metadata: Arc<PgStatementMetadata>;
let format = if let Some(mut arguments) = arguments {
let (statement, metadata_) = self
.get_or_prepare(query, &arguments.types, persistent, metadata_opt)
.await?;
metadata = metadata_;
arguments.apply_patches(self, &metadata.parameters).await?;
self.stream.write(Bind {
portal: None,
statement,
formats: &[PgValueFormat::Binary],
num_params: arguments.types.len() as i16,
params: &*arguments.buffer,
result_formats: &[PgValueFormat::Binary],
});
self.stream.write(message::Execute {
portal: None,
limit: limit.into(),
});
self.write_sync();
PgValueFormat::Binary
} else {
self.stream.write(Query(query));
self.pending_ready_for_query_count += 1;
metadata = Arc::new(PgStatementMetadata::default());
PgValueFormat::Text
};
self.stream.flush().await?;
Ok(try_stream! {
loop {
let message = self.stream.recv().await?;
match message.format {
MessageFormat::BindComplete
| MessageFormat::ParseComplete
| MessageFormat::ParameterDescription
| MessageFormat::NoData => {
}
MessageFormat::CommandComplete => {
let cc: CommandComplete = message.decode()?;
r#yield!(Either::Left(PgQueryResult {
rows_affected: cc.rows_affected(),
}));
}
MessageFormat::EmptyQueryResponse => {
}
MessageFormat::RowDescription => {
let (columns, column_names) = self
.handle_row_description(Some(message.decode()?), false)
.await?;
metadata = Arc::new(PgStatementMetadata {
column_names,
columns,
parameters: Vec::default(),
});
}
MessageFormat::DataRow => {
logger.increment_rows();
let data: DataRow = message.decode()?;
let row = PgRow {
data,
format,
metadata: Arc::clone(&metadata),
};
r#yield!(Either::Right(row));
}
MessageFormat::ReadyForQuery => {
self.handle_ready_for_query(message)?;
break;
}
_ => {
return Err(err_protocol!(
"execute: unexpected message: {:?}",
message.format
));
}
}
}
Ok(())
})
}
}
impl<'c> Executor<'c> for &'c mut PgConnection {
type Database = Postgres;
fn fetch_many<'e, 'q: 'e, E: 'q>(
self,
mut query: E,
) -> BoxStream<'e, Result<Either<PgQueryResult, PgRow>, Error>>
where
'c: 'e,
E: Execute<'q, Self::Database>,
{
let sql = query.sql();
let metadata = query.statement().map(|s| Arc::clone(&s.metadata));
let arguments = query.take_arguments();
let persistent = query.persistent();
Box::pin(try_stream! {
let s = self.run(sql, arguments, 0, persistent, metadata).await?;
pin_mut!(s);
while let Some(v) = s.try_next().await? {
r#yield!(v);
}
Ok(())
})
}
fn fetch_optional<'e, 'q: 'e, E: 'q>(
self,
mut query: E,
) -> BoxFuture<'e, Result<Option<PgRow>, Error>>
where
'c: 'e,
E: Execute<'q, Self::Database>,
{
let sql = query.sql();
let metadata = query.statement().map(|s| Arc::clone(&s.metadata));
let arguments = query.take_arguments();
let persistent = query.persistent();
Box::pin(async move {
let s = self.run(sql, arguments, 1, persistent, metadata).await?;
pin_mut!(s);
while let Some(s) = s.try_next().await? {
if let Either::Right(r) = s {
return Ok(Some(r));
}
}
Ok(None)
})
}
fn prepare_with<'e, 'q: 'e>(
self,
sql: &'q str,
parameters: &'e [PgTypeInfo],
) -> BoxFuture<'e, Result<PgStatement<'q>, Error>>
where
'c: 'e,
{
Box::pin(async move {
self.wait_until_ready().await?;
let (_, metadata) = self.get_or_prepare(sql, parameters, true, None).await?;
Ok(PgStatement {
sql: Cow::Borrowed(sql),
metadata,
})
})
}
fn describe<'e, 'q: 'e>(
self,
sql: &'q str,
) -> BoxFuture<'e, Result<Describe<Self::Database>, Error>>
where
'c: 'e,
{
Box::pin(async move {
self.wait_until_ready().await?;
let (stmt_id, metadata) = self.get_or_prepare(sql, &[], true, None).await?;
let nullable = self.get_nullable_for_columns(stmt_id, &metadata).await?;
Ok(Describe {
columns: metadata.columns.clone(),
nullable,
parameters: Some(Either::Left(metadata.parameters.clone())),
})
})
}
}