Merge pull request #12 from nirupama-dev/main

Bug fix on pagination and execution history

Author: harsha-accenture

pyproject.toml

@@ -34,6 +34,7 @@ dependencies = [
     "google-cloud-compute",
     "google-cloud-iam",
     "cron-descriptor>=1.4.5",
+    "gcs-jupyter-plugin",
 ]
 dynamic = ["version", "description", "authors", "urls", "keywords"]
 

scheduler_jupyter_plugin/controllers/composer.py

@@ -39,3 +39,21 @@ async def get(self):
         except Exception as e:
             self.log.exception(f"Error fetching composer environments: {str(e)}")
             self.finish({"error": str(e)})
+
+class EnvironmentGetController(APIHandler):
+    @tornado.web.authenticated
+    async def get(self):
+        """Returns details of composer environment"""
+        try:
+            env_name = self.get_argument("env_name")
+            async with aiohttp.ClientSession() as client_session:
+                client = composer.Client(
+                    await credentials.get_cached(), self.log, client_session
+                )
+                environment = await client.get_environment(env_name)
+                self.set_header("Content-Type", "application/json")
+                self.finish(json.dumps(environment, default=lambda x: x.dict()))
+        except Exception as e:
+            self.log.exception(f"Error fetching composer environment: {str(e)}")
+            self.finish({"error": str(e)})
+

scheduler_jupyter_plugin/dagTemplates/localPythonTemplate-v1.txt

@@ -64,7 +64,7 @@ def run_notebook_task(**kwargs):
     parameters = '''
     {{parameters}}
     '''
-    if len(parameters)>0:
+    if parameters.strip():
         parameters_dict = convert_parameters(parameters)
     else:
         parameters_dict = {}

scheduler_jupyter_plugin/handlers.py

@@ -156,6 +156,7 @@ def full_path(name):
         "getGcpServiceUrls": UrlHandler,
         "log": LogHandler,
         "composerList": composer.EnvironmentListController,
+        "getComposerEnvironment": composer.EnvironmentGetController,
         "dagRun": airflow.DagRunController,
         "dagRunTask": airflow.DagRunTaskController,
         "dagRunTaskLogs": airflow.DagRunTaskLogsController,

scheduler_jupyter_plugin/models/models.py

@@ -27,8 +27,10 @@ class ComposerEnvironment(BaseModel):
     name: str
     label: str
     description: str
+    state: str
     file_extensions: List[str]  # Supported input file types
     metadata: Optional[Dict[str, str]]  # Optional metadata
+    pypi_packages: Optional[Dict[str, str]] = None
 
     def __str__(self):
         return self.json()

scheduler_jupyter_plugin/services/composer.py

@@ -67,19 +67,20 @@ async def list_environments(
                         return environments
                     else:
                         environment = resp.get("environments", [])
-                        # Extract the 'name' values from the 'environments' list
-                        names = [env["name"] for env in environment]
-                        # Extract the last value after the last slash for each 'name'
-                        last_values = [name.split("/")[-1] for name in names]
-                        for env in last_values:
-                            name = env
+                        for env in environment:
+                            path = env["name"]
+                            name = env["name"].split("/")[-1]
+                            state = env["state"]
+                            pypi_packages = env.get("config", {}).get("softwareConfig", {}).get("pypiPackages", None)
                             environments.append(
                                 ComposerEnvironment(
                                     name=name,
                                     label=name,
                                     description=f"Environment: {name}",
+                                    state=state,
                                     file_extensions=["ipynb"],
-                                    metadata={"path": env},
+                                    metadata={"path": path},
+                                    pypi_packages=pypi_packages
                                 )
                             )
                         return environments
@@ -94,3 +95,47 @@ async def list_environments(
         except Exception as e:
             self.log.exception(f"Error fetching environments list: {str(e)}")
             return {"Error fetching environments list": str(e)}
+        
+    async def get_environment(
+        self, env_name
+    ) -> ComposerEnvironment:
+        try:
+            environment = {}
+            composer_url = await urls.gcp_service_url(COMPOSER_SERVICE_NAME)
+            api_endpoint = f"{composer_url}/v1/{env_name}"
+
+            headers = self.create_headers()
+            async with self.client_session.get(
+                api_endpoint, headers=headers
+            ) as response:
+                if response.status == HTTP_STATUS_OK:
+                    resp = await response.json()
+                    if not resp:
+                        return environment
+                    else:
+                        path = resp.get("name")
+                        name = resp.get("name").split("/")[-1]
+                        state = resp.get("state")
+                        pypi_packages = resp.get("config", {}).get("softwareConfig", {}).get("pypiPackages", None)
+                        environment = ComposerEnvironment(
+                            name=name,
+                            label=name,
+                            description=f"Environment: {name}",
+                            state=state,
+                            file_extensions=["ipynb"],
+                            metadata={"path": path},
+                            pypi_packages=pypi_packages
+                        )
+                            
+                        return environment
+                elif response.status == HTTP_STATUS_FORBIDDEN:
+                    resp = await response.json()
+                    return resp
+                else:
+                    self.log.exception("Error fetching environment")
+                    raise Exception(
+                        f"Error getting composer: {response.reason} {await response.text()}"
+                    )
+        except Exception as e:
+            self.log.exception(f"Error fetching environment: {str(e)}")
+            return {"Error fetching environment": str(e)}

scheduler_jupyter_plugin/tests/test_composer.py

@@ -31,10 +31,12 @@ def get(self, api_endpoint, headers=None):
             {
                 "environments": [
                     {
-                        "name": "projects/mock-project/locations/mock-location/environments/env1"
+                        "name": "projects/mock-project/locations/mock-location/environments/env1",
+                        "state": "RUNNING",
                     },
                     {
-                        "name": "projects/mock-project/locations/mock-location/environments/env2"
+                        "name": "projects/mock-project/locations/mock-location/environments/env2",
+                        "state": "UPDATING",
                     },
                 ]
             }

src/scheduler/common/SchedulerInteface.ts

@@ -69,11 +69,13 @@ export interface IClusterAPIResponse {
 }
 
 export interface IComposerAPIResponse {
-  name: string;
+  name: string | '';
   label: string;
   description: string;
+  state: string;
   file_extensions: [];
   metadata: Record<string, never>;
+  pypi_packages: Record<string, string> | undefined;
 }
 
 export interface IDagRunList {