Merge pull request #130 from nightly-labs/db-tables-update

Db tables update, ip geolocation
nightly-labs · Mar 19, 2024 · 8e96a50 · 8e96a50
2 parents 70de604 + 61f828e
commit 8e96a50
Show file tree

Hide file tree

Showing 34 changed files with 391 additions and 90 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -27,7 +27,7 @@ tower = { version = "0.4.13", features = [
   "limit",
 ] }
 tower-http = { version = "0.5.1", features = ["cors", "trace"] }
-reqwest = "0.11.24"
+reqwest = {version = "0.11.24", features = ["json"]}
 tokio = { version = "1.35.1", features = ["full"] }
 async-trait = "0.1.77"
 
@@ -39,4 +39,4 @@ sqlx = { version = "0.7.3", features = [
   "postgres",
   "chrono",
 ] }
-chrono = { version = "0.4.22", features = ["serde"] }
+chrono = { version = "0.4.22", features = ["serde"] }
diff --git a/database/migrations/0001_types.sql b/database/migrations/0001_types.sql
@@ -23,4 +23,11 @@ CREATE TYPE client_data AS (
     wallet_name TEXT,
     wallet_type TEXT,
     connected_at TIMESTAMPTZ
+);
+
+CREATE TYPE geo_location AS (
+    country TEXT,
+    city TEXT,
+    lat FLOAT8,
+    lon FLOAT8
 );
