Line data    Source code

       1              : use std::net::IpAddr;
       2              : 
       3              : use postgres_protocol2::message::backend::Message;
       4              : use tokio::net::TcpStream;
       5              : use tokio::sync::mpsc;
       6              : 
       7              : use crate::client::SocketConfig;
       8              : use crate::codec::BackendMessage;
       9              : use crate::config::Host;
      10              : use crate::connect_raw::connect_raw;
      11              : use crate::connect_socket::connect_socket;
      12              : use crate::tls::{MakeTlsConnect, TlsConnect};
      13              : use crate::{Client, Config, Connection, Error, RawConnection};
      14              : 
      15            0 : pub async fn connect<T>(
      16            0 :     mut tls: T,
      17            0 :     config: &Config,
      18            0 : ) -> Result<(Client, Connection<TcpStream, T::Stream>), Error>
      19            0 : where
      20            0 :     T: MakeTlsConnect<TcpStream>,
      21            0 : {
      22            0 :     let hostname = match &config.host {
      23            0 :         Host::Tcp(host) => host.as_str(),
      24              :     };
      25              : 
      26            0 :     let tls = tls
      27            0 :         .make_tls_connect(hostname)
      28            0 :         .map_err(|e| Error::tls(e.into()))?;
      29              : 
      30            0 :     match connect_once(config.host_addr, &config.host, config.port, tls, config).await {
      31            0 :         Ok((client, connection)) => Ok((client, connection)),
      32            0 :         Err(e) => Err(e),
      33              :     }
      34            0 : }
      35              : 
      36            0 : async fn connect_once<T>(
      37            0 :     host_addr: Option<IpAddr>,
      38            0 :     host: &Host,
      39            0 :     port: u16,
      40            0 :     tls: T,
      41            0 :     config: &Config,
      42            0 : ) -> Result<(Client, Connection<TcpStream, T::Stream>), Error>
      43            0 : where
      44            0 :     T: TlsConnect<TcpStream>,
      45            0 : {
      46            0 :     let socket = connect_socket(host_addr, host, port, config.connect_timeout).await?;
      47              :     let RawConnection {
      48            0 :         stream,
      49            0 :         parameters,
      50            0 :         delayed_notice,
      51            0 :         process_id,
      52            0 :         secret_key,
      53            0 :     } = connect_raw(socket, tls, config).await?;
      54              : 
      55            0 :     let socket_config = SocketConfig {
      56            0 :         host_addr,
      57            0 :         host: host.clone(),
      58            0 :         port,
      59            0 :         connect_timeout: config.connect_timeout,
      60            0 :     };
      61            0 : 
      62            0 :     let (sender, receiver) = mpsc::unbounded_channel();
      63            0 :     let client = Client::new(
      64            0 :         sender,
      65            0 :         socket_config,
      66            0 :         config.ssl_mode,
      67            0 :         process_id,
      68            0 :         secret_key,
      69            0 :     );
      70            0 : 
      71            0 :     // delayed notices are always sent as "Async" messages.
      72            0 :     let delayed = delayed_notice
      73            0 :         .into_iter()
      74            0 :         .map(|m| BackendMessage::Async(Message::NoticeResponse(m)))
      75            0 :         .collect();
      76            0 : 
      77            0 :     let connection = Connection::new(stream, delayed, parameters, receiver);
      78            0 : 
      79            0 :     Ok((client, connection))
      80            0 : }