WIP: allow users to reference tables with schema.table

Author: ajshedivy

src/chat/chat.ts

@@ -65,6 +65,7 @@ export function activateChat(context: vscode.ExtensionContext) {
             `Getting information from ${Statement.prettyName(usingSchema)}...`
           );
           let refs = await findPossibleTables(
+            stream,
             usingSchema,
             request.prompt.split(` `)
           );
@@ -77,17 +78,17 @@ export function activateChat(context: vscode.ExtensionContext) {
 
           if (Object.keys(refs).length === 0) {
             stream.progress(`No references found. Doing bigger lookup...`);
-            refs = await findPossibleTables(usingSchema, []);
+            refs = await findPossibleTables(stream, usingSchema, []);
           }
 
           if (Object.keys(refs).length > 0) {
             stream.progress(`Building response...`);
             messages.push(
               vscode.LanguageModelChatMessage.User(
-                `Give the developer an SQL statement or information based on the prompt and following table references. Always include code examples where is makes sense. Do not make suggestions for reference you do not have.`
+                `Provide the developer with SQL statements or relevant information based on the user's prompt and referenced table structures. Always include practical code examples where applicable. Ensure all suggestions are directly applicable to the structures and data provided and avoid making suggestions outside the scope of the available information.`
               ),
               vscode.LanguageModelChatMessage.User(
-                `Here are the table references for current schema ${usingSchema}\n${refsToMarkdown(
+                `Here are the table references ${refsToMarkdown(
                   refs
                 )}`
               ),

src/chat/context.ts

@@ -1,4 +1,5 @@
 import { JobManager } from "../config";
+import * as vscode from "vscode";
 import Statement from "../database/statement";
 
 export function canTalkToDb() {
@@ -14,21 +15,60 @@ export function getDefaultSchema(): string {
 
 export type TableRefs = { [key: string]: TableColumn[] };
 
-export async function findPossibleTables(schema: string, words: string[]) {
-  words = words.map((word) =>
-    word.replace(/[.,\/#!?$%\^&\*;:{}=\-_`~()]/g, "")
-  );
+export async function getTableMetaData(schema: string, tableName: string): Promise<TableColumn[]> {
+  const objectFindStatement = [
+    `SELECT `,
+    `  column.TABLE_NAME,`,
+    `  column.COLUMN_NAME,`,
+    `  key.CONSTRAINT_NAME,`,
+    `  column.DATA_TYPE, `,
+    `  column.CHARACTER_MAXIMUM_LENGTH,`,
+    `  column.NUMERIC_SCALE, `,
+    `  column.NUMERIC_PRECISION,`,
+    `  column.IS_NULLABLE, `,
+    // `  column.HAS_DEFAULT, `,
+    // `  column.COLUMN_DEFAULT, `,
+    `  column.COLUMN_TEXT, `,
+    `  column.IS_IDENTITY`,
+    `FROM QSYS2.SYSCOLUMNS2 as column`,
+    `LEFT JOIN QSYS2.syskeycst as key`,
+    `  on `,
+    `    column.table_schema = key.table_schema and`,
+    `    column.table_name = key.table_name and`,
+    `    column.column_name = key.column_name`,
+    `WHERE column.TABLE_SCHEMA = '${Statement.delimName(schema, true)}'`,
+    `AND column.TABLE_NAME = '${Statement.delimName(tableName, true)}'`,
+    `ORDER BY column.ORDINAL_POSITION`,
+  ].join(` `);
 
-  // Add extra words for words with S at the end, to ignore possible plurals
-  words.forEach((item) => {
-    if (item.endsWith(`s`)) {
-      words.push(item.slice(0, -1));
+  return await JobManager.runSQL(objectFindStatement);
+}
+
+export async function parsePromptForRefs(stream: vscode.ChatResponseStream, prompt: string[]): Promise<TableRefs> {
+  const tables: TableRefs = {};
+  for (const word of prompt) {
+    const [schema, table] = word.split(`.`);
+    if (schema && table) {
+      stream.progress(`looking up information for ${schema}.${table}`)
+      const data = await getTableMetaData(schema, table);
+      tables[table] = tables[table] || [];
+      tables[table].push(...data);
     }
-  });
+  }
+  return tables;
+}
+
+export async function findPossibleTables(stream: vscode.ChatResponseStream, schema: string, words: string[]) {
 
-  const validWords = words
-    .filter((item) => item.length > 2 && !item.includes(`'`))
-    .map((item) => `'${Statement.delimName(item, true)}'`);
+  let tables: TableRefs = {}
+
+  // parse all SCHEMA.TABLE references first
+  tables = await parsePromptForRefs(stream, words.filter(word => word.includes('.')));
+
+  // filter prompt for possible refs to tables
+  const validWords = words.map(word => word.replace(/[.,\/#!?$%\^&\*;:{}=\-_`~()]/g, ""))
+    .filter(word => word.length > 2 && !word.endsWith('s') && !word.includes(`'`))
+    .map(word => `'${Statement.delimName(word, true)}'`);
 
   const objectFindStatement = [
     `SELECT `,
@@ -62,16 +102,10 @@ export async function findPossibleTables(schema: string, words: string[]) {
   // TODO
   const result: TableColumn[] = await JobManager.runSQL(objectFindStatement);
 
-  const tables: TableRefs = {};
-
-  for (const row of result) {
-    if (!tables[row.TABLE_NAME]) {
-      tables[row.TABLE_NAME] = [];
-    }
-
+  result.forEach(row => {
+    if (!tables[row.TABLE_NAME]) tables[row.TABLE_NAME] = [];
     tables[row.TABLE_NAME].push(row);
-  }
-
+  });
   return tables;
 }