spark-sql-language-server 0.0.1-beta.6 → 0.2.0-beta.0

package/out-tsc/lsp-server.js

@@ -15,65 +15,128 @@ var __setModuleDefault = (this && this.__setModuleDefault) || (Object.create ? (
 }) : function(o, v) {
     o["default"] = v;
 });
-var __importStar = (this && this.__importStar) || function (mod) {
-    if (mod && mod.__esModule) return mod;
-    var result = {};
-    if (mod != null) for (var k in mod) if (k !== "default" && Object.prototype.hasOwnProperty.call(mod, k)) __createBinding(result, mod, k);
-    __setModuleDefault(result, mod);
-    return result;
-};
+var __importStar = (this && this.__importStar) || (function () {
+    var ownKeys = function(o) {
+        ownKeys = Object.getOwnPropertyNames || function (o) {
+            var ar = [];
+            for (var k in o) if (Object.prototype.hasOwnProperty.call(o, k)) ar[ar.length] = k;
+            return ar;
+        };
+        return ownKeys(o);
+    };
+    return function (mod) {
+        if (mod && mod.__esModule) return mod;
+        var result = {};
+        if (mod != null) for (var k = ownKeys(mod), i = 0; i < k.length; i++) if (k[i] !== "default") __createBinding(result, mod, k[i]);
+        __setModuleDefault(result, mod);
+        return result;
+    };
+})();
 Object.defineProperty(exports, "__esModule", { value: true });
 exports.LSPServer = void 0;
+exports.isSemanticTokenReference = isSemanticTokenReference;
 const antlr4_c3_1 = require("antlr4-c3");
+const antlr4ts_1 = require("antlr4ts");
 const tree_1 = require("antlr4ts/tree");
-const sql_formatter_1 = require("sql-formatter");
 const lsp = __importStar(require("vscode-languageserver"));
-const vscode_languageserver_1 = require("vscode-languageserver");
+const vscode_languageserver_textdocument_1 = require("vscode-languageserver-textdocument");
 const built_in_functions_1 = require("./assets/built-in-functions");
-const SparkSqlLexer_1 = require("./lib/SparkSqlLexer");
-const SparkSqlParser_1 = require("./lib/SparkSqlParser");
-const schema_listener_1 = require("./listeners/schema.listener");
+const completion_parser_1 = require("./completion-parser");
+const cursor_1 = require("./cursor");
+const formatter_adapter_1 = require("./formatter/formatter-adapter");
+const SparkSQLLexer_1 = require("./lib/SparkSQLLexer");
+const SparkSQLParser_1 = require("./lib/SparkSQLParser");
 const statement_listener_1 = require("./listeners/statement.listener");
+const structure_listener_1 = require("./listeners/structure.listener");
+const tokens_collector_listener_1 = require("./listeners/tokens-collector.listener");
 const parsing_warehouse_1 = require("./parsing-warehouse");
 const protocol_translation_1 = require("./protocol-translation");
 const schema_registry_1 = require("./schema-registry");
+const typings_1 = require("./typings");
 const utils_1 = require("./utils");
-const ignoredTokens = [
-    SparkSqlLexer_1.SparkSqlLexer.EOF,
-    SparkSqlLexer_1.SparkSqlLexer.DOT,
-    SparkSqlLexer_1.SparkSqlLexer.SEMI,
-    SparkSqlLexer_1.SparkSqlLexer.COMMA,
-    SparkSqlLexer_1.SparkSqlLexer.LEFT_BRACKET,
-    SparkSqlLexer_1.SparkSqlLexer.RIGHT_BRACKET,
-    SparkSqlLexer_1.SparkSqlLexer.EQ,
-    SparkSqlLexer_1.SparkSqlLexer.DOUBLEQUOTED_STRING,
-    SparkSqlLexer_1.SparkSqlLexer.BACKQUOTED_IDENTIFIER,
-    SparkSqlLexer_1.SparkSqlLexer.AT,
-    SparkSqlLexer_1.SparkSqlLexer.DIV,
-    SparkSqlLexer_1.SparkSqlLexer.NSEQ,
-    SparkSqlLexer_1.SparkSqlLexer.NEQ,
-    SparkSqlLexer_1.SparkSqlLexer.NEQJ,
-    SparkSqlLexer_1.SparkSqlLexer.LT,
-    SparkSqlLexer_1.SparkSqlLexer.LTE,
-    SparkSqlLexer_1.SparkSqlLexer.GT,
-    SparkSqlLexer_1.SparkSqlLexer.GTE,
-    SparkSqlLexer_1.SparkSqlLexer.PLUS,
-    SparkSqlLexer_1.SparkSqlLexer.MINUS,
-    SparkSqlLexer_1.SparkSqlLexer.IDENTIFIER
-];
+const lineage_visitor_1 = require("./visitors/lineage.visitor");
 const preferredRules = [
-    SparkSqlParser_1.SparkSqlParser.RULE_functionName,
-    SparkSqlParser_1.SparkSqlParser.RULE_identifierReference,
-    SparkSqlParser_1.SparkSqlParser.RULE_tableIdentifier,
-    SparkSqlParser_1.SparkSqlParser.RULE_catalogNameDef,
-    SparkSqlParser_1.SparkSqlParser.RULE_databaseNameDef,
-    SparkSqlParser_1.SparkSqlParser.RULE_tableNameDef,
-    SparkSqlParser_1.SparkSqlParser.RULE_viewNameDef,
-    SparkSqlParser_1.SparkSqlParser.RULE_columnNameDef,
-    SparkSqlParser_1.SparkSqlParser.RULE_identifier
+    SparkSQLParser_1.SparkSQLParser.RULE_functionName,
+    SparkSQLParser_1.SparkSQLParser.RULE_identifierReference,
+    SparkSQLParser_1.SparkSQLParser.RULE_multipartIdentifier,
+    SparkSQLParser_1.SparkSQLParser.RULE_tableIdentifier,
+    SparkSQLParser_1.SparkSQLParser.RULE_identifier
 ];
