use chrono::{DateTime, Utc};
use mz_audit_log::ObjectType;
use mz_sql::catalog::EnvironmentId;
use mz_sql_parser::ast::StatementKind;
use serde::{Deserialize, Serialize};
use serde_json::json;
use uuid::Uuid;
#[derive(Debug, Clone, Default)]
pub struct EventDetails<'a> {
pub user_id: Option<Uuid>,
pub application_name: Option<&'a str>,
pub timestamp: Option<DateTime<Utc>>,
}
pub trait SegmentClientExt {
fn environment_track<S>(
&self,
environment_id: &EnvironmentId,
event: S,
properties: serde_json::Value,
details: EventDetails<'_>,
) where
S: Into<String>;
}
impl SegmentClientExt for mz_segment::Client {
fn environment_track<S>(
&self,
environment_id: &EnvironmentId,
event: S,
mut properties: serde_json::Value,
EventDetails {
user_id,
timestamp,
application_name,
}: EventDetails<'_>,
) where
S: Into<String>,
{
{
let properties = match &mut properties {
serde_json::Value::Object(map) => map,
_ => {
panic!("SegmentClientExt::environment_track called with non-object properties")
}
};
properties.insert("event_source".into(), json!("environment"));
properties.insert(
"cloud_provider".into(),
json!(environment_id.cloud_provider().to_string()),
);
properties.insert(
"cloud_provider_region".into(),
json!(environment_id.cloud_provider_region()),
);
if let Some(application_name) = application_name {
properties.insert("application_name".into(), json!(application_name));
}
}
let context = json!({
"group_id": environment_id.organization_id()
});
let user_id = user_id.unwrap_or_else(|| environment_id.organization_id());
self.track(user_id, event, properties, Some(context), timestamp);
}
}
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(rename_all = "kebab-case")]
pub enum StatementFailureType {
ParseFailure,
}
impl StatementFailureType {
pub fn as_title_case(&self) -> &'static str {
match self {
StatementFailureType::ParseFailure => "Parse Failed",
}
}
}
serde_plain::derive_display_from_serialize!(StatementFailureType);
#[derive(Serialize, Deserialize)]
pub enum StatementAction {
Alter,
Create,
}
impl StatementAction {
pub fn as_title_case(&self) -> &'static str {
match self {
StatementAction::Alter => "Alter",
StatementAction::Create => "Create",
}
}
}
pub fn analyze_audited_statement(
statement: StatementKind,
) -> Option<(StatementAction, ObjectType)> {
match statement {
StatementKind::AlterConnection => Some((StatementAction::Alter, ObjectType::Connection)),
StatementKind::AlterIndex => Some((StatementAction::Alter, ObjectType::Index)),
StatementKind::AlterRole => Some((StatementAction::Alter, ObjectType::Role)),
StatementKind::AlterSecret => Some((StatementAction::Alter, ObjectType::Secret)),
StatementKind::AlterSource => Some((StatementAction::Alter, ObjectType::Source)),
StatementKind::CreateCluster => Some((StatementAction::Create, ObjectType::Cluster)),
StatementKind::CreateClusterReplica => {
Some((StatementAction::Create, ObjectType::ClusterReplica))
}
StatementKind::CreateConnection => Some((StatementAction::Create, ObjectType::Connection)),
StatementKind::CreateDatabase => Some((StatementAction::Create, ObjectType::Database)),
StatementKind::CreateIndex => Some((StatementAction::Create, ObjectType::Index)),
StatementKind::CreateMaterializedView => {
Some((StatementAction::Create, ObjectType::MaterializedView))
}
StatementKind::CreateRole => Some((StatementAction::Create, ObjectType::Role)),
StatementKind::CreateSchema => Some((StatementAction::Create, ObjectType::Schema)),
StatementKind::CreateSecret => Some((StatementAction::Create, ObjectType::Secret)),
StatementKind::CreateSink => Some((StatementAction::Create, ObjectType::Sink)),
StatementKind::CreateSource => Some((StatementAction::Create, ObjectType::Source)),
StatementKind::CreateTable => Some((StatementAction::Create, ObjectType::Table)),
StatementKind::CreateView => Some((StatementAction::Create, ObjectType::View)),
_ => None,
}
}