Table Viewer optimization in ITC

client/src/connection/itc/CodeRunner.ts

@@ -2,25 +2,55 @@
 // SPDX-License-Identifier: Apache-2.0
 import { commands } from "vscode";
 
+import { v4 } from "uuid";
+
 import { ITCSession } from ".";
 import { LogLine, getSession } from "..";
 import { useRunStore } from "../../store";
+import { Session } from "../session";
+import { extractTextBetweenTags } from "../util";
 
 let wait: Promise<string> | undefined;
 
+export async function executeRawCode(code: string): Promise<string> {
+  const randomId = v4();
+  const startTag = `<${randomId}>`;
+  const endTag = `</${randomId}>`;
+  const task = () =>
+    _runCode(
+      async (session) => {
+        // eslint-disable-next-line @typescript-eslint/consistent-type-assertions
+        await (session as ITCSession).execute(
+          `Write-Host "${startTag}"\n${code}\nWrite-Host "${endTag}"\n`,
+        );
+      },
+      startTag,
+      endTag,
+    );
+
+  wait = wait ? wait.then(task) : task();
+  return wait;
+}
+
 export async function runCode(
   code: string,
   startTag: string = "",
   endTag: string = "",
 ): Promise<string> {
-  const task = () => _runCode(code, startTag, endTag);
+  const task = () =>
+    _runCode(
+      // eslint-disable-next-line @typescript-eslint/consistent-type-assertions
+      async (session) => await (session as ITCSession).run(code, true),
+      startTag,
+      endTag,
+    );
 
   wait = wait ? wait.then(task) : task();
   return wait;
 }
 
 async function _runCode(
-  code: string,
+  runCallback: (session: Session) => void,
   startTag: string = "",
   endTag: string = "",
 ): Promise<string> {
@@ -45,30 +75,20 @@ async function _runCode(
   const onExecutionLogFn = session.onExecutionLogFn;
   const outputLines = [];
 
-  const addLine = (logLines: LogLine[]) =>
+  const addLine = (logLines: LogLine[]) => {
     outputLines.push(...logLines.map(({ line }) => line));
+  };
 
   try {
     await session.setup(true);
 
     // Lets capture output to use it on
     session.onExecutionLogFn = addLine;
 
-    // eslint-disable-next-line @typescript-eslint/consistent-type-assertions
-    await (session as ITCSession).run(code, true);
+    await runCallback(session);
 
     const logOutput = outputLines.filter((line) => line.trim()).join("");
-
-    logText =
-      startTag && endTag
-        ? logOutput
-            .slice(
-              logOutput.lastIndexOf(startTag),
-              logOutput.lastIndexOf(endTag),
-            )
-            .replace(startTag, "")
-            .replace(endTag, "")
-        : logOutput;
+    logText = extractTextBetweenTags(logOutput, startTag, endTag);
   } finally {
     unsubscribe && unsubscribe();
     // Lets update our session to write to the log

client/src/connection/itc/ItcLibraryAdapter.ts

@@ -11,8 +11,9 @@ import {
   TableData,
   TableRow,
 } from "../../components/LibraryNavigator/types";
-import { Column, ColumnCollection } from "../rest/api/compute";
-import { runCode } from "./CodeRunner";
+import { ColumnCollection } from "../rest/api/compute";
+import { getColumnIconType } from "../util";
+import { executeRawCode, runCode } from "./CodeRunner";
 import { Config } from "./types";
 
 class ItcLibraryAdapter implements LibraryAdapter {
@@ -45,31 +46,14 @@ class ItcLibraryAdapter implements LibraryAdapter {
   }
 
   public async getColumns(item: LibraryItem): Promise<ColumnCollection> {
-    const sql = `
-      %let OUTPUT;
-      proc sql;
-        select catx(',', name, type, varnum) as column into: OUTPUT separated by '~'
-        from sashelp.vcolumn
-        where libname='${item.library}' and memname='${item.name}'
-        order by varnum;
-      quit;
-      %put <COLOUTPUT> &OUTPUT; %put </COLOUTPUT>;
+    const code = `
+      $runner.GetColumns("${item.library}", "${item.name}")
     `;
-
-    const columnLines = processQueryRows(
-      await this.runCode(sql, "<COLOUTPUT>", "</COLOUTPUT>"),
-    );
-
-    const columns = columnLines.map((lineText): Column => {
-      const [name, type, index] = lineText.split(",");
-
-      return {
-        name,
-        type,
-        index: parseInt(index, 10),
-      };
-    });
-
+    const output = await executeRawCode(code);
+    const columns = JSON.parse(output).map((column) => ({
+      ...column,
+      type: getColumnIconType(column),
+    }));
     return {
       items: columns,
       count: -1,
@@ -213,48 +197,16 @@ class ItcLibraryAdapter implements LibraryAdapter {
     start: number,
     limit: number,
   ): Promise<{ rows: Array<string[]>; count: number }> {
-    const maxTableNameLength = 32;
-    const tempTable = `${item.name}${hms()}${start}`.substring(
-      0,
-      maxTableNameLength,
-    );
+    const fullTableName = `${item.library}.${item.name}`;
     const code = `
-      options nonotes nosource nodate nonumber;
-      %let COUNT;
-      proc sql;
-        SELECT COUNT(1) into: COUNT FROM  ${item.library}.${item.name};
-      quit;
-      data work.${tempTable};
-        set ${item.library}.${item.name};
-        if ${start + 1} <= _N_ <= ${start + limit} then output;
-      run;
-
-      filename out temp;
-      proc json nokeys out=out pretty; export work.${tempTable}; run;
-
-      %put <TABLEDATA>;
-      %put <Count>&COUNT</Count>;
-      data _null_; infile out; input; put _infile_; run;
-      %put </TABLEDATA>;
-      proc datasets library=work nolist nodetails; delete ${tempTable}; run;
-      options notes source date number;
+      $runner.GetDatasetRecords("${fullTableName}", ${start}, ${limit})
     `;
-
-    let output = await this.runCode(code, "<TABLEDATA>", "</TABLEDATA>");
-
-    // Extract result count
-    const countRegex = /<Count>(.*)<\/Count>/;
-    const countMatches = output.match(countRegex);
-    const count = parseInt(countMatches[1].replace(/\s|\n/gm, ""), 10);
-    output = output.replace(countRegex, "");
-
-    const rows = output.replace(/\n|\t/gm, "").slice(output.indexOf("{"));
+    const output = await executeRawCode(code);
     try {
-      const tableData = JSON.parse(rows);
-      return { rows: tableData[`SASTableData+${tempTable}`], count };
+      return JSON.parse(output);
     } catch (e) {
       console.warn("Failed to load table data with error", e);
-      console.warn("Raw output", rows);
+      console.warn("Raw output", output);
       throw new Error(
         l10n.t(
           "An error was encountered when loading table data. This usually happens when a table is too large or the data couldn't be processed. See console for more details.",
@@ -263,19 +215,29 @@ class ItcLibraryAdapter implements LibraryAdapter {
     }
   }
 
-  protected async runCode(
-    code: string,
-    startTag: string = "",
-    endTag: string = "",
+  protected async executionHandler(
+    callback: () => Promise<string>,
   ): Promise<string> {
     try {
-      return await runCode(code, startTag, endTag);
+      return await callback();
     } catch (e) {
       onRunError(e);
       commands.executeCommand("setContext", "SAS.librariesDisplayed", false);
       return "";
     }
   }
+
+  protected async runCode(
+    code: string,
+    startTag: string = "",
+    endTag: string = "",
+  ): Promise<string> {
+    return this.executionHandler(() => runCode(code, startTag, endTag));
+  }
+
+  protected async executeRawCode(code: string): Promise<string> {
+    return this.executionHandler(() => executeRawCode(code));
+  }
 }
 
 const processQueryRows = (response: string): string[] => {
@@ -289,9 +251,4 @@ const processQueryRows = (response: string): string[] => {
     .filter((value, index, array) => array.indexOf(value) === index);
 };
 
-const hms = () => {
-  const date = new Date();
-  return `${date.getHours()}${date.getMinutes()}${date.getSeconds()}`;
-};
-
 export default ItcLibraryAdapter;

client/src/connection/itc/index.ts

@@ -27,7 +27,7 @@ import {
   WORK_DIR_END_TAG,
   WORK_DIR_START_TAG,
 } from "./const";
-import { scriptContent } from "./script";
+import { getScript } from "./script";
 import { Config, ITCProtocol, LineCodes } from "./types";
 import { decodeEntities, escapePowershellString } from "./util";
 
@@ -110,9 +110,16 @@ export class ITCSession extends Session {
     );
     this._shellProcess.stdout.on("data", this.onShellStdOut);
     this._shellProcess.stderr.on("data", this.onShellStdErr);
-    this._shellProcess.stdin.write(scriptContent + "\n", this.onWriteComplete);
     this._shellProcess.stdin.write(
-      `$runner = New-Object -TypeName SASRunner -ArgumentList "${escapePowershellString(interopLibraryFolderPath || "")}"\n`,
+      getScript({
+        interopLibraryFolderPath: escapePowershellString(
+          interopLibraryFolderPath || "",
+        ),
+      }) + "\n",
+      this.onWriteComplete,
+    );
+    this._shellProcess.stdin.write(
+      `$runner = New-Object -TypeName SASRunner\n`,
       this.onWriteComplete,
     );
 
@@ -246,6 +253,24 @@ export class ITCSession extends Session {
     return runPromise;
   };
 
+  public execute = async (code: string): Promise<RunResult> => {
+    const runPromise = new Promise<RunResult>((resolve, reject) => {
+      this._runResolve = resolve;
+      this._runReject = reject;
+    });
+
+    this._html5FileName = "";
+    this._pollingForLogResults = true;
+    const codeToExecute = `${code}\nWrite-Host "${LineCodes.RunEndCode}"\n`;
+    this._shellProcess.stdin.write(codeToExecute, async (error) => {
+      if (error) {
+        this._runReject(error);
+      }
+    });
+
+    return runPromise;
+  };
+
   /**
    * Cleans up resources for the given SAS session.
    * @returns void promise.