Grab objects based on the library list

worksofliam · worksofliam · commit ac6f693b896a · 2025-02-25T17:56:39.000-05:00
Signed-off-by: worksofliam &lt;mrliamallan@live.co.uk&gt;
diff --git a/global.d.ts b/global.d.ts
@@ -33,6 +33,12 @@ interface SQLParm {
   ROW_TYPE: "P" | "R",
 }
 
+interface ResolvedSqlObject {
+  schema: string;
+  name: string;
+  sqlType: string;
+}
+
 interface BasicSQLObject {
   type: string;
   tableType: string;
diff --git a/package.json b/package.json
@@ -1296,14 +1296,6 @@
           {
             "name": "build",
             "description": "Build an SQL statement"
-          },
-          {
-            "name": "activity",
-            "description": "Summarize the activity on the system"
-          },
-          {
-            "name": "set-schema",
-            "description": "Set the current schema"
           }
         ]
       }
diff --git a/src/aiProviders/context.ts b/src/aiProviders/context.ts
@@ -25,7 +25,7 @@ interface MarkdownRef {
 
 interface ContextDefinition {
   id: string;
-  type: SQLType;
+  type: string;
   content: any;
 }
 
@@ -112,58 +112,106 @@ export async function buildSchemaDefinition(schema: string): Promise<Partial<Bas
   return compressedData;
 }
 
+// hello my.world how/are you? -> [hello, my, ., world, how, /, are, you]
+function splitUpUserInput(input: string): string[] {
+  input = input.replace(/[,!?$%\^&\*;:{}=\-`~()]/g, "")
+
+  let parts: string[] = [];
+
+  // Split the input string by spaces, dots and forward slash
+
+  let cPart = ``;
+  let char: string;
+
+  const addPart = () => {
+    if (cPart) {
+      parts.push(cPart);
+      cPart = ``;
+    }
+  }
+
+  for (let i = 0; i < input.length; i++) {
+    char = input[i];
+
+    switch (char) {
+      case ` `:
+        addPart();
+        break;
+
+      case `/`:
+      case `.`:
+        addPart();
+        parts.push(char);
+        break;
+
+      default:
+        if ([`/`, `.`].includes(cPart)) {
+          addPart();
+        }
+
+        cPart += char;
+        break;
+    }
+  }
+  
+  addPart();
+
+  return parts;
+}
+
 /**
- * Generates the SQL table definitions for the given schema and input words.
+ * Generates the SQL object definitions for the given input string.
  *
- * This function parses the input words to identify potential table references,
- * filters them based on the schema's available tables, and generates the SQL
- * definitions for the identified tables.
+ * This function parses the input words to identify potential SQL object references,
+ * and generates the SQL definitions for the identified objects based on the library list.
  *
- * @param {string} schema - The schema name to search for tables.
- * @param {string[]} input - An array of words that may contain table references.
- * @returns {Promise<{[key: string]: string | undefined}>} A promise that resolves to an object
- * where the keys are table names and the values are their corresponding SQL definitions.
+ * @param {string} input - A string that may contain table references.
  */
-export async function generateTableDefinition(schema: string, input: string[]) {
-  let tables: ContextDefinition[] = [];
+export async function getSqlContextItems(input: string): Promise<ContextDefinition[]> {
+  let contextItems: ContextDefinition[] = [];
+
   // Parse all SCHEMA.TABLE references first
-  const schemaTableRefs = input.filter((word) => word.includes("."));
-  const justWords = input.map((word) =>
-    word.replace(/[,\/#!?$%\^&\*;:{}=\-_`~()]/g, "")
-  );
+  const tokens = splitUpUserInput(input);
 
   // Remove plurals from words
