feat: plpgsql syntax errors

crates/pgt_lsp/tests/server.rs

@@ -1678,3 +1678,84 @@ ALTER TABLE ONLY "public"."campaign_contact_list"
 
     Ok(())
 }
+
+#[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")]
+async fn test_plpgsql(test_db: PgPool) -> Result<()> {
+    let factory = ServerFactory::default();
+    let mut fs = MemoryFileSystem::default();
+
+    let mut conf = PartialConfiguration::init();
+    conf.merge_with(PartialConfiguration {
+        db: Some(PartialDatabaseConfiguration {
+            database: Some(
+                test_db
+                    .connect_options()
+                    .get_database()
+                    .unwrap()
+                    .to_string(),
+            ),
+            ..Default::default()
+        }),
+        ..Default::default()
+    });
+    fs.insert(
+        url!("postgrestools.jsonc").to_file_path().unwrap(),
+        serde_json::to_string_pretty(&conf).unwrap(),
+    );
+
+    let (service, client) = factory
+        .create_with_fs(None, DynRef::Owned(Box::new(fs)))
+        .into_inner();
+
+    let (stream, sink) = client.split();
+    let mut server = Server::new(service);
+
+    let (sender, mut receiver) = channel(CHANNEL_BUFFER_SIZE);
+    let reader = tokio::spawn(client_handler(stream, sink, sender));
+
+    server.initialize().await?;
+    server.initialized().await?;
+
+    server.load_configuration().await?;
+
+    let initial_content = r#"
+create function test_organisation_id ()
+    returns setof text
+    language plpgsql
+    security invoker
+    as $$
+    declre
+        v_organisation_id uuid;
+begin
+    return next is(private.organisation_id(), v_organisation_id, 'should return organisation_id of token');
+end
+$$;
+"#;
+
+    server.open_document(initial_content).await?;
+
+    let notification = tokio::time::timeout(Duration::from_secs(5), async {
+        loop {
+            match receiver.next().await {
+                Some(ServerNotification::PublishDiagnostics(msg)) => {
+                    if msg.diagnostics.iter().any(|d| {
+                        d.message
+                            .contains("Invalid statement: syntax error at or near \"declre\"")
+                    }) {
+                        return true;
+                    }
+                }
+                _ => continue,
+            }
+        }
+    })
+    .await
+    .is_ok();
+
+    assert!(notification, "expected diagnostics for unknown column");
+
+    server.shutdown().await?;
+    reader.abort();
+
+    Ok(())
+}

crates/pgt_query_ext/src/diagnostics.rs

@@ -15,6 +15,16 @@ pub struct SyntaxDiagnostic {
     pub message: MessageAndDescription,
 }
 
+impl SyntaxDiagnostic {
+    /// Create a new syntax diagnostic with the given message and optional span.
+    pub fn new(message: impl Into<String>, span: Option<TextRange>) -> Self {
+        SyntaxDiagnostic {
+            span,
+            message: MessageAndDescription::from(message.into()),
+        }
+    }
+}
+
 impl From<pg_query::Error> for SyntaxDiagnostic {
     fn from(err: pg_query::Error) -> Self {
         SyntaxDiagnostic {

crates/pgt_query_ext/src/lib.rs

@@ -25,3 +25,38 @@ pub fn parse(sql: &str) -> Result<NodeEnum> {
             .ok_or_else(|| Error::Parse("Unable to find root node".to_string()))
     })?
 }
+
+/// This function parses a PL/pgSQL function.
+///
+/// It expects the entire `CREATE FUNCTION` statement.
+pub fn parse_plpgsql(sql: &str) -> Result<()> {
+    // we swallow the error until we have a proper binding
+    let _ = pg_query::parse_plpgsql(sql)?;
+
+    Ok(())
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_parse_plpgsql_err() {
+        let input = "
+create function test_organisation_id ()
+    returns setof text
+    language plpgsql
+    security invoker
+    as $$
+    -- syntax error here
+    decare
+        v_organisation_id uuid;
+begin
+    select 1;
+end
+$$;
+        ";
+
+        assert!(parse_plpgsql(input).is_err());
+    }
+}

crates/pgt_workspace/src/workspace/server.rs

@@ -53,6 +53,7 @@ mod async_helper;
 mod connection_key;
 mod connection_manager;
 pub(crate) mod document;
+mod function_utils;
 mod migration;
 mod pg_query;
 mod schema_cache_manager;
@@ -528,7 +529,7 @@ impl Workspace for WorkspaceServer {
 
         diagnostics.extend(
             doc.iter(SyncDiagnosticsMapper)
-                .flat_map(|(_id, range, ast, diag)| {
+                .flat_map(|(range, ast, diag)| {
                     let mut errors: Vec<Error> = vec![];
 
                     if let Some(diag) = diag {
@@ -560,9 +561,12 @@ impl Workspace for WorkspaceServer {
                                     },
                                 );
 
+                            // adjust the span of the diagnostics to the statement (if it has one)
+                            let span = d.location().span.map(|s| s + range.start());
+
                             SDiagnostic::new(
                                 d.with_file_path(params.path.as_path().display().to_string())
-                                    .with_file_span(range)
+                                    .with_file_span(span.unwrap_or(range))
                                     .with_severity(severity),
                             )
                         })