list tables context provider

src/aiProviders/continue/continueContextProvider.ts

@@ -151,7 +151,7 @@ export class db2ContextProvider implements IContextProvider {
                   columnData.length > 0 ? columnData[0].TABLE_SCHEMA : null;
 
                 // create context item
-                let prompt = `Db2 for i Table meta data for schema ${tableSchema} table ${table}\n`;
+                let prompt = `Db2 for i table Assistant: The following information is based on the ${table} table within the ${tableSchema} schema. Utilize the provided schema and table metadata to assist the user:\n`;
                 prompt += `Column Info: ${JSON.stringify(columnData)}\n\n`;
 
                 contextItems.push({

src/aiProviders/continue/listTablesContextProvider.ts

@@ -0,0 +1,127 @@
+import { ContextItem, ContextProviderDescription, ContextProviderExtras, ContextSubmenuItem, IContextProvider, LoadSubmenuItemsArgs } from "@continuedev/core";
+import * as fs from 'fs';
+import * as os from 'os';
+import * as path from 'path';
+import * as vscode from "vscode";
+import { JobManager } from "../../config";
+import { JobInfo } from "../../connection/manager";
+import Schemas from "../../database/schemas";
+import Table from "../../database/table";
+import { findPossibleTables } from "../context";
+
+const listDb2Table: ContextProviderDescription = {
+  title: "list Db2i Tables",
+  displayTitle: "Db2i-tables",
+  description: "Add Db2i Table info to Context",
+  type: "submenu"
+}
+
+export let listDb2TableContextProvider: Boolean = false;
+
+class ListDb2iTables implements IContextProvider {
+  get description(): ContextProviderDescription {
+    return listDb2Table;
+  }
+
+  getCurrentJob(): JobInfo {
+    const currentJob: JobInfo = JobManager.getSelection();
+    return currentJob;
+  }
+
+  private getDefaultSchema = (): string => {
+    const currentJob: JobInfo = this.getCurrentJob();
+    return currentJob?.job.options.libraries[0] || "QGPL";
+  };
+
+  async getColumnInfoForAllTables(schema: string) {
+    const items: TableColumn[] = await Table.getItems(schema);
+
+    return items.map((column) => ({
+      table_name: column.TABLE_NAME,
+      schema: column.TABLE_SCHEMA,
+      column_name: column.COLUMN_NAME,
+      column_data_type: column.DATA_TYPE,
+    }));
+  }
+
+  async getContextItems(
+    query: string,
+    extras: ContextProviderExtras
+  ): Promise<ContextItem[]> {
+    let contextitems: ContextItem[] = [];
+    const schema = this.getDefaultSchema();
+    if (query.toUpperCase() === schema.toUpperCase()) {
+      const tableInfo = await this.getColumnInfoForAllTables(schema);
+      contextitems.push({
+        name: `Info for all tables in ${schema}`,
+        content: `Db2 for i table Assistant: The following table and column information is from the ${query} schema. Utilize the provided schema and table metadata to assist the user:\n${JSON.stringify(tableInfo)}`,
+        description: "table metadata",
+      });
+    } else {
+      const tableInfo = await findPossibleTables(
+        null,
+        schema,
+        query.split(` `)
+      );
+      contextitems.push({
+        name: `${query}`,
+        content: `Db2 for i table Assistant: The following information is based on the ${query} table within the ${schema} schema. Utilize the provided schema and table metadata to assist the user:\n${JSON.stringify(tableInfo)}`,
+        description: "table metadata",
+      });
+    }
+    return contextitems;
+  }
+
+  async loadSubmenuItems(
+    args: LoadSubmenuItemsArgs
+  ): Promise<ContextSubmenuItem[]> {
+    const schema = this.getDefaultSchema();
+    const tables: BasicSQLObject[] = await Schemas.getObjects(schema, [
+      `tables`,
+    ]);
+
+    const schemaSubmenuItem: ContextSubmenuItem = {
+      id: schema,
+      title: schema,
+      description: `All table info in schema: ${schema}`,
+    };
+
+    const tableSubmenuItems: ContextSubmenuItem[] = tables.map((table) => ({
+      id: table.name,
+      title: table.name,
+      description: `${table.schema}-${table.name}`,
+    }));
+
+    return [schemaSubmenuItem, ...tableSubmenuItems];
+  }
+}
+
+export async function registerDb2iTablesProvider() {
+  const provider = new ListDb2iTables();
+  const continueID = `Continue.continue`;
+  const continueEx = vscode.extensions.getExtension(continueID);
+  if (continueEx) {
+    if (!continueEx.isActive) {
+      await continueEx.activate();
+    }
+
+    const continueAPI = continueEx?.exports;
+    if (listDb2TableContextProvider) {
+
+      // HACK: re register context provider work around
+      // save continue config file to trigger a config reload to update list tables provider
+      const configFile = path.join(os.homedir(), `.continue`, `config.json`);
+      const now = new Date();
+      fs.utimes(configFile, now, now, (err) => {
+        if (err) {
+          console.error('Error saving Continue config file:', err);
+          return;
+        }
+        vscode.window.showInformationMessage('Updated @Db2-Tables!');
+      });
+    } else {
+      continueAPI?.registerCustomContextProvider(provider);
+      listDb2TableContextProvider = true;
+    }
+  }
+}

src/database/table.ts

@@ -8,7 +8,8 @@ export default class Table {
    * @param {string} name Not user input
    * @returns {Promise<TableColumn[]>}
    */
-  static async getItems(schema: string, name: string): Promise<TableColumn[]> {
+  static async getItems(schema: string, name?: string): Promise<TableColumn[]> {
+    const params = name ? [schema, name] : [schema];
     const sql = [
       `SELECT `,
       `  column.TABLE_SCHEMA,`,
@@ -30,11 +31,14 @@ export default class Table {
       `    column.table_schema = key.table_schema and`,
       `    column.table_name = key.table_name and`,
       `    column.column_name = key.column_name`,
-      `WHERE column.TABLE_SCHEMA = ? AND column.TABLE_NAME = ?`,
+      `WHERE column.TABLE_SCHEMA = ?`,
+      ...[
+        name ? `AND column.TABLE_NAME = ${name}` : ``,
+      ],
       `ORDER BY column.ORDINAL_POSITION`,
     ].join(` `);
 
-    return JobManager.runSQL(sql, {parameters: [schema, name]});
+    return JobManager.runSQL(sql, {parameters: params});
   }
 
   /**

src/extension.ts

@@ -24,6 +24,7 @@ import { SelfTreeDecorationProvider, selfCodesResultsView } from "./views/jobMan
 import { registerContinueProvider } from "./aiProviders/continue/continueContextProvider";
 import { queryHistory } from "./views/queryHistoryView";
 import { activateChat, registerCopilotProvider } from "./aiProviders/copilot";
+import { registerDb2iTablesProvider } from "./aiProviders/continue/listTablesContextProvider";
 
 export interface Db2i {
   sqlJobManager: SQLJobManager,
@@ -98,6 +99,8 @@ export function activate(context: vscode.ExtensionContext): Db2i {
     onConnectOrServerInstall().then(() => {
       exampleBrowser.refresh();
       selfCodesView.setRefreshEnabled(Configuration.get(`jobSelfViewAutoRefresh`) || false);
+      // register list tables
+      registerDb2iTablesProvider();
       if (devMode && runTests) {
         runTests();
       }
@@ -110,6 +113,8 @@ export function activate(context: vscode.ExtensionContext): Db2i {
   // register continue provider
   registerContinueProvider();
 
+
+
   instance.subscribe(context, `disconnected`, `db2i-disconnected`, () => ServerComponent.reset());
 
   return { sqlJobManager: JobManager, sqlJob: (options?: JDBCOptions) => new OldSQLJob(options) };

src/views/jobManager/jobManagerView.ts

@@ -12,6 +12,7 @@ import { SelfCodesQuickPickItem } from "./selfCodes/selfCodesBrowser";
 import { updateStatusBar } from "./statusBar";
 import { setCancelButtonVisibility } from "../results";
 import { JDBCOptions } from "@ibm/mapepire-js/dist/src/types";
+import { registerDb2iTablesProvider } from "../../aiProviders/continue/listTablesContextProvider";
 
 const selectJobCommand = `vscode-db2i.jobManager.selectJob`;
 const activeColor = new vscode.ThemeColor(`minimapGutter.addedBackground`);
@@ -124,6 +125,9 @@ export class JobManagerView implements TreeDataProvider<any> {
 
                 try {
                   await selected.job.connect();
+
+                  // re register tables provider with potential new Schema
+                  await registerDb2iTablesProvider();
                 } catch (e) {
                   window.showErrorMessage(`Failed to start new job with updated properties.`);
                 }