ok

Author: juleswritescode

crates/pgt_completions/src/providers/roles.rs

@@ -27,7 +27,7 @@ pub fn complete_roles<'a>(ctx: &CompletionContext<'a>, builder: &mut CompletionB
 
 #[cfg(test)]
 mod tests {
-    use sqlx::PgPool;
+    use sqlx::{Executor, PgPool};
 
     use crate::test_helper::{CURSOR_POS, CompletionAssertion, assert_complete_results};
 
@@ -81,9 +81,45 @@ mod tests {
         .await;
     }
 
-    async fn works_in_set_statement() {
-        // set role ROLE;
-        // set session authorization ROLE;
+    #[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")]
+    async fn works_in_set_statement(pool: PgPool) {
+        pool.execute(SETUP).await.unwrap();
+
+        assert_complete_results(
+            format!("set role {}", CURSOR_POS).as_str(),
+            vec![
+                CompletionAssertion::LabelAndKind("admin".into(), crate::CompletionItemKind::Role),
+                CompletionAssertion::LabelAndKind(
+                    "test_login".into(),
+                    crate::CompletionItemKind::Role,
+                ),
+                CompletionAssertion::LabelAndKind(
+                    "test_nologin".into(),
+                    crate::CompletionItemKind::Role,
+                ),
+            ],
+            None,
+            &pool,
+        )
+        .await;
+
+        assert_complete_results(
+            format!("set session authorization {}", CURSOR_POS).as_str(),
+            vec![
+                CompletionAssertion::LabelAndKind("admin".into(), crate::CompletionItemKind::Role),
+                CompletionAssertion::LabelAndKind(
+                    "test_login".into(),
+                    crate::CompletionItemKind::Role,
+                ),
+                CompletionAssertion::LabelAndKind(
+                    "test_nologin".into(),
+                    crate::CompletionItemKind::Role,
+                ),
+            ],
+            None,
+            &pool,
+        )
+        .await;
     }
 
     async fn works_in_policies() {}

crates/pgt_completions/src/relevance/filtering.rs

@@ -76,7 +76,8 @@ impl CompletionFilter<'_> {
                     CompletionRelevanceData::Table(_) => match clause {
                         WrappingClause::Select
                         | WrappingClause::Where
-                        | WrappingClause::ColumnDefinitions => false,
+                        | WrappingClause::ColumnDefinitions
+                        | WrappingClause::SetStatement => false,
 
                         WrappingClause::Insert => {
                             ctx.wrapping_node_kind
@@ -101,6 +102,7 @@ impl CompletionFilter<'_> {
                         match clause {
                             WrappingClause::From
                             | WrappingClause::ColumnDefinitions
+                            | WrappingClause::SetStatement
                             | WrappingClause::AlterTable
                             | WrappingClause::DropTable => false,
 
@@ -170,9 +172,12 @@ impl CompletionFilter<'_> {
                         matches!(clause, WrappingClause::PolicyName)
                     }
 
-                    CompletionRelevanceData::Role(_) => {
-                        matches!(clause, WrappingClause::DropRole | WrappingClause::AlterRole)
-                    }
+                    CompletionRelevanceData::Role(_) => match clause {
+                        WrappingClause::DropRole | WrappingClause::AlterRole => true,
+                        WrappingClause::SetStatement => ctx
+                            .before_cursor_matches_kind(&["keyword_role", "keyword_authorization"]),
+                        _ => false,
+                    },
                 }
             })
             .and_then(|is_ok| if is_ok { Some(()) } else { None })