so far!

Author: juleswritescode

crates/pgt_completions/src/context/mod.rs

@@ -33,6 +33,9 @@ pub enum WrappingClause<'a> {
     DropTable,
     PolicyName,
     ToRoleAssignment,
+    SetStatement,
+    AlterRole,
+    DropRole,
 }
 
 #[derive(PartialEq, Eq, Hash, Debug, Clone)]
@@ -424,7 +427,7 @@ impl<'a> CompletionContext<'a> {
             }
 
             "where" | "update" | "select" | "delete" | "from" | "join" | "column_definitions"
-            | "drop_table" | "alter_table" => {
+            | "drop_table" | "alter_table" | "alter_role" | "drop_role" | "set_statement" => {
                 self.wrapping_clause_type =
                     self.get_wrapping_clause_from_current_node(current_node, &mut cursor);
             }
@@ -628,7 +631,10 @@ impl<'a> CompletionContext<'a> {
             "delete" => Some(WrappingClause::Delete),
             "from" => Some(WrappingClause::From),
             "drop_table" => Some(WrappingClause::DropTable),
+            "alter_role" => Some(WrappingClause::AlterRole),
+            "drop_role" => Some(WrappingClause::DropRole),
             "alter_table" => Some(WrappingClause::AlterTable),
+            "set_statement" => Some(WrappingClause::SetStatement),
             "column_definitions" => Some(WrappingClause::ColumnDefinitions),
             "insert" => Some(WrappingClause::Insert),
             "join" => {

crates/pgt_completions/src/providers/roles.rs

@@ -24,3 +24,70 @@ pub fn complete_roles<'a>(ctx: &CompletionContext<'a>, builder: &mut CompletionB
         builder.add_item(item);
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use crate::test_helper::{CURSOR_POS, CompletionAssertion, assert_complete_results};
+
+    const SETUP: &'static str = r#"
+            do $$
+            begin
+                if not exists (
+                    select from pg_catalog.pg_roles
+                    where rolname = 'test'
+                ) then
+                    create role test;
+                end if;
+            end $$;
+
+            create table users (
+              id serial primary key,
+              email varchar,
+              address text
+            );
+        "#;
+
+    #[tokio::test]
+    async fn works_in_drop_role() {
+        assert_complete_results(
+            format!("drop role {}", CURSOR_POS).as_str(),
+            vec![CompletionAssertion::LabelAndKind(
+                "test".into(),
+                crate::CompletionItemKind::Role,
+            )],
+            SETUP,
+        )
+        .await;
+    }
+
+    #[tokio::test]
+    async fn works_in_alter_role() {
+        assert_complete_results(
+            format!("alter role {}", CURSOR_POS).as_str(),
+            vec![CompletionAssertion::LabelAndKind(
+                "test".into(),
+                crate::CompletionItemKind::Role,
+            )],
+            SETUP,
+        )
+        .await;
+    }
+
+    async fn works_in_set_statement() {
+        // set role ROLE;
+        // set session authorization ROLE;
+    }
+
+    async fn works_in_policies() {}
+
+    async fn works_in_grant_statements() {
+        // grant select on my_table to ROLE;
+        // grant ROLE to OTHER_ROLE with admin option;
+    }
+
+    async fn works_in_revoke_statements() {
+        // revoke select on my_table from ROLE;
+        // revoke ROLE from OTHER_ROLE;
+        // revoke admin option for ROLE from OTHER_ROLE;
+    }
+}

crates/pgt_completions/src/relevance/filtering.rs

@@ -170,7 +170,9 @@ impl CompletionFilter<'_> {
                         matches!(clause, WrappingClause::PolicyName)
                     }
 
-                    CompletionRelevanceData::Role(_) => false,
+                    CompletionRelevanceData::Role(_) => {
+                        matches!(clause, WrappingClause::DropRole | WrappingClause::AlterRole)
+                    }
                 }
             })
             .and_then(|is_ok| if is_ok { Some(()) } else { None })

crates/pgt_completions/src/relevance/scoring.rs

@@ -128,7 +128,10 @@ impl CompletionScore<'_> {
                 _ => -50,
             },
 
-            CompletionRelevanceData::Role(_) => 0,
+            CompletionRelevanceData::Role(_) => match clause_type {
+                WrappingClause::DropRole | WrappingClause::AlterRole => 25,
+                _ => -50,
+            },
         }
     }
 
@@ -201,7 +204,7 @@ impl CompletionScore<'_> {
             CompletionRelevanceData::Column(c) => Some(c.schema_name.as_str()),
             CompletionRelevanceData::Schema(s) => Some(s.name.as_str()),
             CompletionRelevanceData::Policy(p) => Some(p.schema_name.as_str()),
-            CompletionRelevanceData::Role(p) => None,
+            CompletionRelevanceData::Role(_) => None,
         }
     }
 
@@ -245,6 +248,30 @@ impl CompletionScore<'_> {
     }
 
     fn check_is_user_defined(&mut self) {
+        if let CompletionRelevanceData::Role(r) = self.data {
+            match r.name.as_str() {
+                "pg_read_all_data"
+                | "pg_write_all_data"
+                | "pg_read_all_settings"
+                | "pg_read_all_stats"
+                | "pg_stat_scan_tables"
+                | "pg_monitor"
+                | "pg_database_owner"
+                | "pg_signal_backend"
+                | "pg_read_server_files"
+                | "pg_write_server_files"
+                | "pg_execute_server_program"
+                | "pg_checkpoint"
+                | "pg_maintain"
+                | "pg_use_reserved_connections"
+                | "pg_create_subscription"
+                | "postgres" => self.score -= 20,
+                _ => {}
+            };
+
+            return;
+        }
+
         let schema = match self.get_schema_name() {
             Some(s) => s.to_string(),
             None => return,

crates/pgt_schema_cache/src/roles.rs

@@ -27,8 +27,6 @@ mod tests {
 
     #[tokio::test]
     async fn loads_roles() {
-        let test_db = get_new_test_db().await;
-
         let setup = r#"
             do $$
             begin
@@ -53,6 +51,8 @@ mod tests {
             end $$;
         "#;
 
+        let test_db = get_new_test_db().await;
+
         test_db
             .execute(setup)
             .await