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::{CopyInHandlerRef, InfluxdbLineProtocolHandlerRef};
48
49pub(crate) struct GreptimeDBStartupParameters {
50    version: String,
51}
52
53impl GreptimeDBStartupParameters {
54    fn new() -> GreptimeDBStartupParameters {
55        GreptimeDBStartupParameters {
56            version: format!(
57                "16.3-{}-{}",
58                common_version::product_name(),
59                common_version::version()
60            ),
61        }
62    }
63}
64
65impl ServerParameterProvider for GreptimeDBStartupParameters {
66    fn server_parameters<C>(&self, _client: &C) -> Option<HashMap<String, String>>
67    where
68        C: ClientInfo,
69    {
70        Some(HashMap::from([
71            ("server_version".to_owned(), self.version.clone()),
72            ("server_encoding".to_owned(), "UTF8".to_owned()),
73            ("client_encoding".to_owned(), "UTF8".to_owned()),
74            ("DateStyle".to_owned(), "ISO YMD".to_owned()),
75            ("integer_datetimes".to_owned(), "on".to_owned()),
76        ]))
77    }
78}
79
80pub struct PostgresServerHandlerInner {
81    query_handler: ServerSqlQueryHandlerRef,
82    auth_source: Option<Arc<PgUserProviderSource>>,
83    copy_in_handler: CopyInHandlerRef,
84    /// Executes influxdb line-protocol batches flushed by a table-less
85    /// COPY FROM STDIN — the same handler the HTTP write endpoints use.
86    /// Absent on servers not wired for it; such connections reject the
87    /// influxdb format when a copy starts.
88    influxdb_handler: Option<InfluxdbLineProtocolHandlerRef>,
89    force_tls: bool,
90    auth_driver: AuthDriver<Arc<GreptimeDBStartupParameters>, UserInfoRef>,
91    /// SCRAM selected; failures get the uniform metric label.
92    is_scram_auth: AtomicBool,
93
94    batching_enabled: bool,
95    session: Arc<Session>,
96    query_parser: Arc<DefaultQueryParser>,
97
98    /// State of an in-progress `COPY ... FROM STDIN`; set when the copy
99    /// starts and taken when it finishes, fails or the connection drops.
100    copy_in_state: tokio::sync::Mutex<Option<copy_in::CopyInState>>,
101}
102
103#[derive(Builder)]
104pub(crate) struct MakePostgresServerHandler {
105    query_handler: ServerSqlQueryHandlerRef,
106    copy_in_handler: CopyInHandlerRef,
107    #[builder(default)]
108    influxdb_handler: Option<InfluxdbLineProtocolHandlerRef>,
109    user_provider: Option<UserProviderRef>,
110    #[builder(default = "Arc::new(GreptimeDBStartupParameters::new())")]
111    param_provider: Arc<GreptimeDBStartupParameters>,
112    force_tls: bool,
113}
114
115pub(crate) struct PostgresServerHandler(Arc<PostgresServerHandlerInner>);
116
117impl PgWireServerHandlers for PostgresServerHandler {
118    fn simple_query_handler(&self) -> Arc<impl SimpleQueryHandler> {
119        self.0.clone()
120    }
121
122    fn extended_query_handler(&self) -> Arc<impl ExtendedQueryHandler> {
123        self.0.clone()
124    }
125
126    fn startup_handler(&self) -> Arc<impl StartupHandler> {
127        self.0.clone()
128    }
129
130    fn copy_handler(&self) -> Arc<impl pgwire::api::copy::CopyHandler> {
131        self.0.clone()
132    }
133
134    fn error_handler(&self) -> Arc<impl ErrorHandler> {
135        self.0.clone()
136    }
137}
138
139impl MakePostgresServerHandler {
140    fn make(
141        &self,
142        addr: Option<SocketAddr>,
143        process_id: u32,
144        batching_enabled: bool,
145    ) -> PostgresServerHandler {
146        let session = Arc::new(Session::new(
147            addr,
148            Channel::Postgres,
149            Default::default(),
150            process_id,
151        ));
152        let handler = PostgresServerHandlerInner {
153            query_handler: self.query_handler.clone(),
154            auth_source: self
155                .user_provider
156                .clone()
157                .map(|p| Arc::new(PgUserProviderSource::new(p))),
158            copy_in_handler: self.copy_in_handler.clone(),
159            influxdb_handler: self.influxdb_handler.clone(),
160            force_tls: self.force_tls,
161            auth_driver: AuthDriver::new(self.param_provider.clone())
162                .with_pid_secret_key_generator(Arc::new(ZeroPidSecretKeyGenerator)),
163            is_scram_auth: AtomicBool::new(false),
164
165            batching_enabled,
166            session: session.clone(),
167            query_parser: Arc::new(DefaultQueryParser::new(self.query_handler.clone(), session)),
168            copy_in_state: tokio::sync::Mutex::new(None),
169        };
170        PostgresServerHandler(Arc::new(handler))
171    }
172}
servers/postgres.rs

servers/
postgres.rs
