diff --git a/src/aiProviders/continue/continueContextProvider.ts b/src/aiProviders/continue/continueContextProvider.ts index 6f08f8a5..e2464b88 100644 --- a/src/aiProviders/continue/continueContextProvider.ts +++ b/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({ diff --git a/src/aiProviders/continue/listTablesContextProvider.ts b/src/aiProviders/continue/listTablesContextProvider.ts new file mode 100644 index 00000000..9d6ab22e --- /dev/null +++ b/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 { + 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 { + 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; + } + } +} \ No newline at end of file diff --git a/src/database/table.ts b/src/database/table.ts index 06517510..91d3cf08 100644 --- a/src/database/table.ts +++ b/src/database/table.ts @@ -8,7 +8,8 @@ export default class Table { * @param {string} name Not user input * @returns {Promise} */ - static async getItems(schema: string, name: string): Promise { + static async getItems(schema: string, name?: string): Promise { + 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}); } /** diff --git a/src/extension.ts b/src/extension.ts index 2138f7c7..ed5782a9 100644 --- a/src/extension.ts +++ b/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) }; diff --git a/src/views/jobManager/jobManagerView.ts b/src/views/jobManager/jobManagerView.ts index e967a3a5..71e7caf8 100644 --- a/src/views/jobManager/jobManagerView.ts +++ b/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 { 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.`); }