sqlx_core_guts/net/
socket.rs1#![allow(dead_code)]
2
3use std::io;
4use std::path::Path;
5use std::pin::Pin;
6use std::task::{Context, Poll};
7
8use sqlx_rt::{AsyncRead, AsyncWrite, TcpStream};
9
10#[derive(Debug)]
11pub enum Socket {
12 Tcp(TcpStream),
13
14 #[cfg(unix)]
15 Unix(sqlx_rt::UnixStream),
16}
17
18impl Socket {
19 pub async fn connect_tcp(host: &str, port: u16) -> io::Result<Self> {
20 TcpStream::connect((host.trim_matches(|c| c == '[' || c == ']'), port))
22 .await
23 .map(Socket::Tcp)
24 }
25
26 #[cfg(unix)]
27 pub async fn connect_uds(path: impl AsRef<Path>) -> io::Result<Self> {
28 sqlx_rt::UnixStream::connect(path.as_ref())
29 .await
30 .map(Socket::Unix)
31 }
32
33 #[cfg(not(unix))]
34 pub async fn connect_uds(_: impl AsRef<Path>) -> io::Result<Self> {
35 Err(io::Error::new(
36 io::ErrorKind::Other,
37 "Unix domain sockets are not supported outside Unix platforms.",
38 ))
39 }
40
41 pub async fn shutdown(&mut self) -> io::Result<()> {
42 #[cfg(feature = "_rt-async-std")]
43 {
44 use std::net::Shutdown;
45
46 match self {
47 Socket::Tcp(s) => s.shutdown(Shutdown::Both),
48
49 #[cfg(unix)]
50 Socket::Unix(s) => s.shutdown(Shutdown::Both),
51 }
52 }
53
54 #[cfg(any(feature = "_rt-actix", feature = "_rt-tokio"))]
55 {
56 use sqlx_rt::AsyncWriteExt;
57
58 match self {
59 Socket::Tcp(s) => s.shutdown().await,
60
61 #[cfg(unix)]
62 Socket::Unix(s) => s.shutdown().await,
63 }
64 }
65 }
66}
67
68impl AsyncRead for Socket {
69 fn poll_read(
70 mut self: Pin<&mut Self>,
71 cx: &mut Context<'_>,
72 buf: &mut super::PollReadBuf<'_>,
73 ) -> Poll<io::Result<super::PollReadOut>> {
74 match &mut *self {
75 Socket::Tcp(s) => Pin::new(s).poll_read(cx, buf),
76
77 #[cfg(unix)]
78 Socket::Unix(s) => Pin::new(s).poll_read(cx, buf),
79 }
80 }
81}
82
83impl AsyncWrite for Socket {
84 fn poll_write(
85 mut self: Pin<&mut Self>,
86 cx: &mut Context<'_>,
87 buf: &[u8],
88 ) -> Poll<io::Result<usize>> {
89 match &mut *self {
90 Socket::Tcp(s) => Pin::new(s).poll_write(cx, buf),
91
92 #[cfg(unix)]
93 Socket::Unix(s) => Pin::new(s).poll_write(cx, buf),
94 }
95 }
96
97 fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
98 match &mut *self {
99 Socket::Tcp(s) => Pin::new(s).poll_flush(cx),
100
101 #[cfg(unix)]
102 Socket::Unix(s) => Pin::new(s).poll_flush(cx),
103 }
104 }
105
106 #[cfg(any(feature = "_rt-actix", feature = "_rt-tokio"))]
107 fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
108 match &mut *self {
109 Socket::Tcp(s) => Pin::new(s).poll_shutdown(cx),
110
111 #[cfg(unix)]
112 Socket::Unix(s) => Pin::new(s).poll_shutdown(cx),
113 }
114 }
115
116 #[cfg(feature = "_rt-async-std")]
117 fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
118 match &mut *self {
119 Socket::Tcp(s) => Pin::new(s).poll_close(cx),
120
121 #[cfg(unix)]
122 Socket::Unix(s) => Pin::new(s).poll_close(cx),
123 }
124 }
125}