Merge pull request #195 from Shubha-accenture/notebook-kernel-status

Notebook kernel status, Create runtime network and API disable error message changes

dataproc_jupyter_plugin/controllers/composer.py

@@ -32,11 +32,7 @@ async def get(self):
                     await credentials.get_cached(), self.log, client_session
                 )
                 environments = await client.list_environments()
-                response = []
-                for environment in environments:
-                    env = environment.dict()
-                    response.append(env)
-                self.finish(json.dumps(response))
+                self.finish(json.dumps(environments, default=lambda x: x.dict()))
         except Exception as e:
             self.log.exception(f"Error fetching composer environments: {str(e)}")
             self.finish({"error": str(e)})

src/bigQuery/bigQueryService.tsx

@@ -420,4 +420,29 @@ export class BigQueryService {
       );
     }
   };
+
+  static listBigQueryDatasetsAPIService = async (projectId: string) => {
+    const pageToken = '';
+    try {
+      const data: any = await requestAPI(
+        `bigQueryDataset?project_id=${projectId}&pageToken=${pageToken}`
+      );
+      return data;
+    } catch (reason) {
+      return reason;
+    }
+  };
+  static getBigQuerySearchCatalogAPIService = async () => {
+    try {
+      const data: any = await requestAPI(
+        `bigQuerySearch?search_string=''&type=(table|dataset)&system=bigquery`,
+        {
+          method: 'POST'
+        }
+      );
+      return data;
+    } catch (reason) {
+      return reason;
+    }
+  };
 }

src/controls/NotebookButtonExtension.tsx

@@ -33,6 +33,8 @@ import notebookSchedulerIcon from '../../style/icons/scheduler_calendar_month.sv
 import { NotebookScheduler } from '../scheduler/notebookScheduler';
 import { ISettingRegistry } from '@jupyterlab/settingregistry';
 
