GoogleCloudDataproc · ptwng · May 20, 2024 · Mar 5, 2024 · Mar 5, 2024 · Mar 5, 2024
diff --git a/dataproc_jupyter_plugin/contollers/bigqueryController.py b/dataproc_jupyter_plugin/contollers/bigqueryController.py
@@ -14,9 +14,9 @@
 
 
 import json
+
 from jupyter_server.base.handlers import APIHandler
 import tornado
-from dataproc_jupyter_plugin import handlers
 from dataproc_jupyter_plugin.services.bigqueryService import (
     BigQueryDatasetInfoService,
     BigQueryDatasetListService,
@@ -28,6 +28,7 @@
 )
 from dataproc_jupyter_plugin.utils.constants import bq_public_dataset_project_id
 from google.cloud.jupyter_config import gcp_project
+from dataproc_jupyter_plugin.utils.credentials import GetCachedCredentials
 
 
 class BigqueryDatasetController(APIHandler):
@@ -37,7 +38,7 @@ def get(self):
             page_token = self.get_argument("pageToken")
             project_id = self.get_argument("project_id")
             bigquery_dataset = BigQueryDatasetListService()
-            credentials = handlers.get_cached_credentials(self.log)
+            credentials = GetCachedCredentials.get_cached_credentials(self.log)
             dataset_list = bigquery_dataset.list_datasets(
                 credentials, page_token, project_id, self.log
             )
@@ -55,7 +56,7 @@ def get(self):
             dataset_id = self.get_argument("dataset_id")
             project_id = self.get_argument("project_id")
             bigquery_dataset = BigQueryTableListService()
-            credentials = handlers.get_cached_credentials(self.log)
+            credentials = GetCachedCredentials.get_cached_credentials(self.log)
             table_list = bigquery_dataset.list_table(
                 credentials, dataset_id, page_token, project_id, self.log
             )
@@ -72,7 +73,7 @@ def get(self):
             dataset_id = self.get_argument("dataset_id")
             project_id = self.get_argument("project_id")
             bq_dataset = BigQueryDatasetInfoService()
-            credentials = handlers.get_cached_credentials(self.log)
+            credentials = GetCachedCredentials.get_cached_credentials(self.log)
             dataset_info = bq_dataset.list_dataset_info(
                 credentials, dataset_id, project_id, self.log
             )
@@ -90,7 +91,7 @@ def get(self):
             table_id = self.get_argument("table_id")
             project_id = self.get_argument("project_id")
             bq_table = BigQueryTableInfoService()
-            credentials = handlers.get_cached_credentials(self.log)
+            credentials = GetCachedCredentials.get_cached_credentials(self.log)
             table_info = bq_table.list_table_info(
                 credentials, dataset_id, table_id, project_id, self.log
             )
@@ -106,12 +107,19 @@ def get(self):
         try:
             dataset_id = self.get_argument("dataset_id")
             table_id = self.get_argument("table_id")
-            page_token = self.get_argument("pageToken")
+            max_results = self.get_argument("max_results")
+            start_index = self.get_argument("start_index")
             project_id = self.get_argument("project_id")
             bq_preview = BigQueryPreviewService()
-            credentials = handlers.get_cached_credentials(self.log)
+            credentials = GetCachedCredentials.get_cached_credentials(self.log)
             preview_data = bq_preview.bigquery_preview_data(
-                credentials, dataset_id, table_id, page_token, project_id, self.log
+                credentials,
+                dataset_id,
+                table_id,
+                max_results,
+                start_index,
+                project_id,
+                self.log,
             )
             self.finish(json.dumps(preview_data))
         except Exception as e:
@@ -139,7 +147,7 @@ def get(self):
             system = self.get_argument("system")
             projects = [gcp_project(), bq_public_dataset_project_id]
             bq_search = BigQuerySearchService()
-            credentials = handlers.get_cached_credentials(self.log)
+            credentials = GetCachedCredentials.get_cached_credentials(self.log)
             search_data = bq_search.bigquery_search(
                 credentials, search_string, type, system, projects, self.log
             )

diff --git a/dataproc_jupyter_plugin/contollers/clusterController.py b/dataproc_jupyter_plugin/contollers/clusterController.py
@@ -14,9 +14,9 @@
 
 
 import json
