use crate::model::*;
use crate::util::Result;
use regex::Regex;
use std::collections::BTreeSet;
use super::util::unquote_ident;
fn parse_privileges(privileges_str: &str, object_type: Option<&str>) -> BTreeSet<Privilege> {
let mut privileges = BTreeSet::new();
for priv_str in privileges_str.split(',') {
let priv_trimmed = priv_str.trim().to_uppercase();
match priv_trimmed.as_str() {
"SELECT" => {
privileges.insert(Privilege::Select);
}
"INSERT" => {
privileges.insert(Privilege::Insert);
}
"UPDATE" => {
privileges.insert(Privilege::Update);
}
"DELETE" => {
privileges.insert(Privilege::Delete);
}
"TRUNCATE" => {
privileges.insert(Privilege::Truncate);
}
"REFERENCES" => {
privileges.insert(Privilege::References);
}
"TRIGGER" => {
privileges.insert(Privilege::Trigger);
}
"USAGE" => {
privileges.insert(Privilege::Usage);
}
"EXECUTE" => {
privileges.insert(Privilege::Execute);
}
"CREATE" => {
privileges.insert(Privilege::Create);
}
"ALL" | "ALL PRIVILEGES" => match object_type {
Some("SCHEMA") => {
privileges.insert(Privilege::Usage);
privileges.insert(Privilege::Create);
}
Some("SEQUENCE") => {
privileges.insert(Privilege::Usage);
privileges.insert(Privilege::Select);
privileges.insert(Privilege::Update);
}
Some("FUNCTION") => {
privileges.insert(Privilege::Execute);
}
Some("TYPE") => {
privileges.insert(Privilege::Usage);
}
_ => {
privileges.insert(Privilege::Select);
privileges.insert(Privilege::Insert);
privileges.insert(Privilege::Update);
privileges.insert(Privilege::Delete);
privileges.insert(Privilege::Truncate);
privileges.insert(Privilege::References);
privileges.insert(Privilege::Trigger);
}
},
_ => continue,
}
}
privileges
}
fn parse_grant_all_in_schema(sql: &str, schema: &mut Schema) {
let re = Regex::new(
r#"(?i)GRANT\s+(.+?)\s+ON\s+ALL\s+(TABLES|SEQUENCES|FUNCTIONS)\s+IN\s+SCHEMA\s+("[^"]+"|\w+)\s+TO\s+("[^"]+"|\w+|PUBLIC)\s*(WITH\s+GRANT\s+OPTION)?\s*;"#
).unwrap();
for cap in re.captures_iter(sql) {
let privileges_str = cap.get(1).unwrap().as_str();
let object_kind = cap.get(2).unwrap().as_str().to_uppercase();
let schema_name = unquote_ident(cap.get(3).unwrap().as_str());
let grantee = unquote_ident(cap.get(4).unwrap().as_str());
let with_grant_option = cap.get(5).is_some();
let inferred_type = match object_kind.as_str() {
"TABLES" => "TABLE",
"SEQUENCES" => "SEQUENCE",
"FUNCTIONS" => "FUNCTION",
_ => continue,
};
let privileges = parse_privileges(privileges_str, Some(inferred_type));
if privileges.is_empty() {
continue;
}
let grant = Grant {
grantee: grantee.to_string(),
privileges,
with_grant_option,
};
match object_kind.as_str() {
"TABLES" => {
let keys: Vec<String> = schema
.tables
.iter()
.filter(|(_, t)| t.schema == schema_name)
.map(|(k, _)| k.clone())
.collect();
for key in keys {
if let Some(table) = schema.tables.get_mut(&key) {
table.grants.push(grant.clone());
}
}
let view_keys: Vec<String> = schema
.views
.iter()
.filter(|(_, v)| v.schema == schema_name)
.map(|(k, _)| k.clone())
.collect();
for key in view_keys {
if let Some(view) = schema.views.get_mut(&key) {
view.grants.push(grant.clone());
}
}
}
"SEQUENCES" => {
let keys: Vec<String> = schema
.sequences
.iter()
.filter(|(_, s)| s.schema == schema_name)
.map(|(k, _)| k.clone())
.collect();
for key in keys {
if let Some(seq) = schema.sequences.get_mut(&key) {
seq.grants.push(grant.clone());
}
}
}
"FUNCTIONS" => {
let keys: Vec<String> = schema
.functions
.iter()
.filter(|(_, f)| f.schema == schema_name)
.map(|(k, _)| k.clone())
.collect();
for key in keys {
if let Some(func) = schema.functions.get_mut(&key) {
func.grants.push(grant.clone());
}
}
}
_ => {}
}
}
}
pub(super) fn parse_grant_statements(sql: &str, schema: &mut Schema) -> Result<()> {
parse_grant_all_in_schema(sql, schema);
let grant_re = Regex::new(
r#"(?i)GRANT\s+(.+?)\s+ON\s+(?:(TABLE|VIEW|SEQUENCE|FUNCTION|SCHEMA|TYPE)\s+)?(.+?)\s+TO\s+("[^"]+"|\w+|PUBLIC)\s*(WITH\s+GRANT\s+OPTION)?"#
).unwrap();
for cap in grant_re.captures_iter(sql) {
let privileges_str = cap.get(1).unwrap().as_str();
let object_type = cap.get(2).map(|m| m.as_str().to_uppercase());
let object_name_raw = cap.get(3).unwrap().as_str();
let grantee = unquote_ident(cap.get(4).unwrap().as_str());
let with_grant_option = cap.get(5).is_some();
if object_name_raw.to_uppercase().starts_with("ALL ") {
continue;
}
let inferred_type = object_type.as_deref().unwrap_or("TABLE");
let privileges = parse_privileges(privileges_str, Some(inferred_type));
if privileges.is_empty() {
continue;
}
let grant = Grant {
grantee: grantee.to_string(),
privileges,
with_grant_option,
};
match inferred_type {
"TABLE" | "VIEW" => {
let (obj_schema, obj_name) = parse_object_name(object_name_raw);
let key = qualified_name(&obj_schema, &obj_name);
if let Some(table) = schema.tables.get_mut(&key) {
table.grants.push(grant);
} else if let Some(view) = schema.views.get_mut(&key) {
view.grants.push(grant);
} else {
let pending_type = if inferred_type == "VIEW" {
PendingGrantObjectType::View
} else {
PendingGrantObjectType::Table
};
schema.pending_grants.push(PendingGrant {
object_type: pending_type,
object_key: key,
grant,
});
}
}
"SEQUENCE" => {
let (obj_schema, obj_name) = parse_object_name(object_name_raw);
let key = qualified_name(&obj_schema, &obj_name);
if let Some(sequence) = schema.sequences.get_mut(&key) {
sequence.grants.push(grant);
} else {
schema.pending_grants.push(PendingGrant {
object_type: PendingGrantObjectType::Sequence,
object_key: key,
grant,
});
}
}
"FUNCTION" => {
let function_key = parse_function_signature(object_name_raw);
if let Some(func) = schema.functions.get_mut(&function_key) {
func.grants.push(grant);
} else {
schema.pending_grants.push(PendingGrant {
object_type: PendingGrantObjectType::Function,
object_key: function_key,
grant,
});
}
}
"SCHEMA" => {
let schema_name = unquote_ident(object_name_raw.trim());
if let Some(pg_schema) = schema.schemas.get_mut(schema_name) {
pg_schema.grants.push(grant);
} else {
schema.pending_grants.push(PendingGrant {
object_type: PendingGrantObjectType::Schema,
object_key: schema_name.to_string(),
grant,
});
}
}
"TYPE" => {
let (obj_schema, obj_name) = parse_object_name(object_name_raw);
let key = qualified_name(&obj_schema, &obj_name);
if let Some(enum_type) = schema.enums.get_mut(&key) {
enum_type.grants.push(grant);
} else if let Some(domain) = schema.domains.get_mut(&key) {
domain.grants.push(grant);
} else {
schema.pending_grants.push(PendingGrant {
object_type: PendingGrantObjectType::Enum,
object_key: key,
grant,
});
}
}
_ => {}
}
}
Ok(())
}
pub(super) fn parse_revoke_statements(sql: &str, schema: &mut Schema) -> Result<()> {
let revoke_re = Regex::new(
r#"(?i)REVOKE\s+(GRANT\s+OPTION\s+FOR\s+)?(.+?)\s+ON\s+(?:(TABLE|VIEW|SEQUENCE|FUNCTION|SCHEMA|TYPE)\s+)?(.+?)\s+FROM\s+("[^"]+"|\w+|PUBLIC)\s*;"#
).unwrap();
for cap in revoke_re.captures_iter(sql) {
let grant_option_for = cap.get(1).is_some();
let privileges_str = cap.get(2).unwrap().as_str();
let object_type = cap.get(3).map(|m| m.as_str().to_uppercase());
let object_name_raw = cap.get(4).unwrap().as_str();
let grantee = unquote_ident(cap.get(5).unwrap().as_str());
if object_name_raw.to_uppercase().starts_with("ALL ") {
continue;
}
let inferred_type = object_type.as_deref().unwrap_or("TABLE");
let privileges = parse_privileges(privileges_str, Some(inferred_type));
if privileges.is_empty() {
continue;
}
match inferred_type {
"TABLE" | "VIEW" => {
let (obj_schema, obj_name) = parse_object_name(object_name_raw);
let key = qualified_name(&obj_schema, &obj_name);
if let Some(table) = schema.tables.get_mut(&key) {
revoke_from_grants(&mut table.grants, grantee, &privileges, grant_option_for);
} else if let Some(view) = schema.views.get_mut(&key) {
revoke_from_grants(&mut view.grants, grantee, &privileges, grant_option_for);
} else {
let pending_type = if inferred_type == "VIEW" {
PendingGrantObjectType::View
} else {
PendingGrantObjectType::Table
};
schema.pending_revokes.push(PendingRevoke {
object_type: pending_type,
object_key: key,
grantee: grantee.to_string(),
privileges,
grant_option_for,
});
}
}
"SEQUENCE" => {
let (obj_schema, obj_name) = parse_object_name(object_name_raw);
let key = qualified_name(&obj_schema, &obj_name);
if let Some(sequence) = schema.sequences.get_mut(&key) {
revoke_from_grants(
&mut sequence.grants,
grantee,
&privileges,
grant_option_for,
);
} else {
schema.pending_revokes.push(PendingRevoke {
object_type: PendingGrantObjectType::Sequence,
object_key: key,
grantee: grantee.to_string(),
privileges,
grant_option_for,
});
}
}
"FUNCTION" => {
let function_key = parse_function_signature(object_name_raw);
if let Some(func) = schema.functions.get_mut(&function_key) {
revoke_from_grants(&mut func.grants, grantee, &privileges, grant_option_for);
} else {
schema.pending_revokes.push(PendingRevoke {
object_type: PendingGrantObjectType::Function,
object_key: function_key,
grantee: grantee.to_string(),
privileges,
grant_option_for,
});
}
}
"SCHEMA" => {
let schema_name = unquote_ident(object_name_raw.trim());
if let Some(pg_schema) = schema.schemas.get_mut(schema_name) {
revoke_from_grants(
&mut pg_schema.grants,
grantee,
&privileges,
grant_option_for,
);
} else {
schema.pending_revokes.push(PendingRevoke {
object_type: PendingGrantObjectType::Schema,
object_key: schema_name.to_string(),
grantee: grantee.to_string(),
privileges,
grant_option_for,
});
}
}
"TYPE" => {
let (obj_schema, obj_name) = parse_object_name(object_name_raw);
let key = qualified_name(&obj_schema, &obj_name);
if let Some(enum_type) = schema.enums.get_mut(&key) {
revoke_from_grants(
&mut enum_type.grants,
grantee,
&privileges,
grant_option_for,
);
} else if let Some(domain) = schema.domains.get_mut(&key) {
revoke_from_grants(&mut domain.grants, grantee, &privileges, grant_option_for);
} else {
schema.pending_revokes.push(PendingRevoke {
object_type: PendingGrantObjectType::Enum,
object_key: key,
grantee: grantee.to_string(),
privileges,
grant_option_for,
});
}
}
_ => {}
}
}
Ok(())
}
pub(super) fn parse_alter_default_privileges(sql: &str, schema: &mut Schema) -> Result<()> {
let grant_re = Regex::new(
r"(?is)ALTER\s+DEFAULT\s+PRIVILEGES\s+(?:FOR\s+ROLE\s+(\w+)\s+)?(?:IN\s+SCHEMA\s+(\w+)\s+)?GRANT\s+(.+?)\s+ON\s+(TABLES|SEQUENCES|FUNCTIONS|ROUTINES|TYPES|SCHEMAS)\s+TO\s+(\w+|PUBLIC)\s*(WITH\s+GRANT\s+OPTION)?\s*;"
).unwrap();
for cap in grant_re.captures_iter(sql) {
let target_role = cap
.get(1)
.map(|m| m.as_str().to_string())
.unwrap_or_else(|| "CURRENT_ROLE".to_string());
let schema_scope = cap.get(2).map(|m| m.as_str().to_string());
let privileges_str = cap.get(3).unwrap().as_str();
let object_type_str = cap.get(4).unwrap().as_str().to_uppercase();
let grantee = cap.get(5).unwrap().as_str().to_string();
let with_grant_option = cap.get(6).is_some();
let object_type = match DefaultPrivilegeObjectType::from_sql_str(&object_type_str) {
Some(ot) => ot,
None => continue,
};
let mut privileges = BTreeSet::new();
let priv_upper = privileges_str.to_uppercase();
if priv_upper.contains("ALL") {
match object_type {
DefaultPrivilegeObjectType::Tables => {
privileges.insert(Privilege::Select);
privileges.insert(Privilege::Insert);
privileges.insert(Privilege::Update);
privileges.insert(Privilege::Delete);
privileges.insert(Privilege::Truncate);
privileges.insert(Privilege::References);
privileges.insert(Privilege::Trigger);
}
DefaultPrivilegeObjectType::Sequences => {
privileges.insert(Privilege::Usage);
privileges.insert(Privilege::Select);
privileges.insert(Privilege::Update);
}
DefaultPrivilegeObjectType::Functions | DefaultPrivilegeObjectType::Routines => {
privileges.insert(Privilege::Execute);
}
DefaultPrivilegeObjectType::Types => {
privileges.insert(Privilege::Usage);
}
DefaultPrivilegeObjectType::Schemas => {
privileges.insert(Privilege::Usage);
privileges.insert(Privilege::Create);
}
}
} else {
for priv_str in privileges_str.split(',') {
let priv_trimmed = priv_str.trim().to_uppercase();
match priv_trimmed.as_str() {
"SELECT" => privileges.insert(Privilege::Select),
"INSERT" => privileges.insert(Privilege::Insert),
"UPDATE" => privileges.insert(Privilege::Update),
"DELETE" => privileges.insert(Privilege::Delete),
"TRUNCATE" => privileges.insert(Privilege::Truncate),
"REFERENCES" => privileges.insert(Privilege::References),
"TRIGGER" => privileges.insert(Privilege::Trigger),
"USAGE" => privileges.insert(Privilege::Usage),
"EXECUTE" => privileges.insert(Privilege::Execute),
"CREATE" => privileges.insert(Privilege::Create),
_ => continue,
};
}
}
if privileges.is_empty() {
continue;
}
schema.default_privileges.push(DefaultPrivilege {
target_role,
schema: schema_scope,
object_type,
grantee,
privileges,
with_grant_option,
});
}
let revoke_re = Regex::new(
r"(?is)ALTER\s+DEFAULT\s+PRIVILEGES\s+(?:FOR\s+ROLE\s+(\w+)\s+)?(?:IN\s+SCHEMA\s+(\w+)\s+)?REVOKE\s+(.+?)\s+ON\s+(TABLES|SEQUENCES|FUNCTIONS|ROUTINES|TYPES|SCHEMAS)\s+FROM\s+(\w+|PUBLIC)\s*;"
).unwrap();
for cap in revoke_re.captures_iter(sql) {
let target_role = cap
.get(1)
.map(|m| m.as_str().to_string())
.unwrap_or_else(|| "CURRENT_ROLE".to_string());
let schema_scope = cap.get(2).map(|m| m.as_str().to_string());
let privileges_str = cap.get(3).unwrap().as_str();
let object_type_str = cap.get(4).unwrap().as_str().to_uppercase();
let grantee = cap.get(5).unwrap().as_str().to_string();
let object_type = match DefaultPrivilegeObjectType::from_sql_str(&object_type_str) {
Some(ot) => ot,
None => continue,
};
let mut privs_to_revoke = BTreeSet::new();
let priv_upper = privileges_str.to_uppercase();
if priv_upper.contains("ALL") {
match object_type {
DefaultPrivilegeObjectType::Tables => {
privs_to_revoke.insert(Privilege::Select);
privs_to_revoke.insert(Privilege::Insert);
privs_to_revoke.insert(Privilege::Update);
privs_to_revoke.insert(Privilege::Delete);
privs_to_revoke.insert(Privilege::Truncate);
privs_to_revoke.insert(Privilege::References);
privs_to_revoke.insert(Privilege::Trigger);
}
DefaultPrivilegeObjectType::Sequences => {
privs_to_revoke.insert(Privilege::Usage);
privs_to_revoke.insert(Privilege::Select);
privs_to_revoke.insert(Privilege::Update);
}
DefaultPrivilegeObjectType::Functions | DefaultPrivilegeObjectType::Routines => {
privs_to_revoke.insert(Privilege::Execute);
}
DefaultPrivilegeObjectType::Types => {
privs_to_revoke.insert(Privilege::Usage);
}
DefaultPrivilegeObjectType::Schemas => {
privs_to_revoke.insert(Privilege::Usage);
privs_to_revoke.insert(Privilege::Create);
}
}
} else {
for priv_str in privileges_str.split(',') {
let priv_trimmed = priv_str.trim().to_uppercase();
match priv_trimmed.as_str() {
"SELECT" => privs_to_revoke.insert(Privilege::Select),
"INSERT" => privs_to_revoke.insert(Privilege::Insert),
"UPDATE" => privs_to_revoke.insert(Privilege::Update),
"DELETE" => privs_to_revoke.insert(Privilege::Delete),
"TRUNCATE" => privs_to_revoke.insert(Privilege::Truncate),
"REFERENCES" => privs_to_revoke.insert(Privilege::References),
"TRIGGER" => privs_to_revoke.insert(Privilege::Trigger),
"USAGE" => privs_to_revoke.insert(Privilege::Usage),
"EXECUTE" => privs_to_revoke.insert(Privilege::Execute),
"CREATE" => privs_to_revoke.insert(Privilege::Create),
_ => continue,
};
}
}
for dp in &mut schema.default_privileges {
if dp.target_role == target_role
&& dp.schema == schema_scope
&& dp.object_type == object_type
&& dp.grantee == grantee
{
for privilege in &privs_to_revoke {
dp.privileges.remove(privilege);
}
}
}
schema
.default_privileges
.retain(|dp| !dp.privileges.is_empty());
}
Ok(())
}
fn parse_object_name(name: &str) -> (String, String) {
let trimmed = unquote_ident(name.trim());
match trimmed.split_once('.') {
Some((schema, obj)) => (
unquote_ident(schema).to_string(),
unquote_ident(obj).to_string(),
),
None => ("public".to_string(), trimmed.to_string()),
}
}
fn parse_function_signature(sig: &str) -> String {
let trimmed = sig.trim();
if let Some(paren_pos) = trimmed.find('(') {
let before_paren = &trimmed[..paren_pos];
let args_part = &trimmed[paren_pos..];
if let Some(dot_pos) = before_paren.rfind('.') {
let schema_part = unquote_ident(&before_paren[..dot_pos]);
let func_name = unquote_ident(&before_paren[dot_pos + 1..]);
format!("{schema_part}.{func_name}{args_part}")
} else {
let name = unquote_ident(before_paren);
format!("public.{name}{args_part}")
}
} else if trimmed.contains('.') {
trimmed.to_string()
} else {
format!("public.{}", unquote_ident(trimmed))
}
}