Create separate flags for enabling/disabling each preview feature in isolation.

dataproc_jupyter_plugin/__init__.py

@@ -12,12 +12,18 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.from ._version import __version__
 import logging
+import sys
 
 from google.cloud.jupyter_config.tokenrenewer import CommandTokenRenewer
 from jupyter_server.services.sessions.sessionmanager import SessionManager
 from kernels_mixer.kernels import MixingMappingKernelManager
 from kernels_mixer.kernelspecs import MixingKernelSpecManager
 from kernels_mixer.websockets import DelegatingWebsocketConnection
+from traitlets import Bool
+from jupyter_server.serverapp import ServerApp
+
+from dataproc_jupyter_plugin.commons.vertexPlatform import _parse_vertexai_flag
+
 
 from .handlers import DataprocPluginConfig, configure_gateway_client_url, setup_handlers
 
@@ -83,6 +89,11 @@ def _load_jupyter_server_extension(server_app):
     server_app: jupyterlab.labapp.LabApp
         JupyterLab application instance
     """
+    flag_status = _parse_vertexai_flag()
+
+    # Log the verterai flag
+    server_app.log.info(f"Vertexai flag: {flag_status}")
+
     setup_handlers(server_app.web_app)
     name = "dataproc_jupyter_plugin"
     server_app.log.info(f"Registered {name} server extension")

dataproc_jupyter_plugin/commons/vertexPlatform.py

@@ -0,0 +1,22 @@
+# dataproc_jupyter_plugin/flag.py
+
+# This is a simple module to store the verterai flag
+# Initialized with a default value
+VERTEX_PLATFORM_ENABLED = False
+
+
+def _parse_vertexai_flag():
+    """
+    Parse the vertexai flag from command-line arguments
+
+    This function will be called during extension loading
+    """
+    global VERTEX_PLATFORM_ENABLED
+    import sys
+
+    for arg in sys.argv:
+        if arg.startswith("--isVertexPlatform="):
+            VERTEX_PLATFORM_ENABLED = arg.split("=")[1].lower() == "true"
+            break
+
+    return VERTEX_PLATFORM_ENABLED

dataproc_jupyter_plugin/controllers/vertexPlatform.py

@@ -0,0 +1,28 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import json
+import tornado
+from jupyter_server.base.handlers import APIHandler
+from dataproc_jupyter_plugin.commons.vertexPlatform import _parse_vertexai_flag
+
+class VertexPlatformController(APIHandler):
+    @tornado.web.authenticated
+    async def get(self):
+        try:
+            return self.finish({"enabled": _parse_vertexai_flag()})
+
+        except Exception as e:
+            return self.finish({"error": str(e)})

dataproc_jupyter_plugin/handlers.py

@@ -37,6 +37,7 @@
     composer,
     dataproc,
     executor,
+    vertexPlatform,
 )
 
 _region_not_set_error = """GCP region not set in gcloud.
@@ -209,6 +210,7 @@ def full_path(name):
         "bigQueryPreview": bigquery.PreviewController,
         "bigQueryProjectsList": bigquery.ProjectsController,
         "bigQuerySearch": bigquery.SearchController,
+        "api/vertexPlatform": vertexPlatform.VertexPlatformController,
     }
     handlers = [(full_path(name), handler) for name, handler in handlersMap.items()]
     web_app.add_handlers(host_pattern, handlers)

src/index.ts

@@ -78,6 +78,10 @@ const iconPythonLogo = new LabIcon({
   svgstr: pythonLogo
 });
 
+interface VertexResponse {
+  enabled: boolean;
+  error?: string;  // Optional error property for error cases
+}
 const extension: JupyterFrontEndPlugin<void> = {
   id: PLUGIN_ID,
   autoStart: true,
@@ -199,7 +203,7 @@ const extension: JupyterFrontEndPlugin<void> = {
      * previewEnabled flag and hides or shows the GCS browser or DPMS explorer
      * as necessary.
      */
-    const onPreviewEnabledChanged = () => {
+    const onPreviewEnabledChanged = async () => {
       previewEnabled = settings.get('previewEnabled').composite as boolean;
       if (!previewEnabled) {
         // Preview was disabled, tear everything down.
@@ -214,6 +218,8 @@ const extension: JupyterFrontEndPlugin<void> = {
       } else {
         // Preview was enabled, (re)create DPMS and GCS.
         if (!panelDpms && !panelGcs) {
+          const vertexai: VertexResponse = await requestAPI('api/vertexPlatform')
+          if(!vertexai.enabled){
           panelDpms = new Panel();
           panelDpms.id = 'dpms-tab';
           panelDpms.title.caption = 'Dataset Explorer - DPMS';
@@ -247,6 +253,28 @@ const extension: JupyterFrontEndPlugin<void> = {
           }
           app.shell.add(panelDpms, 'left', { rank: 1001 });
         }
+        else{
+          panelDpms = new Panel();
+          panelDpms.id = 'dpms-tab';
+          panelDpms.title.caption = 'Dataset Explorer - DPMS';
+          panelDpms.addWidget(new dpmsWidget(app as JupyterLab, themeManager));
+          if (bqFeature.enable_bigquery_integration && !panelDatasetExplorer) {
+            panelDatasetExplorer = new Panel();
+            panelDatasetExplorer.id = 'dataset-explorer-tab';
+            panelDatasetExplorer.title.caption = 'Dataset Explorer - BigQuery';
+            panelDatasetExplorer.addWidget(
+              new BigQueryWidget(
+                app as JupyterLab,
+                settingRegistry as ISettingRegistry,
+                bqFeature.enable_bigquery_integration as boolean,
+                themeManager
+              )
+            );
+          }
+          onThemeChanged();
+          app.shell.add(panelDpms, 'left', { rank: 1001 });
+        }
+        }
       }
     };
 

src/utils/const.ts

@@ -16,7 +16,8 @@
  */
 
 import { requestAPI } from '../handler/handler';
-export const VERSION_DETAIL = '0.1.78';
+const { version } = require('../../package.json');
+export const VERSION_DETAIL = version;
 export const CREATE_CLUSTER_URL =
   'https://console.cloud.google.com/dataproc/clusters';
 export const CREATE_BATCH_URL =
@@ -170,7 +171,7 @@ export const DCU_HOURS = 3600000;
 export const GB_MONTHS = 2592000;
 export const TITLE_LAUNCHER_CATEGORY = 'Google Cloud Resources';
 export const SPARK_HISTORY_SERVER = 'Spark History Server';
-export const DEFAULT_LABEL_DETAIL = 'client:dataproc-jupyter-plugin';
+export const DEFAULT_LABEL_DETAIL = 'client:bigquery-jupyter-plugin';
 export const JOB_FIELDS_EXCLUDED = ['queryList', 'properties', 'args'];
 export const BATCH_FIELDS_EXCLUDED = ['queryList', 'properties'];
 export const KEY_MESSAGE =