Merge pull request #360 from codefori/fix/lateral_issues

Fix lateral use and improve performance with caching

Author: worksofliam

package-lock.json

@@ -1,3 +1,4 @@
+import Tokenizer from "sql-formatter/lib/src/lexer/Tokenizer";
 import SQLTokeniser, { NameTypes } from "./tokens";
 import { CTEReference, CallableReference, ClauseType, ClauseTypeWord, IRange, ObjectRef, QualifiedObject, StatementType, StatementTypeWord, Token } from "./types";
 
@@ -8,6 +9,15 @@ const tokenIs = (token: Token|undefined, type: string, value?: string) => {
 export default class Statement {
 	public type: StatementType = StatementType.Unknown;
 	private label: string|undefined;
+	private _blockTokens: Token[]|undefined;
+
+	private get blockTokens() {
+		if (!this._blockTokens) {
+			this._blockTokens = SQLTokeniser.createBlocks(this.tokens.slice(0));
+		}
+
+		return this._blockTokens;
+	}
 
   constructor(public tokens: Token[], public range: IRange) {
 		this.tokens = this.tokens.filter(newToken => newToken.type !== `newline`);
@@ -101,62 +111,34 @@ export default class Statement {
 		return currentClause;
 	}
 
-	getBlockRangeAt(offset: number) {
-		let start = -1;
-		let end = -1;
-
-		// Get the current token for the provided offset
-		let i = this.tokens.findIndex((token, i) => (offset >= token.range.start && offset <= token.range.end) || (offset > token.range.end && this.tokens[i+1] && offset < this.tokens[i+1].range.start));
-
-		let depth = 0;
-
-		if (tokenIs(this.tokens[i], `closebracket`)) {
-			i--;
-		}
-
-		if (tokenIs(this.tokens[i], `openbracket`)) {
-			start = i+1;
-			i++;
-		} else {
-			for (let x = i; x >= 0; x--) {
-				if (tokenIs(this.tokens[x], `openbracket`)) {
-					if (depth === 0) {
-						start = x+1;
-						break;
+	getBlockRangeAt(offset: number): IRange|undefined {
+		const blockContainsOffset = (cOffset: number, block: Token[]): IRange|undefined => {
+			const tokenInOffset = block.find(token => cOffset >= token.range.start && cOffset <= token.range.end);
+			if (tokenInOffset) {
+				if (tokenInOffset.type === `block`) {
+					if (tokenInOffset.block!.length > 0) {
+						return blockContainsOffset(cOffset, tokenInOffset.block!);
 					} else {
-						depth--;
+						const rawEnd = this.tokens.findIndex(token => token.range.end === tokenInOffset.range.end);
+						return {
+							start: rawEnd,
+							end: rawEnd
+						}
 					}
-				} else
-				if (tokenIs(this.tokens[x], `closebracket`)) {
-					depth++;
-				}
-			}
-		}
-
-		depth = 0;
-
-		for (let x = i; x <= this.tokens.length; x++) {
-			if (tokenIs(this.tokens[x], `openbracket`)) {
-				depth++;
-			} else
-			if (tokenIs(this.tokens[x], `closebracket`)) {
-				if (depth === 0) {
-					end = x;
-					break;
 				} else {
-					depth--;
+					const rawStart = this.tokens.findIndex(token => token.range.start === block[0].range.start);
+					const rawEnd = this.tokens.findIndex(token => token.range.end === block[block.length-1].range.end);
+					return {
+						start: rawStart,
+						end: rawEnd+1
+					};
 				}
 			}
-		}
 
-		if (start === -1 || end === -1) {
 			return undefined;
-		} else {
-			return {
-				start,
-				end
-			}
 		}
+
+		return blockContainsOffset(offset, this.blockTokens);
 	}
 
 	getCallableDetail(offset: number, withBlocks = false): CallableReference {
@@ -176,13 +158,22 @@ export default class Statement {
 	}
 
 	getBlockAt(offset: number): Token[] {
-		const range = this.getBlockRangeAt(offset);
+		const expandBlock = (tokens: Token[]): Token[] => {
+			const block = tokens.filter(token => token.type === `block`);
+			if (block.length > 0) {
+				return block.reduce((acc, token) => acc.concat(expandBlock(token.block!)), tokens);
+			}
 
-		if (range) {
-			return this.tokens.slice(range.start, range.end)
-		} else {
-			return []
+			return tokens;
+		}
+
+		let blockRange = this.getBlockRangeAt(offset);
+
+		if (blockRange) {
+			return expandBlock(this.tokens.slice(blockRange.start, blockRange.end));
 		}
+
+		return [];
 	}
 
 	getReferenceByOffset(offset: number) {
@@ -227,7 +218,7 @@ export default class Statement {
 	getCTEReferences(): CTEReference[] {
 		if (this.type !== StatementType.With) return [];
 
-		const withBlocks = SQLTokeniser.createBlocks(this.tokens.slice(0));
+		const withBlocks = this.blockTokens;
 
 		let cteList: CTEReference[] = [];
 
@@ -324,11 +315,16 @@ export default class Statement {
 
 	}
 
+	private static readonly BANNED_NAMES = [`VALUES`];
 	getObjectReferences(): ObjectRef[] {
 		let list: ObjectRef[] = [];
 
 		const doAdd = (ref?: ObjectRef) => {
-			if (ref) list.push(ref);
+			if (ref) {
+				if (ref.object.name && !Statement.BANNED_NAMES.includes(ref.object.name.toUpperCase())) {
+					list.push(ref);
+				}
+			}
 		}
 
 		const basicQueryFinder = (startIndex: number): void => {
@@ -365,6 +361,14 @@ export default class Statement {
 							i += 3; //For the brackets
 						}
 					}
+				} else if (currentClause === `from` && tokenIs(this.tokens[i], `function`)) {
+					const sqlObj = this.getRefAtToken(i, {includeParameters: true});
+					if (sqlObj) {
+						i += sqlObj.tokens.length;
+						if (sqlObj.isUDTF || sqlObj.fromLateral) {
+							i += 3; //For the brackets
+						}
+					}
 				}
 			}
 		}
@@ -538,18 +542,17 @@ export default class Statement {
 		return list;
 	}
 
-	private getRefAtToken(i: number, options: {withSystemName?: boolean} = {}): ObjectRef|undefined {
+	private getRefAtToken(i: number, options: {withSystemName?: boolean, includeParameters?: boolean} = {}): ObjectRef|undefined {
 		let sqlObj: ObjectRef;
 
 		let nextIndex = i;
 		let nextToken = this.tokens[i];
 
 		let endIndex = i;
 
-		const isUDTF = tokenIs(nextToken, `function`, `TABLE`);
-		const isLateral = tokenIs(nextToken, `function`, `LATERAL`);
+		const isSubSelect = tokenIs(nextToken, `function`, `TABLE`) || tokenIs(nextToken, `function`, `LATERAL`) || (options.includeParameters && tokenIs(nextToken, `function`));
 
-		if (isUDTF) {
+		if (isSubSelect) {
 			sqlObj = this.getRefAtToken(i+2);
 			if (sqlObj) {
 				sqlObj.isUDTF = true;
@@ -563,14 +566,6 @@ export default class Statement {
 				nextIndex = -1;
 				nextToken = undefined;
 			}
-		} else if (isLateral) {
-			const blockTokens = this.getBlockAt(nextToken.range.end+1);
-			const newStatement = new Statement(blockTokens, {start: nextToken.range.start, end: blockTokens[blockTokens.length-1].range.end});
-			[sqlObj] = newStatement.getObjectReferences();
-
-			sqlObj.fromLateral = true;
-			nextIndex = i + 2 + blockTokens.length;
-			nextToken = this.tokens[nextIndex];
 
 		} else {
 			if (nextToken && NameTypes.includes(nextToken.type)) {
@@ -618,7 +613,7 @@ export default class Statement {
 				}
 			}
 
-			if (!isUDTF) {
+			if (!isSubSelect && !sqlObj.isUDTF) {
 				sqlObj.tokens = this.tokens.slice(i, endIndex+1);
 			}
 

src/language/sql/statement.ts

@@ -1105,7 +1105,29 @@ parserScenarios(`Object references`, ({newDoc}) => {
 
     expect(refs[2].object.name).toBe(`object_statistics`);
     expect(refs[2].object.schema).toBe(`qsys2`);
-    expect(refs[2].alias).toBe(`z`);
+    expect(refs[2].alias).toBe(undefined);
+  });
+
+  test('SELECT FROM LATERAL', () => {
+    const lines = [
+      `SELECT id, id_phone, t.phone_number`,
+      `  FROM testlateral AS s,`,
+      `       LATERAL(VALUES (1, S.phone1),`,
+      `                      (2, S.phone2),`,
+      `                      (3, S.phone3)) AS T(id_phone, phone_number)`,
+    ].join(`\n`);
+
+    const document = new Document(lines);
+
+    expect(document.statements.length).toBe(1);
+
+    const statement = document.statements[0];
+
+    expect(statement.type).toBe(StatementType.Select);
+
+    const refs = statement.getObjectReferences();
+    console.log(refs);
+    expect(refs.length).toBe(1);
   });
 
   test(`Multiple UDTFs`, () => {
@@ -1407,7 +1429,7 @@ parserScenarios(`PL body tests`, ({newDoc}) => {
 
     const refs = statement.getObjectReferences();
     const ctes = statement.getCTEReferences();
-
+    
     expect(refs.length).toBe(10);
     expect(refs[0].object.name).toBe(`shipments`);
     expect(refs[0].alias).toBe(`s`);