Bug tracker and Buganizer fixes

dataproc_jupyter_plugin/commons/constants.py

@@ -12,27 +12,21 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-
-from dataproc_jupyter_plugin.services.gcpUrlService import gcp_service_url
-
 ENVIRONMENT_API = "https://composer.googleapis.com/v1"
 TAGS = "dataproc_jupyter_plugin"
 GCS = "gs://"
 CONTENT_TYPE = "application/json"
-
-
-dataproc_url = gcp_service_url("dataproc")
-dataproc_url = gcp_service_url("dataproc")
-compute_url = gcp_service_url(
-    "compute", default_url="https://compute.googleapis.com/compute/v1"
-)
-metastore_url = gcp_service_url("metastore")
-cloudkms_url = gcp_service_url("cloudkms")
-cloudresourcemanager_url = gcp_service_url("cloudresourcemanager")
-datacatalog_url = gcp_service_url("datacatalog")
-storage_url = gcp_service_url(
-    "storage", default_url="https://storage.googleapis.com/storage/v1/"
-)
-dataplex_url = gcp_service_url(
-    "dataplex", default_url="https://dataplex.googleapis.com/v1"
-)
+PACKAGE_NAME = "dataproc_jupyter_plugin"
+COMPOSER_SERVICE_NAME = "composer"
+DATAPROC_SERVICE_NAME = "dataproc"
+BIGQUERY_SERVICE_NAME = "bigquery"
+DATACATALOG_SERVICE_NAME = "datacatalog"
+CLOUDRESOURCEMANAGER_SERVICE_NAME = "cloudresourcemanager"
+STORAGE_SERVICE_NAME = "storage"
+COMPUTE_SERVICE_NAME = "compute"
+METASTORE_SERVICE_NAME = "metastore"
+CLOUDKMS_SERVICE_NAME = "cloudkms"
+COMPUTE_SERVICE_DEFAULT_URL = "https://compute.googleapis.com/compute/v1"
+STORAGE_SERVICE_DEFAULT_URL = "https://storage.googleapis.com/storage/v1/"
+WRAPPER_PAPPERMILL_FILE = "wrapper_papermill.py"
+LOCAL_DAG_FILE_LOCATION = "./Scheduled-jobs"

dataproc_jupyter_plugin/controllers/airflow.py

