servers/
postgres.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
15mod auth_handler;
16mod copy_in;
17mod fixtures;
18mod handler;
19mod server;
20mod types;
21mod utils;
22
23pub(crate) const METADATA_USER: &str = "user";
24pub(crate) const METADATA_DATABASE: &str = "database";
25/// key to store our parsed catalog
26pub(crate) const METADATA_CATALOG: &str = "catalog";
27/// key to store our parsed schema
28pub(crate) const METADATA_SCHEMA: &str = "schema";
29
30use std::collections::HashMap;
31use std::net::SocketAddr;
32use std::sync::Arc;
33use std::sync::atomic::AtomicBool;
34
35use ::auth::{UserInfoRef, UserProviderRef};
36use derive_builder::Builder;
37use pgwire::api::auth::{AuthDriver, ServerParameterProvider, StartupHandler};
38use pgwire::api::query::{ExtendedQueryHandler, SimpleQueryHandler};
39use pgwire::api::{ClientInfo, ErrorHandler, PgWireServerHandlers};
40pub use server::PostgresServer;
41use session::Session;
42use session::context::Channel;
43
44use self::auth_handler::{PgUserProviderSource, ZeroPidSecretKeyGenerator};
45use self::handler::DefaultQueryParser;
46use crate::query_handler::sql::ServerSqlQueryHandlerRef;
47use crate::query_handler::{
48    CopyInHandlerRef, InfluxdbLineProtocolHandlerRef, OpenTelemetryProtocolHandlerRef,
49};
50
51pub(crate) struct GreptimeDBStartupParameters {
52    version: String,
53}
54
55impl GreptimeDBStartupParameters {
56    fn new() -> GreptimeDBStartupParameters {
57        GreptimeDBStartupParameters {
58            version: format!(
59                "16.3-{}-{}",
60                common_version::product_name(),
61                common_version::version()
62            ),
63        }
64    }
65}
66
67impl ServerParameterProvider for GreptimeDBStartupParameters {
68    fn server_parameters<C>(&self, _client: &C) -> Option<HashMap<String, String>>
69    where
70        C: ClientInfo,
71    {
72        Some(HashMap::from([
73            ("server_version".to_owned(), self.version.clone()),
74            ("server_encoding".to_owned(), "UTF8".to_owned()),
75            ("client_encoding".to_owned(), "UTF8".to_owned()),
76            ("DateStyle".to_owned(), "ISO YMD".to_owned()),
77            ("integer_datetimes".to_owned(), "on".to_owned()),
78        ]))
79    }
80}
81
82pub struct PostgresServerHandlerInner {
83    query_handler: ServerSqlQueryHandlerRef,
84    auth_source: Option<Arc<PgUserProviderSource>>,
85    copy_in_handler: CopyInHandlerRef,
86    /// Executes influxdb line-protocol batches flushed by a table-less
87    /// COPY FROM STDIN — the same handler the HTTP write endpoints use.
88    /// Absent on servers not wired for it; such connections reject the
89    /// influxdb format when a copy starts.
90    influxdb_handler: Option<InfluxdbLineProtocolHandlerRef>,
91    /// Executes the OTLP payload flushed by a table-less
92    /// `COPY FROM STDIN WITH (FORMAT otlp)` — the same handler the
93    /// OTLP/HTTP endpoints use — plus the with-metric-engine flag the
94    /// HTTP wiring passes from its options.
95    otlp_handler: Option<OpenTelemetryProtocolHandlerRef>,
96    otlp_with_metric_engine: bool,
97    force_tls: bool,
98    auth_driver: AuthDriver<Arc<GreptimeDBStartupParameters>, UserInfoRef>,
99    /// SCRAM selected; failures get the uniform metric label.
100    is_scram_auth: AtomicBool,
101
102    batching_enabled: bool,
103    session: Arc<Session>,
104    query_parser: Arc<DefaultQueryParser>,
105
106    /// State of an in-progress `COPY ... FROM STDIN`; set when the copy
107    /// starts and taken when it finishes, fails or the connection drops.
108    copy_in_state: tokio::sync::Mutex<Option<copy_in::CopyInState>>,
109}
110
111#[derive(Builder)]
112pub(crate) struct MakePostgresServerHandler {
113    query_handler: ServerSqlQueryHandlerRef,
114    copy_in_handler: CopyInHandlerRef,
115    #[builder(default)]
116    influxdb_handler: Option<InfluxdbLineProtocolHandlerRef>,
117    #[builder(default)]
118    otlp_handler: Option<OpenTelemetryProtocolHandlerRef>,
119    #[builder(default)]
120    otlp_with_metric_engine: bool,
121    user_provider: Option<UserProviderRef>,
122    #[builder(default = "Arc::new(GreptimeDBStartupParameters::new())")]
123    param_provider: Arc<GreptimeDBStartupParameters>,
124    force_tls: bool,
125}
126
127pub(crate) struct PostgresServerHandler(Arc<PostgresServerHandlerInner>);
128
129impl PgWireServerHandlers for PostgresServerHandler {
130    fn simple_query_handler(&self) -> Arc<impl SimpleQueryHandler> {
131        self.0.clone()
132    }
133
134    fn extended_query_handler(&self) -> Arc<impl ExtendedQueryHandler> {
135        self.0.clone()
136    }
137
138    fn startup_handler(&self) -> Arc<impl StartupHandler> {
139        self.0.clone()
140    }
141
142    fn copy_handler(&self) -> Arc<impl pgwire::api::copy::CopyHandler> {
143        self.0.clone()
144    }
145
146    fn error_handler(&self) -> Arc<impl ErrorHandler> {
147        self.0.clone()
148    }
149}
150
151impl MakePostgresServerHandler {
152    fn make(
153        &self,
154        addr: Option<SocketAddr>,
155        process_id: u32,
156        batching_enabled: bool,
157    ) -> PostgresServerHandler {
158        let session = Arc::new(Session::new(
159            addr,
160            Channel::Postgres,
161            Default::default(),
162            process_id,
163        ));
164        let handler = PostgresServerHandlerInner {
165            query_handler: self.query_handler.clone(),
166            auth_source: self
167                .user_provider
168                .clone()
169                .map(|p| Arc::new(PgUserProviderSource::new(p))),
170            copy_in_handler: self.copy_in_handler.clone(),
171            influxdb_handler: self.influxdb_handler.clone(),
172            otlp_handler: self.otlp_handler.clone(),
173            otlp_with_metric_engine: self.otlp_with_metric_engine,
174            force_tls: self.force_tls,
175            auth_driver: AuthDriver::new(self.param_provider.clone())
176                .with_pid_secret_key_generator(Arc::new(ZeroPidSecretKeyGenerator)),
177            is_scram_auth: AtomicBool::new(false),
178
179            batching_enabled,
180            session: session.clone(),
181            query_parser: Arc::new(DefaultQueryParser::new(self.query_handler.clone(), session)),
182            copy_in_state: tokio::sync::Mutex::new(None),
183        };
184        PostgresServerHandler(Arc::new(handler))
185    }
186}
servers/postgres.rs

servers/
postgres.rs