diff --git a/database/migrations/0007_sessions.sql b/database/migrations/0007_sessions.sql
@@ -1,6 +1,5 @@
 CREATE TABLE sessions (
     session_id TEXT NOT NULL,
-    session_type session_type_enum NOT NULL,
     app_id TEXT NOT NULL,
     app_metadata TEXT NOT NULL,
     persistent BOOLEAN NOT NULL,

diff --git a/database/migrations/0009_connection_events.sql b/database/migrations/0009_connection_events.sql
@@ -5,6 +5,8 @@ CREATE TABLE connection_events (
     entity_id TEXT NOT NULL, -- The ID of the connecting entity (could be the same app_id or client_profile_id or client_id)
     entity_type entity_type_enum NOT NULL, -- Distinguishes between 'client' and 'app'
     ip_address TEXT NOT NULL,
+    session_type session_type_enum, -- Null for app connections
+    geo_location geo_location,
     success BOOLEAN NOT NULL,
     connected_at TIMESTAMPTZ NOT NULL,
     disconnected_at TIMESTAMPTZ

diff --git a/database/migrations/0010_ip_addresses.sql b/database/migrations/0010_ip_addresses.sql
@@ -0,0 +1,10 @@
+CREATE TABLE ip_addresses (
+    ip_addr TEXT NOT NULL UNIQUE,
+    last_updated_at TIMESTAMPTZ NOT NULL,
+    country TEXT,
+    city TEXT,
+    lat FLOAT8,
+    lon FLOAT8
+);
+
+CREATE INDEX ip_addresses_last_updated_at_idx ON ip_addresses (last_updated_at);
diff --git a/database/migrations/0010_requests.sql → database/migrations/0011_requests.sql b/database/migrations/0010_requests.sql → database/migrations/0011_requests.sql
diff --git a/...se/migrations/0011_create_hypertables.sql → ...se/migrations/0012_create_hypertables.sql b/...se/migrations/0011_create_hypertables.sql → ...se/migrations/0012_create_hypertables.sql
diff --git a/database/migrations/0012_requests_stats.sql → database/migrations/0013_requests_stats.sql b/database/migrations/0012_requests_stats.sql → database/migrations/0013_requests_stats.sql
diff --git a/database/migrations/0013_session_stats.sql → database/migrations/0014_session_stats.sql b/database/migrations/0013_session_stats.sql → database/migrations/0014_session_stats.sql
diff --git a/...base/migrations/0014_connection_stats.sql → ...base/migrations/0015_connection_stats.sql b/...base/migrations/0014_connection_stats.sql → ...base/migrations/0015_connection_stats.sql
diff --git a/database/src/aggregated_views_queries/connections_stats.rs b/database/src/aggregated_views_queries/connections_stats.rs
@@ -86,7 +86,6 @@ mod tests {
             let session = DbNcSession {
                 session_id: session_id.clone(),
                 app_id: app_id.clone(),
-                session_type: SessionType::Relay,
                 app_metadata: "test_metadata".to_string(),
                 persistent: true,
                 network: network.to_string(),
@@ -95,7 +94,7 @@ mod tests {
                 session_close_timestamp: None,
             };
 
-            db.handle_new_session(&session).await.unwrap();
+            db.handle_new_session(&session, None).await.unwrap();
 
             // Each time a session is created, means that app has been connected, create 2 more connections
             let mut tx = db.connection_pool.begin().await.unwrap();
@@ -104,6 +103,7 @@ mod tests {
                 &session_id,
                 &app_id,
                 &network.to_string(),
+                None,
             )
             .await
             .unwrap();
@@ -113,6 +113,7 @@ mod tests {
                 &session_id,
                 &app_id,
                 &network.to_string(),
+                None,
             )
             .await
             .unwrap();
@@ -123,7 +124,9 @@ mod tests {
                     &app_id,
                     &session_id,
                     j as i64,
+                    &SessionType::Relay,
                     &network.to_string(),
+                    None,
                 )
                 .await
                 .unwrap();

diff --git a/database/src/aggregated_views_queries/requests_stats.rs b/database/src/aggregated_views_queries/requests_stats.rs
@@ -55,9 +55,7 @@ mod test {
 
     use super::*;
     use crate::{
-        structs::{
-            consts::DAY_IN_SECONDS, request_status::RequestStatus, session_type::SessionType,
-        },
+        structs::{consts::DAY_IN_SECONDS, request_status::RequestStatus},
         tables::{
             registered_app::table_struct::DbRegisteredApp, requests::table_struct::Request,
             sessions::table_struct::DbNcSession, utils::to_microsecond_precision,
@@ -84,7 +82,6 @@ mod test {
         let session = DbNcSession {
             session_id: "test_session_id".to_string(),
             app_id: app_id.to_string(),
-            session_type: SessionType::Relay,
             app_metadata: "test_app_metadata".to_string(),
             persistent: false,
             network: "test_network".to_string(),
@@ -93,7 +90,7 @@ mod test {
             session_close_timestamp: None,
         };
 
-        db.handle_new_session(&session).await.unwrap();
+        db.handle_new_session(&session, None).await.unwrap();
 
         let result = db.get_sessions_by_app_id(&app_id).await.unwrap();
         assert_eq!(result.len(), 1);
@@ -187,17 +184,15 @@ mod test {
         let session = DbNcSession {
             session_id: "test_session_id".to_string(),
             app_id: "test_app_id".to_string(),
-            session_type: SessionType::Relay,
             app_metadata: "test_app_metadata".to_string(),
-
             persistent: false,
             network: "test_network".to_string(),
             client_data: None,
             session_open_timestamp: to_microsecond_precision(&Utc::now()),
             session_close_timestamp: None,
         };
 
-        db.handle_new_session(&session).await.unwrap();
+        db.handle_new_session(&session, None).await.unwrap();
 
         let result = db.get_sessions_by_app_id(&app_id).await.unwrap();
         assert_eq!(result.len(), 1);
@@ -323,7 +318,6 @@ mod test {
         let session = DbNcSession {
             session_id: "test_session_id".to_string(),
             app_id: second_app_id.to_string(),
-            session_type: SessionType::Relay,
             app_metadata: "test_app_metadata".to_string(),
             persistent: false,
             network: "test_network".to_string(),
@@ -332,7 +326,7 @@ mod test {
             session_close_timestamp: None,
         };
 
-        db_arc.handle_new_session(&session).await.unwrap();
+        db_arc.handle_new_session(&session, None).await.unwrap();
 
         let mut tasks = Vec::new();
         for i in 0..10 {

diff --git a/database/src/aggregated_views_queries/session_average_time.rs b/database/src/aggregated_views_queries/session_average_time.rs
@@ -26,7 +26,7 @@ impl Db {
 #[cfg(test)]
 mod test {
 
-    use crate::{structs::session_type::SessionType, tables::sessions::table_struct::DbNcSession};
+    use crate::tables::sessions::table_struct::DbNcSession;
     use sqlx::types::chrono::Utc;
     use std::time::Duration;
 
@@ -64,17 +64,15 @@ mod test {
             let session = DbNcSession {
                 session_id: format!("session_id_{}", start.timestamp()),
                 app_id: "test_app_id".to_string(),
-                session_type: SessionType::Relay,
                 app_metadata: "test_app_metadata".to_string(),
-
                 persistent: false,
                 network: "test_network".to_string(),
                 client_data: None,
                 session_open_timestamp: *start,
                 session_close_timestamp: None,
             };
 
-            db.handle_new_session(&session).await.unwrap();
+            db.handle_new_session(&session, None).await.unwrap();
             db.close_session(&session.session_id, *end).await.unwrap();
         }
 

diff --git a/database/src/aggregated_views_queries/sessions_stats.rs b/database/src/aggregated_views_queries/sessions_stats.rs
@@ -52,10 +52,7 @@ impl Db {
 
 #[cfg(test)]
 mod tests {
-    use crate::{
-        structs::{session_type::SessionType, time_filters::TimeFilter},
-        tables::sessions::table_struct::DbNcSession,
-    };
+    use crate::{structs::time_filters::TimeFilter, tables::sessions::table_struct::DbNcSession};
     use sqlx::types::chrono::Utc;
     use std::time::Duration;
 
@@ -87,7 +84,6 @@ mod tests {
             let session = DbNcSession {
                 session_id: format!("session_{}_{}", app_id, i),
                 app_id: app_id.clone(),
-                session_type: SessionType::Relay,
                 app_metadata: "test_metadata".to_string(),
                 persistent: false,
                 network: "test_network".to_string(),
@@ -96,7 +92,7 @@ mod tests {
                 session_close_timestamp: None,
             };
 
-            db.handle_new_session(&session).await.unwrap();
+            db.handle_new_session(&session, None).await.unwrap();
             db.close_session(&session.session_id, session_end)
                 .await
                 .unwrap();
@@ -148,7 +144,6 @@ mod tests {
             let session = DbNcSession {
                 session_id: format!("session_{}_{}", app_id, i),
                 app_id: app_id.clone(),
-                session_type: SessionType::Relay,
                 app_metadata: "test_metadata".to_string(),
                 persistent: false,
                 network: "test_network".to_string(),
@@ -157,7 +152,7 @@ mod tests {
                 session_close_timestamp: Some(session_end),
             };
 
-            db.handle_new_session(&session).await.unwrap();
+            db.handle_new_session(&session, None).await.unwrap();
             db.close_session(&session.session_id, session_end)
                 .await
                 .unwrap();
@@ -172,7 +167,6 @@ mod tests {
             let session = DbNcSession {
                 session_id: format!("session_{}_{}_2nd", app_id, i), // Ensure unique session IDs for the second period
                 app_id: app_id.clone(),
-                session_type: SessionType::Relay,
                 app_metadata: "test_metadata".to_string(),
                 persistent: false,
                 network: "test_network".to_string(),
@@ -181,7 +175,7 @@ mod tests {
                 session_close_timestamp: Some(session_end),
             };
 
-            db.handle_new_session(&session).await.unwrap();
+            db.handle_new_session(&session, None).await.unwrap();
             db.close_session(&session.session_id, session_end)
                 .await
                 .unwrap();

diff --git a/database/src/structs/geo_location.rs b/database/src/structs/geo_location.rs
@@ -0,0 +1,10 @@
+use sqlx::Type;
+
+#[derive(Clone, Debug, PartialEq, Type)]
+#[sqlx(type_name = "geo_location")]
+pub struct GeoLocation {
+    pub country: Option<String>,
+    pub city: Option<String>,
+    pub lat: Option<f64>,
+    pub lon: Option<f64>,
+}
diff --git a/database/src/structs/mod.rs b/database/src/structs/mod.rs
@@ -3,6 +3,7 @@ pub mod consts;
 pub mod db_error;
 pub mod entity_type;
 pub mod filter_requests;
+pub mod geo_location;
 pub mod privilege_level;
 pub mod request_status;
 pub mod session_type;

diff --git a/database/src/tables/connection_events/table_struct.rs b/database/src/tables/connection_events/table_struct.rs
@@ -1,4 +1,6 @@
-use crate::structs::entity_type::EntityType;
+use crate::structs::{
+    entity_type::EntityType, geo_location::GeoLocation, session_type::SessionType,
+};
 use sqlx::{
     postgres::PgRow,
     types::chrono::{DateTime, Utc},
@@ -7,16 +9,18 @@ use sqlx::{
 
 pub const CONNECTION_EVENTS_TABLE_NAME: &str = "connection_events";
 pub const CONNECTION_EVENTS_KEYS_KEYS: &str =
-    "event_id, app_id, session_id, entity_id, entity_type, ip_address, success, connected_at, disconnected_at";
+    "event_id, app_id, session_id, entity_id, entity_type, ip_address, session_type, geo_location, success, connected_at, disconnected_at";
 
-#[derive(Clone, Debug, Eq, PartialEq)]
+#[derive(Clone, Debug, PartialEq)]
 pub struct ConnectionEvent {
     pub event_id: i64,
     pub app_id: String,
     pub session_id: String,
     pub entity_id: String,
     pub entity_type: EntityType,
     pub ip_address: String,
+    pub session_type: Option<SessionType>,
+    pub geo_location: Option<GeoLocation>,
     pub success: bool,
     pub connected_at: DateTime<Utc>,
     pub disconnected_at: Option<DateTime<Utc>>,
@@ -31,6 +35,8 @@ impl FromRow<'_, PgRow> for ConnectionEvent {
             entity_id: row.get("entity_id"),
             entity_type: row.get("entity_type"),
             ip_address: row.get("ip_address"),
+            session_type: row.get("session_type"),
+            geo_location: row.get("geo_location"),
             success: row.get("success"),
             connected_at: row.get("connected_at"),
             disconnected_at: row.get("disconnected_at"),