From 5ba9b374f075c4a5852281ca9a05acb59b929329 Mon Sep 17 00:00:00 2001
From: Adam Shedivy <ajshedivyaj@gmail.com>
Date: Tue, 14 Jan 2025 12:06:49 -0600
Subject: [PATCH] improve checks for current job and current schema

---
 .../continue/listTablesContextProvider.ts     |  5 ++++-
 src/extension.ts                              |  4 +++-
 src/views/jobManager/jobManagerView.ts        | 21 +++++++++++++++----
 3 files changed, 24 insertions(+), 6 deletions(-)

diff --git a/src/aiProviders/continue/listTablesContextProvider.ts b/src/aiProviders/continue/listTablesContextProvider.ts
index fff02842..47969a14 100644
--- a/src/aiProviders/continue/listTablesContextProvider.ts
+++ b/src/aiProviders/continue/listTablesContextProvider.ts
@@ -112,7 +112,10 @@ class ListDb2iTables implements IContextProvider {
   }
 }
 
-export async function registerDb2iTablesProvider(schema: string) {
+export async function registerDb2iTablesProvider(schema?: string) {
+  if (!schema) {
+    return;
+  }
   const continueID = `Continue.continue`;
   const continueEx = vscode.extensions.getExtension(continueID);
   if (continueEx) {
diff --git a/src/extension.ts b/src/extension.ts
index 13e5f5f1..219b0a39 100644
--- a/src/extension.ts
+++ b/src/extension.ts
@@ -102,7 +102,9 @@ export function activate(context: vscode.ExtensionContext): Db2i {
       exampleBrowser.refresh();
       selfCodesView.setRefreshEnabled(Configuration.get(`jobSelfViewAutoRefresh`) || false);
       // register list tables
-      registerDb2iTablesProvider(JobManager.getSelection().job.options.libraries[0]);
+      const currentJob = JobManager.getSelection();
+      const currentSchema = currentJob?.job.options.libraries[0];
+      registerDb2iTablesProvider(currentSchema);
       if (devMode && runTests) {
         runTests();
       }
diff --git a/src/views/jobManager/jobManagerView.ts b/src/views/jobManager/jobManagerView.ts
index fdcbc504..e0c97e3c 100644
--- a/src/views/jobManager/jobManagerView.ts
+++ b/src/views/jobManager/jobManagerView.ts
@@ -305,15 +305,28 @@ export class JobManagerView implements TreeDataProvider<any> {
 
     const selectedJob = JobManager.getSelection();
     const selectedSchema = selectedJob.job.options.libraries[0];
-    
+
     // re-register db2i tables context provider with current schema
-    if (provider && provider.getCurrentSchema().toLowerCase() !== selectedSchema.toLowerCase()) {
+    if (
+      provider &&
+      selectedJob !== undefined &&
+      provider.getCurrentSchema() &&
+      selectedSchema &&
+      provider.getCurrentSchema().trim().toLowerCase() !==
+        selectedSchema.trim().toLowerCase()
+    ) {
       registerDb2iTablesProvider(selectedSchema);
     }
 
-    setCancelButtonVisibility(selectedJob && selectedJob.job.getStatus() === "busy");
+    setCancelButtonVisibility(
+      selectedJob && selectedJob.job.getStatus() === "busy"
+    );
     sqlLanguageStatus.setState(selectedJob !== undefined);
-    commands.executeCommand(`setContext`, `vscode-db2i:jobManager.hasJob`, selectedJob !== undefined);
+    commands.executeCommand(
+      `setContext`,
+      `vscode-db2i:jobManager.hasJob`,
+      selectedJob !== undefined
+    );
   }
 
   getTreeItem(element: vscode.TreeItem) {