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 fixtures;
17mod handler;
18mod server;
19mod types;
20mod utils;
21
22pub(crate) const METADATA_USER: &str = "user";
23pub(crate) const METADATA_DATABASE: &str = "database";
24/// key to store our parsed catalog
25pub(crate) const METADATA_CATALOG: &str = "catalog";
26/// key to store our parsed schema
27pub(crate) const METADATA_SCHEMA: &str = "schema";
28
29use std::collections::HashMap;
30use std::net::SocketAddr;
31use std::sync::Arc;
32use std::sync::atomic::AtomicBool;
33
34use ::auth::{UserInfoRef, UserProviderRef};
35use derive_builder::Builder;
36use pgwire::api::auth::{AuthDriver, ServerParameterProvider, StartupHandler};
37use pgwire::api::query::{ExtendedQueryHandler, SimpleQueryHandler};
38use pgwire::api::{ClientInfo, ErrorHandler, PgWireServerHandlers};
39pub use server::PostgresServer;
40use session::Session;
41use session::context::Channel;
42
43use self::auth_handler::{PgUserProviderSource, ZeroPidSecretKeyGenerator};
44use self::handler::DefaultQueryParser;
45use crate::query_handler::sql::ServerSqlQueryHandlerRef;
46
47pub(crate) struct GreptimeDBStartupParameters {
48    version: String,
49}
50
51impl GreptimeDBStartupParameters {
52    fn new() -> GreptimeDBStartupParameters {
53        GreptimeDBStartupParameters {
54            version: format!(
55                "16.3-{}-{}",
56                common_version::product_name(),
57                common_version::version()
58            ),
59        }
60    }
61}
62
63impl ServerParameterProvider for GreptimeDBStartupParameters {
64    fn server_parameters<C>(&self, _client: &C) -> Option<HashMap<String, String>>
65    where
66        C: ClientInfo,
67    {
68        Some(HashMap::from([
69            ("server_version".to_owned(), self.version.clone()),
70            ("server_encoding".to_owned(), "UTF8".to_owned()),
71            ("client_encoding".to_owned(), "UTF8".to_owned()),
72            ("DateStyle".to_owned(), "ISO YMD".to_owned()),
73            ("integer_datetimes".to_owned(), "on".to_owned()),
74        ]))
75    }
76}
77
78pub struct PostgresServerHandlerInner {
79    query_handler: ServerSqlQueryHandlerRef,
80    auth_source: Option<Arc<PgUserProviderSource>>,
81    force_tls: bool,
82    auth_driver: AuthDriver<Arc<GreptimeDBStartupParameters>, UserInfoRef>,
83    /// SCRAM selected; failures get the uniform metric label.
84    is_scram_auth: AtomicBool,
85
86    batching_enabled: bool,
87    session: Arc<Session>,
88    query_parser: Arc<DefaultQueryParser>,
89}
90
91#[derive(Builder)]
92pub(crate) struct MakePostgresServerHandler {
93    query_handler: ServerSqlQueryHandlerRef,
94    user_provider: Option<UserProviderRef>,
95    #[builder(default = "Arc::new(GreptimeDBStartupParameters::new())")]
96    param_provider: Arc<GreptimeDBStartupParameters>,
97    force_tls: bool,
98}
99
100pub(crate) struct PostgresServerHandler(Arc<PostgresServerHandlerInner>);
101
102impl PgWireServerHandlers for PostgresServerHandler {
103    fn simple_query_handler(&self) -> Arc<impl SimpleQueryHandler> {
104        self.0.clone()
105    }
106
107    fn extended_query_handler(&self) -> Arc<impl ExtendedQueryHandler> {
108        self.0.clone()
109    }
110
111    fn startup_handler(&self) -> Arc<impl StartupHandler> {
112        self.0.clone()
113    }
114
115    fn error_handler(&self) -> Arc<impl ErrorHandler> {
116        self.0.clone()
117    }
118}
119
120impl MakePostgresServerHandler {
121    fn make(
122        &self,
123        addr: Option<SocketAddr>,
124        process_id: u32,
125        batching_enabled: bool,
126    ) -> PostgresServerHandler {
127        let session = Arc::new(Session::new(
128            addr,
129            Channel::Postgres,
130            Default::default(),
131            process_id,
132        ));
133        let handler = PostgresServerHandlerInner {
134            query_handler: self.query_handler.clone(),
135            auth_source: self
136                .user_provider
137                .clone()
138                .map(|p| Arc::new(PgUserProviderSource::new(p))),
139            force_tls: self.force_tls,
140            auth_driver: AuthDriver::new(self.param_provider.clone())
141                .with_pid_secret_key_generator(Arc::new(ZeroPidSecretKeyGenerator)),
142            is_scram_auth: AtomicBool::new(false),
143
144            batching_enabled,
145            session: session.clone(),
146            query_parser: Arc::new(DefaultQueryParser::new(self.query_handler.clone(), session)),
147        };
148        PostgresServerHandler(Arc::new(handler))
149    }
150}
servers/postgres.rs

servers/
postgres.rs
