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::CopyInHandlerRef;
47use crate::query_handler::sql::ServerSqlQueryHandlerRef;
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    force_tls: bool,
85    auth_driver: AuthDriver<Arc<GreptimeDBStartupParameters>, UserInfoRef>,
86    /// SCRAM selected; failures get the uniform metric label.
87    is_scram_auth: AtomicBool,
88
89    batching_enabled: bool,
90    session: Arc<Session>,
91    query_parser: Arc<DefaultQueryParser>,
92
93    /// State of an in-progress `COPY ... FROM STDIN`; set when the copy
94    /// starts and taken when it finishes, fails or the connection drops.
95    copy_in_state: tokio::sync::Mutex<Option<copy_in::CopyInState>>,
96}
97
98#[derive(Builder)]
99pub(crate) struct MakePostgresServerHandler {
100    query_handler: ServerSqlQueryHandlerRef,
101    copy_in_handler: CopyInHandlerRef,
102    user_provider: Option<UserProviderRef>,
103    #[builder(default = "Arc::new(GreptimeDBStartupParameters::new())")]
104    param_provider: Arc<GreptimeDBStartupParameters>,
105    force_tls: bool,
106}
107
108pub(crate) struct PostgresServerHandler(Arc<PostgresServerHandlerInner>);
109
110impl PgWireServerHandlers for PostgresServerHandler {
111    fn simple_query_handler(&self) -> Arc<impl SimpleQueryHandler> {
112        self.0.clone()
113    }
114
115    fn extended_query_handler(&self) -> Arc<impl ExtendedQueryHandler> {
116        self.0.clone()
117    }
118
119    fn startup_handler(&self) -> Arc<impl StartupHandler> {
120        self.0.clone()
121    }
122
123    fn copy_handler(&self) -> Arc<impl pgwire::api::copy::CopyHandler> {
124        self.0.clone()
125    }
126
127    fn error_handler(&self) -> Arc<impl ErrorHandler> {
128        self.0.clone()
129    }
130}
131
132impl MakePostgresServerHandler {
133    fn make(
134        &self,
135        addr: Option<SocketAddr>,
136        process_id: u32,
137        batching_enabled: bool,
138    ) -> PostgresServerHandler {
139        let session = Arc::new(Session::new(
140            addr,
141            Channel::Postgres,
142            Default::default(),
143            process_id,
144        ));
145        let handler = PostgresServerHandlerInner {
146            query_handler: self.query_handler.clone(),
147            auth_source: self
148                .user_provider
149                .clone()
150                .map(|p| Arc::new(PgUserProviderSource::new(p))),
151            copy_in_handler: self.copy_in_handler.clone(),
152            force_tls: self.force_tls,
153            auth_driver: AuthDriver::new(self.param_provider.clone())
154                .with_pid_secret_key_generator(Arc::new(ZeroPidSecretKeyGenerator)),
155            is_scram_auth: AtomicBool::new(false),
156
157            batching_enabled,
158            session: session.clone(),
159            query_parser: Arc::new(DefaultQueryParser::new(self.query_handler.clone(), session)),
160            copy_in_state: tokio::sync::Mutex::new(None),
161        };
162        PostgresServerHandler(Arc::new(handler))
163    }
164}
servers/postgres.rs

servers/
postgres.rs
