feat: expose new workspace function splitStatements with ranges (#730)

Author: lneves12

crates/pgls_wasm/build-wasm.sh

@@ -85,6 +85,7 @@ EXPORTED_FUNCTIONS="[
   '_pgls_complete',
   '_pgls_hover',
   '_pgls_parse',
+  '_pgls_split_statements',
   '_pgls_version',
   '_pgls_lsp_handle_message'
 ]"

crates/pgls_wasm/src/ffi.rs

@@ -256,6 +256,27 @@ pub unsafe extern "C" fn pgls_parse(sql: *const c_char) -> *mut c_char {
     })
 }
 
+/// Split SQL into individual statements.
+/// Returns a JSON array of `Statement` objects, each with `sql`, `start`, and `end`.
+/// The returned string must be freed with `pgls_free_string`.
+///
+/// # Safety
+/// The sql pointer must be valid and point to a null-terminated UTF-8 string.
+#[unsafe(no_mangle)]
+pub unsafe extern "C" fn pgls_split_statements(sql: *const c_char) -> *mut c_char {
+    // SAFETY: Caller guarantees sql is valid
+    let sql_str = match unsafe { c_str_to_str(sql) } {
+        Some(s) => s,
+        None => return str_to_c_string("ERROR: Invalid SQL pointer"),
+    };
+
+    with_workspace(|ws| {
+        let statements = ws.split_statements(sql_str);
+        let json = serde_json::to_string(&statements).unwrap_or_else(|e| format!("ERROR: {e}"));
+        str_to_c_string(&json)
+    })
+}
+
 /// Get the version of the library.
 /// The returned string must be freed with `pgls_free_string`.
 #[unsafe(no_mangle)]
@@ -371,4 +392,23 @@ mod tests {
             pgls_free_string(result);
         }
     }
+
+    #[test]
+    fn test_ffi_split_statements() {
+        pgls_init();
+
+        unsafe {
+            let sql = CString::new("SELECT 1; SELECT 2;").unwrap();
+            let result = pgls_split_statements(sql.as_ptr());
+            assert!(!result.is_null());
+            let result_str = CStr::from_ptr(result).to_str().unwrap();
+            let parsed: Vec<serde_json::Value> = serde_json::from_str(result_str).unwrap();
+            assert_eq!(parsed.len(), 2);
+            assert_eq!(parsed[0]["sql"], "SELECT 1;");
+            assert_eq!(parsed[0]["start"], 0);
+            assert_eq!(parsed[0]["end"], 9);
+            assert_eq!(parsed[1]["sql"], "SELECT 2;");
+            pgls_free_string(result);
+        }
+    }
 }

crates/pgls_wasm/src/lib.rs

@@ -66,6 +66,18 @@ pub struct CompletionItem {
     pub detail: Option<String>,
 }
 
+/// A single SQL statement extracted from a multi-statement source.
+#[derive(Debug, Serialize, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct Statement {
+    /// The SQL text of the statement
+    pub sql: String,
+    /// Start byte offset in the original SQL string
+    pub start: u32,
+    /// End byte offset in the original SQL string
+    pub end: u32,
+}
+
 /// In-memory file system for storing SQL files.
 ///
 /// This wraps `pgls_fs::MemoryFileSystem` with a simpler API
@@ -367,6 +379,22 @@ impl Workspace {
             Err(e) => vec![e.to_string()],
         }
     }
+
+    /// Split SQL into individual statements.
+    ///
+    /// Each returned `Statement` contains the SQL text and its byte offsets
+    /// in the original string.
+    pub fn split_statements(&self, sql: &str) -> Vec<Statement> {
+        self.inner
+            .split_statements(sql)
+            .into_iter()
+            .map(|range| Statement {
+                sql: sql[range].to_string(),
+                start: u32::from(range.start()),
+                end: u32::from(range.end()),
+            })
+            .collect()
+    }
 }
 
 #[cfg(test)]
@@ -397,6 +425,30 @@ mod tests {
         assert!(!errors.is_empty());
     }
 
