feat: improve server parameter apis

sunng87 · Jun 24, 2023 · c9e5adf · c9e5adf
1 parent 8203b95
commit c9e5adf
Show file tree

Hide file tree

Showing 4 changed files with 47 additions and 40 deletions.
diff --git a/examples/sqlite.rs b/examples/sqlite.rs
@@ -45,33 +45,6 @@ impl AuthSource for DummyAuthSource {
     }
 }
 
-struct SqliteParameters {
-    version: &'static str,
-}
-
-impl SqliteParameters {
-    fn new() -> SqliteParameters {
-        SqliteParameters {
-            version: rusqlite::version(),
-        }
-    }
-}
-
-impl ServerParameterProvider for SqliteParameters {
-    fn server_parameters<C>(&self, client: &C) -> Option<HashMap<String, String>>
-    where
-        C: ClientInfo,
-    {
-        let provider = DefaultServerParameterProvider;
-        if let Some(mut params) = provider.server_parameters(client) {
-            params.insert("server_version".to_owned(), self.version.to_owned());
-            Some(params)
-        } else {
-            None
-        }
-    }
-}
-
 #[async_trait]
 impl SimpleQueryHandler for SqliteBackend {
     async fn do_query<'a, C>(&self, _client: &C, query: &'a str) -> PgWireResult<Vec<Response<'a>>>
@@ -323,9 +296,12 @@ impl MakeHandler for MakeSqliteBackend {
 
 #[tokio::main]
 pub async fn main() {
+    let mut parameters = DefaultServerParameterProvider::default();
+    parameters.set_server_version(rusqlite::version().to_owned());
+
     let authenticator = Arc::new(MakeMd5PasswordAuthStartupHandler::new(
         Arc::new(DummyAuthSource),
-        Arc::new(SqliteParameters::new()),
+        Arc::new(parameters),
     ));
     let processor = Arc::new(MakeSqliteBackend::new());
 

diff --git a/src/api/auth/mod.rs b/src/api/auth/mod.rs
@@ -33,23 +33,53 @@ pub trait ServerParameterProvider: Send + Sync {
         C: ClientInfo;
 }
 
-/// Default noop parameter provider
-pub struct DefaultServerParameterProvider;
+/// Default noop parameter provider.
+///
+/// This provider responds frontend with default parameters:
+///
+/// - `DateStyle: ISO YMD`: the default text serialization in this library is
+/// using `YMD` style date. If you override this, or use your own serialization
+/// for date types, remember to update this as well.
+/// - `server_encoding: UTF8`
+/// - `client_encoding: UTF8`
+/// - `integer_datetimes: on`:
+///
+#[derive(Debug, Getters, Setters)]
+#[getset(get = "pub", set = "pub")]
+pub struct DefaultServerParameterProvider {
+    server_version: String,
+    server_encoding: String,
+    client_encoding: String,
+    date_style: String,
+    integer_datetimes: String,
+}
+
+impl Default for DefaultServerParameterProvider {
+    fn default() -> Self {
+        Self {
+            server_version: env!("CARGO_PKG_VERSION").to_owned(),
+            server_encoding: "UTF8".to_owned(),
+            client_encoding: "UTF8".to_owned(),
+            date_style: "ISO YMD".to_owned(),
+            integer_datetimes: "on".to_owned(),
+        }
+    }
+}
 
 impl ServerParameterProvider for DefaultServerParameterProvider {
     fn server_parameters<C>(&self, _client: &C) -> Option<HashMap<String, String>>
     where
         C: ClientInfo,
     {
-        let mut params = HashMap::with_capacity(4);
+        let mut params = HashMap::with_capacity(5);
+        params.insert("server_version".to_owned(), self.server_version.clone());
+        params.insert("server_encoding".to_owned(), self.server_encoding.clone());
+        params.insert("client_encoding".to_owned(), self.client_encoding.clone());
+        params.insert("DateStyle".to_owned(), self.date_style.clone());
         params.insert(
-            "server_version".to_owned(),
-            env!("CARGO_PKG_VERSION").to_owned(),
+            "integer_datetimes".to_owned(),
+            self.integer_datetimes.clone(),
         );
-        params.insert("server_encoding".to_owned(), "UTF8".to_owned());
-        params.insert("client_encoding".to_owned(), "UTF8".to_owned());
-        params.insert("DateStyle".to_owned(), "ISO YMD".to_owned());
-        params.insert("integer_datetimes".to_owned(), "on".to_owned());
 
         Some(params)
     }

diff --git a/src/api/auth/noop.rs b/src/api/auth/noop.rs
@@ -23,7 +23,7 @@ impl StartupHandler for NoopStartupHandler {
     {
         if let PgWireFrontendMessage::Startup(ref startup) = message {
             super::save_startup_parameters_to_metadata(client, startup);
-            super::finish_authentication(client, &DefaultServerParameterProvider).await;
+            super::finish_authentication(client, &DefaultServerParameterProvider::default()).await;
         }
         Ok(())
     }

diff --git a/src/api/query.rs b/src/api/query.rs
@@ -230,7 +230,8 @@ pub trait ExtendedQueryHandler: Send + Sync {
 
     /// Called when client sends `sync` command.
     ///
-    /// The default implementation flushes client buffer.
+    /// The default implementation flushes client buffer and sends
+    /// `READY_FOR_QUERY` response to client
     async fn on_sync<C>(&self, client: &mut C, _message: PgSync) -> PgWireResult<()>
     where
         C: ClientInfo + Sink<PgWireBackendMessage> + Unpin + Send + Sync,
@@ -271,7 +272,7 @@ pub trait ExtendedQueryHandler: Send + Sync {
         Ok(())
     }
 
-    /// Return resultset metadata without actually execute statement or portal
+    /// Return resultset metadata without actually executing statement or portal
     async fn do_describe<C>(
         &self,
         client: &mut C,