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::context::Channel;
40use session::Session;
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!("16.3-greptimedb-{}", common_version::version()),
54        }
55    }
56}
57
58impl ServerParameterProvider for GreptimeDBStartupParameters {
59    fn server_parameters<C>(&self, _client: &C) -> Option<HashMap<String, String>>
60    where
61        C: ClientInfo,
62    {
63        Some(HashMap::from([
64            ("server_version".to_owned(), self.version.clone()),
65            ("server_encoding".to_owned(), "UTF8".to_owned()),
66            ("client_encoding".to_owned(), "UTF8".to_owned()),
67            ("DateStyle".to_owned(), "ISO YMD".to_owned()),
68            ("integer_datetimes".to_owned(), "on".to_owned()),
69        ]))
70    }
71}
72
73pub struct PostgresServerHandlerInner {
74    query_handler: ServerSqlQueryHandlerRef,
75    login_verifier: PgLoginVerifier,
76    force_tls: bool,
77    param_provider: Arc<GreptimeDBStartupParameters>,
78
79    session: Arc<Session>,
80    query_parser: Arc<DefaultQueryParser>,
81}
82
83#[derive(Builder)]
84pub(crate) struct MakePostgresServerHandler {
85    query_handler: ServerSqlQueryHandlerRef,
86    user_provider: Option<UserProviderRef>,
87    #[builder(default = "Arc::new(GreptimeDBStartupParameters::new())")]
88    param_provider: Arc<GreptimeDBStartupParameters>,
89    force_tls: bool,
90}
91
92pub(crate) struct PostgresServerHandler(Arc<PostgresServerHandlerInner>);
93
94impl PgWireServerHandlers for PostgresServerHandler {
95    fn simple_query_handler(&self) -> Arc<impl SimpleQueryHandler> {
96        self.0.clone()
97    }
98
99    fn extended_query_handler(&self) -> Arc<impl ExtendedQueryHandler> {
100        self.0.clone()
101    }
102
103    fn startup_handler(&self) -> Arc<impl StartupHandler> {
104        self.0.clone()
105    }
106
107    fn error_handler(&self) -> Arc<impl ErrorHandler> {
108        self.0.clone()
109    }
110}
111
112impl MakePostgresServerHandler {
113    fn make(&self, addr: Option<SocketAddr>, process_id: u32) -> PostgresServerHandler {
114        let session = Arc::new(Session::new(
115            addr,
116            Channel::Postgres,
117            Default::default(),
118            process_id,
119        ));
120        let handler = PostgresServerHandlerInner {
121            query_handler: self.query_handler.clone(),
122            login_verifier: PgLoginVerifier::new(self.user_provider.clone()),
123            force_tls: self.force_tls,
124            param_provider: self.param_provider.clone(),
125
126            session: session.clone(),
127            query_parser: Arc::new(DefaultQueryParser::new(self.query_handler.clone(), session)),
128        };
129        PostgresServerHandler(Arc::new(handler))
130    }
131}