+import { Widget } from '@lumino/widgets';
+
 const iconLogs = new LabIcon({
   name: 'launcher:logs-icon',
   svgstr: logsIcon
@@ -56,6 +58,7 @@ const iconNotebookScheduler = new LabIcon({
 class NotebookButtonExtensionPoint implements IDisposable {
   // IDisposable required.
   isDisposed: boolean;
+  private kernelStatusWidget: Widget;
   private readonly sparkLogsButton: ToolbarButton;
   private readonly sessionDetailsButton: ToolbarButton;
   private readonly sessionDetailsButtonDisable: ToolbarButton;
@@ -78,6 +81,37 @@ class NotebookButtonExtensionPoint implements IDisposable {
     this.isDisposed = false;
     this.context.sessionContext.sessionChanged.connect(this.onSessionChanged);
 
+    this.context.sessionContext.statusChanged.connect(this.updateKernelStatus);
+
+    // Create the kernel status widget
+    this.kernelStatusWidget = new Widget();
+    this.kernelStatusWidget.node.className = 'kernel-status-toolbar';
+    this.kernelStatusWidget.node.style.display = 'flex';
+    this.kernelStatusWidget.node.style.alignItems = 'center';
+
+    const kernelStatusLabel = document.createElement('span');
+    kernelStatusLabel.textContent =
+      this.context.sessionContext.kernelDisplayName.includes('Local') ||
+      this.context.sessionContext.kernelDisplayName.includes('No Kernel')
+        ? ''
+        : '|';
+    kernelStatusLabel.style.marginRight = '5px';
+    kernelStatusLabel.className = 'kernel-status-label';
+
+    const kernelStatusValue = document.createElement('span');
+    kernelStatusValue.textContent = 'Loading...';
+    kernelStatusValue.style.color = 'gray';
+    kernelStatusValue.className = 'kernel-status-value';
+
+    this.kernelStatusWidget.node.appendChild(kernelStatusLabel);
+    this.kernelStatusWidget.node.appendChild(kernelStatusValue);
+
+    // Add the widget to the toolbar
+    this.panel.toolbar.insertItem(12, 'kernel-status', this.kernelStatusWidget);
+
+    // Initial fetch of kernel status
+    this.fetchAndUpdateKernelStatus();
+
     this.sparkLogsButton = new ToolbarButton({
       icon: iconLogs,
       onClick: this.onSparkLogsClick,
@@ -125,6 +159,102 @@ class NotebookButtonExtensionPoint implements IDisposable {
     );
   }
 
+  /**
+   * Fetch and update kernel status from the KernelAPI.
+   */
+  private async fetchAndUpdateKernelStatus() {
+    try {
+      const currentKernelId = this.context.sessionContext.session?.kernel?.id;
+      if (!currentKernelId) {
+        this.setKernelStatus('Initializing', '#455A64');
+        return;
+      }
+
+      const runningKernels = await KernelAPI.listRunning();
+      const currentKernel = runningKernels.find(
+        kernel => kernel.id === currentKernelId
+      );
+
+      if (currentKernel) {
+        const kernelStatus = currentKernel?.execution_state ?? 'Unknown';
+        this.setKernelStatus(kernelStatus, this.getStatusColor(kernelStatus));
+      } else {
+        this.setKernelStatus('Failed', '#D93025');
+      }
+    } catch (error) {
+      console.error('Error fetching kernel status:', error);
+      this.setKernelStatus('Failed', '#D93025');
+    }
+  }
+
+  /**
+   * Update the kernel status label and color.
+   */
+  private setKernelStatus(status: string, color: string) {
+    const kernelStatusValue = this.kernelStatusWidget.node.querySelector(
+      '.kernel-status-value'
+    ) as HTMLElement;
+    const kernelStatusLabel = this.kernelStatusWidget.node.querySelector(
+      '.kernel-status-label'
+    ) as HTMLElement;
+
+    if (kernelStatusValue && kernelStatusLabel) {
+      kernelStatusLabel.textContent =
+        this.context.sessionContext.kernelDisplayName.includes('Local') ||
+        this.context.sessionContext.kernelDisplayName.includes('No Kernel')
+          ? ''
+          : '|';
+
+      let displayStatus = '';
+      if (!this.context.sessionContext.kernelDisplayName.includes('Local')) {
+        switch (status) {
+          case 'idle':
+            displayStatus = 'Idle (Ready)';
+            break;
+          case 'busy':
+            displayStatus = 'Busy (Executing)';
+            break;
+          default:
+            displayStatus = status.charAt(0).toUpperCase() + status.slice(1);
+            break;
+        }
+      }
+
+      kernelStatusValue.textContent = displayStatus;
+      kernelStatusValue.style.color = color;
+      kernelStatusValue.style.marginRight = '5px';
+    }
+  }
+
+  /**
+   * Get color based on kernel status.
+   */
+  private getStatusColor(status: string): string {
+    switch (status.toLowerCase()) {
+      case 'idle':
+        return '#188038';
+      case 'busy':
+        return '#3A8DFF';
+      case 'starting':
+        return '#FFB300';
+      case 'initializing':
+        return ' #455A64';
+      case 'restarting':
+        return '#9747FF';
+      case 'Unknown':
+        return '#C9C9C9';
+      default:
+        return '#FF9800';
+    }
+  }
+
+  /**
+   * Event handler to update kernel status when the session status changes.
+   */
+  private updateKernelStatus = async () => {
+    await this.fetchAndUpdateKernelStatus();
+  };
+
   private onNotebookSchedulerClick = () => {
     const content = new NotebookScheduler(
       this.app as JupyterLab,
@@ -234,10 +364,15 @@ class NotebookButtonExtensionPoint implements IDisposable {
    * Event handler for when the session changes, we need to reattach kernel change events.
    */
   private onSessionChanged = async (session: ISessionContext) => {
+    session.connectionStatusChanged.connect(this.updateKernelStatus);
     session.kernelChanged.connect(this.onKernelChanged);
   };
 
   dispose() {
+    this.context.sessionContext.statusChanged.disconnect(
+      this.updateKernelStatus
+    );
+    this.kernelStatusWidget.dispose();
     this.context.sessionContext.sessionChanged.disconnect(
       this.onSessionChanged
     );

src/gcs/gcsDrive.ts

@@ -70,6 +70,8 @@ export class GCSDrive implements Contents.IDrive {
       const para = document.createElement('p');
       para.id = 'gcs-list-bucket-error';
       para.style.color = '#ff0000';
+      para.style.maxWidth= '100%';
+      para.style.whiteSpace='normal';
       para.textContent = content?.error?.message;
       paragraph = document.getElementById('filter-buckets-objects');
       paragraph?.after(para);

src/index.ts

@@ -64,6 +64,10 @@ import * as path from 'path';
 import { requestAPI } from './handler/handler';
 import { eventEmitter } from './utils/signalEmitter';
 import { BigQueryWidget } from './bigQuery/bigQueryWidget';
+import { RunTimeSerive } from './runtime/runtimeService';
+import { Notification } from '@jupyterlab/apputils';
+import { BigQueryService } from './bigQuery/bigQueryService';
+import { SchedulerService } from './scheduler/schedulerServices';
 
 const iconDpms = new LabIcon({
   name: 'launcher:dpms-icon',
@@ -165,11 +169,95 @@ const extension: JupyterFrontEndPlugin<void> = {
 
     // Capture the signal
     eventEmitter.on('dataprocConfigChange', (message: string) => {
+      checkAllApisEnabled();
       if (bqFeature.enable_bigquery_integration) {
         loadBigQueryWidget('');
       }
     });
 
+    const checkAllApisEnabled = async () => {
+      const dataprocClusterResponse =
+        await RunTimeSerive.listClustersDataprocAPIService();
+
+      let bigqueryDatasetsResponse;
+      const credentials = await authApi();
+      if (credentials?.project_id) {
+        bigqueryDatasetsResponse =
+          await BigQueryService.listBigQueryDatasetsAPIService(
+            credentials.project_id
+          );
+      }
+
+      const composerListResponse =
+        await SchedulerService.listComposersAPICheckService();
+      const dataCatalogResponse =
+        await BigQueryService.getBigQuerySearchCatalogAPIService();
+
+      const apiChecks = [
+        {
+          response: dataprocClusterResponse,
+          errorKey: 'error.message',
+          errorMessage: 'Cloud Dataproc API has not been used in project',
+          notificationMessage: 'The Cloud Dataproc API is not enabled.',
+          enableLink:
+            'https://console.cloud.google.com/apis/library/dataproc.googleapis.com'
+        },
+        {
+          response: bigqueryDatasetsResponse,
+          errorKey: 'error',
+          errorMessage: 'has not enabled BigQuery',
+          notificationMessage: 'The BigQuery API is not enabled.',
+          enableLink:
+            'https://console.cloud.google.com/apis/library/bigquery.googleapis.com'
+        },
+        {
+          response: composerListResponse,
+          errorKey: 'Error fetching environments list',
+          errorMessage: 'Cloud Composer API has not been used in project',
+          notificationMessage: 'The Cloud Composer API is not enabled.',
+          enableLink:
+            'https://console.cloud.google.com/apis/library/composer.googleapis.com'
+        },
+        {
+          response: dataCatalogResponse,
+          errorKey: 'error',
+          errorMessage:
+            'Google Cloud Data Catalog API has not been used in project',
+          notificationMessage: 'Google Cloud Data Catalog API is not enabled.',
+          enableLink:
+            'https://console.cloud.google.com/apis/library/datacatalog.googleapis.com'
+        }
+      ];
+
+      apiChecks.forEach(
+        ({
+          response,
+          errorKey,
+          errorMessage,
+          notificationMessage,
+          enableLink
+        }) => {
+          const errorValue = errorKey
+            .split('.')
+            .reduce((acc, key) => acc?.[key], response);
+          if (errorValue && errorValue.includes(errorMessage)) {
+            Notification.error(notificationMessage, {
+              actions: [
+                {
+                  label: 'Enable',
+                  callback: () => window.open(enableLink, '_blank'),
+                  displayType: 'link'
+                }
+              ],
+              autoClose: false
+            });
+          }
+        }
+      );
+    };
+
+    await checkAllApisEnabled();
+
     /**
      * Handler for when the Jupyter Lab theme changes.
      */

src/runtime/createRunTime.tsx

@@ -1447,7 +1447,7 @@ function CreateRunTime({
                               handleSubNetworkChange(val)
                             }
                             renderInput={params => (
-                              <TextField {...params} label="subnetwork" />
+                              <TextField {...params} label="subnetwork*" />
                             )}
                           />
                         </div>