Skip to main content

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;
32
33use ::auth::UserProviderRef;
34use derive_builder::Builder;
35use pgwire::api::auth::{ServerParameterProvider, StartupHandler};
36use pgwire::api::query::{ExtendedQueryHandler, SimpleQueryHandler};
37use pgwire::api::{ClientInfo, ErrorHandler, PgWireServerHandlers};
38pub use server::PostgresServer;
39use session::Session;
40use session::context::Channel;
41
42use self::auth_handler::PgLoginVerifier;
43use self::handler::DefaultQueryParser;
44use crate::query_handler::sql::ServerSqlQueryHandlerRef;
45
46pub(crate) struct GreptimeDBStartupParameters {
47    version: String,
48}
49
50impl GreptimeDBStartupParameters {
51    fn new() -> GreptimeDBStartupParameters {
52        GreptimeDBStartupParameters {
53            version: format!(
54                "16.3-{}-{}",
55                common_version::product_name(),
56                common_version::version()
57            ),
58        }
59    }
60}
61
62impl ServerParameterProvider for GreptimeDBStartupParameters {
63    fn server_parameters<C>(&self, _client: &C) -> Option<HashMap<String, String>>
64    where
65        C: ClientInfo,
66    {
67        Some(HashMap::from([
68            ("server_version".to_owned(), self.version.clone()),
69            ("server_encoding".to_owned(), "UTF8".to_owned()),
70            ("client_encoding".to_owned(), "UTF8".to_owned()),
71            ("DateStyle".to_owned(), "ISO YMD".to_owned()),
72            ("integer_datetimes".to_owned(), "on".to_owned()),
73        ]))
74    }
75}
76
77pub struct PostgresServerHandlerInner {
78    query_handler: ServerSqlQueryHandlerRef,
79    login_verifier: PgLoginVerifier,
80    force_tls: bool,
81    param_provider: Arc<GreptimeDBStartupParameters>,
82
83    batching_enabled: bool,
84    session: Arc<Session>,
85    query_parser: Arc<DefaultQueryParser>,
86}
87
88#[derive(Builder)]
89pub(crate) struct MakePostgresServerHandler {
90    query_handler: ServerSqlQueryHandlerRef,
91    user_provider: Option<UserProviderRef>,
92    #[builder(default = "Arc::new(GreptimeDBStartupParameters::new())")]
93    param_provider: Arc<GreptimeDBStartupParameters>,
94    force_tls: bool,
95}
96
97pub(crate) struct PostgresServerHandler(Arc<PostgresServerHandlerInner>);
98
99impl PgWireServerHandlers for PostgresServerHandler {
100    fn simple_query_handler(&self) -> Arc<impl SimpleQueryHandler> {
101        self.0.clone()
102    }
103
104    fn extended_query_handler(&self) -> Arc<impl ExtendedQueryHandler> {
105        self.0.clone()
106    }
107
108    fn startup_handler(&self) -> Arc<impl StartupHandler> {
109        self.0.clone()
110    }
111
112    fn error_handler(&self) -> Arc<impl ErrorHandler> {
113        self.0.clone()
114    }
115}
116
117impl MakePostgresServerHandler {
118    fn make(
119        &self,
120        addr: Option<SocketAddr>,
121        process_id: u32,
122        batching_enabled: bool,
123    ) -> PostgresServerHandler {
124        let session = Arc::new(Session::new(
125            addr,
126            Channel::Postgres,
127            Default::default(),
128            process_id,
129        ));
130        let handler = PostgresServerHandlerInner {
131            query_handler: self.query_handler.clone(),
132            login_verifier: PgLoginVerifier::new(self.user_provider.clone()),
133            force_tls: self.force_tls,
134            param_provider: self.param_provider.clone(),
135
136            batching_enabled,
137            session: session.clone(),
138            query_parser: Arc::new(DefaultQueryParser::new(self.query_handler.clone(), session)),
139        };
140        PostgresServerHandler(Arc::new(handler))
141    }
142}