maybe

Author: juleswritescode

crates/pg_completions/Cargo.toml

@@ -16,6 +16,8 @@ async-std = "1.12.0"
 
 text-size.workspace = true
 
+serde                     = { workspace = true, features = ["derive"] }
+serde_json                = { workspace = true }
 pg_schema_cache.workspace = true
 tree-sitter.workspace     = true
 tree_sitter_sql.workspace = true

crates/pg_completions/src/complete.rs

@@ -1,3 +1,4 @@
+use serde::{Deserialize, Serialize};
 use text_size::TextSize;
 
 use crate::{
@@ -17,9 +18,9 @@ pub struct CompletionParams<'a> {
     pub tree: Option<&'a tree_sitter::Tree>,
 }
 
-#[derive(Debug, Default)]
+#[derive(Debug, Default, Serialize, Deserialize)]
 pub struct CompletionResult {
-    pub items: Vec<CompletionItem>,
+    pub(crate) items: Vec<CompletionItem>,
 }
 
 impl IntoIterator for CompletionResult {

crates/pg_completions/src/item.rs

@@ -1,10 +1,12 @@
-#[derive(Debug, PartialEq, Eq)]
+use serde::{Deserialize, Serialize};
+
+#[derive(Debug, PartialEq, Eq, Serialize, Deserialize)]
 pub enum CompletionItemKind {
     Table,
     Function,
 }
 
-#[derive(Debug)]
+#[derive(Debug, Serialize, Deserialize)]
 pub struct CompletionItem {
     pub label: String,
     pub(crate) score: i32,

crates/pg_lsp_new/Cargo.toml

@@ -10,6 +10,10 @@ name                 = "pg_lsp_new"
 repository.workspace = true
 version              = "0.0.0"
 
+[[bin]]
+name = "pglsp_new"
+path = "src/main.rs"
+test = false
 
 [dependencies]
 anyhow              = { workspace = true }
@@ -18,6 +22,7 @@ futures             = "0.3.31"
 pg_configuration    = { workspace = true }
 pg_console          = { workspace = true }
 pg_diagnostics      = { workspace = true }
+pg_completions      = { workspace = true }
 pg_fs               = { workspace = true }
 pg_lsp_converters   = { workspace = true }
 pg_text_edit        = { workspace = true }
@@ -29,6 +34,7 @@ text-size.workspace = true
 tokio               = { workspace = true, features = ["rt", "io-std"] }
 tower-lsp           = { version = "0.20.0" }
 tracing             = { workspace = true, features = ["attributes"] }
+tracing-subscriber  = { workspace = true }
 
 [dev-dependencies]
 

crates/pg_lsp_new/src/handlers.rs

@@ -1 +1,2 @@
+pub(crate) mod completions;
 pub(crate) mod text_document;

crates/pg_lsp_new/src/handlers/completions.rs

@@ -0,0 +1,45 @@
+use crate::session::Session;
+use anyhow::Result;
+use pg_workspace_new::workspace::CompletionParams;
+use text_size::TextSize;
+use tower_lsp::lsp_types::{self, CompletionItem, CompletionItemLabelDetails};
+
+pub fn get_completions(
+    session: &Session,
+    params: lsp_types::CompletionParams,
+) -> Result<lsp_types::CompletionResponse> {
+    let pos = params.text_document_position.position;
+    let url = params.text_document_position.text_document.uri;
+
+    let path = session.file_path(&url)?;
+
+    let completion_result = session.workspace.get_completions(CompletionParams {
+        path,
+        position: TextSize::from(pos.character),
+    })?;
+
+    let items: Vec<CompletionItem> = completion_result
+        .into_iter()
+        .map(|i| CompletionItem {
+            label: i.label,
+            label_details: Some(CompletionItemLabelDetails {
+                description: Some(i.description),
+                detail: None,
+            }),
+            preselect: Some(i.preselected),
+            kind: Some(to_lsp_types_completion_item_kind(i.kind)),
+            ..CompletionItem::default()
+        })
+        .collect();
+
+    Ok(lsp_types::CompletionResponse::Array(items))
+}
+
+fn to_lsp_types_completion_item_kind(
+    pg_comp_kind: pg_completions::CompletionItemKind,
+) -> lsp_types::CompletionItemKind {
+    match pg_comp_kind {
+        pg_completions::CompletionItemKind::Function
+        | pg_completions::CompletionItemKind::Table => lsp_types::CompletionItemKind::CLASS,
+    }
+}

crates/pg_lsp_new/src/handlers/text_document.rs

@@ -24,11 +24,11 @@ pub(crate) async fn did_open(
     let version = params.text_document.version;
     let content = params.text_document.text;
 
-    let biome_path = session.file_path(&url)?;
+    let path = session.file_path(&url)?;
     let doc = Document::new(version, &content);
 
     session.workspace.open_file(OpenFileParams {
-        path: biome_path,
+        path,
         version,
         content,
     })?;

crates/pg_lsp_new/src/main.rs

@@ -0,0 +1,36 @@
+use std::{fs::File, path::PathBuf, str::FromStr};
+
+use pg_lsp_new::ServerFactory;
+use tracing_subscriber::fmt::format::FmtSpan;
+
+#[tokio::main]
+async fn main() -> anyhow::Result<()> {
+    // This is currently only used for local debugging.
+    // The `pglsp.toml` and `pglsp.log` filepaths should be dynamic.
+
+    let log_path = PathBuf::from_str("pglsp.log").expect("Opened the log file.");
+    let log_file = File::create(log_path).expect("Could not open the file.");
+
+    let subscriber = tracing_subscriber::FmtSubscriber::builder()
+        .with_span_events(FmtSpan::ENTER)
+        .with_span_events(FmtSpan::CLOSE)
+        .with_ansi(false)
+        .with_writer(log_file)
+        .finish();
+
+    tracing::subscriber::set_global_default(subscriber)?;
+
+    let stdin = tokio::io::stdin();
+    let stdout = tokio::io::stdout();
+
+    let config_path = PathBuf::from_str("pglsp.toml").expect("Could not find the pglsp.toml file.");
+
+    tracing::info!("Starting server.");
+
+    ServerFactory::new(true)
+        .create(Some(config_path))
+        .accept(stdin, stdout)
+        .await;
+
+    Ok(())
+}

crates/pg_lsp_new/src/server.rs

@@ -232,6 +232,14 @@ impl LanguageServer for LSPServer {
             .await
             .ok();
     }
+
+    #[tracing::instrument(level = "trace", skip(self))]
+    async fn completion(&self, params: CompletionParams) -> LspResult<Option<CompletionResponse>> {
+        match handlers::completions::get_completions(&self.session, params) {
+            Ok(result) => LspResult::Ok(Some(result)),
+            Err(e) => LspResult::Err(into_lsp_error(e)),
+        }
+    }
 }
 
 impl Drop for LSPServer {
@@ -379,6 +387,7 @@ impl ServerFactory {
         workspace_method!(builder, change_file);
         workspace_method!(builder, close_file);
         workspace_method!(builder, pull_diagnostics);
+        workspace_method!(builder, get_completions);
 
         let (service, socket) = builder.finish();
         ServerConnection { socket, service }

crates/pg_workspace_new/Cargo.toml

@@ -19,6 +19,7 @@ ignore                    = { workspace = true }
 pg_configuration          = { workspace = true }
 pg_console                = { workspace = true }
 pg_diagnostics            = { workspace = true }
+pg_completions            = { workspace = true }
 pg_fs                     = { workspace = true, features = ["serde"] }
 pg_query_ext              = { workspace = true }
 pg_schema_cache           = { workspace = true }

crates/pg_workspace_new/src/diagnostics.rs

@@ -60,6 +60,12 @@ impl WorkspaceError {
     pub fn vcs_disabled() -> Self {
         Self::Vcs(VcsDiagnostic::DisabledVcs(DisabledVcs {}))
     }
+
+    pub fn runtime(msg: &str) -> Self {
+        Self::RuntimeError(RuntimeError {
+            message: msg.into(),
+        })
+    }
 }
 
 impl Error for WorkspaceError {}

crates/pg_workspace_new/src/workspace.rs

@@ -4,7 +4,7 @@ pub use self::client::{TransportRequest, WorkspaceClient, WorkspaceTransport};
 use pg_configuration::PartialConfiguration;
 use pg_fs::PgLspPath;
 use serde::{Deserialize, Serialize};
-use text_size::TextRange;
+use text_size::{TextRange, TextSize};
 
 use crate::WorkspaceError;
 
@@ -39,6 +39,14 @@ pub struct PullDiagnosticsParams {
     // pub skip: Vec<RuleSelector>,
 }
 
+#[derive(Debug, serde::Serialize, serde::Deserialize)]
+pub struct CompletionParams {
+    /// The File for which a completion is requested.
+    pub path: PgLspPath,
+    /// The Cursor position in the file for which a completion is requested.
+    pub position: TextSize,
+}
+
 #[derive(Debug, serde::Serialize, serde::Deserialize)]
 pub struct PullDiagnosticsResult {
     pub diagnostics: Vec<pg_diagnostics::serde::Diagnostic>,
@@ -105,6 +113,11 @@ pub trait Workspace: Send + Sync + RefUnwindSafe {
         params: PullDiagnosticsParams,
     ) -> Result<PullDiagnosticsResult, WorkspaceError>;
 
+    fn get_completions(
+        &self,
+        params: CompletionParams,
+    ) -> Result<pg_completions::CompletionResult, WorkspaceError>;
+
     /// Refresh the schema cache for this workspace
     fn refresh_schema_cache(&self) -> Result<(), WorkspaceError>;
 

crates/pg_workspace_new/src/workspace/client.rs

@@ -127,4 +127,11 @@ where
     ) -> Result<super::PullDiagnosticsResult, WorkspaceError> {
         self.request("pglsp/pull_diagnostics", params)
     }
+
+    fn get_completions(
+        &self,
+        params: super::CompletionParams,
+    ) -> Result<pg_completions::CompletionResult, WorkspaceError> {
+        self.request("pglsp/get_completions", params)
+    }
 }

crates/pg_workspace_new/src/workspace/server.rs

@@ -342,6 +342,45 @@ impl Workspace for WorkspaceServer {
             skipped_diagnostics: 0,
         })
     }
+
+    #[tracing::instrument(level = "info", skip(self))]
+    fn get_completions(
+        &self,
+        params: super::CompletionParams,
+    ) -> Result<pg_completions::CompletionResult, WorkspaceError> {
+        let doc = self
+            .documents
+            .get(&params.path)
+            .ok_or(WorkspaceError::not_found())?;
+
+        tracing::info!("Found the document.");
+        tracing::info!("Looking for statement at position: {:?}", &params.position);
+
+        let statement_ref = match doc.statement_ref_at_offset(&params.position) {
+            Some(s) => s,
+            None => return Ok(pg_completions::CompletionResult::default()),
+        };
+
+        let tree = self.tree_sitter.fetch(&statement_ref);
+        let text = doc
+            .statement_by_id(statement_ref.id)
+            .expect("Found statement_ref but no matching statement")
+            .text;
+
+        let schema_cache = self
+            .schema_cache
+            .read()
+            .map_err(|_| WorkspaceError::runtime("Unable to load SchemaCache"))?;
+
+        let result = pg_completions::complete(pg_completions::CompletionParams {
+            position: params.position,
+            schema: &schema_cache,
+            tree: tree.as_deref(),
+            text,
+        });
+
+        Ok(result)
+    }
 }
 
 /// Returns `true` if `path` is a directory or

crates/pg_workspace_new/src/workspace/server/document.rs

@@ -129,6 +129,13 @@ impl Document {
             .collect()
     }
 
+    pub(super) fn statement_by_id(&self, id: StatementId) -> Option<Statement> {
+        self.statements
+            .iter()
+            .find(|s| s.0 == id)
+            .map(|pos| self.statement(pos))
+    }
+
     pub(super) fn statement_ref(&self, inner_ref: &StatementPosition) -> StatementRef {
         StatementRef {
             id: inner_ref.0,