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::InfluxdbLineProtocolHandlerRef;
34use crate::query_handler::sql::ServerSqlQueryHandlerRef;
35use crate::server::{AbortableStream, BaseTcpServer, Server};
36use crate::tls::ReloadableTlsServerConfig;
37
38pub struct PostgresServer {
39    base_server: BaseTcpServer,
40    make_handler: Arc<MakePostgresServerHandler>,
41    tls_server_config: Arc<ReloadableTlsServerConfig>,
42    keep_alive_secs: u64,
43    batching_enabled: bool,
44    influxdb_handler: Option<InfluxdbLineProtocolHandlerRef>,
45    bind_addr: Option<SocketAddr>,
46    process_manager: Option<ProcessManagerRef>,
47}
48
49impl PostgresServer {
50    /// Creates a new Postgres server with provided query_handler and async runtime.
51    ///
52    /// `query_handler` and `copy_in_handler` are expected to be handles to
53    /// the same frontend instance (the SQL and COPY FROM STDIN facets of
54    /// one object); they are separate parameters only because they are
55    /// separate traits.
56    #[allow(clippy::too_many_arguments)]
57    pub fn new(
58        query_handler: ServerSqlQueryHandlerRef,
59        copy_in_handler: CopyInHandlerRef,
60        force_tls: bool,
61        tls_server_config: Arc<ReloadableTlsServerConfig>,
62        keep_alive_secs: u64,
63        io_runtime: Runtime,
64        user_provider: Option<UserProviderRef>,
65        process_manager: Option<ProcessManagerRef>,
66    ) -> PostgresServer {
67        let make_handler = Arc::new(
68            MakePostgresServerHandlerBuilder::default()
69                .query_handler(query_handler.clone())
70                .copy_in_handler(copy_in_handler)
71                .user_provider(user_provider.clone())
72                .force_tls(force_tls)
73                .build()
74                .unwrap(),
75        );
76        PostgresServer {
77            base_server: BaseTcpServer::create_server("Postgres", io_runtime),
78            make_handler,
79            tls_server_config,
80            keep_alive_secs,
81            batching_enabled: false,
82            influxdb_handler: None,
83            bind_addr: None,
84            process_manager,
85        }
86    }
87
88    /// Enables ordinary-table batching for connections accepted by this server.
89    pub fn with_batching_enabled(mut self, enabled: bool) -> Self {
90        self.batching_enabled = enabled;
91        self
92    }
93
94    /// Wires the influxdb line-protocol handler used by table-less
95    /// `COPY FROM STDIN WITH (FORMAT influxdb)` flushes — the same
96    /// handler the HTTP influxdb write endpoints use, so a flushed batch
97    /// runs through the identical pipeline.
98    pub fn with_influxdb_handler(mut self, handler: InfluxdbLineProtocolHandlerRef) -> Self {
99        self.influxdb_handler = Some(handler);
100        self
101    }
102
103    fn accept(
104        &self,
105        io_runtime: Runtime,
106        accepting_stream: AbortableStream,
107    ) -> impl Future<Output = ()> + use<> {
108        // Assemble the per-connection handler directly, including any
109        // `with_*` wiring applied after construction (the influxdb
110        // line-protocol handler). Constructing the struct — instead of
111        // going through the builder — keeps this infallible: no
112        // `.build().unwrap()` on the hot accept path.
113        let make_handler = &self.make_handler;
114        let handler_maker = Arc::new(MakePostgresServerHandler {
115            query_handler: make_handler.query_handler.clone(),
116            copy_in_handler: make_handler.copy_in_handler.clone(),
117            influxdb_handler: self.influxdb_handler.clone(),
118            user_provider: make_handler.user_provider.clone(),
119            param_provider: make_handler.param_provider.clone(),
120            force_tls: make_handler.force_tls,
121        });
122        let batching_enabled = self.batching_enabled;
123        let tls_server_config = self.tls_server_config.clone();
124        let process_manager = self.process_manager.clone();
125        accepting_stream.for_each(move |tcp_stream| {
126            let io_runtime = io_runtime.clone();
127            let tls_acceptor = tls_server_config.get_config().map(TlsAcceptor::from);
128            let handler_maker = handler_maker.clone();
129            let process_id = process_manager.as_ref().map(|p| p.next_id()).unwrap_or(0);
130
131            async move {
132                match tcp_stream {
133                    Err(error) => debug!("Broken pipe: {}", error), // IoError doesn't impl ErrorExt.
134                    Ok(io_stream) => {
135                        let addr = match io_stream.peer_addr() {
136                            Ok(addr) => {
137                                debug!("PostgreSQL client coming from {}", addr);
138                                Some(addr)
139                            }
140                            Err(e) => {
141                                warn!(e; "Failed to get PostgreSQL client addr");
142                                None
143                            }
144                        };
145
146                        let _handle = io_runtime.spawn(async move {
147                            crate::metrics::METRIC_POSTGRES_CONNECTIONS.inc();
148                            let pg_handler =
149                                Arc::new(handler_maker.make(addr, process_id, batching_enabled));
150                            let r =
151                                process_socket(io_stream, tls_acceptor.clone(), pg_handler).await;
152                            crate::metrics::METRIC_POSTGRES_CONNECTIONS.dec();
153                            r
154                        });
155                    }
156                };
157            }
158        })
159    }
160}
161
162pub const POSTGRES_SERVER: &str = "POSTGRES_SERVER";
163
164#[async_trait]
165impl Server for PostgresServer {
166    async fn shutdown(&self) -> Result<()> {
167        self.base_server.shutdown().await
168    }
169
170    async fn start(&mut self, listening: SocketAddr) -> Result<()> {
171        let (stream, addr) = self
172            .base_server
173            .bind(listening, self.keep_alive_secs)
174            .await?;
175
176        let io_runtime = self.base_server.io_runtime();
177        let join_handle = common_runtime::spawn_global(self.accept(io_runtime, stream));
178
179        self.base_server.start_with(join_handle).await?;
180
181        self.bind_addr = Some(addr);
182        Ok(())
183    }
184
185    fn name(&self) -> &str {
186        POSTGRES_SERVER
187    }
188
189    fn bind_addr(&self) -> Option<SocketAddr> {
190        self.bind_addr
191    }
192
193    fn as_any(&self) -> &dyn Any {
194        self
195    }
196}
servers/postgres/server.rs

servers/postgres/
server.rs
