Skip to content

Commit

Permalink
Merge pull request #219 from nightly-labs/grafana_restore
Browse files Browse the repository at this point in the history
Grafana restore
  • Loading branch information
dzlk17 authored Nov 4, 2024
2 parents c245751 + f61d142 commit 809a6d4
Show file tree
Hide file tree
Showing 24 changed files with 226 additions and 51 deletions.
1 change: 1 addition & 0 deletions database/migrations/0002_team.sql
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
CREATE TABLE team(
team_id TEXT NOT NULL UNIQUE,
grafana_id TEXT UNIQUE,
team_name TEXT NOT NULL,
personal BOOLEAN NOT NULL,
subscription subscription,
Expand Down
17 changes: 15 additions & 2 deletions database/src/tables/team/select.rs
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,9 @@ impl Db {
tx: Option<&mut Transaction<'_, sqlx::Postgres>>,
team_id: &String,
) -> Result<Option<Team>, DbError> {
let query = format!("SELECT * FROM {TEAM_TABLE_NAME} WHERE team_id = $1 AND deactivated_at IS NULL");
let query = format!(
"SELECT * FROM {TEAM_TABLE_NAME} WHERE team_id = $1 AND deactivated_at IS NULL"
);
let typed_query = query_as::<_, Team>(&query);

match tx {
Expand Down Expand Up @@ -100,7 +102,9 @@ impl Db {
}

pub async fn get_team_by_admin_id(&self, admin_id: &String) -> Result<Option<Team>, DbError> {
let query = format!("SELECT * FROM {TEAM_TABLE_NAME} WHERE team_admin_id = $1 AND deactivated_at IS NULL");
let query = format!(
"SELECT * FROM {TEAM_TABLE_NAME} WHERE team_admin_id = $1 AND deactivated_at IS NULL"
);
let typed_query = query_as::<_, Team>(&query);

return typed_query
Expand All @@ -109,4 +113,13 @@ impl Db {
.await
.map_err(|e| e.into());
}
pub async fn get_all_teams(&self) -> Result<Vec<Team>, DbError> {
let query = format!("SELECT * FROM {TEAM_TABLE_NAME} WHERE deactivated_at IS NULL");
let typed_query = query_as::<_, Team>(&query);

return typed_query
.fetch_all(&self.connection_pool)
.await
.map_err(|e| e.into());
}
}
6 changes: 5 additions & 1 deletion database/src/tables/team/table_struct.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7,11 +7,12 @@ use sqlx::{

pub const TEAM_TABLE_NAME: &str = "team";
pub const TEAM_KEYS: &str =
"team_id, team_name, personal, subscription, team_admin_id, registration_timestamp, deactivated_at";
"team_id, grafana_id, team_name, personal, subscription, team_admin_id, registration_timestamp, deactivated_at";

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct Team {
pub team_id: String,
pub grafana_id: String,
pub personal: bool,
pub team_name: String,
// Subscription is required to get access to the statistics
Expand All @@ -25,6 +26,9 @@ impl FromRow<'_, PgRow> for Team {
fn from_row(row: &sqlx::postgres::PgRow) -> std::result::Result<Self, sqlx::Error> {
Ok(Team {
team_id: row.get("team_id"),
grafana_id: row
.try_get::<Option<String>, _>("grafana_id")?
.unwrap_or_default(),
team_name: row.get("team_name"),
personal: row.get("personal"),
subscription: row.get("subscription"),
Expand Down
35 changes: 33 additions & 2 deletions database/src/tables/team/update.rs
Original file line number Diff line number Diff line change
Expand Up @@ -16,11 +16,12 @@ impl Db {
team: &Team,
) -> Result<(), DbError> {
let query_body = format!(
"INSERT INTO {TEAM_TABLE_NAME} ({TEAM_KEYS}) VALUES ($1, $2, $3, $4, $5, $6, NULL)"
"INSERT INTO {TEAM_TABLE_NAME} ({TEAM_KEYS}) VALUES ($1, $2, $3, $4, $5, $6, $7, NULL)"
);

let query_result = query(&query_body)
.bind(&team.team_id)
.bind(&team.grafana_id)
.bind(&team.team_name)
.bind(&team.personal)
.bind(&team.subscription)
Expand All @@ -37,11 +38,12 @@ impl Db {

pub async fn create_new_team(&self, team: &Team) -> Result<(), DbError> {
let query_body = format!(
"INSERT INTO {TEAM_TABLE_NAME} ({TEAM_KEYS}) VALUES ($1, $2, $3, $4, $5, $6, NULL)"
"INSERT INTO {TEAM_TABLE_NAME} ({TEAM_KEYS}) VALUES ($1, $2, $3, $4, $5, $6, $7, NULL)"
);

let query_result = query(&query_body)
.bind(&team.team_id)
.bind(&team.grafana_id)
.bind(&team.team_name)
.bind(&team.personal)
.bind(&team.subscription)
Expand Down Expand Up @@ -96,6 +98,34 @@ impl Db {
Err(e) => Err(e).map_err(|e| e.into()),
}
}

pub async fn update_grafana_id(&self, team_id: &str, grafana_id: &str) -> Result<(), DbError> {
let query_body = format!(
"UPDATE {TEAM_TABLE_NAME} SET grafana_id = $1 WHERE team_id = $2 AND deactivated_at IS NULL",
);

let query_result = query(&query_body)
.bind(grafana_id)
.bind(team_id)
.execute(&self.connection_pool)
.await;

match query_result {
Ok(_) => Ok(()),
Err(e) => Err(e).map_err(|e| e.into()),
}
}

pub async fn clear_all_grafana_ids(&self) -> Result<(), DbError> {
let query_body = format!("UPDATE {TEAM_TABLE_NAME} SET grafana_id = NULL",);

let query_result = query(&query_body).execute(&self.connection_pool).await;

match query_result {
Ok(_) => Ok(()),
Err(e) => Err(e).map_err(|e| e.into()),
}
}
}

#[cfg(feature = "cloud_integration_tests")]
Expand All @@ -121,6 +151,7 @@ mod tests {
// Create team and register app
let team = Team {
team_id: "test_team_id".to_string(),
grafana_id: "test_grafana_id".to_string(),
team_name: "test_team_name".to_string(),
personal: false,
subscription: None,
Expand Down
1 change: 1 addition & 0 deletions database/src/tables/test_utils.rs
Original file line number Diff line number Diff line change
Expand Up @@ -97,6 +97,7 @@ pub mod test_utils {

let team = Team {
team_id: team_id.clone(),
grafana_id: "test_grafana_id".to_string(),
team_name: "test_team_name".to_string(),
personal: false,
subscription: None,
Expand Down
5 changes: 3 additions & 2 deletions database/src/tables/user_app_privileges/select.rs
Original file line number Diff line number Diff line change
Expand Up @@ -114,7 +114,7 @@ impl Db {
// 17.10.2024 Hubert: "If this ever breaks, I will write comments."
let query = format!(
"WITH RelevantTeams AS (
SELECT DISTINCT t.team_id, t.team_name, t.personal, t.subscription,
SELECT DISTINCT t.team_id, t.grafana_id, t.team_name, t.personal, t.subscription,
t.registration_timestamp, gu.email AS team_admin_email,
gu.user_id AS team_admin_id, t.deactivated_at,
CASE
Expand All @@ -127,7 +127,7 @@ impl Db {
JOIN {USERS_TABLE_NAME} gu ON t.team_admin_id = gu.user_id
WHERE (t.team_admin_id = $1 OR uap.user_id = $1) AND t.deactivated_at IS NULL
)
SELECT rt.team_id, rt.team_name, rt.personal, rt.subscription, rt.registration_timestamp,
SELECT rt.team_id, rt.grafana_id, rt.team_name, rt.personal, rt.subscription, rt.registration_timestamp,
rt.team_admin_email, rt.team_admin_id, ra.app_id, ra.app_name, ra.whitelisted_domains,
ra.ack_public_keys, ra.registration_timestamp AS app_registration_timestamp,
uap.user_id, uap.privilege_level, uap.creation_timestamp AS privilege_creation_timestamp,
Expand Down Expand Up @@ -163,6 +163,7 @@ impl Db {
for row in rows {
let team = Team {
team_id: row.get("team_id"),
grafana_id: row.get("grafana_id"),
personal: row.get("personal"),
team_name: row.get("team_name"),
subscription: row.get("subscription"),
Expand Down
1 change: 1 addition & 0 deletions database/src/tables/user_app_privileges/update.rs
Original file line number Diff line number Diff line change
Expand Up @@ -371,6 +371,7 @@ mod tests {

let team = Team {
team_id: team_id.clone(),
grafana_id: "test_grafana_id_2".to_string(),
team_name: "test_team_name".to_string(),
personal: false,
subscription: None,
Expand Down
114 changes: 114 additions & 0 deletions server/src/bin/grafana_restore.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,114 @@
use core::panic;
use database::db::Db;
use openapi::apis::configuration::Configuration;
use server::env::{GF_SECURITY_ADMIN_PASSWORD, GF_SECURITY_ADMIN_USER, GRAFANA_BASE_PATH};
use server::http::cloud::grafana_utils::add_user_to_team::handle_grafana_add_user_to_team;
use server::http::cloud::grafana_utils::create_new_app::handle_grafana_create_new_app;
use server::http::cloud::grafana_utils::create_new_team::handle_grafana_create_new_team;
use server::utils::import_template_dashboards;
use std::collections::HashSet;
use std::sync::Arc;

// This script is used to restore the state of the Grafana instance
// Before running this script, clear the contents of the grafana-data folder
#[tokio::main]
async fn main() {
let db = Db::connect_to_the_pool().await;
let mut conf: Configuration = Configuration::new();
conf.base_path = GRAFANA_BASE_PATH().to_string();
conf.basic_auth = Some((
GF_SECURITY_ADMIN_USER().to_string(),
Some(GF_SECURITY_ADMIN_PASSWORD().to_string()),
));

let grafana_client_conf = Arc::new(conf);
// Setup template dashboards
import_template_dashboards(&grafana_client_conf).await;

if let Err(err) = db.clear_all_grafana_ids().await {
panic!("Failed to clear grafana ids in database: {:?}", err);
}

let teams = match db.get_all_teams().await {
Ok(teams) => teams,
Err(e) => {
panic!("Failed to get teams. Error: {:?}", e);
}
};
for team in teams {
// create teams
let grafana_id = match handle_grafana_create_new_team(
&grafana_client_conf,
&team.team_admin_id,
&team.team_name,
)
.await
{
Ok(id) => {
if let Err(err) = db.update_grafana_id(&team.team_id, &id.to_string()).await {
panic!("Failed to update grafana id in database: {:?}", err);
}
id.to_string()
}
Err(err) => {
panic!("Failed to create team in grafana: {:?}", err);
}
};

let privileges = match db.get_privileges_by_team_id(&team.team_id).await {
Ok(privileges) => privileges,
Err(e) => {
panic!("Failed to get privileges. Error: {:?}", e);
}
};
let unique_user_ids: Vec<String> = privileges
.into_iter()
.map(|privilege| privilege.user_id)
.collect::<HashSet<_>>()
.into_iter()
.collect();

let users_emails = match db.get_users_emails_by_ids(&unique_user_ids).await {
Ok(emails) => emails,
Err(e) => {
panic!("Failed to get users emails. Error: {:?}", e);
}
};
for (_, user_email) in users_emails {
// add users to teams
match handle_grafana_add_user_to_team(&grafana_client_conf, &grafana_id, &user_email)
.await
{
Ok(id) => id,
Err(err) => {
panic!("Failed to add user to team in grafana: {:?}", err);
}
};
}
let apps = match db.get_registered_apps_by_team_id(&team.team_id).await {
Ok(apps) => apps,
Err(e) => {
panic!("Failed to get apps. Error: {:?}", e);
}
};
for app in apps {
let app_id = app.app_id.clone();
let app_name = app.app_name.clone();
// create apps
match handle_grafana_create_new_app(
&grafana_client_conf,
&app_id,
&app_name,
&grafana_id,
)
.await
{
Ok(id) => id,
Err(err) => {
panic!("Failed to create app in grafana: {:?}", err);
}
};
}
}
println!("Got it! Exiting...");
}
22 changes: 19 additions & 3 deletions server/src/http/cloud/accept_team_invite.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
use super::{
grafana_utils::add_user_to_team::handle_grafana_add_user_to_team,
utils::{custom_validate_team_id, validate_request},
utils::{custom_validate_uuid, validate_request},
};
use crate::{
env::is_env_production, middlewares::auth_middleware::UserId,
Expand All @@ -19,7 +19,7 @@ use ts_rs::TS;
#[ts(export)]
#[serde(rename_all = "camelCase")]
pub struct HttpAcceptTeamInviteRequest {
#[garde(custom(custom_validate_team_id))]
#[garde(custom(custom_validate_uuid))]
pub team_id: String,
}

Expand Down Expand Up @@ -101,10 +101,26 @@ pub async fn accept_team_invite(
}
}

let grafana_team_id = match db.get_team_by_team_id(None, &request.team_id).await {
Ok(Some(team)) => team.grafana_id,
Ok(None) => {
return Err((
StatusCode::BAD_REQUEST,
CloudApiErrors::TeamDoesNotExist.to_string(),
));
}
Err(err) => {
error!("Failed to get team by team_id: {:?}", err);
return Err((
StatusCode::INTERNAL_SERVER_ERROR,
CloudApiErrors::DatabaseError.to_string(),
));
}
};
// Grafana add user to the team
if is_env_production() {
if let Err(err) =
handle_grafana_add_user_to_team(&grafana_conf, &request.team_id, &user.email).await
handle_grafana_add_user_to_team(&grafana_conf, &grafana_team_id, &user.email).await
{
error!("Failed to add user to the team in grafana: {:?}", err);
return Err((
Expand Down
4 changes: 2 additions & 2 deletions server/src/http/cloud/cancel_team_user_invite.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
use super::utils::{custom_validate_team_id, validate_request};
use super::utils::{custom_validate_uuid, validate_request};
use crate::{
middlewares::auth_middleware::UserId,
structs::cloud::{api_cloud_errors::CloudApiErrors, team_invite::TeamInvite},
Expand All @@ -15,7 +15,7 @@ use ts_rs::TS;
#[ts(export)]
#[serde(rename_all = "camelCase")]
pub struct HttpCancelTeamUserInviteRequest {
#[garde(custom(custom_validate_team_id))]
#[garde(custom(custom_validate_uuid))]
pub team_id: String,
#[garde(email)]
pub user_email: String,
Expand Down
4 changes: 2 additions & 2 deletions server/src/http/cloud/cancel_user_team_invite.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
use super::utils::{custom_validate_team_id, validate_request};
use super::utils::{custom_validate_uuid, validate_request};
use crate::{
middlewares::auth_middleware::UserId, structs::cloud::api_cloud_errors::CloudApiErrors,
};
Expand All @@ -14,7 +14,7 @@ use ts_rs::TS;
#[ts(export)]
#[serde(rename_all = "camelCase")]
pub struct HttpCancelUserTeamInviteRequest {
#[garde(custom(custom_validate_team_id))]
#[garde(custom(custom_validate_uuid))]
pub team_id: String,
}

Expand Down
4 changes: 2 additions & 2 deletions server/src/http/cloud/change_user_privileges.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
use super::utils::{custom_validate_team_id, custom_validate_uuid, validate_request};
use super::utils::{custom_validate_uuid, validate_request};
use crate::{
middlewares::auth_middleware::UserId,
structs::cloud::{
Expand Down Expand Up @@ -36,7 +36,7 @@ pub struct PrivilegeChange {
#[ts(export)]
#[serde(rename_all = "camelCase")]
pub struct HttpChangeUsersPrivilegesRequest {
#[garde(custom(custom_validate_team_id))]
#[garde(custom(custom_validate_uuid))]
pub team_id: String,
#[garde(dive)]
pub privileges_changes: Vec<PrivilegeChange>,
Expand Down
Loading

0 comments on commit 809a6d4

Please sign in to comment.