Fetch schema info when referenced

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

Author: worksofliam

src/aiProviders/context.ts

@@ -209,7 +209,15 @@ export async function getContentItemsForRefs(allObjects: ResolvedSqlObject[]): P
     allObjects.map(async (o) => {
       try {
         if (o.sqlType === `SCHEMA`) {
-          // TODO: maybe we want to include info about a schema here?
+          const schemaSemantic = await buildSchemaDefinition(o.name);
+          if (schemaSemantic) {
+            return {
+              name: `SCHEMA Definition`,
+              description: `${o.name} definition`,
+              content: JSON.stringify(schemaSemantic)
+            };
+          }
+
           return undefined;
 
         } else {

src/aiProviders/prompt.ts

@@ -36,21 +36,8 @@ export async function getContextItems(input: string, options: PromptOptions = {}
 
   if (currentJob) {
     const currentSchema = currentJob?.job.options.libraries[0] || "QGPL";
-
     const useSchemaDef: boolean = Configuration.get<boolean>(`ai.useSchemaDefinition`);
 
-    if (useSchemaDef) {
-      progress(`Building schema definition for ${currentSchema}...`);
-      const schemaSemantic = await buildSchemaDefinition(currentSchema);
-      if (schemaSemantic) {
-        contextItems.push({
-          name: `SCHEMA Definition`,
-          description: `${currentSchema} definition`,
-          content: JSON.stringify(schemaSemantic)
-        });
-      }
-    }
-
     // TODO: self?
 
     progress(`Finding objects to work with...`);
@@ -73,24 +60,42 @@ export async function getContextItems(input: string, options: PromptOptions = {}
     if (context.refs.length === 1) {
       const ref = context.refs[0];
       const prettyNameRef = Statement.prettyName(ref.name);
-      progress(`Finding objects related to ${prettyNameRef}...`);
 
-      const relatedObjects = await Schemas.getRelatedObjects(ref);
-      const contentItems = await getContentItemsForRefs(relatedObjects);
+      if (ref.sqlType === `SCHEMA`) {
+        followUps.push(
+          `What are some objects in that schema?`,
+          `What is the difference between a schema and a library?`,
+        );
+      } else {
+        progress(`Finding objects related to ${prettyNameRef}...`);
+
+        const relatedObjects = await Schemas.getRelatedObjects(ref);
+        const contentItems = await getContentItemsForRefs(relatedObjects);
 
-      contextItems.push(...contentItems);
+        contextItems.push(...contentItems);
 
-      if (relatedObjects.length === 1) {
-        followUps.push(`How is ${prettyNameRef} related to ${Statement.prettyName(relatedObjects[0].name)}?`);
-      } else if (ref.sqlType === `TABLE`) {
-        followUps.push(`What are some objects related to that table?`);
+        if (relatedObjects.length === 1) {
+          followUps.push(`How is ${prettyNameRef} related to ${Statement.prettyName(relatedObjects[0].name)}?`);
+        } else if (ref.sqlType === `TABLE`) {
+          followUps.push(`What are some objects related to that table?`);
+        }
       }
 
     } else if (context.refs.length > 1) {
       const randomRef = context.refs[Math.floor(Math.random() * context.refs.length)];
       const prettyNameRef = Statement.prettyName(randomRef.name);
 
       followUps.push(`What are some objects related to ${prettyNameRef}?`);
+    } else if (useSchemaDef) {
+      progress(`Getting info for schema ${currentSchema}...`);
+      const schemaSemantic = await buildSchemaDefinition(currentSchema);
+      if (schemaSemantic) {
+        contextItems.push({
+          name: `SCHEMA Definition`,
+          description: `${currentSchema} definition`,
+          content: JSON.stringify(schemaSemantic)
+        });
+      }
     }
 
     if (options.withDb2Prompt) {

src/database/schemas.ts

@@ -71,12 +71,12 @@ export default class Schemas {
     for (const obj of sqlObjects) {
       if (obj.schema) {
         statements.push(
-          `${BASE_RESOLVE_SELECT} from table(qsys2.object_statistics(?, '*ALL', object_name => ?)) where SQL_OBJECT_TYPE IS NOT NULL`
+          `${BASE_RESOLVE_SELECT} from table(qsys2.object_statistics(?, '*ALL', object_name => ?))`
         );
         parameters.push(obj.schema, obj.name);
       } else {
         statements.push(
-          `${BASE_RESOLVE_SELECT} from table(qsys2.object_statistics('*LIBL', '*ALL', object_name => ?)) where SQL_OBJECT_TYPE IS NOT NULL`
+          `${BASE_RESOLVE_SELECT} from table(qsys2.object_statistics('*LIBL', '*ALL', object_name => ?))`
         );
         parameters.push(obj.name);
       }
@@ -117,7 +117,7 @@ export default class Schemas {
       name: object.NAME,
       schema: object.SCHEMA,
       sqlType: object.SQLTYPE
-    }));
+    })).filter(o => o.sqlType);
 
     return resolvedObjects;
   }