+from dataproc_jupyter_plugin.utils.credentials import GetCachedCredentials
 from jupyter_server.base.handlers import APIHandler
 import tornado
-from dataproc_jupyter_plugin import handlers
 from dataproc_jupyter_plugin.services.clusterListService import ClusterListService
 
 
@@ -27,7 +27,7 @@ def get(self):
             page_token = self.get_argument("pageToken")
             page_size = self.get_argument("pageSize")
             cluster = ClusterListService()
-            credentials = handlers.get_cached_credentials(self.log)
+            credentials = GetCachedCredentials.get_cached_credentials(self.log)
             cluster_list = cluster.list_clusters(
                 credentials, page_size, page_token, self.log
             )

diff --git a/dataproc_jupyter_plugin/contollers/composerController.py b/dataproc_jupyter_plugin/contollers/composerController.py
@@ -13,9 +13,9 @@
 # limitations under the License.
 
 import json
+from dataproc_jupyter_plugin.utils.credentials import GetCachedCredentials
 from jupyter_server.base.handlers import APIHandler
 import tornado
-from dataproc_jupyter_plugin import handlers
 from dataproc_jupyter_plugin.services.composerService import ComposerService
 from requests import HTTPError
 
@@ -26,7 +26,7 @@ def get(self):
         """Returns names of available composer environments"""
         try:
             environments_manager = ComposerService()
-            credentials = handlers.get_cached_credentials(self.log)
+            credentials = GetCachedCredentials.get_cached_credentials(self.log)
             environments = environments_manager.list_environments(credentials, self.log)
 
         except Exception as e:

diff --git a/dataproc_jupyter_plugin/contollers/dagController.py b/dataproc_jupyter_plugin/contollers/dagController.py
@@ -15,9 +15,9 @@
 
 import json
 import subprocess
+from dataproc_jupyter_plugin.utils.credentials import GetCachedCredentials
 from jupyter_server.base.handlers import APIHandler
 import tornado