-  justWords.push(
-    ...justWords
+  tokens.push(
+    ...tokens
       .filter((word) => word.endsWith("s"))
       .map((word) => word.slice(0, -1))
   );
 
-  // Filter prompt for possible refs to tables
-  const validWords = justWords
-    .filter(
-      (word) => word.length > 2 && !word.endsWith("s") && !word.includes(`'`)
-    )
-    .map((word) => `${Statement.delimName(word, true)}`);
-
-  const allTables: BasicSQLObject[] = await Schemas.getObjects(schema, [
-    `tables`,
-  ]);
-
-  const filteredTables = Array.from(
-    new Set(
-      validWords.filter((word) => allTables.some((table) => table.name == word))
-    )
-  );
+  let possibleRefs: {name: string, schema?: string}[] = [];
+  for (let i = 0; i < tokens.length; i++) {
+    const token = tokens[i];
+    
+    if (token[i+1] && [`.`, `/`].includes(token[i+1]) && tokens[i + 2]) {
+      const nextToken = tokens[i + 2];
+
+      possibleRefs.push({
+        name: Statement.delimName(nextToken, true),
+        schema: Statement.delimName(token, true),
+      });
+
+      i += 2; // Skip the next token as it's already processed
+
+    } else {
+      possibleRefs.push({
+        name: Statement.delimName(token, true),
+      });
+    }
+  }
+
+  const allObjects = await Schemas.resolveObjects(possibleRefs);
 
   await Promise.all(
-    filteredTables.map(async (token) => {
+    allObjects.map(async (o) => {
       try {
-        const content = await Schemas.generateSQL(schema, token, `TABLE`);
+        const content = await Schemas.generateSQL(o.schema, o.name, o.sqlType);
+
         if (content) {
-          tables.push({
-            id: token,
-            type: `tables`,
+          contextItems.push({
+            id: o.name,
+            type: o.sqlType,
             content: content,
           });
         }
@@ -173,19 +221,7 @@ export async function generateTableDefinition(schema: string, input: string[]) {
     })
   );
 
-  // check for QSYS2
-  // for (const item of schemaTableRefs) {
-  //   const [curSchema, table] = item.split(`.`);
-  //   if (curSchema.toUpperCase() === `QSYS2`) {
-  //     try {
-  //       const content = await Schemas.getObjects
-  //     } catch (e) {
-  //       continue
-  //     }
-  //   }
-  // }
-
-  return tables;
+  return contextItems;
 }
 
 /**
diff --git a/src/aiProviders/continue/continueContextProvider.ts b/src/aiProviders/continue/continueContextProvider.ts
@@ -13,7 +13,7 @@ import { SelfCodeNode } from "../../views/jobManager/selfCodes/nodes";
 import {
   buildSchemaDefinition,
   canTalkToDb,
-  generateTableDefinition
+  getSqlContextItems
 } from "../context";
 import { DB2_SELF_PROMPT, DB2_SYSTEM_PROMPT } from "../prompts";
 import Configuration from "../../configuration";
@@ -179,10 +179,7 @@ export class db2ContextProvider implements IContextProvider {
             break;
           default:
             // 2. TABLE References
-            const tablesRefs = await generateTableDefinition(
-              schema,
-              fullInput.split(` `)
-            );
+            const tablesRefs = await getSqlContextItems(fullInput);
             for (const table of tablesRefs) {
               contextItems.push({
                 name: `table definition for ${table.id}`,
diff --git a/src/aiProviders/continue/listTablesContextProvider.ts b/src/aiProviders/continue/listTablesContextProvider.ts
@@ -11,7 +11,7 @@ import Schemas from "../../database/schemas";
 import Table from "../../database/table";
 import {
   buildSchemaDefinition,
-  generateTableDefinition
+  getSqlContextItems
 } from "../context";
 import Configuration from "../../configuration";
 
@@ -82,10 +82,7 @@ class ListDb2iTables implements IContextProvider {
       }
 
     } else {
-      const tablesRefs = await generateTableDefinition(
-        this.schema,
-        extras.fullInput.split(` `)
-      );
+      const tablesRefs = await getSqlContextItems(extras.fullInput);
       for (const table of tablesRefs) {
         contextItems.push({
           name: `table definition for ${table.id}`,
diff --git a/src/aiProviders/prompt.ts b/src/aiProviders/prompt.ts
@@ -1,7 +1,7 @@
 import { JobManager } from "../config";
 import Configuration from "../configuration";
 import { JobInfo } from "../connection/manager";
-import { buildSchemaDefinition, canTalkToDb, generateTableDefinition } from "./context";
+import { buildSchemaDefinition, canTalkToDb, getSqlContextItems } from "./context";
 import { DB2_SYSTEM_PROMPT } from "./prompts";
 
 export interface PromptOptions {
@@ -48,10 +48,7 @@ export async function buildPrompt(input: string, options: PromptOptions = {}): P
     // TODO: self?
 
     progress(`Building table definition for ${currentSchema}...`);
-    const refs = await generateTableDefinition(
-      currentSchema,
-      input.split(` `)
-    );
+    const refs = await getSqlContextItems(input);
 
     if (options.history) {
       contextItems.push(...options.history);
diff --git a/src/database/schemas.ts b/src/database/schemas.ts
@@ -47,6 +47,69 @@ function getFilterClause(againstColumn: string, filter: string, noAnd?: boolean)
 }
 
 export default class Schemas {
+  /**
+   * Resolves to the following SQL types: TABLE, VIEW, ALIAS, INDEX, FUNCTION and PROCEDURE
+   */
+  static async resolveObjects(sqlObjects: {name: string, schema?: string}[]): Promise<ResolvedSqlObject[]> {
+    let statements: string[] = [];
+    let parameters: BasicColumnType[] = [];
+
+    // First, we use OBJECT_STATISTICS to resolve the object based on the library list.
+    // But, if the object is qualified with a schema, we need to use that schema to get the correct object.
+    for (const obj of sqlObjects) {
+      if (obj.schema) {
+        statements.push(
+          `select OBJLONGNAME as name, OBJLONGSCHEMA as schema, SQL_OBJECT_TYPE as sqlType from table(qsys2.object_statistics(?, '*ALL', object_name => ?)) where SQL_OBJECT_TYPE IS NOT NULL`
+        );
+        parameters.push(obj.schema, obj.name);
+      } else {
+        statements.push(
+          `select OBJLONGNAME as name, OBJLONGSCHEMA as schema, SQL_OBJECT_TYPE as sqlType from table(qsys2.object_statistics('*LIBL', '*ALL', object_name => ?)) where SQL_OBJECT_TYPE IS NOT NULL`
+        );
+        parameters.push(obj.name);
+      }
+    }
+
+    // We have to do a little bit more logic for routines, because they are not included properly in OBJECT_STATISTICS.
+    // So we do a join against the library list view and SYSROUTINES (which has a list of all routines in a given schema)
+    // to get the correct schema and name.
+    const unqualified = sqlObjects.filter(obj => !obj.schema).map(obj => obj.name);
+    const qualified = sqlObjects.filter(obj => obj.schema);
+    const qualifiedClause = qualified.map(obj => `(s.routine_name = ? AND s.routine_schema = ?)`).join(` OR `);
+
+    let baseStatement = [
+      `select s.routine_name as name, l.schema_name as schema, s.ROUTINE_TYPE as sqlType`,
+      `from qsys2.library_list_info as l`,
+      `right join qsys2.sysroutines as s on l.schema_name = s.routine_schema`,
+      `where `,
+      `  l.schema_name is not null and`,
+      `  s.routine_name in (${sqlObjects.map(() => `?`).join(`, `)})`,
+    ].join(` `);
+    parameters.push(...unqualified);
+
+    if (qualified.length > 0) {
+      baseStatement += ` and (${qualifiedClause})`;
+      parameters.push(...qualified.flatMap(obj => [obj.name, obj.schema]));
+    }
+
+    statements.push(baseStatement);
+
+    if (statements.length === 0) {
+      return [];
+    }
+
+    const query = `${statements.join(" UNION ALL ")}`;
+    const objects: any[] = await JobManager.runSQL(query, { parameters });
+    
+    const resolvedObjects: ResolvedSqlObject[] = objects.map(object => ({
+      name: object.NAME,
+      schema: object.SCHEMA,
+      sqlType: object.SQLTYPE
+    }));
+
+    return resolvedObjects;
+  }
+
   /**
    * @param schema Not user input
    */

Original file line number	Diff line number	Diff line change
`@@ -1296,14 +1296,6 @@`
`1296`	`1296`	`{`
`1297`	`1297`	`"name": "build",`
`1298`	`1298`	`"description": "Build an SQL statement"`
`1299`		`- },`
`1300`		`- {`
`1301`		`- "name": "activity",`
`1302`		`- "description": "Summarize the activity on the system"`
`1303`		`- },`
`1304`		`- {`
`1305`		`- "name": "set-schema",`
`1306`		`- "description": "Set the current schema"`
`1307`	`1299`	`}`
`1308`	`1300`	`]`
`1309`	`1301`	`}`