@@ -0,0 +1,167 @@
+# 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 subprocess
+from dataproc_jupyter_plugin import credentials
+from dataproc_jupyter_plugin.services import airflow
+from jupyter_server.base.handlers import APIHandler
+import tornado
+
+
+class DagListController(APIHandler):
+    @tornado.web.authenticated
+    async def get(self):
+        try:
+            client = airflow.Client(await credentials.get_cached(), self.log)
+            composer_name = self.get_argument("composer")
+            dag_list = await client.list_jobs(composer_name)
+            self.finish(json.dumps(dag_list))
+        except Exception as e:
+            self.log.exception(f"Error fetching cluster list")
+            self.finish({"error": str(e)})
+
+
+class DagDeleteController(APIHandler):
+    @tornado.web.authenticated
+    async def get(self):
+        try:
+            client = airflow.Client(await credentials.get_cached(), self.log)
+            composer = self.get_argument("composer")
+            dag_id = self.get_argument("dag_id")
+            from_page = self.get_argument("from_page", default=None)
+            delete_response = await client.delete_job(composer, dag_id, from_page)
+            if delete_response == 0:
+                self.finish(json.dumps({"status": delete_response}))
+            else:
+                self.log.exception(f"Error deleting dag file")
+                self.finish(json.dumps({"status": delete_response}))
+        except Exception as e:
+            self.log.exception(f"Error deleting dag file: {str(e)}")
+            self.finish({"error": str(e)})
+
+
+class DagUpdateController(APIHandler):
+    @tornado.web.authenticated
+    async def get(self):
+        try:
+            client = airflow.Client(await credentials.get_cached(), self.log)
+            composer = self.get_argument("composer")
+            dag_id = self.get_argument("dag_id")
+            status = self.get_argument("status")
+            update_response = await client.update_job(composer, dag_id, status)
+            if update_response == 0:
+                self.finish({"status": 0})
+            else:
+                self.log.exception(f"Error updating status")
+                self.finish(json.dumps(update_response))
+        except Exception as e:
+            self.log.exception(f"Error updating status: {str(e)}")
+            self.finish({"error": str(e)})
+
+
+class DagRunController(APIHandler):
+    @tornado.web.authenticated
+    async def get(self):
+        try:
+            client = airflow.Client(await credentials.get_cached(), self.log)
+            composer_name = self.get_argument("composer")
+            dag_id = self.get_argument("dag_id")
+            start_date = self.get_argument("start_date")
+            offset = self.get_argument("offset")
+            end_date = self.get_argument("end_date")
+            dag_run_list = await client.list_dag_runs(
+                composer_name, dag_id, start_date, end_date, offset
+            )
+            self.finish(json.dumps(dag_run_list))
+        except Exception as e:
+            self.log.exception(f"Error fetching dag run list {str(e)}")
+            self.finish({"error": str(e)})
+
+
+class DagRunTaskController(APIHandler):
+    @tornado.web.authenticated
+    async def get(self):
+        try:
+            client = airflow.Client(await credentials.get_cached(), self.log)
+            composer_name = self.get_argument("composer")
+            dag_id = self.get_argument("dag_id")
+            dag_run_id = self.get_argument("dag_run_id")
+            dag_run_list = await client.list_dag_run_task(
+                composer_name, dag_id, dag_run_id
+            )
+            self.finish(json.dumps(dag_run_list))
+        except Exception as e:
+            self.log.exception(f"Error fetching dag run tasks: {str(e)}")
+            self.finish({"error": str(e)})
+
+
+class DagRunTaskLogsController(APIHandler):
+    @tornado.web.authenticated
+    async def get(self):
+        try:
+            client = airflow.Client(await credentials.get_cached(), self.log)
+            composer_name = self.get_argument("composer")
+            dag_id = self.get_argument("dag_id")
+            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")
+            dag_run_list = await client.list_dag_run_task_logs(
+                composer_name, dag_id, dag_run_id, task_id, task_try_number
+            )
+            self.finish(json.dumps(dag_run_list))
+        except Exception as e:
+            self.log.exception(f"Error fetching dag run task logs: {str(e)}")
+            self.finish({"error": str(e)})
+
+
+class EditDagController(APIHandler):
+    @tornado.web.authenticated
+    async def get(self):
+        try:
+            client = airflow.Client(await credentials.get_cached(), self.log)
+            bucket_name = self.get_argument("bucket_name")
+            dag_id = self.get_argument("dag_id")
+            dag_details = await client.edit_jobs(dag_id, bucket_name)
+            self.finish(json.dumps(dag_details))
+        except Exception as e:
+            self.log.exception(f"Error getting dag details")
+            self.finish({"error": str(e)})
+
+
+class ImportErrorController(APIHandler):
+    @tornado.web.authenticated
+    async def get(self):
+        try:
+            client = airflow.Client(await credentials.get_cached(), self.log)
+            composer_name = self.get_argument("composer")
+            import_errors_list = await client.list_import_errors(composer_name)
+            self.finish(json.dumps(import_errors_list))
+        except Exception as e:
+            self.log.exception(f"Error fetching import error list")
+            self.finish({"error": str(e)})
+
+
+class TriggerDagController(APIHandler):
+    @tornado.web.authenticated
+    async def get(self):
+        try:
+            client = airflow.Client(await credentials.get_cached(), self.log)
+            dag_id = self.get_argument("dag_id")
+            composer = self.get_argument("composer")
+            trigger = await client.dag_trigger(dag_id, composer)
+            self.finish(json.dumps(trigger))
+        except Exception as e:
+            self.log.exception(f"Error triggering dag")
+            self.finish({"error": str(e)})

dataproc_jupyter_plugin/controllers/bigquery.py

@@ -23,7 +23,6 @@
 from dataproc_jupyter_plugin import credentials
 from dataproc_jupyter_plugin.services import bigquery
 
-
 # GCP project holding BigQuery public datasets.
 BQ_PUBLIC_DATASET_PROJECT_ID = "bigquery-public-data"
 

dataproc_jupyter_plugin/controllers/composerController.py → dataproc_jupyter_plugin/controllers/composer.py

@@ -1,4 +1,4 @@
-# Copyright 2023 Google LLC
+# 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.
@@ -13,22 +13,19 @@
 # limitations under the License.
 
 import json
-from dataproc_jupyter_plugin.commons.gcloudOperations import GetCachedCredentials
+from dataproc_jupyter_plugin import credentials
 from jupyter_server.base.handlers import APIHandler
 import tornado
-from dataproc_jupyter_plugin.services.composerService import ComposerService
-from requests import HTTPError
+from dataproc_jupyter_plugin.services import composer
 
 
-class ComposerListController(APIHandler):
+class EnvironmentListController(APIHandler):
     @tornado.web.authenticated
-    def get(self):
+    async def get(self):
         """Returns names of available composer environments"""
         try:
-            environments_manager = ComposerService()
-            credentials = GetCachedCredentials.get_cached_credentials(self.log)
-            environments = environments_manager.list_environments(credentials, self.log)
-
+            client = composer.Client(await credentials.get_cached(), self.log)
+            environments = await client.list_environments()
         except Exception as e:
             self.log.exception(f"Error fetching composer environments: {str(e)}")
             self.finish({"error": str(e)})