laying the infra…

Author: juleswritescode

crates/pgt_completions/src/item.rs

@@ -11,6 +11,7 @@ pub enum CompletionItemKind {
     Function,
     Column,
     Schema,
+    Policy,
 }
 
 impl Display for CompletionItemKind {
@@ -20,6 +21,7 @@ impl Display for CompletionItemKind {
             CompletionItemKind::Function => "Function",
             CompletionItemKind::Column => "Column",
             CompletionItemKind::Schema => "Schema",
+            CompletionItemKind::Policy => "Policy",
         };
 
         write!(f, "{txt}")

crates/pgt_completions/src/providers/mod.rs

@@ -1,10 +1,12 @@
 mod columns;
 mod functions;
 mod helper;
+mod policies;
 mod schemas;
 mod tables;
 
 pub use columns::*;
 pub use functions::*;
+pub use policies::*;
 pub use schemas::*;
 pub use tables::*;

crates/pgt_completions/src/providers/policies.rs

@@ -0,0 +1,35 @@
+use crate::{
+    CompletionItemKind,
+    builder::{CompletionBuilder, PossibleCompletionItem},
+    context::{CompletionContext, WrappingClause},
+    relevance::{CompletionRelevanceData, filtering::CompletionFilter, scoring::CompletionScore},
+};
+
+use super::helper::{find_matching_alias_for_table, get_completion_text_with_schema_or_alias};
+
+pub fn complete_policies<'a>(ctx: &CompletionContext<'a>, builder: &mut CompletionBuilder<'a>) {
+    let available_policies = &ctx.schema_cache.policies;
+
+    for pol in available_policies {
+        let relevance = CompletionRelevanceData::Policy(pol);
+
+        let mut item = PossibleCompletionItem {
+            label: pol.name.clone(),
+            score: CompletionScore::from(relevance.clone()),
+            filter: CompletionFilter::from(relevance),
+            description: format!("Table: {}.{}", pol.schema_name, pol.table_name),
+            kind: CompletionItemKind::Column,
+            completion_text: None,
+        };
+
+        // autocomplete with the alias in a join clause if we find one
+        if matches!(ctx.wrapping_clause_type, Some(WrappingClause::Join { .. })) {
+            item.completion_text = find_matching_alias_for_table(ctx, pol.table_name.as_str())
+                .and_then(|alias| {
+                    get_completion_text_with_schema_or_alias(ctx, pol.name.as_str(), alias.as_str())
+                });
+        }
+
+        builder.add_item(item);
+    }
+}

crates/pgt_completions/src/relevance.rs

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

crates/pgt_completions/src/relevance/filtering.rs

@@ -130,6 +130,9 @@ impl CompletionFilter<'_> {
                 // we should never allow schema suggestions if there already was one.
                 false
             }
+
+            // no aliases and schemas for policies
+            CompletionRelevanceData::Policy(_) => false,
         };
 
         if !matches {

crates/pgt_completions/src/relevance/scoring.rs

@@ -45,6 +45,7 @@ impl CompletionScore<'_> {
             CompletionRelevanceData::Table(t) => t.name.as_str(),
             CompletionRelevanceData::Column(c) => c.name.as_str(),
             CompletionRelevanceData::Schema(s) => s.name.as_str(),
+            CompletionRelevanceData::Policy(p) => p.name.as_str(),
         };
 
         let fz_matcher = SkimMatcherV2::default();
@@ -116,6 +117,7 @@ impl CompletionScore<'_> {
                 WrappingClause::Delete if !has_mentioned_schema => 15,
                 _ => -50,
             },
+            CompletionRelevanceData::Policy(_) => 0,
         }
     }
 
@@ -149,6 +151,7 @@ impl CompletionScore<'_> {
                 WrappingNode::Relation if !has_mentioned_schema && has_node_text => 0,
                 _ => -50,
             },
+            CompletionRelevanceData::Policy(_) => 0,
         }
     }
 
@@ -182,13 +185,15 @@ impl CompletionScore<'_> {
             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.name.as_str(),
         }
     }
 
     fn get_table_name(&self) -> Option<&str> {
         match self.data {
             CompletionRelevanceData::Column(c) => Some(c.table_name.as_str()),
             CompletionRelevanceData::Table(t) => Some(t.name.as_str()),
+            CompletionRelevanceData::Policy(p) => Some(p.table_name.as_str()),
             _ => None,
         }
     }

crates/pgt_lsp/src/handlers/completions.rs

@@ -65,5 +65,6 @@ fn to_lsp_types_completion_item_kind(
         pgt_completions::CompletionItemKind::Table => lsp_types::CompletionItemKind::CLASS,
         pgt_completions::CompletionItemKind::Column => lsp_types::CompletionItemKind::FIELD,
         pgt_completions::CompletionItemKind::Schema => lsp_types::CompletionItemKind::CLASS,
+        pgt_completions::CompletionItemKind::Policy => lsp_types::CompletionItemKind::VALUE,
     }
 }

crates/pgt_schema_cache/src/lib.rs

@@ -13,6 +13,7 @@ mod versions;
 
 pub use columns::*;
 pub use functions::{Behavior, Function, FunctionArg, FunctionArgs};
+pub use policies::{Policy, PolicyCommand};
 pub use schema_cache::SchemaCache;
 pub use schemas::Schema;
 pub use tables::{ReplicaIdentity, Table};

crates/pgt_schema_cache/src/policies.rs

@@ -56,14 +56,14 @@ impl From<PolicyQueried> for Policy {
 
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct Policy {
-    name: String,
-    table_name: String,
-    schema_name: String,
-    is_permissive: bool,
-    command: PolicyCommand,
-    role_names: Vec<String>,
-    security_qualification: Option<String>,
-    with_check: Option<String>,
+    pub name: String,
+    pub table_name: String,
+    pub schema_name: String,
+    pub is_permissive: bool,
+    pub command: PolicyCommand,
+    pub role_names: Vec<String>,
+    pub security_qualification: Option<String>,
+    pub with_check: Option<String>,
 }
 
 impl SchemaCacheItem for Policy {