+const ignoredTokens = [
+    SparkSQLLexer_1.SparkSQLLexer.EOF,
+    SparkSQLLexer_1.SparkSQLLexer.SEMICOLON,
+    SparkSQLLexer_1.SparkSQLLexer.COMMA,
+    SparkSQLLexer_1.SparkSQLLexer.DOT,
+    SparkSQLLexer_1.SparkSQLLexer.ASTERISK,
+    SparkSQLLexer_1.SparkSQLLexer.COLON,
+    SparkSQLLexer_1.SparkSQLLexer.LEFT_PAREN,
+    SparkSQLLexer_1.SparkSQLLexer.RIGHT_PAREN,
+    SparkSQLLexer_1.SparkSQLLexer.LEFT_BRACKET,
+    SparkSQLLexer_1.SparkSQLLexer.RIGHT_BRACKET,
+    SparkSQLLexer_1.SparkSQLLexer.STRING_LITERAL,
+    SparkSQLLexer_1.SparkSQLLexer.DECIMAL_VALUE,
+    SparkSQLLexer_1.SparkSQLLexer.IDENTIFIER
+];
+/**
+ * Simple LRU Cache for folding ranges
+ * Limits memory usage by evicting least recently used entries
+ */
+class LRUFoldingCache {
+    constructor(maxSize = 50) {
+        this.cache = new Map();
+        this.maxSize = maxSize;
+    }
+    get(key) {
+        const entry = this.cache.get(key);
+        if (entry) {
+            // Update last access time
+            entry.lastAccess = Date.now();
+            // Move to end (most recently used)
+            this.cache.delete(key);
+            this.cache.set(key, entry);
+        }
+        return entry;
+    }
+    set(key, entry) {
+        // Remove if exists (to update position)
+        if (this.cache.has(key)) {
+            this.cache.delete(key);
+        }
+        // Evict oldest if at capacity
+        else if (this.cache.size >= this.maxSize) {
+            // The first key is the least recently used
+            const oldestKey = this.cache.keys().next().value;
+            if (oldestKey) {
+                this.cache.delete(oldestKey);
+            }
+        }
+        entry.lastAccess = Date.now();
+        this.cache.set(key, entry);
+    }
+    delete(key) {
+        return this.cache.delete(key);
+    }
+    deleteByPrefix(prefix) {
+        for (const key of this.cache.keys()) {
+            if (key.startsWith(prefix)) {
+                this.cache.delete(key);
+            }
+        }
+    }
+    clear() {
+        this.cache.clear();
+    }
+    get size() {
+        return this.cache.size;
+    }
+    keys() {
+        return this.cache.keys();
+    }
+}
 class LSPServer {
     constructor() {
+        this._builtInFunctionItems = null;
+        this._foldingCache = new LRUFoldingCache(50); // Max 50 cached documents
         this._warehouse = new parsing_warehouse_1.ParsingWarehouse();
         this._schemaRegistry = new schema_registry_1.SchemaRegistry();
     }
@@ -83,235 +146,607 @@ class LSPServer {
     get schemaRegistry() {
         return this._schemaRegistry;
     }
+    /**
+     * Get cached built-in function completion items
+     */
+    getBuiltInFunctionItems() {
+        if (!this._builtInFunctionItems) {
+            this._builtInFunctionItems = built_in_functions_1.BuiltInFunctions.map(fn => {
+                const funcName = fn.name.split('(')[0];
+                return (0, protocol_translation_1.toCompletionItem)('built-in', funcName, fn.expr, fn.funcDesc, (0, built_in_functions_1.generateDescription)(fn));
+            });
+        }
+        return this._builtInFunctionItems;
+    }
+    getTokenByPosition(document, position) {
+        const parsingValue = this._warehouse.getParsingValue(document);
+        const offset = document.offsetAt(position);
+        return parsingValue.tokenStream.getTokens().find(t => t.startIndex <= offset && offset <= t.stopIndex + 1);
+    }
     doValidation(document) {
         const parsingValue = this._warehouse.getParsingValue(document);
         return parsingValue.errors;
     }
-    doCompletion(document, position) {
+    doFormatting(document, options, returnResult) {
         const parsingValue = this._warehouse.getParsingValue(document);
-        let offset = document.offsetAt(position);
-        const token = parsingValue.tokenStream.getTokens().find(t => t.startIndex <= offset && offset <= t.stopIndex + 1);
-        if (!token) {
-            return [];
-        }
-        const core = new antlr4_c3_1.CodeCompletionCore(parsingValue.parser);
-        core.ignoredTokens = new Set(ignoredTokens);
-        core.preferredRules = new Set(preferredRules);
-        let index = token.tokenIndex;
-        let bias = 0;
-        if (token.type === SparkSqlLexer_1.SparkSqlLexer.DOT) {
-            index += 1;
-            offset -= 1;
-            bias = 1;
-        }
-        const completionItems = [];
+        // Collect parse errors for reporting (but don't block formatting)
+        const parseErrors = parsingValue.errors.map(e => ({
+            message: e.message,
+            line: e.range.start.line
+        }));
         try {
-            const candidates = core.collectCandidates(index);
-            const identifiers = [];
-            if (candidates.rules.size > 0) {
-                const listener = new schema_listener_1.SchemaListener();
-                tree_1.ParseTreeWalker.DEFAULT.walk(listener, parsingValue.tree);
-                const schemaContexts = listener.getSchemaContexts();
-                const context = schemaContexts.find(c => c.startOffset <= offset && c.endOffset >= offset);
-                const schemaParts = context?.text.replace(/[`'"]/g, '').split('.');
-                if (candidates.rules.has(SparkSqlParser_1.SparkSqlParser.RULE_catalogNameDef)) {
-                    const catalogsInDocument = (0, utils_1.unique)(schemaContexts.filter(c => c !== context && c.kind === schema_listener_1.SchemaContextKindEnum.CATALOG).map(c => c.catalogName));
-                    completionItems.push(...catalogsInDocument.map(c => (0, protocol_translation_1.toCompletionItem)('catalog', c, `\`${c}\``, 'Catalog in document')), ...this._schemaRegistry
-                        .getSchemaCatalogs()
-                        .map(c => (0, protocol_translation_1.toCompletionItem)('catalog', c.label, `\`${c.label}\``, 'Schema catalog', c.description)));
-                    identifiers.push('catalog');
-                }
-                if (candidates.rules.has(SparkSqlParser_1.SparkSqlParser.RULE_databaseNameDef)) {
-                    const databasesInDocument = (0, utils_1.unique)(schemaContexts.filter(d => d !== context && d.kind === schema_listener_1.SchemaContextKindEnum.DATABASE).map(d => d.databaseName));
-                    completionItems.push(...databasesInDocument.map(d => (0, protocol_translation_1.toCompletionItem)('database', d, d, 'Database in document')));
-                    if (!!schemaParts && schemaParts.length <= 1 - bias) {
-                        completionItems.push(...this._schemaRegistry
-                            .getSchemaDatabases({})
-                            .map(d => (0, protocol_translation_1.toCompletionItem)('database', d.label, d.label, `Schema database: ${d.catalog}.*`, d.description)));
-                    }
-                    else if (schemaParts) {
-                        completionItems.push(...this._schemaRegistry
-                            .getSchemaDatabases({ catalog: schemaParts?.[0] })
-                            .map(d => (0, protocol_translation_1.toCompletionItem)('database', d.label, d.label, `Schema database: ${d.catalog}.*`, d.description)));
+            // Use formatSQLFromText which handles normalization internally
+            // Returns a single edit that replaces the entire document
+            const formattingOutput = (0, formatter_adapter_1.formatSQLFromText)(document.getText(), options, formatter_adapter_1.DEFAULT_FORMATTER_CONFIG);
+            // Merge parse errors with formatting errors
+            const allErrors = [...parseErrors, ...formattingOutput.errors];
+            // Check if we should use fallback formatting
+            // Trigger fallback if there are parse errors and we detect that many statements were skipped
+            const shouldUseFallback = parseErrors.length > 0 && this.shouldUseFallbackFormatting(document, formattingOutput.edits, parsingValue.tree);
+            if (shouldUseFallback) {
+                try {
+                    const fallbackResult = this.doStatementByStatementFormatting(document, options);
+                    if (fallbackResult.edits.length > formattingOutput.edits.length) {
+                        // Fallback produced more edits, use it
+                        const result = {
+                            edits: fallbackResult.edits,
+                            errors: [
+                                ...parseErrors,
+                                ...fallbackResult.errors,
+                                {
+                                    message: `Used fallback formatting due to parse errors (recovered ${fallbackResult.edits.length} edits vs ${formattingOutput.edits.length})`,
+                                    line: 0
+                                }
+                            ],
+                            success: false
+                        };
+                        return returnResult ? result : fallbackResult.edits;
                     }
-                    identifiers.push('database');
                 }
-                if (candidates.rules.has(SparkSqlParser_1.SparkSqlParser.RULE_tableNameDef)) {
-                    let tablesInDocument = [];
-                    if (!!schemaParts && schemaParts.length <= 1 - bias) {
-                        completionItems.push(...this._schemaRegistry
-                            .getSchemaTables({})
-                            .map(t => (0, protocol_translation_1.toCompletionItem)('table', t.label, t.label, `Schema table: ${t.catalog}.${t.database}.*`, t.description)));
-                        tablesInDocument = (0, utils_1.unique)(schemaContexts.filter(t => t !== context && t.kind === schema_listener_1.SchemaContextKindEnum.TABLE).map(t => t.tableName));
-                    }
-                    else if (!!schemaParts && schemaParts.length === 2 - bias) {
-                        completionItems.push(...this._schemaRegistry
-                            .getSchemaTables({ database: schemaParts[0] })
-                            .map(t => (0, protocol_translation_1.toCompletionItem)('table', t.label, t.label, `Schema table: ${t.catalog}.${t.database}.*`, t.description)));
-                        tablesInDocument = (0, utils_1.unique)(schemaContexts.filter(t => t !== context &&
-                            t.kind === schema_listener_1.SchemaContextKindEnum.TABLE &&
-                            (!schemaParts[0] || t.databaseName === schemaParts[0])).map(t => t.tableName));
-                    }
-                    else {
-                        completionItems.push(...this._schemaRegistry
-                            .getSchemaTables({ catalog: schemaParts?.[0], database: schemaParts?.[1] })
-                            .map(t => (0, protocol_translation_1.toCompletionItem)('table', t.label, t.label, `Schema table: ${t.catalog}.${t.database}.*`, t.description)));
-                        tablesInDocument = (0, utils_1.unique)(schemaContexts.filter(t => t !== context &&
-                            t.kind === schema_listener_1.SchemaContextKindEnum.TABLE &&
-                            (!schemaParts?.[0] || t.catalogName === schemaParts?.[0]) &&
-                            (!schemaParts?.[1] || t.databaseName === schemaParts?.[1])).map(t => t.tableName));
-                    }
-                    completionItems.push(...tablesInDocument.map(t => (0, protocol_translation_1.toCompletionItem)('table', t, t, 'Table in document')));
-                    identifiers.push('table');
+                catch (fallbackError) {
+                    // Fallback failed, continue with primary results
                 }
-                if (candidates.rules.has(SparkSqlParser_1.SparkSqlParser.RULE_columnNameDef)) {
-                    let nearestTable;
-                    let columnDefinitionMap = undefined;
-                    const [startIndex, endIndex] = this.getSemicolonRange(document.getText(), offset);
-                    if (startIndex && endIndex) {
-                        nearestTable = schemaContexts.find(c => c.kind === schema_listener_1.SchemaContextKindEnum.TABLE && c.startOffset >= startIndex && c.endOffset <= endIndex);
-                    }
-                    const { catalogName, databaseName, tableName } = nearestTable || {};
-                    let columns;
-                    if (!!schemaParts && schemaParts.length <= 1) {
-                        if (bias > 0) {
-                            columnDefinitionMap = {
-                                table: schemaParts[0]
-                            };
-                        }
-                        columns = this._schemaRegistry.getSchemaColumns({
-                            catalog: catalogName,
-                            database: databaseName,
-                            table: bias > 0 ? schemaParts[0] : tableName
-                        });
-                    }
-                    else if (!!schemaParts && schemaParts.length === 2) {
-                        columnDefinitionMap = {
-                            table: schemaParts[0]
-                        };
-                        columns = this._schemaRegistry.getSchemaColumns({ table: schemaParts[0] || tableName });
-                    }
-                    else if (!!schemaParts && schemaParts.length === 3) {
-                        columnDefinitionMap = {
-                            database: schemaParts[0],
-                            table: schemaParts[1]
-                        };
-                        columns = this._schemaRegistry.getSchemaColumns({
-                            database: schemaParts[0] || databaseName,
-                            table: schemaParts[1] || tableName
-                        });
-                    }
-                    else {
-                        columnDefinitionMap = {
-                            catalog: schemaParts?.[0],
-                            database: schemaParts?.[1],
-                            table: schemaParts?.[2]
+            }
+            // formatSQLFromText returns a single edit that replaces the entire document
+            // No need to normalize blank lines since normalization is done during preprocessing
+            const result = {
+                edits: formattingOutput.edits,
+                errors: allErrors,
+                success: allErrors.length === 0
+            };
+            return returnResult ? result : result.edits;
+        }
+        catch (e) {
+            const errorMessage = e instanceof Error ? e.message : String(e);
+            // If formatting failed and there are parse errors, try statement-by-statement formatting
+            if (parseErrors.length > 0) {
+                try {
+                    const fallbackResult = this.doStatementByStatementFormatting(document, options);
+                    if (fallbackResult.edits.length > 0) {
+                        // Fallback formatting succeeded, return those results
+                        const result = {
+                            edits: fallbackResult.edits,
+                            errors: [
+                                ...parseErrors,
+                                ...fallbackResult.errors,
+                                {
+                                    message: `Used fallback formatting due to parse errors`,
+                                    line: 0
+                                }
+                            ],
+                            success: false
                         };
-                        columns = this._schemaRegistry.getSchemaColumns({
-                            catalog: schemaParts?.[0] || catalogName,
-                            database: schemaParts?.[1] || databaseName,
-                            table: schemaParts?.[2] || tableName
-                        });
+                        return returnResult ? result : fallbackResult.edits;
                     }
-                    const columnsInDocument = (0, utils_1.unique)(schemaContexts.filter(t => {
-                        return (t !== context &&
-                            t.kind === schema_listener_1.SchemaContextKindEnum.COLUMN &&
-                            (columnDefinitionMap
-                                ? (!columnDefinitionMap.catalog || columnDefinitionMap.catalog === t.catalogName) &&
-                                    (!columnDefinitionMap.database || columnDefinitionMap.database === t.databaseName) &&
-                                    (!columnDefinitionMap.table || columnDefinitionMap.table === t.tableName)
-                                : (!t.tableName && !nearestTable) ||
-                                    !nearestTable?.tableName ||
-                                    t.tableName === nearestTable.tableName));
-                    }));
-                    completionItems.push(...columnsInDocument.map(d => (0, protocol_translation_1.toCompletionItem)('column', d.text, d.text, `Column ${d.tableName ? `of ${d.tableName}` : ''} in document`)));
-                    if (columns.length) {
-                        completionItems.push(...columns.map(c => (0, protocol_translation_1.toCompletionItem)('column', c.label, c.label, `Schema column: ${c.catalog}.${c.database}.${c.table}.*`, c.description)));
-                    }
-                    identifiers.push('column');
                 }
-                if (candidates.rules.has(SparkSqlParser_1.SparkSqlParser.RULE_functionName)) {
-                    let functions;
-                    if (!!schemaParts && schemaParts.length <= 1 - bias) {
-                        functions = this._schemaRegistry.getSchemaFunctions({});
+                catch (fallbackError) {
+                    // Fallback also failed, continue to original error handling
+                }
+            }
+            const result = {
+                edits: [],
+                errors: [
+                    ...parseErrors,
+                    {
+                        message: `Formatting failed: ${errorMessage}`,
+                        line: 0
+                    }
+                ],
+                success: false
+            };
+            return returnResult ? result : [];
+        }
+    }
+    /**
+     * Format with a specific formatter configuration (legacy or strategy)
+     * Used for comparing formatter outputs in playground
+     */
+    doFormattingWithConfig(document, options, config) {
+        const parsingValue = this._warehouse.getParsingValue(document);
+        // Collect parse errors for reporting
+        const parseErrors = parsingValue.errors.map(e => ({
+            message: e.message,
+            line: e.range.start.line
+        }));
+        try {
+            // Use formatSQLFromText which handles normalization internally
+            // Returns a single edit that replaces the entire document
+            const formattingOutput = (0, formatter_adapter_1.formatSQLFromText)(document.getText(), options, config);
+            const allErrors = [...parseErrors, ...formattingOutput.errors];
+            return {
+                edits: formattingOutput.edits,
+                errors: allErrors,
+                success: allErrors.length === 0
+            };
+        }
+        catch (e) {
+            const errorMessage = e instanceof Error ? e.message : String(e);
+            return {
+                edits: [],
+                errors: [
+                    ...parseErrors,
+                    {
+                        message: `Formatting failed: ${errorMessage}`,
+                        line: 0
                     }
-                    else if (!!schemaParts && schemaParts.length === 2 - bias) {
-                        functions = this._schemaRegistry.getSchemaFunctions({ database: schemaParts[0] });
+                ],
+                success: false
+            };
+        }
+    }
+    /**
+     * Fallback formatting: split document by semicolons and format each statement independently
+     * This is used when the primary formatting fails due to parse errors
+     * IMPORTANT: Preserves original separators (blank lines and comments) between statements
+     */
+    doStatementByStatementFormatting(document, options) {
+        const text = document.getText();
+        const lines = text.split('\n');
+        const allEdits = [];
+        const allErrors = [];
+        // Parse document into statement segments with separators preserved
+        const segments = this.splitIntoStatementSegments(lines);
+        let currentLineOffset = 0;
+        for (const segment of segments) {
+            if (segment.type === 'statement') {
+                // Format the statement
+                const statementText = segment.lines.join('\n');
+                const statementDoc = vscode_languageserver_textdocument_1.TextDocument.create(`${document.uri}#statement-${currentLineOffset}`, document.languageId, document.version, statementText);
+                try {
+                    const stmtParsingValue = this._warehouse.parseDocument(statementDoc);
+                    // Only format if parsing succeeded (no errors)
+                    if (stmtParsingValue.errors.length === 0) {
+                        // Use formatSQLFromText which handles normalization internally
+                        const formattingOutput = (0, formatter_adapter_1.formatSQLFromText)(statementDoc.getText(), options, formatter_adapter_1.DEFAULT_FORMATTER_CONFIG);
+                        // Adjust edit positions to match the original document
+                        for (const edit of formattingOutput.edits) {
+                            allEdits.push({
+                                range: {
+                                    start: {
+                                        line: edit.range.start.line + currentLineOffset,
+                                        character: edit.range.start.character
+                                    },
+                                    end: {
+                                        line: edit.range.end.line + currentLineOffset,
+                                        character: edit.range.end.character
+                                    }
+                                },
+                                newText: edit.newText
+                            });
+                        }
+                        // Collect formatting errors
+                        for (const error of formattingOutput.errors) {
+                            allErrors.push({
+                                message: error.message,
+                                line: (error.line ?? 0) + currentLineOffset
+                            });
+                        }
                     }
                     else {
-                        functions = this._schemaRegistry.getSchemaFunctions({
-                            catalog: schemaParts?.[0],
-                            database: schemaParts?.[1]
+                        // Skip statements with parse errors
+                        allErrors.push({
+                            message: `Skipped statement due to parse errors`,
+                            line: currentLineOffset
                         });
                     }
-                    completionItems.push(...functions.map(f => (0, protocol_translation_1.toCompletionItem)('function', f.label, f.label, `UDF: ${f.catalog}.${f.database}.*`, f.description)), ...built_in_functions_1.BuiltInFunctions.map(f => (0, protocol_translation_1.toCompletionItem)('built-in', f.name, f.name, 'built-in function', f.description)));
-                    identifiers.push('function');
+                }
+                catch (e) {
+                    // Skip this statement if formatting failed
+                    allErrors.push({
+                        message: `Statement formatting failed: ${e instanceof Error ? e.message : String(e)}`,
+                        line: currentLineOffset
+                    });
                 }
             }
-            for (const candidate of Array.from(candidates.tokens)) {
-                const keyword = parsingValue.parser.vocabulary.getDisplayName(candidate[0]).replace(/'/g, '');
-                completionItems.push({
-                    kind: vscode_languageserver_1.CompletionItemKind.Keyword,
-                    label: keyword,
-                    insertText: keyword,
-                    insertTextFormat: vscode_languageserver_1.InsertTextFormat.Snippet
-                });
+            // For 'separator' type segments, we don't format them - they are preserved as-is
+            // Move line offset forward
+            currentLineOffset += segment.lines.length;
+        }
+        return { edits: allEdits, errors: allErrors };
+    }
+    /**
+     * Split lines into statement segments and separator segments
+     * Statement segments contain SQL statements (ending with semicolon)
+     * Separator segments contain blank lines and comments between statements
+     *
+     * NOTE: This method does NOT normalize separators anymore
+     * Normalization is done later via normalizeBlankLinesInEdits
+     */
+    splitIntoStatementSegments(lines) {
+        const segments = [];
+        let currentStatementLines = [];
+        let currentSeparatorLines = [];
+        let inStatement = false;
+        for (let i = 0; i < lines.length; i++) {
+            const line = lines[i];
+            const trimmedLine = line.trim();
+            // Check if this is a comment or blank line
+            const isCommentOrBlank = trimmedLine.startsWith('--') || trimmedLine.length === 0;
+            if (!inStatement && !isCommentOrBlank) {
+                // Start of a new statement
+                // First, save any accumulated separator lines (preserve original)
+                if (currentSeparatorLines.length > 0) {
+                    segments.push({ type: 'separator', lines: currentSeparatorLines });
+                    currentSeparatorLines = [];
+                }
+                inStatement = true;
+                currentStatementLines.push(line);
+            }
+            else if (inStatement) {
+                // Continue building the current statement
+                currentStatementLines.push(line);
+                // Check if this line ends the statement (contains semicolon)
+                if (line.includes(';')) {
+                    // Statement complete
+                    segments.push({ type: 'statement', lines: currentStatementLines });
+                    currentStatementLines = [];
+                    inStatement = false;
+                }
+            }
+            else {
+                // Not in a statement and this is a comment or blank line
+                currentSeparatorLines.push(line);
             }
         }
-        catch (e) {
+        // Handle remaining lines
+        if (currentSeparatorLines.length > 0) {
+            segments.push({ type: 'separator', lines: currentSeparatorLines });
+        }
+        if (currentStatementLines.length > 0 && currentStatementLines.some(l => l.trim().length > 0)) {
+            // Last statement without semicolon
+            segments.push({ type: 'statement', lines: currentStatementLines });
         }
-        return completionItems;
+        return segments;
     }
-    doFormatting(document, _options) {
-        const parsingValue = this._warehouse.getParsingValue(document);
-        try {
-            const formattedCode = (0, sql_formatter_1.formatDialect)(parsingValue.origin, {
-                dialect: sql_formatter_1.spark,
-                tabWidth: 2,
-                useTabs: false,
-                keywordCase: 'preserve',
-                linesBetweenQueries: 1,
-                newlineBeforeSemicolon: true
-            });
-            const textEdit = lsp.TextEdit.replace(lsp.Range.create(0, 0, document.lineCount, 0), formattedCode);
-            return [textEdit];
+    /**
+     * Normalize separator lines (blank lines + comments between statements)
+     * NOTE: This method is no longer used as normalization is done via normalizeBlankLinesInEdits
+     * Kept for potential future use
+     */
+    normalizeSeparator(lines) {
+        const result = ['']; // Always start with 1 blank line
+        let lastWasBlank = true; // We just added a blank line
+        for (const line of lines) {
+            const trimmedLine = line.trim();
+            const isBlank = trimmedLine.length === 0;
+            const isComment = trimmedLine.startsWith('--');
+            if (isComment) {
+                // Always preserve comments
+                result.push(line);
+                lastWasBlank = false;
+            }
+            else if (isBlank) {
+                // Only add blank line if the last line wasn't blank (prevent consecutive blanks)
+                if (!lastWasBlank) {
+                    result.push('');
+                    lastWasBlank = true;
+                }
+                // Otherwise skip this blank line
+            }
+            else {
+                // Non-comment, non-blank content (shouldn't happen in separators, but handle it)
+                result.push(line);
+                lastWasBlank = false;
+            }
         }
-        catch (e) {
-            return [];
+        return result;
+    }
+    /**
+     * Normalize blank lines in the formatted result
+     * This is a simplified safe version that works on text directly
+     * Rules:
+     * - Remove excessive consecutive blank lines (2+ become 1)
+     * - Preserves comments and statement structure
+     */
+    normalizeBlankLinesInEdits(document, edits) {
+        // If no edits, no need to normalize
+        if (edits.length === 0) {
+            return edits;
+        }
+        // Create a new document with edits applied
+        const text = document.getText();
+        const lines = text.split('\n');
+        // Apply edits using TextDocument.applyEdits logic
+        // Sort edits from bottom to top to avoid position shifting
+        const sortedEdits = [...edits].sort((a, b) => {
+            if (a.range.start.line !== b.range.start.line) {
+                return b.range.start.line - a.range.start.line;
+            }
+            return b.range.start.character - a.range.start.character;
+        });
+        // Build formatted text by applying edits
+        const result = lines.map((line, index) => ({ line, index }));
+        for (const edit of sortedEdits) {
+            const startLine = edit.range.start.line;
+            const endLine = edit.range.end.line;
+            if (startLine >= result.length)
+                continue;
+            // Handle single line edit
+            if (startLine === endLine) {
+                const lineObj = result[startLine];
+                const line = lineObj.line;
+                const before = line.substring(0, edit.range.start.character);
+                const after = line.substring(edit.range.end.character);
+                lineObj.line = before + edit.newText + after;
+            }
+            else {
+                // Multi-line edit: remove lines and merge
+                if (endLine >= result.length)
+                    continue;
+                const firstLine = result[startLine].line.substring(0, edit.range.start.character);
+                const lastLine = result[endLine].line.substring(edit.range.end.character);
+                result[startLine].line = firstLine + edit.newText + lastLine;
+                result.splice(startLine + 1, endLine - startLine);
+            }
+        }
+        const formattedLines = result.map(r => r.line);
+        // Now normalize blank lines
+        const normalizedLines = [];
+        let consecutiveBlankCount = 0;
+        let hasChanges = false;
+        for (let i = 0; i < formattedLines.length; i++) {
+            const line = formattedLines[i];
+            const isBlank = line.trim().length === 0;
+            if (isBlank) {
+                consecutiveBlankCount++;
+                // Only keep first blank line
+                if (consecutiveBlankCount === 1) {
+                    normalizedLines.push(line);
+                }
+                else {
+                    // Skip this blank line
+                    hasChanges = true;
+                }
+            }
+            else {
+                consecutiveBlankCount = 0;
+                normalizedLines.push(line);
+            }
         }
+        // Only return new edit if we actually removed blank lines
+        if (!hasChanges) {
+            return edits;
+        }
+        // Create a single edit that replaces entire document
+        const normalizedText = normalizedLines.join('\n');
+        return [
+            {
+                range: {
+                    start: { line: 0, character: 0 },
+                    end: { line: lines.length - 1, character: lines[lines.length - 1].length }
+                },
+                newText: normalizedText
+            }
+        ];
     }
-    doFoldingRanges(document) {
+    /**
+     * Determine if we should use fallback formatting
+     * Criteria: Parse errors exist and the coverage is below 60%
+     */
+    shouldUseFallbackFormatting(document, edits, _tree) {
+        const text = document.getText();
+        const totalLines = text.split('\n').length;
+        // Count how many lines have edits
+        const editedLines = new Set();
+        for (const edit of edits) {
+            for (let line = edit.range.start.line; line <= edit.range.end.line; line++) {
+                editedLines.add(line);
+            }
+        }
+        const coverage = editedLines.size / totalLines;
+        // Use fallback if coverage is less than 60%
+        // This indicates that many statements were not formatted due to parse errors
+        // We use a conservative threshold to avoid triggering fallback unnecessarily
+        return coverage < 0.6;
+    }
+    doFoldingRanges(document, options = typings_1.DEFAULT_FOLDING_OPTIONS) {
+        const optionsHash = options.level || 'all';
+        // 使用 uri + optionsHash 作为缓存 key，支持不同选项独立缓存
+        const cacheKey = `${document.uri}#${optionsHash}`;
+        const currentVersion = document.version;
+        const currentText = document.getText();
+        const currentHash = this._hashText(currentText);
+        // 检查缓存 (包含 options)
+        const cached = this._foldingCache.get(cacheKey);
+        if (cached &&
+            cached.version === currentVersion &&
+            cached.textHash === currentHash &&
+            cached.optionsHash === optionsHash) {
+            return cached.ranges;
+        }
+        // 缓存未命中，重新计算
         const parsingValue = this._warehouse.getParsingValue(document);
-        const listener = new statement_listener_1.StatementListener(document);
+        const listener = new statement_listener_1.StatementListener(document, options.level || 'all');
         tree_1.ParseTreeWalker.DEFAULT.walk(listener, parsingValue.tree);
-        const ranges = listener.getRanges();
-        return ranges.map(range => (0, protocol_translation_1.toFoldingRange)(range));
+        const foldingRanges = listener.getFoldingRanges();
+        const result = foldingRanges.map(info => (0, protocol_translation_1.toFoldingRange)(info.range, info.kind));
+        // 更新缓存 (lastAccess 由 LRUFoldingCache.set 自动设置)
+        this._foldingCache.set(cacheKey, {
+            ranges: result,
+            version: currentVersion,
+            textHash: currentHash,
+            optionsHash: optionsHash,
+            lastAccess: 0 // Will be overwritten by LRUFoldingCache.set
+        });
+        return result;
     }
-    getTokenByPosition(document, position) {
+    /**
+     * 清除指定文档的折叠缓存
+     * 清除所有与该文档 URI 相关的缓存（包括不同 options）
+     */
+    clearFoldingCache(uri) {
+        this._foldingCache.deleteByPrefix(uri);
+    }
+    /**
+     * 简单文本哈希算法
+     */
+    _hashText(text) {
+        let hash = 0;
+        for (let i = 0; i < text.length; i++) {
+            const char = text.charCodeAt(i);
+            hash = (hash << 5) - hash + char;
+            hash = hash & hash; // Convert to 32bit integer
+        }
+        return hash;
+    }
+    doReferences(document, position, skipSelf = true) {
         const parsingValue = this._warehouse.getParsingValue(document);
-        const offset = document.offsetAt(position);
-        return parsingValue.tokenStream.getTokens().find(t => t.startIndex <= offset && offset <= t.stopIndex + 1);
+        const tokens = parsingValue.tokenStream.getTokens();
+        const currentOffset = document.offsetAt(position);
+        const currentToken = tokens.find(token => {
+            return token.startIndex <= currentOffset && currentOffset <= token.stopIndex;
+        });
+        if (currentToken?.text === undefined) {
+            return [];
+        }
+        const listener = new tokens_collector_listener_1.TokensCollectorListener(document);
+        tree_1.ParseTreeWalker.DEFAULT.walk(listener, parsingValue.tree);
+        const semanticTokens = listener.getSemanticTokens();
+        const references = [];
+        const compareText = (str1, str2) => (0, utils_1.sanitizeText)(str1).localeCompare((0, utils_1.sanitizeText)(str2)) === 0;
+        const findReferences = (token) => {
+            if (currentToken.text && compareText(currentToken.text, token.name)) {
+                references.push(token.range);
+            }
+        };
+        semanticTokens.forEach(findReferences);
+        if (skipSelf) {
+            return references.filter(range => !(0, protocol_translation_1.positionInRange)(document, position, range));
+        }
+        return references;
     }
-    getSemicolonRange(str, target) {
-        if (target < 0 || target >= str.length) {
-            return [null, null];
-        }
-        let prevSemicolonIndex = null;
-        let nextSemicolonIndex = null;
-        for (let i = target - 1; i >= 0; i--) {
-            if (str[i] === ';') {
-                prevSemicolonIndex = i;
-                break;
+    doCompletion(document, position) {
+        const sql = cursor_1.defaultCursor.insertAt(document.getText(), position);
+        const charStream = antlr4ts_1.CharStreams.fromString(sql);
+        const lexer = new SparkSQLLexer_1.SparkSQLLexer(charStream);
+        const tokenStream = new antlr4ts_1.CommonTokenStream(lexer);
+        const parser = new SparkSQLParser_1.SparkSQLParser(tokenStream);
+        parser.buildParseTree = true;
+        parser.removeErrorListeners();
+        const tree = parser.program();
+        const completionItems = [];
+        // Use SparkSQLCompletionParser for context-aware completions
+        const sqlCompletionParser = new completion_parser_1.SparkSQLCompletionParser(tree, cursor_1.defaultCursor);
+        const suggestions = sqlCompletionParser.getSuggestions({
+            getCatalogs: () => {
+                return this.schemaRegistry.getSchemaCatalogs();
+            },
+            getDatabases: args => {
+                return this.schemaRegistry.getSchemaDatabases(args);
+            },
+            getTables: args => {
+                return this.schemaRegistry.getSchemaTables(args);
+            },
+            getColumns: args => {
+                return this.schemaRegistry.getSchemaColumns(args);
             }
+        });
+        completionItems.push(...suggestions);
+        const parsingValue = this._warehouse.getParsingValue(document);
+        const offset = document.offsetAt(position);
+        const token = parsingValue.tokenStream.getTokens().find(t => t.startIndex <= offset && offset <= t.stopIndex + 1);
+        if (!token) {
+            return { completionItems: [], identifiers: [] };
         }
-        for (let i = target + 1; i < str.length; i++) {
-            if (str[i] === ';') {
-                nextSemicolonIndex = i;
-                break;
+        const core = new antlr4_c3_1.CodeCompletionCore(parsingValue.parser);
+        core.ignoredTokens = new Set(ignoredTokens);
+        core.preferredRules = new Set(preferredRules);
+        let index = token.tokenIndex;
+        if (token.type === SparkSQLLexer_1.SparkSQLLexer.DOT) {
+            index += 1;
+        }
+        const candidates = core.collectCandidates(index);
+        const identifiers = [];
+        if (candidates.rules.size > 0) {
+            if (candidates.rules.has(SparkSQLParser_1.SparkSQLParser.RULE_functionName)) {
+                completionItems.push(...this.getBuiltInFunctionItems());
             }
         }
-        return [prevSemicolonIndex, nextSemicolonIndex || str.length];
+        for (const candidate of Array.from(candidates.tokens)) {
+            const keyword = parsingValue.parser.vocabulary.getDisplayName(candidate[0]).replace(/'/g, '');
+            completionItems.push({
+                kind: lsp.CompletionItemKind.Keyword,
+                label: keyword,
+                insertText: keyword,
+                insertTextFormat: lsp.InsertTextFormat.Snippet
+            });
+        }
+        return { completionItems, identifiers };
+    }
+    parseStructure(document, options = { ignoreInsertSource: false, addCreateContext: false }) {
+        const parsingValue = this._warehouse.getParsingValue(document);
+        const listener = new structure_listener_1.SparkStructureListener(document);
+        tree_1.ParseTreeWalker.DEFAULT.walk(listener, parsingValue.tree);
+        let nodes = listener.getNodes();
+        let edges = listener.getEdges();
+        if (options.ignoreInsertSource) {
+            const tmp = (0, structure_listener_1.ignoreInsertSource)(nodes, edges);
+            nodes = tmp.nodes;
+            edges = tmp.edges;
+        }
+        if (options.addCreateContext) {
+            const tmp = (0, structure_listener_1.addFromCreateContext)(nodes, edges);
+            nodes = tmp.nodes;
+            edges = tmp.edges;
+        }
+        return { nodes, edges };
+    }
+    collectTokens(document, range) {
+        const parsingValue = range
+            ? this._warehouse.parseDocument(document, range)
+            : this._warehouse.getParsingValue(document);
+        if (parsingValue.tokenCollection) {
+            return parsingValue.tokenCollection;
+        }
+        const listener = new tokens_collector_listener_1.TokensCollectorListener(document);
+        tree_1.ParseTreeWalker.DEFAULT.walk(listener, parsingValue.tree);
+        if (!range) {
+            this._warehouse.setTokenCollection(document, {
+                semanticTokens: listener.getSemanticTokens(),
+                stringTokens: listener.getStringTokens()
+            });
+        }
+        return { semanticTokens: listener.getSemanticTokens(), stringTokens: listener.getStringTokens() };
+    }
+    /**
+     * Lineage
+     */
+    getLineage(document, options = {}) {
+        const parsingValue = this._warehouse.getParsingValue(document);
+        const visitor = new lineage_visitor_1.LineageVisitor(options);
+        try {
+            parsingValue.tree.accept(visitor);
+            return visitor.getLineage();
+        }
+        catch (e) {
+            console.error(`Language Server: Lineage visitor traversal failed.\n\n${e}`);
+            return { nodes: [], edges: [] };
+        }
     }
 }
 exports.LSPServer = LSPServer;
+function isSemanticTokenReference(token1, token2) {
+    return (token1.name.replace(/['"`]/g, '').localeCompare(token2.name.replace(/['"`]/g, '')) === 0 &&
+        token1.type === token2.type);
+}
 //# sourceMappingURL=lsp-server.js.map