Better support for THEN and LOOP

Signed-off-by: worksofliam <[email protected]>

Author: worksofliam

src/language/sql/document.ts

@@ -50,8 +50,14 @@ export default class Document {
 
         case `keyword`:
           switch (tokens[i].value?.toUpperCase()) {
+            case `LOOP`:
+              // This handles the case that 'END LOOP' is supported.
+              if (currentStatementType === StatementType.End) {
+                break;
+              } 
             case `BEGIN`:
             case `DO`:
+            case `THEN`:
               // We include BEGIN in the current statement
               // then the next statement beings
               const statementTokens = tokens.slice(statementStart, i+1);
@@ -96,6 +102,7 @@ export default class Document {
     let depth = 0;
 
     for (const statement of this.statements) {
+      console.log(currentGroup);
         if (statement.isBlockEnder()) {
           if (depth > 0) {
             currentGroup.push(statement);

src/language/sql/statement.ts

@@ -43,7 +43,7 @@ export default class Statement {
 		// These statements can end with BEGIN, which signifies a block starter
 		if ([StatementType.Create, StatementType.Declare].includes(this.type)) {
 			const last = this.tokens[this.tokens.length-1];
-			if (tokenIs(last, `keyword`, `BEGIN`) || tokenIs(last, `keyword`, `DO`)) {
+			if (tokenIs(last, `keyword`, `BEGIN`)) {
 				return true;
 			}
 		}

src/language/sql/tests/statements.test.ts

@@ -1013,6 +1013,76 @@ describe(`Object references`, () => {
     expect(refs[1].object.name).toBe(`OBJECT_STATISTICS`);
     expect(refs[1].object.schema).toBe(`QSYS2`);
     expect(refs[1].alias).toBe(`b`);
+  });
+
+  test('LOOP statements', () => {
+    const lines = [
+      `CREATE OR REPLACE PROCEDURE KRAKEN917.Wait_For_Kraken(kraken_job_name varchar(10), delay_time bigint default 30)`,
+      `BEGIN`,
+      `    DECLARE v_sql_stmt CLOB(1M) CCSID 37;`,
+      ``,
+      `    DECLARE number_of_active_jobs INT;`,
+      ``,
+      `    CALL systools.lprintf('Waiting for job to finish...');`,
+      ``,
+      `    fetch_loop: LOOP`,
+      `        SET v_sql_stmt ='values(SELECT COUNT(*) FROM TABLE (qsys2.active_job_info(subsystem_list_filter => ''QBATCH'')) WHERE JOB_NAME_SHORT LIKE ''' CONCAT kraken_job_name CONCAT ''') into ?';`,
+      ``,
+      `        PREPARE values_st FROM v_sql_stmt;`,
+      ``,
+      `        EXECUTE values_st USING number_of_active_jobs;`,
+      ``,
+      `        IF number_of_active_jobs = 0 THEN`,
+      `            CALL SYSTOOLS.LPRINTF(kraken_job_name CONCAT ' JOB DONE');`,
+      ``,
+      `            LEAVE fetch_loop;`,
+      ``,
+      `        END IF;`,
+      ``,
+      `        CALL qsys2.qcmdexc('DLYJOB ' CONCAT delay_time);`,
+      ``,
+      `    END LOOP fetch_loop;`,
+      ``,
+      `END;`,
+    ].join(`\n`);
+
+    const document = new Document(lines);
+
+    const groups = document.getStatementGroups();
+    expect(groups.length).toBe(1);
+
+    const group = groups[0];
+
+    // console.log(group.statements.map((s, so) => `${so}  ` + s.type.padEnd(10) + ` ` + s.tokens.map(t => t.value).join(' ')));
+
+    expect(group.statements.length).toBe(16);
+    expect(group.statements.map(s => s.type)).toEqual([
+      'Create',  'Declare',
+      'Declare', 'Call',
+      'Unknown', 'Unknown',
+      'Unknown', 'Unknown',
+      'Unknown', 'Call',
+      'Unknown', 'End',
+      'Call',    'End',
+      'Unknown', 'End'
+    ]);
+
+    let refs;
+
+    const firstCall = group.statements[3];
+    refs = firstCall.getObjectReferences();
+    expect(refs.length).toBe(1);
+    expect(refs[0].object.name).toBe(`lprintf`);
+
+    const secondCall = group.statements[9];
+    refs = secondCall.getObjectReferences();
+    expect(refs.length).toBe(1);
+    expect(refs[0].object.name).toBe(`LPRINTF`);
+
+    const thirdCall = group.statements[12];
+    refs = thirdCall.getObjectReferences();
+    expect(refs.length).toBe(1);
+    expect(refs[0].object.name).toBe(`qcmdexc`);
   })
 });
 

src/language/sql/tokens.ts

@@ -74,7 +74,7 @@ export default class SQLTokeniser {
     {
       name: `KEYWORD`,
       match: [{ type: `word`, match: (value: string) => {
-        return [`AS`, `FOR`, `OR`, `REPLACE`, `BEGIN`, `DO`, `END`, `CURSOR`, `DEFAULT`, `HANDLER`, `REFERENCES`, `ON`, `UNIQUE`, `SPECIFIC`, `EXTERNAL`].includes(value.toUpperCase());
+        return [`AS`, `FOR`, `OR`, `REPLACE`, `BEGIN`, `DO`, `THEN`, `LOOP`, `END`, `CURSOR`, `DEFAULT`, `HANDLER`, `REFERENCES`, `ON`, `UNIQUE`, `SPECIFIC`, `EXTERNAL`].includes(value.toUpperCase());
       } }],
       becomes: `keyword`,
     },