more

PeerDB-io · Nov 22, 2023 · 7ab8398 · 7ab8398
1 parent 0ef50fd
commit 7ab8398
Show file tree

Hide file tree

Showing 6 changed files with 25 additions and 32 deletions.
diff --git a/nexus/catalog/src/lib.rs b/nexus/catalog/src/lib.rs
@@ -211,7 +211,7 @@ impl Catalog {
 
         let rows = self.pg.query(&stmt, &[]).await?;
 
-        let mut peers = HashMap::new();
+        let mut peers = HashMap::with_capacity(rows.len());
 
         for row in rows {
             let name: &str = row.get(1);

diff --git a/nexus/peer-bigquery/src/stream.rs b/nexus/peer-bigquery/src/stream.rs
@@ -101,7 +101,7 @@ impl BqRecordStream {
     }
 
     pub fn convert_result_set_item(&self, result_set: &ResultSet) -> anyhow::Result<Record> {
-        let mut values = Vec::new();
+        let mut values = Vec::with_capacity(self.schema.fields.len());
         for field in &self.schema.fields {
             let field_type = &field.r#type;
             let field_name = &field.name;

diff --git a/nexus/peer-snowflake/src/lib.rs b/nexus/peer-snowflake/src/lib.rs
@@ -7,7 +7,7 @@ use pgwire::error::{ErrorInfo, PgWireError, PgWireResult};
 use sqlparser::dialect::GenericDialect;
 use sqlparser::parser;
 use std::cmp::min;
-use std::{collections::HashMap, time::Duration};
+use std::time::Duration;
 use stream::SnowflakeDataType;
 
 use auth::SnowflakeAuth;
@@ -36,14 +36,23 @@ const TIME_OUTPUT_FORMAT: &str = "HH:MI:SS.FF";
 const TIMESTAMP_OUTPUT_FORMAT: &str = "YYYY-MM-DDTHH24:MI:SS.FF";
 const TIMESTAMP_TZ_OUTPUT_FORMAT: &str = "YYYY-MM-DDTHH24:MI:SS.FFTZHTZM";
 
+#[derive(Debug, Serialize)]
+struct SQLStatementParameters<'a> {
+    pub date_output_format: &'a str,
+    pub time_output_format: &'a str,
+    pub timestamp_ltz_output_format: &'a str,
+    pub timestamp_ntz_output_format: &'a str,
+    pub timestamp_tz_output_format: &'a str,
+}
+
 #[derive(Debug, Serialize)]
 struct SQLStatement<'a> {
     statement: &'a str,
     timeout: u64,
     database: &'a str,
     warehouse: &'a str,
     role: &'a str,
-    parameters: HashMap<String, String>,
+    parameters: SQLStatementParameters<'a>,
 }
 
 #[allow(non_snake_case)]
@@ -147,28 +156,6 @@ impl SnowflakeQueryExecutor {
     #[async_recursion]
     #[tracing::instrument(name = "peer_sflake::process_query", skip_all)]
     async fn process_query(&self, query_str: &str) -> anyhow::Result<ResultSet> {
-        let mut parameters = HashMap::new();
-        parameters.insert(
-            "date_output_format".to_string(),
-            DATE_OUTPUT_FORMAT.to_string(),
-        );
-        parameters.insert(
-            "time_output_format".to_string(),
-            TIME_OUTPUT_FORMAT.to_string(),
-        );
-        parameters.insert(
-            "timestamp_ltz_output_format".to_string(),
-            TIMESTAMP_TZ_OUTPUT_FORMAT.to_string(),
-        );
-        parameters.insert(
-            "timestamp_ntz_output_format".to_string(),
-            TIMESTAMP_OUTPUT_FORMAT.to_string(),
-        );
-        parameters.insert(
-            "timestamp_tz_output_format".to_string(),
-            TIMESTAMP_TZ_OUTPUT_FORMAT.to_string(),
-        );
-
         let mut auth = self.auth.clone();
         let jwt = auth.get_jwt()?;
         let secret = jwt.expose_secret().clone();
@@ -186,7 +173,13 @@ impl SnowflakeQueryExecutor {
                 database: &self.config.database,
                 warehouse: &self.config.warehouse,
                 role: &self.config.role,
-                parameters,
+                parameters: SQLStatementParameters {
+                    date_output_format: DATE_OUTPUT_FORMAT,
+                    time_output_format: TIME_OUTPUT_FORMAT,
+                    timestamp_ltz_output_format: TIMESTAMP_TZ_OUTPUT_FORMAT,
+                    timestamp_ntz_output_format: TIMESTAMP_OUTPUT_FORMAT,
+                    timestamp_tz_output_format: TIMESTAMP_TZ_OUTPUT_FORMAT,
+                },
             })
             .send()
             .await

diff --git a/nexus/server/src/cursor.rs b/nexus/server/src/cursor.rs
@@ -20,8 +20,8 @@ impl PeerCursors {
         self.cursors.insert(name, peer);
     }
 
-    pub fn remove_cursor(&mut self, name: String) {
-        self.cursors.remove(&name);
+    pub fn remove_cursor(&mut self, name: &str) {
+        self.cursors.remove(name);
     }
 
     pub fn get_peer(&self, name: &str) -> Option<&Peer> {

diff --git a/nexus/server/src/main.rs b/nexus/server/src/main.rs
@@ -141,7 +141,7 @@ impl NexusBackend {
                     }
                     peer_cursor::CursorModification::Closed(cursors) => {
                         for cursor_name in cursors {
-                            peer_cursors.remove_cursor(cursor_name);
+                            peer_cursors.remove_cursor(&cursor_name);
                         }
                         Ok(vec![Response::Execution(Tag::new_for_execution(
                             "CLOSE CURSOR",

diff --git a/nexus/value/src/lib.rs b/nexus/value/src/lib.rs
@@ -203,7 +203,7 @@ impl Value {
                 }
             }
             serde_json::Value::Object(map) => {
-                let mut hstore = HashMap::new();
+                let mut hstore = HashMap::with_capacity(map.len());
                 for (key, value) in map {
                     hstore.insert(key.clone(), value.to_string());
                 }
@@ -253,7 +253,7 @@ impl Value {
             Value::Uuid(u) => serde_json::Value::String(u.to_string()),
             Value::Enum(s) => serde_json::Value::String(s.clone()),
             Value::Hstore(map) => {
-                let mut object = serde_json::Map::new();
+                let mut object = serde_json::Map::with_capacity(map.len());
                 for (key, value) in map {
                     object.insert(key.clone(), serde_json::Value::String(value.clone()));
                 }