fix(cubesql): Catch error on TcpStream.peer_addr() (#6300)

Author: ovr

rust/cubesql/cubesql/src/compile/engine/information_schema/postgres/pg_stat_activity.rs

@@ -78,11 +78,13 @@ impl PgStatActivityBuilder {
         self.application_name
             .append_option(session.application_name)
             .unwrap();
-        self.client_addr.append_option(session.client_addr).unwrap();
+        self.client_addr.append_value(session.client_addr).unwrap();
         self.client_hostname
             .append_option(session.client_hostname)
             .unwrap();
-        self.client_port.append_option(session.client_port).unwrap();
+        self.client_port
+            .append_value(session.client_port.to_string())
+            .unwrap();
         self.backend_start.append_null().unwrap();
         self.xact_start.append_null().unwrap();
         self.query_start.append_null().unwrap();

rust/cubesql/cubesql/src/compile/mod.rs

@@ -7322,6 +7322,20 @@ ORDER BY \"COUNT(count)\" DESC"
         Ok(())
     }
 
+    #[tokio::test]
+    async fn test_pgcatalog_pg_stat_activity_postgres() -> Result<(), CubeError> {
+        insta::assert_snapshot!(
+            "pgcatalog_pg_stat_activity_postgres",
+            execute_query(
+                "SELECT * FROM pg_catalog.pg_stat_activity".to_string(),
+                DatabaseProtocol::PostgreSQL
+            )
+            .await?
+        );
+
+        Ok(())
+    }
+
     #[tokio::test]
     async fn test_pgcatalog_pguser_postgres() -> Result<(), CubeError> {
         insta::assert_snapshot!(

rust/cubesql/cubesql/src/compile/snapshots/cubesql__compile__tests__information_schema_processlist.snap

@@ -1,8 +1,7 @@
 ---
 source: cubesql/src/compile/mod.rs
-assertion_line: 4118
-expression: "execute_query(\"SELECT * FROM information_schema.processlist\".to_string()).await?"
-
+assertion_line: 5710
+expression: "execute_query(\"SELECT * FROM information_schema.processlist\".to_string(),\n            DatabaseProtocol::MySQL).await?"
 ---
 +----+------+-----------+----+---------+------+------------------------+------+
 | ID | USER | HOST      | DB | COMMAND | TIME | STATE                  | INFO |

rust/cubesql/cubesql/src/compile/snapshots/cubesql__compile__tests__pgcatalog_pg_stat_activity_postgres.snap

@@ -0,0 +1,10 @@
+---
+source: cubesql/src/compile/mod.rs
+assertion_line: 7327
+expression: "execute_query(\"SELECT * FROM pg_catalog.pg_stat_activity\".to_string(),\n            DatabaseProtocol::PostgreSQL).await?"
+---
++-----+---------+-----+------------+----------+---------+------------------+-------------+-----------------+-------------+---------------+------------+-------------+--------------+-----------------+------------+-------+-------------+--------------+----------+-------+----------------+
+| oid | datname | pid | leader_pid | usesysid | usename | application_name | client_addr | client_hostname | client_port | backend_start | xact_start | query_start | state_change | wait_event_type | wait_event | state | backend_xid | backend_xmin | query_id | query | backend_type   |
++-----+---------+-----+------------+----------+---------+------------------+-------------+-----------------+-------------+---------------+------------+-------------+--------------+-----------------+------------+-------+-------------+--------------+----------+-------+----------------+
+| 1   | cubedb  | 1   | NULL       | NULL     | ovr     | NULL             | 127.0.0.1   | NULL            | 1234        | NULL          | NULL       | NULL        | NULL         | NULL            | NULL       | NULL  | NULL        | NULL         | NULL     | NULL  | client backend |
++-----+---------+-----+------------+----------+---------+------------------+-------------+-----------------+-------------+---------------+------------+-------------+--------------+-----------------+------------+-------+-------------+--------------+----------+-------+----------------+

rust/cubesql/cubesql/src/compile/snapshots/cubesql__compile__tests__show_processlist.snap

@@ -1,8 +1,7 @@
 ---
 source: cubesql/src/compile/mod.rs
-assertion_line: 4116
-expression: "execute_query(\"SHOW processlist\".to_string()).await?"
-
+assertion_line: 5676
+expression: "execute_query(\"SHOW processlist\".to_string(), DatabaseProtocol::MySQL).await?"
 ---
 +----+------+-----------+----+---------+------+------------------------+------+
 | ID | USER | HOST      | DB | COMMAND | TIME | STATE                  | INFO |

rust/cubesql/cubesql/src/compile/test/mod.rs

@@ -189,7 +189,7 @@ pub async fn get_test_session(protocol: DatabaseProtocol) -> Arc<Session> {
     };
     let session_manager = Arc::new(SessionManager::new(server.clone()));
     let session = session_manager
-        .create_session(protocol, "127.0.0.1".to_string())
+        .create_session(protocol, "127.0.0.1".to_string(), 1234)
         .await;
 
     // Populate like shims

rust/cubesql/cubesql/src/sql/mysql/service.rs

@@ -501,12 +501,21 @@ impl ProcessingLoop for MySqlServer {
                 }
             };
 
+            let (client_addr, client_port) = match socket.peer_addr() {
+                Ok(peer_addr) => (peer_addr.ip().to_string(), peer_addr.port()),
+                Err(e) => {
+                    error!(
+                        "[mysql] Error while calling peer_addr() on TcpStream: {}",
+                        e
+                    );
+
+                    ("127.0.0.1".to_string(), 0000_u16)
+                }
+            };
+
             let session = self
                 .session_manager
-                .create_session(
-                    DatabaseProtocol::MySQL,
-                    socket.peer_addr().unwrap().to_string(),
-                )
+                .create_session(DatabaseProtocol::MySQL, client_addr, client_port)
                 .await;
 
             let logger = Arc::new(SessionLogger::new(session.state.clone()));

rust/cubesql/cubesql/src/sql/postgres/service.rs

@@ -56,12 +56,18 @@ impl ProcessingLoop for PostgresServer {
                 }
             };
 
+            let (client_addr, client_port) = match socket.peer_addr() {
+                Ok(peer_addr) => (peer_addr.ip().to_string(), peer_addr.port()),
+                Err(e) => {
+                    error!("[pg] Error while calling peer_addr() on TcpStream: {}", e);
+
+                    ("127.0.0.1".to_string(), 0000_u16)
+                }
+            };
+
             let session = self
                 .session_manager
-                .create_session(
-                    DatabaseProtocol::PostgreSQL,
-                    socket.peer_addr().unwrap().to_string(),
-                )
+                .create_session(DatabaseProtocol::PostgreSQL, client_addr, client_port)
                 .await;
             let logger = Arc::new(SessionLogger::new(session.state.clone()));
 
@@ -74,7 +80,7 @@ impl ProcessingLoop for PostgresServer {
             tokio::spawn(async move {
                 tx.closed().await;
 
-                trace!("[postgres] Removing connection {}", connection_id);
+                trace!("[pg] Removing connection {}", connection_id);
 
                 session_manager.drop_session(connection_id).await;
             });

rust/cubesql/cubesql/src/sql/session.rs

@@ -80,8 +80,10 @@ pub struct SessionState {
     pub connection_id: u32,
     // secret for this session
     pub secret: u32,
-    // client address, immutable
-    pub host: String,
+    // client ip, immutable
+    pub client_ip: String,
+    // client port, immutable
+    pub client_port: u16,
     // client protocol, mysql/postgresql, immutable
     pub protocol: DatabaseProtocol,
 
@@ -104,7 +106,8 @@ pub struct SessionState {
 impl SessionState {
     pub fn new(
         connection_id: u32,
-        host: String,
+        client_ip: String,
+        client_port: u16,
         protocol: DatabaseProtocol,
         auth_context: Option<AuthContextRef>,
     ) -> Self {
@@ -113,7 +116,8 @@ impl SessionState {
         Self {
             connection_id,
             secret: rng.gen(),
-            host,
+            client_ip,
+            client_port,
             protocol,
             variables: RwLockSync::new(None),
             properties: RwLockSync::new(SessionProperties::new(None, None)),
@@ -395,9 +399,9 @@ impl Session {
             usesysid: 0,
             usename: self.state.user(),
             application_name,
-            client_addr: None,
+            client_addr: self.state.client_ip.clone(),
             client_hostname: None,
-            client_port: None,
+            client_port: self.state.client_port.clone(),
             query,
         }
     }
@@ -406,7 +410,7 @@ impl Session {
     pub fn to_process_list(self: &Arc<Self>) -> SessionProcessList {
         SessionProcessList {
             id: self.state.connection_id,
-            host: self.state.host.clone(),
+            host: self.state.client_ip.clone(),
             user: self.state.user(),
             database: self.state.database(),
         }
@@ -430,8 +434,8 @@ pub struct SessionStatActivity {
     pub usesysid: u32,
     pub usename: Option<String>,
     pub application_name: Option<String>,
-    pub client_addr: Option<String>,
+    pub client_addr: String,
     pub client_hostname: Option<String>,
-    pub client_port: Option<String>,
+    pub client_port: u16,
     pub query: Option<String>,
 }

rust/cubesql/cubesql/src/sql/session_manager.rs

@@ -35,14 +35,21 @@ impl SessionManager {
     pub async fn create_session(
         self: &Arc<Self>,
         protocol: DatabaseProtocol,
-        host: String,
+        client_addr: String,
+        client_port: u16,
     ) -> Arc<Session> {
         let connection_id = self.last_id.fetch_add(1, Ordering::SeqCst);
 
         let sess = Session {
             session_manager: self.clone(),
             server: self.server.clone(),
-            state: Arc::new(SessionState::new(connection_id, host, protocol, None)),
+            state: Arc::new(SessionState::new(
+                connection_id,
+                client_addr,
+                client_port,
+                protocol,
+                None,
+            )),
         };
 
         let session_ref = Arc::new(sess);