Code review comments fix

GoogleCloudDataproc · Sep 5, 2024 · 4d70971 · 4d70971
1 parent 68c21db
commit 4d70971
Show file tree

Hide file tree

Showing 10 changed files with 244 additions and 338 deletions.
diff --git a/dataproc_jupyter_plugin/controllers/cluster.py b/dataproc_jupyter_plugin/controllers/cluster.py
diff --git a/dataproc_jupyter_plugin/controllers/dataproc.py b/dataproc_jupyter_plugin/controllers/dataproc.py
@@ -22,7 +22,7 @@
 from dataproc_jupyter_plugin.services import dataproc
 
 
-class ClusterListController(APIHandler):
+class RuntimeController(APIHandler):
     @tornado.web.authenticated
     async def get(self):
         try:
@@ -32,25 +32,74 @@ async def get(self):
                 client = dataproc.Client(
                     await credentials.get_cached(), self.log, client_session
                 )
-                cluster_list = await client.list_clusters(page_size, page_token)
-            self.finish(json.dumps(cluster_list))
+                runtime_list = await client.list_runtime(page_size, page_token)
+            self.finish(json.dumps(runtime_list))
         except Exception as e:
-            self.log.exception("Error fetching cluster list")
+            self.log.exception(f"Error fetching runtime template list: {str(e)}")
             self.finish({"error": str(e)})
 
 
-class RuntimeController(APIHandler):
+class ClusterListController(APIHandler):
     @tornado.web.authenticated
     async def get(self):
         try:
             page_token = self.get_argument("pageToken")
             page_size = self.get_argument("pageSize")
-            async with aiohttp.ClientSession() as client_session:
-                client = dataproc.Client(
-                    await credentials.get_cached(), self.log, client_session
-                )
-                runtime_list = await client.list_runtime(page_size, page_token)
-            self.finish(json.dumps(runtime_list))
+            client = dataproc.Client(await credentials.get_cached(), self.log)
+            cluster_list = await client.list_clusters(page_size, page_token)
+            self.finish(json.dumps(cluster_list))
         except Exception as e:
-            self.log.exception(f"Error fetching runtime template list: {str(e)}")
+            self.log.exception(f"Error fetching cluster list")
+            self.finish({"error": str(e)})
+
+
+class ClusterDetailController(APIHandler):
+    @tornado.web.authenticated
+    async def get(self):
+        try:
+            cluster = self.get_argument("cluster")
+            client = dataproc.Client(await credentials.get_cached(), self.log)
+            get_cluster = await client.get_cluster_detail(cluster)
+            self.finish(json.dumps(get_cluster))
+        except Exception as e:
+            self.log.exception(f"Error fetching get cluster")
+            self.finish({"error": str(e)})
+
+
+class StopClusterController(APIHandler):
+    @tornado.web.authenticated
+    async def post(self):
+        try:
+            cluster = self.get_argument("cluster")
+            client = dataproc.Client(await credentials.get_cached(), self.log)
+            stop_cluster = await client.stop_cluster(cluster)
+            self.finish(json.dumps(stop_cluster))
+        except Exception as e:
+            self.log.exception(f"Error fetching stop cluster")
+            self.finish({"error": str(e)})
+
+
+class StartClusterController(APIHandler):
+    @tornado.web.authenticated
+    async def post(self):
+        try:
+            cluster = self.get_argument("cluster")
+            client = dataproc.Client(await credentials.get_cached(), self.log)
+            start_cluster = await client.start_cluster(cluster)
+            self.finish(json.dumps(start_cluster))
+        except Exception as e:
+            self.log.exception(f"Error fetching start cluster")
+            self.finish({"error": str(e)})
+
+
+class DeleteClusterController(APIHandler):
+    @tornado.web.authenticated
+    async def delete(self):
+        try:
+            cluster = self.get_argument("cluster")
+            client = dataproc.Client(await credentials.get_cached(), self.log)
+            delete_cluster = await client.delete_cluster(cluster)
+            self.finish(json.dumps(delete_cluster))
+        except Exception as e:
+            self.log.exception(f"Error deleting cluster")
             self.finish({"error": str(e)})
diff --git a/dataproc_jupyter_plugin/handlers.py b/dataproc_jupyter_plugin/handlers.py
@@ -34,7 +34,6 @@
 from dataproc_jupyter_plugin.controllers import (
     airflow,
     bigquery,
-    cluster,
     composer,
     dataproc,
     executor,
@@ -194,11 +193,10 @@ def full_path(name):
         "dagRunTask": airflow.DagRunTaskController,
         "dagRunTaskLogs": airflow.DagRunTaskLogsController,
         "clusterList": dataproc.ClusterListController,
-        "clusterListPage": cluster.ClusterListPageController,
-        "clusterDetail": cluster.ClusterDetailController,
-        "stopCluster": cluster.StopClusterController,
-        "startCluster": cluster.StartClusterController,
-        "deleteCluster": cluster.DeleteClusterController,
+        "clusterDetail": dataproc.ClusterDetailController,
+        "stopCluster": dataproc.StopClusterController,
+        "startCluster": dataproc.StartClusterController,
+        "deleteCluster": dataproc.DeleteClusterController,
         "runtimeList": dataproc.RuntimeController,
         "createJobScheduler": executor.ExecutorController,
         "dagList": airflow.DagListController,

diff --git a/dataproc_jupyter_plugin/services/cluster.py b/dataproc_jupyter_plugin/services/cluster.py