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

servers/postgres/
server.rs