-from dataproc_jupyter_plugin import handlers
 from dataproc_jupyter_plugin.services.dagListService import (
     DagListService,
     DagDeleteService,
@@ -32,7 +32,7 @@ def get(self):
         try:
             dag = DagListService()
             composer_name = self.get_argument("composer")
-            credentials = handlers.get_cached_credentials(self.log)
+            credentials = GetCachedCredentials.get_cached_credentials(self.log)
             dag_list = dag.list_jobs(credentials, composer_name, TAGS, self.log)
             self.finish(json.dumps(dag_list))
         except Exception as e:
@@ -70,7 +70,7 @@ def get(self):
             composer = self.get_argument("composer")
             dag_id = self.get_argument("dag_id")
             from_page = self.get_argument("from_page", default=None)
-            credentials = handlers.get_cached_credentials(self.log)
+            credentials = GetCachedCredentials.get_cached_credentials(self.log)
             delete_response = dag.delete_job(
                 credentials, composer, dag_id, from_page, self.log
             )
@@ -92,7 +92,7 @@ def get(self):
             composer = self.get_argument("composer")
             dag_id = self.get_argument("dag_id")
             status = self.get_argument("status")
-            credentials = handlers.get_cached_credentials(self.log)
+            credentials = GetCachedCredentials.get_cached_credentials(self.log)
             update_response = dag.update_job(
                 credentials, composer, dag_id, status, self.log
             )

diff --git a/dataproc_jupyter_plugin/contollers/dagRunController.py b/dataproc_jupyter_plugin/contollers/dagRunController.py
@@ -14,9 +14,9 @@
 
 
 import json
+from dataproc_jupyter_plugin.utils.credentials import GetCachedCredentials
 from jupyter_server.base.handlers import APIHandler
 import tornado
-from dataproc_jupyter_plugin import handlers
 from dataproc_jupyter_plugin.services.dagRunService import (
     DagRunListService,
     DagRunTaskListService,
@@ -34,7 +34,7 @@ def get(self):
             start_date = self.get_argument("start_date")
             offset = self.get_argument("offset")
             end_date = self.get_argument("end_date")
-            credentials = handlers.get_cached_credentials(self.log)
+            credentials = GetCachedCredentials.get_cached_credentials(self.log)
             dag_run_list = dag_run.list_dag_runs(
                 credentials,
                 composer_name,
@@ -58,7 +58,7 @@ def get(self):
             composer_name = self.get_argument("composer")
             dag_id = self.get_argument("dag_id")
             dag_run_id = self.get_argument("dag_run_id")
-            credentials = handlers.get_cached_credentials(self.log)
+            credentials = GetCachedCredentials.get_cached_credentials(self.log)
             dag_run_list = dag_run.list_dag_run_task(
                 credentials, composer_name, dag_id, dag_run_id, self.log
             )
@@ -78,7 +78,7 @@ def get(self):
             dag_run_id = self.get_argument("dag_run_id")
             task_id = self.get_argument("task_id")
             task_try_number = self.get_argument("task_try_number")
-            credentials = handlers.get_cached_credentials(self.log)
+            credentials = GetCachedCredentials.get_cached_credentials(self.log)
             dag_run_list = dag_run.list_dag_run_task_logs(
                 credentials,
                 composer_name,

diff --git a/dataproc_jupyter_plugin/contollers/editDagController.py b/dataproc_jupyter_plugin/contollers/editDagController.py
@@ -13,9 +13,9 @@
 # limitations under the License.
 
 import json
+from dataproc_jupyter_plugin.utils.credentials import GetCachedCredentials
 from jupyter_server.base.handlers import APIHandler
 import tornado
-from dataproc_jupyter_plugin import handlers
 from dataproc_jupyter_plugin.services.editDagService import DagEditService
 from dataproc_jupyter_plugin.utils.constants import TAGS
 
@@ -27,7 +27,7 @@ def get(self):
             dag = DagEditService()
             bucket_name = self.get_argument("bucket_name")
             dag_id = self.get_argument("dag_id")
-            credentials = handlers.get_cached_credentials(self.log)
+            credentials = GetCachedCredentials.get_cached_credentials(self.log)
             dag_details = dag.edit_jobs(dag_id, bucket_name, credentials, self.log)
             self.finish(json.dumps(dag_details))
         except Exception as e:

diff --git a/dataproc_jupyter_plugin/contollers/executorController.py b/dataproc_jupyter_plugin/contollers/executorController.py
@@ -13,9 +13,9 @@
 # limitations under the License.
 
 
+from dataproc_jupyter_plugin.utils.credentials import GetCachedCredentials
 from jupyter_server.base.handlers import APIHandler
 import tornado
-from dataproc_jupyter_plugin import handlers
 from dataproc_jupyter_plugin.services.executorService import ExecutorService
 
 
@@ -25,7 +25,7 @@ def post(self):
         try:
             input_data = self.get_json_body()
             execute = ExecutorService()
-            credentials = handlers.get_cached_credentials(self.log)
+            credentials = GetCachedCredentials.get_cached_credentials(self.log)
             execute.execute(credentials, input_data, self.log)
         except Exception as e:
             self.log.exception(f"Error creating dag schedule: {str(e)}")

diff --git a/dataproc_jupyter_plugin/contollers/importErrorController.py b/dataproc_jupyter_plugin/contollers/importErrorController.py
@@ -14,9 +14,9 @@
 
 
 import json
+from dataproc_jupyter_plugin.utils.credentials import GetCachedCredentials
 from jupyter_server.base.handlers import APIHandler
 import tornado
-from dataproc_jupyter_plugin import handlers
 from dataproc_jupyter_plugin.services.importErrorService import ImportErrorService
 
 
@@ -26,7 +26,7 @@ def get(self):
         try:
             import_errors = ImportErrorService()
             composer_name = self.get_argument("composer")
-            credentials = handlers.get_cached_credentials(self.log)
+            credentials = GetCachedCredentials.get_cached_credentials(self.log)
             import_errors_list = import_errors.list_import_errors(
                 credentials, composer_name, self.log
             )

diff --git a/dataproc_jupyter_plugin/contollers/runtimeController.py b/dataproc_jupyter_plugin/contollers/runtimeController.py
@@ -14,9 +14,9 @@
 
 
 import json
+from dataproc_jupyter_plugin.utils.credentials import GetCachedCredentials
 from jupyter_server.base.handlers import APIHandler
 import tornado
-from dataproc_jupyter_plugin import handlers
 from dataproc_jupyter_plugin.services.runtimeListService import RuntimeListService
 
 
@@ -27,7 +27,7 @@ def get(self):
             page_token = self.get_argument("pageToken")
             page_size = self.get_argument("pageSize")
             runtime = RuntimeListService()
-            credentials = handlers.get_cached_credentials(self.log)
+            credentials = GetCachedCredentials.get_cached_credentials(self.log)
             runtime_list = runtime.list_runtime(
                 credentials, page_size, page_token, self.log
             )

diff --git a/dataproc_jupyter_plugin/contollers/triggerDagController.py b/dataproc_jupyter_plugin/contollers/triggerDagController.py
@@ -14,9 +14,9 @@
 
 
 import json
+from dataproc_jupyter_plugin.utils.credentials import GetCachedCredentials
 from jupyter_server.base.handlers import APIHandler
 import tornado
-from dataproc_jupyter_plugin import handlers
 from dataproc_jupyter_plugin.services.triggerDagService import TriggerDagService
 
 
@@ -27,7 +27,7 @@ def get(self):
             trigger_dag = TriggerDagService()
             dag_id = self.get_argument("dag_id")
             composer = self.get_argument("composer")
-            credentials = handlers.get_cached_credentials(self.log)
+            credentials = GetCachedCredentials.get_cached_credentials(self.log)
             trigger = trigger_dag.dag_trigger(credentials, dag_id, composer, self.log)
             self.finish(json.dumps(trigger))
         except Exception as e:

diff --git a/dataproc_jupyter_plugin/handlers.py b/dataproc_jupyter_plugin/handlers.py
@@ -203,7 +203,7 @@ def post(self):
             run_gcloud_subcommand(f"config set project {project_id}")
             run_gcloud_subcommand(f"config set dataproc/region {region}")
             clear_gcloud_cache()
-            update_gateway_client_url(self.config, self.log)
+            configure_gateway_client_url(self.config, self.log)
             self.finish({"config": ERROR_MESSAGE + "successful"})
         except subprocess.CalledProcessError as er:
             self.finish({"config": ERROR_MESSAGE + "failed"})

diff --git a/dataproc_jupyter_plugin/services/bigqueryService.py b/dataproc_jupyter_plugin/services/bigqueryService.py
@@ -129,7 +129,14 @@ def list_table_info(self, credentials, dataset_id, table_id, project_id, log):
 
 class BigQueryPreviewService:
     def bigquery_preview_data(
-        self, credentials, dataset_id, table_id, page_token, project_id, log
+        self,
+        credentials,
+        dataset_id,
+        table_id,
+        max_results,
+        start_index,
+        project_id,
+        log,
     ):
         try:
             if (
@@ -138,7 +145,7 @@ def bigquery_preview_data(
                 and ("region_id" in credentials)
             ):
                 access_token = credentials["access_token"]
-                api_endpoint = f"https://bigquery.googleapis.com/bigquery/v2/projects/{project_id}/datasets/{dataset_id}/tables/{table_id}/data?pageToken={page_token}"
+                api_endpoint = f"https://bigquery.googleapis.com/bigquery/v2/projects/{project_id}/datasets/{dataset_id}/tables/{table_id}/data?maxResults={max_results}&startIndex={start_index}"
                 headers = ApiHeaders.create_headers(access_token)
                 response = requests.get(api_endpoint, headers=headers)
                 if response.status_code == 200:

diff --git a/dataproc_jupyter_plugin/utils/credentials.py b/dataproc_jupyter_plugin/utils/credentials.py
@@ -0,0 +1,45 @@
+# 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.
+
+
+from google.cloud.jupyter_config.config import (
+    gcp_credentials,
+    gcp_project,
+    gcp_project_number,
+    gcp_region,
+)
+
+
+class GetCachedCredentials:
+    def get_cached_credentials(self):
+        credentials = {
+            "project_id": "",
+            "project_number": 0,
+            "region_id": "",
+            "access_token": "",
+            "config_error": 0,
+            "login_error": 0,
+        }
+
+        try:
+            credentials["project_id"] = gcp_project()
+            credentials["region_id"] = gcp_region()
+            credentials["config_error"] = 0
+            credentials["access_token"] = gcp_credentials()
+            credentials["project_number"] = gcp_project_number()
+            return credentials
+        except Exception as ex:
+            self.log.exception(f"Error fetching credentials from gcloud ", ex)
+            credentials["config_error"] = 1
+            return credentials