+    #[test]
+    fn test_workspace_split_statements() {
+        let workspace = Workspace::new();
+
+        let statements = workspace.split_statements("SELECT 1; SELECT 2;");
+        assert_eq!(statements.len(), 2);
+        assert_eq!(statements[0].sql, "SELECT 1;");
+        assert_eq!(statements[0].start, 0);
+        assert_eq!(statements[0].end, 9);
+        assert_eq!(statements[1].sql, "SELECT 2;");
+        assert_eq!(statements[1].start, 10);
+        assert_eq!(statements[1].end, 19);
+    }
+
+    #[test]
+    fn test_workspace_split_statements_ignores_comments() {
+        let workspace = Workspace::new();
+
+        let statements = workspace.split_statements("-- comment\nSELECT 1;\n\nSELECT 2;");
+        assert_eq!(statements.len(), 2);
+        assert_eq!(statements[0].sql, "SELECT 1;");
+        assert_eq!(statements[1].sql, "SELECT 2;");
+    }
+
     #[test]
     fn test_memory_fs() {
         let mut fs = MemoryFs::new();

crates/pgls_workspace/src/workspace/server.rs

@@ -26,6 +26,7 @@ use pgls_diagnostics::{
 };
 use pgls_fs::{ConfigName, PgLSPath};
 use pgls_schema_cache::SchemaCache;
+use pgls_text_size::TextRange;
 #[cfg(feature = "db")]
 use pgls_typecheck::{IdentifierType, TypecheckParams, TypedIdentifier};
 use pgls_workspace_macros::ignored_path;
@@ -151,6 +152,12 @@ impl WorkspaceServer {
         self.schema_cache.clear();
     }
 
+    /// Split raw SQL into byte ranges of individual statements.
+    #[allow(clippy::unused_self)]
+    pub fn split_statements(&self, sql: &str) -> Vec<TextRange> {
+        pgls_statement_splitter::split(sql).ranges
+    }
+
     /// Get a clone of the current schema.
     pub fn get_schema(&self) -> Option<Arc<SchemaCache>> {
         self.schema_cache.get()

packages/@postgres-language-server/wasm/README.md

@@ -34,6 +34,10 @@ const workspace = await createWorkspace();
 const errors = workspace.parse("SELECT * FROM users;");
 console.log(errors); // []
 
+// Split SQL into individual statements
+const statements = workspace.splitStatements("SELECT 1; SELECT 2;");
+// [{ sql: "SELECT 1;", start: 0, end: 9 }, { sql: "SELECT 2;", start: 10, end: 19 }]
+
 // Insert a file and lint it
 workspace.insertFile("/query.sql", "SELECT * FROM users;");
 const diagnostics = workspace.lint("/query.sql");
@@ -132,17 +136,18 @@ lsp.handleMessage({
 
 ### Workspace
 
-| Method                      | Description                                |
-| --------------------------- | ------------------------------------------ |
-| `parse(sql: string)`        | Parse SQL, returns array of error messages |
-| `insertFile(path, content)` | Add or update a file in the workspace      |
-| `removeFile(path)`          | Remove a file from the workspace           |
-| `lint(path)`                | Get diagnostics for a file                 |
-| `complete(path, offset)`    | Get completions at position                |
-| `hover(path, offset)`       | Get hover info at position                 |
-| `setSchema(json)`           | Set database schema                        |
-| `clearSchema()`             | Clear the current schema                   |
-| `version()`                 | Get library version                        |
+| Method                      | Description                                 |
+| --------------------------- | ------------------------------------------- |
+| `parse(sql: string)`        | Parse SQL, returns array of error messages  |
+| `splitStatements(sql)`      | Split SQL into statements with byte offsets |
+| `insertFile(path, content)` | Add or update a file in the workspace       |
+| `removeFile(path)`          | Remove a file from the workspace            |
+| `lint(path)`                | Get diagnostics for a file                  |
+| `complete(path, offset)`    | Get completions at position                 |
+| `hover(path, offset)`       | Get hover info at position                  |
+| `setSchema(json)`           | Set database schema                         |
+| `clearSchema()`             | Clear the current schema                    |
+| `version()`                 | Get library version                         |
 
 ### LanguageServer
 

packages/@postgres-language-server/wasm/src/index.ts

@@ -26,6 +26,7 @@ export type {
   Diagnostic,
   CompletionItem,
   SchemaCache,
+  Statement,
   WorkspaceOptions,
   PGLSModule,
   JsonRpcMessage,

packages/@postgres-language-server/wasm/src/types.ts

@@ -49,6 +49,18 @@ export interface WorkspaceOptions {
   schema?: SchemaCache | string;
 }
 
+/**
+ * A single SQL statement extracted from a multi-statement source.
+ */
+export interface Statement {
+  /** The SQL text of the statement */
+  sql: string;
+  /** Start byte offset in the original SQL string */
+  start: number;
+  /** End byte offset in the original SQL string */
+  end: number;
+}
+
 /**
  * A JSON-RPC message (request, response, or notification).
  * This is the standard LSP message format.
@@ -86,6 +98,7 @@ export interface PGLSModule {
   _pgls_complete(pathPtr: number, offset: number): number;
   _pgls_hover(pathPtr: number, offset: number): number;
   _pgls_parse(sqlPtr: number): number;
+  _pgls_split_statements(sqlPtr: number): number;
   _pgls_version(): number;
 
   // Language Server API

packages/@postgres-language-server/wasm/src/workspace.ts

@@ -22,10 +22,11 @@ import type {
   Diagnostic,
   PGLSModule,
   SchemaCache,
+  Statement,
   WorkspaceOptions,
 } from "./types.js";
 
-export type { Diagnostic, CompletionItem, SchemaCache, WorkspaceOptions, PGLSModule };
+export type { Diagnostic, CompletionItem, SchemaCache, Statement, WorkspaceOptions, PGLSModule };
 
 /**
  * The Workspace class provides a direct API for SQL parsing, linting,
@@ -154,6 +155,23 @@ export class Workspace {
     }
   }
 
+  /**
+   * Split SQL into individual statements.
+   *
+   * Each returned `Statement` contains the SQL text and its byte offsets
+   * in the original string.
+   */
+  splitStatements(sql: string): Statement[] {
+    const sqlPtr = allocateString(this.module, sql);
+    try {
+      const resultPtr = this.module._pgls_split_statements(sqlPtr);
+      const result = readAndFreeString(this.module, resultPtr);
+      return parseResult<Statement[]>(result) ?? [];
+    } finally {
+      this.module._free(sqlPtr);
+    }
+  }
+
   /**
    * Get the version of the library.
    */