servers/postgres/
server.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::any::Any;
16use std::future::Future;
17use std::net::SocketAddr;
18use std::sync::Arc;
19
20use ::auth::UserProviderRef;
21use async_trait::async_trait;
22use catalog::process_manager::ProcessManagerRef;
23use common_runtime::Runtime;
24use common_runtime::runtime::RuntimeTrait;
25use common_telemetry::{debug, warn};
26use futures::StreamExt;
27use pgwire::tokio::process_socket;
28use tokio_rustls::TlsAcceptor;
29
30use crate::error::Result;
31use crate::postgres::{MakePostgresServerHandler, MakePostgresServerHandlerBuilder};
32use crate::query_handler::CopyInHandlerRef;
33use crate::query_handler::sql::ServerSqlQueryHandlerRef;
34use crate::server::{AbortableStream, BaseTcpServer, Server};
35use crate::tls::ReloadableTlsServerConfig;
36
37pub struct PostgresServer {
38    base_server: BaseTcpServer,
39    make_handler: Arc<MakePostgresServerHandler>,
40    tls_server_config: Arc<ReloadableTlsServerConfig>,
41    keep_alive_secs: u64,
42    batching_enabled: bool,
43    bind_addr: Option<SocketAddr>,
44    process_manager: Option<ProcessManagerRef>,
45}
46
47impl PostgresServer {
48    /// Creates a new Postgres server with provided query_handler and async runtime.
49    ///
50    /// `query_handler` and `copy_in_handler` are expected to be handles to
51    /// the same frontend instance (the SQL and COPY FROM STDIN facets of
52    /// one object); they are separate parameters only because they are
53    /// separate traits.
54    #[allow(clippy::too_many_arguments)]
55    pub fn new(
56        query_handler: ServerSqlQueryHandlerRef,
57        copy_in_handler: CopyInHandlerRef,
58        force_tls: bool,
59        tls_server_config: Arc<ReloadableTlsServerConfig>,
60        keep_alive_secs: u64,
61        io_runtime: Runtime,
62        user_provider: Option<UserProviderRef>,
63        process_manager: Option<ProcessManagerRef>,
64    ) -> PostgresServer {
65        let make_handler = Arc::new(
66            MakePostgresServerHandlerBuilder::default()
67                .query_handler(query_handler.clone())
68                .copy_in_handler(copy_in_handler)
69                .user_provider(user_provider.clone())
70                .force_tls(force_tls)
71                .build()
72                .unwrap(),
73        );
74        PostgresServer {
75            base_server: BaseTcpServer::create_server("Postgres", io_runtime),
76            make_handler,
77            tls_server_config,
78            keep_alive_secs,
79            batching_enabled: false,
80            bind_addr: None,
81            process_manager,
82        }
83    }
84
85    /// Enables ordinary-table batching for connections accepted by this server.
86    pub fn with_batching_enabled(mut self, enabled: bool) -> Self {
87        self.batching_enabled = enabled;
88        self
89    }
90
91    fn accept(
92        &self,
93        io_runtime: Runtime,
94        accepting_stream: AbortableStream,
95    ) -> impl Future<Output = ()> + use<> {
96        let handler_maker = self.make_handler.clone();
97        let batching_enabled = self.batching_enabled;
98        let tls_server_config = self.tls_server_config.clone();
99        let process_manager = self.process_manager.clone();
100        accepting_stream.for_each(move |tcp_stream| {
101            let io_runtime = io_runtime.clone();
102            let tls_acceptor = tls_server_config.get_config().map(TlsAcceptor::from);
103            let handler_maker = handler_maker.clone();
104            let process_id = process_manager.as_ref().map(|p| p.next_id()).unwrap_or(0);
105
106            async move {
107                match tcp_stream {
108                    Err(error) => debug!("Broken pipe: {}", error), // IoError doesn't impl ErrorExt.
109                    Ok(io_stream) => {
110                        let addr = match io_stream.peer_addr() {
111                            Ok(addr) => {
112                                debug!("PostgreSQL client coming from {}", addr);
113                                Some(addr)
114                            }
115                            Err(e) => {
116                                warn!(e; "Failed to get PostgreSQL client addr");
117                                None
118                            }
119                        };
120
121                        let _handle = io_runtime.spawn(async move {
122                            crate::metrics::METRIC_POSTGRES_CONNECTIONS.inc();
123                            let pg_handler =
124                                Arc::new(handler_maker.make(addr, process_id, batching_enabled));
125                            let r =
126                                process_socket(io_stream, tls_acceptor.clone(), pg_handler).await;
127                            crate::metrics::METRIC_POSTGRES_CONNECTIONS.dec();
128                            r
129                        });
130                    }
131                };
132            }
133        })
134    }
135}
136
137pub const POSTGRES_SERVER: &str = "POSTGRES_SERVER";
138
139#[async_trait]
140impl Server for PostgresServer {
141    async fn shutdown(&self) -> Result<()> {
142        self.base_server.shutdown().await
143    }
144
145    async fn start(&mut self, listening: SocketAddr) -> Result<()> {
146        let (stream, addr) = self
147            .base_server
148            .bind(listening, self.keep_alive_secs)
149            .await?;
150
151        let io_runtime = self.base_server.io_runtime();
152        let join_handle = common_runtime::spawn_global(self.accept(io_runtime, stream));
153
154        self.base_server.start_with(join_handle).await?;
155
156        self.bind_addr = Some(addr);
157        Ok(())
158    }
159
160    fn name(&self) -> &str {
161        POSTGRES_SERVER
162    }
163
164    fn bind_addr(&self) -> Option<SocketAddr> {
165        self.bind_addr
166    }
167
168    fn as_any(&self) -> &dyn Any {
169        self
170    }
171}
servers/postgres/server.rs

servers/postgres/
server.rs
