use std::collections::BTreeMap;
use std::ops::{Deref, DerefMut};
use std::str::FromStr;
use futures_channel::mpsc::UnboundedSender;
use futures_util::SinkExt;
use log::Level;
use sqlx_core::bytes::{Buf, Bytes};
use crate::connection::tls::MaybeUpgradeTls;
use crate::error::Error;
use crate::io::{Decode, Encode};
use crate::message::{Message, MessageFormat, Notice, Notification, ParameterStatus};
use crate::net::{self, BufferedSocket, Socket};
use crate::{PgConnectOptions, PgDatabaseError, PgSeverity};
pub struct PgStream {
inner: BufferedSocket<Box<dyn Socket>>,
pub(crate) notifications: Option<UnboundedSender<Notification>>,
pub(crate) parameter_statuses: BTreeMap<String, String>,
pub(crate) server_version_num: Option<u32>,
}
impl PgStream {
pub(super) async fn connect(options: &PgConnectOptions) -> Result<Self, Error> {
let socket_future = match options.fetch_socket() {
Some(ref path) => net::connect_uds(path, MaybeUpgradeTls(options)).await?,
None => net::connect_tcp(&options.host, options.port, MaybeUpgradeTls(options)).await?,
};
let socket = socket_future.await?;
Ok(Self {
inner: BufferedSocket::new(socket),
notifications: None,
parameter_statuses: BTreeMap::default(),
server_version_num: None,
})
}
pub(crate) async fn send<'en, T>(&mut self, message: T) -> Result<(), Error>
where
T: Encode<'en>,
{
self.write(message);
self.flush().await?;
Ok(())
}
pub(crate) async fn recv_expect<'de, T: Decode<'de>>(
&mut self,
format: MessageFormat,
) -> Result<T, Error> {
let message = self.recv().await?;
if message.format != format {
return Err(err_protocol!(
"expecting {:?} but received {:?}",
format,
message.format
));
}
message.decode()
}
pub(crate) async fn recv_unchecked(&mut self) -> Result<Message, Error> {
let mut header: Bytes = self.inner.read(5).await?;
let format = MessageFormat::try_from_u8(header.get_u8())?;
let size = (header.get_u32() - 4) as usize;
let contents = self.inner.read(size).await?;
Ok(Message { format, contents })
}
pub(crate) async fn recv(&mut self) -> Result<Message, Error> {
loop {
let message = self.recv_unchecked().await?;
match message.format {
MessageFormat::ErrorResponse => {
return Err(PgDatabaseError(message.decode()?).into());
}
MessageFormat::NotificationResponse => {
if let Some(buffer) = &mut self.notifications {
let notification: Notification = message.decode()?;
let _ = buffer.send(notification).await;
continue;
}
}
MessageFormat::ParameterStatus => {
let ParameterStatus { name, value } = message.decode()?;
match name.as_str() {
"server_version" => {
self.server_version_num = parse_server_version(&value);
}
_ => {
self.parameter_statuses.insert(name, value);
}
}
continue;
}
MessageFormat::NoticeResponse => {
let notice: Notice = message.decode()?;
let (log_level, tracing_level) = match notice.severity() {
PgSeverity::Fatal | PgSeverity::Panic | PgSeverity::Error => {
(Level::Error, tracing::Level::ERROR)
}
PgSeverity::Warning => (Level::Warn, tracing::Level::WARN),
PgSeverity::Notice => (Level::Info, tracing::Level::INFO),
PgSeverity::Debug => (Level::Debug, tracing::Level::DEBUG),
PgSeverity::Info | PgSeverity::Log => (Level::Trace, tracing::Level::TRACE),
};
let log_is_enabled = log::log_enabled!(
target: "sqlx::postgres::notice",
log_level
) || sqlx_core::private_tracing_dynamic_enabled!(
target: "sqlx::postgres::notice",
tracing_level
);
if log_is_enabled {
let message = format!("{}", notice.message());
sqlx_core::private_tracing_dynamic_event!(
target: "sqlx::postgres::notice",
tracing_level,
message
);
}
continue;
}
_ => {}
}
return Ok(message);
}
}
}
impl Deref for PgStream {
type Target = BufferedSocket<Box<dyn Socket>>;
#[inline]
fn deref(&self) -> &Self::Target {
&self.inner
}
}
impl DerefMut for PgStream {
#[inline]
fn deref_mut(&mut self) -> &mut Self::Target {
&mut self.inner
}
}
fn parse_server_version(s: &str) -> Option<u32> {
let mut parts = Vec::<u32>::with_capacity(3);
let mut from = 0;
let mut chs = s.char_indices().peekable();
while let Some((i, ch)) = chs.next() {
match ch {
'.' => {
if let Ok(num) = u32::from_str(&s[from..i]) {
parts.push(num);
from = i + 1;
} else {
break;
}
}
_ if ch.is_digit(10) => {
if chs.peek().is_none() {
if let Ok(num) = u32::from_str(&s[from..]) {
parts.push(num);
}
break;
}
}
_ => {
if let Ok(num) = u32::from_str(&s[from..i]) {
parts.push(num);
}
break;
}
};
}
let version_num = match parts.as_slice() {
[major, minor, rev] => (100 * major + minor) * 100 + rev,
[major, minor] if *major >= 10 => 100 * 100 * major + minor,
[major, minor] => (100 * major + minor) * 100,
[major] => 100 * 100 * major,
_ => return None,
};
Some(version_num)
}
#[cfg(test)]
mod tests {
use super::parse_server_version;
#[test]
fn test_parse_server_version_num() {
assert_eq!(parse_server_version("9.6.1"), Some(90601));
assert_eq!(parse_server_version("10.1"), Some(100001));
assert_eq!(parse_server_version("9.6devel"), Some(90600));
assert_eq!(parse_server_version("10devel"), Some(100000));
assert_eq!(parse_server_version("13devel87"), Some(130000));
assert_eq!(parse_server_version("unknown"), None);
}
}