sure

Author: juleswritescode

crates/pgt_completions/src/complete.rs

@@ -5,7 +5,8 @@ use crate::{
     context::CompletionContext,
     item::CompletionItem,
     providers::{
-        complete_columns, complete_functions, complete_policies, complete_schemas, complete_tables,
+        complete_columns, complete_functions, complete_policies, complete_roles, complete_schemas,
+        complete_tables,
     },
     sanitization::SanitizedCompletionParams,
 };
@@ -36,6 +37,7 @@ pub fn complete(params: CompletionParams) -> Vec<CompletionItem> {
     complete_columns(&ctx, &mut builder);
     complete_schemas(&ctx, &mut builder);
     complete_policies(&ctx, &mut builder);
+    complete_roles(&ctx, &mut builder);
 
     builder.finish()
 }

crates/pgt_completions/src/item.rs

@@ -12,6 +12,7 @@ pub enum CompletionItemKind {
     Column,
     Schema,
     Policy,
+    Role,
 }
 
 impl Display for CompletionItemKind {
@@ -22,6 +23,7 @@ impl Display for CompletionItemKind {
             CompletionItemKind::Column => "Column",
             CompletionItemKind::Schema => "Schema",
             CompletionItemKind::Policy => "Policy",
+            CompletionItemKind::Role => "Role",
         };
 
         write!(f, "{txt}")

crates/pgt_completions/src/providers/mod.rs

@@ -2,11 +2,13 @@ mod columns;
 mod functions;
 mod helper;
 mod policies;
+mod roles;
 mod schemas;
 mod tables;
 
 pub use columns::*;
 pub use functions::*;
 pub use policies::*;
+pub use roles::*;
 pub use schemas::*;
 pub use tables::*;

crates/pgt_completions/src/providers/roles.rs

@@ -0,0 +1,26 @@
+use crate::{
+    CompletionItemKind,
+    builder::{CompletionBuilder, PossibleCompletionItem},
+    context::CompletionContext,
+    relevance::{CompletionRelevanceData, filtering::CompletionFilter, scoring::CompletionScore},
+};
+
+pub fn complete_roles<'a>(ctx: &CompletionContext<'a>, builder: &mut CompletionBuilder<'a>) {
+    let available_roles = &ctx.schema_cache.roles;
+
+    for role in available_roles {
+        let relevance = CompletionRelevanceData::Role(role);
+
+        let item = PossibleCompletionItem {
+            label: role.name.chars().take(35).collect::<String>(),
+            score: CompletionScore::from(relevance.clone()),
+            filter: CompletionFilter::from(relevance),
+            description: role.name.clone(),
+            kind: CompletionItemKind::Role,
+            completion_text: None,
+            detail: None,
+        };
+
+        builder.add_item(item);
+    }
+}

crates/pgt_completions/src/relevance.rs

@@ -8,4 +8,5 @@ pub(crate) enum CompletionRelevanceData<'a> {
     Column(&'a pgt_schema_cache::Column),
     Schema(&'a pgt_schema_cache::Schema),
     Policy(&'a pgt_schema_cache::Policy),
+    Role(&'a pgt_schema_cache::Role),
 }

crates/pgt_completions/src/relevance/filtering.rs

@@ -169,6 +169,8 @@ impl CompletionFilter<'_> {
                     CompletionRelevanceData::Policy(_) => {
                         matches!(clause, WrappingClause::PolicyName)
                     }
+
+                    CompletionRelevanceData::Role(_) => false,
                 }
             })
             .and_then(|is_ok| if is_ok { Some(()) } else { None })
@@ -204,8 +206,8 @@ impl CompletionFilter<'_> {
 
             // we should never allow schema suggestions if there already was one.
             CompletionRelevanceData::Schema(_) => false,
-            // no policy comletion if user typed a schema node first.
-            CompletionRelevanceData::Policy(_) => false,
+            // no policy or row completion if user typed a schema node first.
+            CompletionRelevanceData::Policy(_) | CompletionRelevanceData::Role(_) => false,
         };
 
         if !matches {

crates/pgt_completions/src/relevance/scoring.rs

@@ -47,6 +47,7 @@ impl CompletionScore<'_> {
             CompletionRelevanceData::Column(c) => c.name.as_str().to_ascii_lowercase(),
             CompletionRelevanceData::Schema(s) => s.name.as_str().to_ascii_lowercase(),
             CompletionRelevanceData::Policy(p) => p.name.as_str().to_ascii_lowercase(),
+            CompletionRelevanceData::Role(r) => r.name.as_str().to_ascii_lowercase(),
         };
 
         let fz_matcher = SkimMatcherV2::default();
@@ -126,6 +127,8 @@ impl CompletionScore<'_> {
                 WrappingClause::PolicyName => 25,
                 _ => -50,
             },
+
+            CompletionRelevanceData::Role(_) => 0,
         }
     }
 
@@ -160,6 +163,7 @@ impl CompletionScore<'_> {
                 _ => -50,
             },
             CompletionRelevanceData::Policy(_) => 0,
+            CompletionRelevanceData::Role(_) => 0,
         }
     }
 
@@ -178,7 +182,10 @@ impl CompletionScore<'_> {
             Some(n) => n,
         };
 
-        let data_schema = self.get_schema_name();
+        let data_schema = match self.get_schema_name() {
+            Some(s) => s,
+            None => return,
+        };
 
         if schema_name == data_schema {
             self.score += 25;
@@ -187,13 +194,14 @@ impl CompletionScore<'_> {
         }
     }
 
-    fn get_schema_name(&self) -> &str {
+    fn get_schema_name(&self) -> Option<&str> {
         match self.data {
-            CompletionRelevanceData::Function(f) => f.schema.as_str(),
-            CompletionRelevanceData::Table(t) => t.schema.as_str(),
-            CompletionRelevanceData::Column(c) => c.schema_name.as_str(),
-            CompletionRelevanceData::Schema(s) => s.name.as_str(),
-            CompletionRelevanceData::Policy(p) => p.schema_name.as_str(),
+            CompletionRelevanceData::Function(f) => Some(f.schema.as_str()),
+            CompletionRelevanceData::Table(t) => Some(t.schema.as_str()),
+            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,
         }
     }
 
@@ -212,7 +220,10 @@ impl CompletionScore<'_> {
             _ => {}
         }
 
-        let schema = self.get_schema_name().to_string();
+        let schema = match self.get_schema_name() {
+            Some(s) => s.to_string(),
+            None => return,
+        };
         let table_name = match self.get_table_name() {
             Some(t) => t,
             None => return,
@@ -234,7 +245,10 @@ impl CompletionScore<'_> {
     }
 
     fn check_is_user_defined(&mut self) {
-        let schema = self.get_schema_name().to_string();
+        let schema = match self.get_schema_name() {
+            Some(s) => s.to_string(),
+            None => return,
+        };
 
         let system_schemas = ["pg_catalog", "information_schema", "pg_toast"];
 

crates/pgt_lsp/src/handlers/completions.rs

@@ -76,5 +76,6 @@ fn to_lsp_types_completion_item_kind(
         pgt_completions::CompletionItemKind::Column => lsp_types::CompletionItemKind::FIELD,
         pgt_completions::CompletionItemKind::Schema => lsp_types::CompletionItemKind::CLASS,
         pgt_completions::CompletionItemKind::Policy => lsp_types::CompletionItemKind::CONSTANT,
+        pgt_completions::CompletionItemKind::Role => lsp_types::CompletionItemKind::CONSTANT,
     }
 }