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

servers/postgres/
server.rs
