Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Unwraps removed #224

Merged
merged 5 commits into from
Dec 3, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 0 additions & 1 deletion database/src/tables/client_profiles/update.rs
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,6 @@ impl Db {
#[cfg(feature = "cloud_integration_tests")]
#[cfg(test)]
mod tests {
use crate::tables::utils::to_microsecond_precision;
use sqlx::types::chrono::Utc;

#[tokio::test]
Expand Down
6 changes: 4 additions & 2 deletions database/src/tables/requests/table_struct.rs
Original file line number Diff line number Diff line change
Expand Up @@ -26,9 +26,11 @@ impl FromRow<'_, PgRow> for Request {
Ok(Request {
request_id: row.get("request_id"),
app_id: row.get("app_id"),
request_type: RequestType::from_str(row.get("request_type")).unwrap(),
request_type: RequestType::from_str(row.get("request_type"))
.map_err(|_| sqlx::Error::Decode(format!("Invalid request_type")))?,
session_id: row.get("session_id"),
request_status: row.get("request_status"),
request_status: RequestStatus::from_str(row.get("request_status"))
.map_err(|_| sqlx::Error::Decode(format!("Invalid request_status")))?,
network: row.get("network"),
creation_timestamp: row.get("creation_timestamp"),
})
Expand Down
10 changes: 8 additions & 2 deletions database/src/tables/sessions/update.rs
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,10 @@ impl Db {
ip_address: &String,
current_time: &DateTime<Utc>,
) -> Result<(), DbError> {
let mut tx = self.connection_pool.begin().await.unwrap();
let mut tx = match self.connection_pool.begin().await {
Ok(tx) => tx,
Err(err) => return Err(err.into()),
};

// 1. Save the new session
if let Err(err) = self.save_new_session(&mut tx, &session).await {
Expand Down Expand Up @@ -116,7 +119,10 @@ impl Db {
session_id: &String,
) -> Result<(), DbError> {
// Start a new transaction
let mut tx = self.connection_pool.begin().await.unwrap();
let mut tx = match self.connection_pool.begin().await {
Ok(tx) => tx,
Err(err) => return Err(err.into()),
};

// User can't connect to the session without any connected keys
if connected_keys.is_empty() {
Expand Down
10 changes: 7 additions & 3 deletions database/src/tables/team/update.rs
Original file line number Diff line number Diff line change
Expand Up @@ -127,9 +127,11 @@ impl Db {
}
}

pub async fn delete_all_user_teams(&self,
pub async fn delete_all_user_teams(
&self,
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
user_id: &str) -> Result<(), DbError> {
user_id: &str,
) -> Result<(), DbError> {
let query_body = format!(
"UPDATE {TEAM_TABLE_NAME} SET deactivated_at = $1 WHERE team_admin_id = $2 AND deactivated_at IS NULL",
);
Expand All @@ -150,7 +152,9 @@ impl Db {
#[cfg(feature = "cloud_integration_tests")]
#[cfg(test)]
mod tests {
use crate::tables::{team::table_struct::Team, utils::to_microsecond_precision};
use crate::tables::{
team::table_struct::Team, test_utils::test_utils::to_microsecond_precision,
};
use sqlx::types::chrono::Utc;

#[tokio::test]
Expand Down
11 changes: 10 additions & 1 deletion database/src/tables/test_utils.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,9 +6,10 @@ pub mod test_utils {
structs::{db_error::DbError, privilege_level::PrivilegeLevel},
tables::{
registered_app::table_struct::DbRegisteredApp, team::table_struct::Team,
user_app_privileges::table_struct::UserAppPrivilege,
user_app_privileges::table_struct::UserAppPrivilege, utils::get_current_datetime,
},
};
use chrono::TimeZone;
use sqlx::{
types::chrono::{DateTime, Utc},
Row, Transaction,
Expand Down Expand Up @@ -159,4 +160,12 @@ pub mod test_utils {
Ok(())
}
}

pub fn to_microsecond_precision(datetime: &DateTime<Utc>) -> DateTime<Utc> {
// Should never fail as we are converting from a valid DateTime<Utc>
match Utc.timestamp_micros(datetime.timestamp_micros()) {
chrono::LocalResult::Single(dt) => dt,
_ => get_current_datetime(),
}
}
}
3 changes: 2 additions & 1 deletion database/src/tables/user_app_privileges/update.rs
Original file line number Diff line number Diff line change
Expand Up @@ -306,7 +306,8 @@ mod tests {
structs::privilege_level::PrivilegeLevel,
tables::{
registered_app::table_struct::DbRegisteredApp, team::table_struct::Team,
user_app_privileges::table_struct::UserAppPrivilege, utils::to_microsecond_precision,
test_utils::test_utils::to_microsecond_precision,
user_app_privileges::table_struct::UserAppPrivilege,
},
};
use sqlx::types::chrono::Utc;
Expand Down
4 changes: 3 additions & 1 deletion database/src/tables/users/update.rs
Original file line number Diff line number Diff line change
Expand Up @@ -114,7 +114,9 @@ impl Db {
#[cfg(feature = "cloud_integration_tests")]
#[cfg(test)]
mod tests {
use crate::tables::{users::table_struct::User, utils::to_microsecond_precision};
use crate::tables::{
test_utils::test_utils::to_microsecond_precision, users::table_struct::User,
};
use sqlx::types::chrono::Utc;

#[tokio::test]
Expand Down
5 changes: 0 additions & 5 deletions database/src/tables/utils.rs
Original file line number Diff line number Diff line change
Expand Up @@ -15,11 +15,6 @@ pub fn get_current_datetime() -> DateTime<Utc> {
Utc::now()
}

pub fn to_microsecond_precision(datetime: &DateTime<Utc>) -> DateTime<Utc> {
// Should never fail as we are converting from a valid DateTime<Utc>
Utc.timestamp_micros(datetime.timestamp_micros()).unwrap()
}

// This function is used to format the keys of a table to be used in a view query
pub fn format_view_keys(prefix: &str, keys: &[(&'static str, bool)]) -> String {
keys.iter()
Expand Down
1 change: 1 addition & 0 deletions server/src/env.rs
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ pub struct ENV {
pub GF_SECURITY_ADMIN_PASSWORD: String,
pub MAILER_ACTIVE: bool,
}

pub fn get_env() -> &'static ENV {
static INSTANCE: OnceCell<ENV> = OnceCell::new();

Expand Down
13 changes: 2 additions & 11 deletions server/src/http/cloud/accept_team_invite.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@ use super::{
};
use crate::{
env::is_env_production, middlewares::auth_middleware::UserId,
structs::cloud::api_cloud_errors::CloudApiErrors,
structs::cloud::api_cloud_errors::CloudApiErrors, utils::start_transaction,
};
use axum::{extract::State, http::StatusCode, Extension, Json};
use database::db::Db;
Expand Down Expand Up @@ -140,16 +140,7 @@ pub async fn accept_team_invite(
};
}
// Accept invite
let mut tx = match db.connection_pool.begin().await {
Ok(tx) => tx,
Err(err) => {
error!("Failed to start transaction: {:?}", err);
return Err((
StatusCode::INTERNAL_SERVER_ERROR,
CloudApiErrors::DatabaseError.to_string(),
));
}
};
let mut tx = start_transaction(&db).await?;

// Accept invite
if let Err(err) = db
Expand Down
27 changes: 15 additions & 12 deletions server/src/http/cloud/change_user_privileges.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ use crate::{
structs::cloud::{
api_cloud_errors::CloudApiErrors, new_user_privilege_level::NewUserPrivilegeLevel,
},
utils::start_transaction,
};
use axum::{extract::State, http::StatusCode, Extension, Json};
use database::{
Expand Down Expand Up @@ -125,18 +126,22 @@ pub async fn change_user_privileges(
.collect();

// Start transaction to update users privileges
let mut tx = db.connection_pool.begin().await.map_err(|err| {
error!("Failed to start transaction: {:?}", err);
(
StatusCode::INTERNAL_SERVER_ERROR,
CloudApiErrors::DatabaseError.to_string(),
)
})?;
let mut tx = start_transaction(&db).await?;

// Update users privileges
for requested_change in request.privileges_changes {
// Determine action
let new_privilege_level = requested_change.new_privilege_level.to_privilege_level();
let new_privilege_level =
match requested_change.new_privilege_level.to_privilege_level() {
Some(privilege) => privilege,
None => {
error!("Failed to convert new privilege level");
return Err((
StatusCode::BAD_REQUEST,
CloudApiErrors::InternalServerError.to_string(),
));
}
};
let user_id = user_ids.get(&requested_change.user_email).ok_or((
StatusCode::BAD_REQUEST,
CloudApiErrors::UserDoesNotExist.to_string(),
Expand All @@ -158,8 +163,7 @@ pub async fn change_user_privileges(
&mut tx,
user_id,
&requested_change.app_id,
// Safe unwrap
new_privilege_level.unwrap(),
new_privilege_level,
)
.await
.map_err(|err| {
Expand Down Expand Up @@ -212,8 +216,7 @@ pub async fn change_user_privileges(
&mut tx,
user_id,
&requested_change.app_id,
// Safe unwrap
new_privilege_level.unwrap(),
new_privilege_level,
)
.await
.map_err(|err| {
Expand Down
9 changes: 2 additions & 7 deletions server/src/http/cloud/delete_account_finish.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ use crate::{
cloud::api_cloud_errors::CloudApiErrors,
session_cache::{ApiSessionsCache, SessionCache, SessionsCacheKey},
},
utils::start_transaction,
};
use axum::{extract::State, http::StatusCode, Extension, Json};
use database::db::Db;
Expand Down Expand Up @@ -82,13 +83,7 @@ pub async fn delete_account_finish(
sessions_cache.remove(&sessions_key);

// Start transaction to update users privileges
let mut tx = db.connection_pool.begin().await.map_err(|err| {
error!("Failed to start transaction: {:?}", err);
(
StatusCode::INTERNAL_SERVER_ERROR,
CloudApiErrors::DatabaseError.to_string(),
)
})?;
let mut tx = start_transaction(&db).await?;

let mut owned_team_grafana_ids = Vec::new();
let mut non_owned_team_grafana_ids = Vec::new();
Expand Down
3 changes: 2 additions & 1 deletion server/src/http/cloud/delete_app.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ use crate::{
env::is_env_production,
http::cloud::grafana_utils::delete_registered_app::handle_grafana_delete_app,
middlewares::auth_middleware::UserId, structs::cloud::api_cloud_errors::CloudApiErrors,
utils::start_transaction,
};
use axum::{extract::State, http::StatusCode, Extension, Json};
use database::{db::Db, structs::privilege_level::PrivilegeLevel};
Expand Down Expand Up @@ -70,7 +71,7 @@ pub async fn delete_app(
}
// Delete the app
// Start a transaction
let mut tx = db.connection_pool.begin().await.unwrap();
let mut tx = start_transaction(&db).await?;

if let Err(err) = db
.remove_privileges_for_inactive_app_within_tx(&mut tx, &request.app_id)
Expand Down
3 changes: 2 additions & 1 deletion server/src/http/cloud/delete_team.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ use crate::{
http::cloud::grafana_utils::delete_team::handle_grafana_delete_team,
middlewares::auth_middleware::UserId,
structs::cloud::{api_cloud_errors::CloudApiErrors, app_info::AppInfo},
utils::start_transaction,
};
use axum::{extract::State, http::StatusCode, Extension, Json};
use database::db::Db;
Expand Down Expand Up @@ -32,7 +33,7 @@ pub async fn delete_team(
validate_request(&request, &())?;
warn!("Delete team request: {:?}", request);
// Start a transaction
let mut tx = db.connection_pool.begin().await.unwrap();
let mut tx = start_transaction(&db).await?;

// First check if team exists
let team = match db.get_team_by_team_id(None, &request.team_id).await {
Expand Down
3 changes: 2 additions & 1 deletion server/src/http/cloud/domains/remove_whitelisted_domain.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ use crate::{
http::cloud::utils::{custom_validate_domain_name, custom_validate_uuid},
middlewares::auth_middleware::UserId,
structs::cloud::api_cloud_errors::CloudApiErrors,
utils::start_transaction,
};
use axum::{extract::State, http::StatusCode, Extension, Json};
use database::{db::Db, structs::privilege_level::PrivilegeLevel};
Expand Down Expand Up @@ -105,7 +106,7 @@ pub async fn remove_whitelisted_domain(
));
}

let mut tx = db.connection_pool.begin().await.unwrap();
let mut tx = start_transaction(&db).await?;

// Remove domain from whitelisted domains
if let Err(err) = db
Expand Down
8 changes: 6 additions & 2 deletions server/src/http/cloud/domains/verify_domain_finish.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ use crate::{
http::cloud::utils::{custom_validate_domain_name, custom_validate_uuid},
middlewares::auth_middleware::UserId,
structs::cloud::api_cloud_errors::CloudApiErrors,
utils::start_transaction,
};
use anyhow::bail;
use axum::{extract::State, http::StatusCode, Extension, Json};
Expand Down Expand Up @@ -183,7 +184,7 @@ pub async fn verify_domain_finish(
}

// Add domain to whitelist
let mut tx = db.connection_pool.begin().await.unwrap();
let mut tx = start_transaction(&db).await?;

if let Err(err) = db
.add_new_whitelisted_domain(&mut tx, &request.app_id, &domain_name)
Expand Down Expand Up @@ -249,7 +250,10 @@ async fn check_verification_code(
let txt_data = txt.txt_data();
// Each TXT record can contain multiple strings, so we iterate through them all
for txt_str in txt_data {
let txt_str = std::str::from_utf8(txt_str).unwrap();
let txt_str = match std::str::from_utf8(txt_str) {
Ok(txt_str) => txt_str,
Err(err) => bail!("Failed to parse TXT record: {:?}", err),
};
// Check if the verification code is present
if txt_str.contains(&code) {
return Ok(());
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -91,7 +91,16 @@ pub async fn process_event_app_connect(
}
} else {
// Reconnection to an existing session
let mut tx = db.connection_pool.begin().await.unwrap();
let mut tx = match db.connection_pool.begin().await {
Ok(tx) => tx,
Err(err) => {
error!(
"Failed to create new transaction to save app connection event, app_id: [{}], event: [{:?}], err: [{}]",
app_id, event, err
);
return;
}
};

// Get the geolocation data
let geo_location_data = get_geolocation_data(&db, &geo_loc_requester, &ip).await;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,16 @@ pub async fn process_event_client_connect_init(
save_event_client_connect(db, app_id, network, event, &current_time).await;

// Save connection attempt by client
let mut tx = db.connection_pool.begin().await.unwrap();
let mut tx = match db.connection_pool.begin().await {
Ok(tx) => tx,
Err(err) => {
error!(
"Failed to create new transaction to save client connect event, app_id: [{}], event: [{:?}], err: [{}]",
app_id, event, err
);
return;
}
};

// Get the geolocation data
let geo_location_data = get_geolocation_data(&db, &geo_loc_requester, &ip).await;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,16 @@ pub async fn process_event_client_disconnect(
save_event_client_disconnect(db, app_id, network, event, &current_timestamp).await;

// Update connection status for user
let mut tx = db.connection_pool.begin().await.unwrap();
let mut tx = match db.connection_pool.begin().await {
Ok(tx) => tx,
Err(err) => {
error!(
"Failed to create new transaction to update client disconnect status, app_id: [{}], event: [{:?}], err: [{}]",
app_id, event, err
);
return;
}
};

match db
.close_client_connection(&mut tx, &app_id, &event.disconnected_session_id)
Expand Down
4 changes: 2 additions & 2 deletions server/src/http/cloud/register_new_app.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@ use super::{
use crate::{
env::is_env_production, middlewares::auth_middleware::UserId,
statics::REGISTERED_APPS_LIMIT_PER_TEAM, structs::cloud::api_cloud_errors::CloudApiErrors,
test_env::is_test_env,
utils::start_transaction,
};
use axum::{extract::State, http::StatusCode, Extension, Json};
use database::{
Expand Down Expand Up @@ -140,7 +140,7 @@ pub async fn register_new_app(

// Register a new app under this team
// Start a transaction
let mut tx = db.connection_pool.begin().await.unwrap();
let mut tx = start_transaction(&db).await?;

// Register a new app
let db_registered_app =
Expand Down
Loading
Loading