added pipeline

Signed-off-by: Rashida Kanchwala <[email protected]>

package/kedro_viz/api/apps.py

@@ -5,6 +5,8 @@
 import time
 from pathlib import Path
 
+import zipfile
+
 from fastapi import FastAPI, HTTPException
 from fastapi.requests import Request
 from fastapi.responses import HTMLResponse, JSONResponse, Response
@@ -110,6 +112,21 @@ async def index():
 
     @app.get("/api/main", response_class=JSONResponse)
     async def main():
-        return json.loads(Path(filepath).read_text(encoding="utf8"))
+        with zipfile.ZipFile(f'{filepath}.zip', 'r') as zip_file:
+            main = zip_file.read('main')
+        return json.loads(main)
+
+    @app.get("/api/nodes/{node_id}", response_class=JSONResponse)
+    async def get_node_metadata(node_id):
+        with zipfile.ZipFile(f'{filepath}.zip', 'r') as zip_file:
+            node_metdata = zip_file.read(f'nodes/{node_id}')
+        return json.loads(node_metdata)
+
+    @app.get("/api/pipelines/{pipeline_id}", response_class=JSONResponse)
+    async def get_registered_pipeline(pipeline_id):
+        with zipfile.ZipFile(f'{filepath}.zip', 'r') as zip_file:
+            pipeline = zip_file.read(f'pipelines/{pipeline_id}')
+        return json.loads(pipeline)
 
     return app
+

package/kedro_viz/api/rest/responses.py

@@ -4,11 +4,21 @@
 from typing import Any, Dict, List, Optional, Union
 
 import orjson
-from fastapi.responses import ORJSONResponse
+from fastapi.responses import ORJSONResponse, JSONResponse
 from pydantic import BaseModel
 
 from kedro_viz.data_access import data_access_manager
 
+from kedro_viz.models.flowchart import (
+    DataNode,
+    DataNodeMetadata,
+    ParametersNodeMetadata,
+    TaskNode,
+    TaskNodeMetadata,
+    TranscodedDataNode,
+    TranscodedDataNodeMetadata,
+)
+
 
 class APIErrorMessage(BaseModel):
     message: str
@@ -241,6 +251,24 @@ class GraphAPIResponse(BaseAPIResponse):
     modular_pipelines: ModularPipelinesTreeAPIResponse
     selected_pipeline: str
 
+class EnhancedORJSONResponse(ORJSONResponse):
+    @staticmethod
+    def encode_to_human_readable(content: Any) -> bytes:
+        """A method to encode the given content to JSON, with the
+        proper formatting to write a human-readable file.
+
+        Returns:
+            A bytes object containing the JSON to write.
+
+        """
+        return orjson.dumps(
+            content,
+            option=orjson.OPT_INDENT_2
+            | orjson.OPT_NON_STR_KEYS
+            | orjson.OPT_SERIALIZE_NUMPY,
+        )
+
+
 
 def get_default_response() -> GraphAPIResponse:
     """Default response for `/api/main`."""
@@ -270,20 +298,53 @@ def get_default_response() -> GraphAPIResponse:
         selected_pipeline=default_selected_pipeline_id,
     )
 
+def get_node_metadata_response(node_id: str):
+
+    node = data_access_manager.nodes.get_node_by_id(node_id)
+    if not node:
+        return JSONResponse(status_code=404, content={"message": "Invalid node ID"})
 
-class EnhancedORJSONResponse(ORJSONResponse):
-    @staticmethod
-    def encode_to_human_readable(content: Any) -> bytes:
-        """A method to encode the given content to JSON, with the
-        proper formatting to write a human-readable file.
+    if not node.has_metadata():
+        return JSONResponse(content={})
 
-        Returns:
-            A bytes object containing the JSON to write.
+    if isinstance(node, TaskNode):
+        return TaskNodeMetadata(node)
 
-        """
-        return orjson.dumps(
-            content,
-            option=orjson.OPT_INDENT_2
-            | orjson.OPT_NON_STR_KEYS
-            | orjson.OPT_SERIALIZE_NUMPY,
+    if isinstance(node, DataNode):
+        return DataNodeMetadata(node)
+
+    if isinstance(node, TranscodedDataNode):
+        return TranscodedDataNodeMetadata(node)
+
+    return ParametersNodeMetadata(node)
+
+
+def get_selected_pipeline_response(registered_pipeline_id: str):
+    if not data_access_manager.registered_pipelines.has_pipeline(
+        registered_pipeline_id
+    ):
+        return JSONResponse(status_code=404, content={"message": "Invalid pipeline ID"})
+
+    modular_pipelines_tree = (
+        data_access_manager.create_modular_pipelines_tree_for_registered_pipeline(
+            registered_pipeline_id
         )
+    )
+
+    return GraphAPIResponse(
+        nodes=data_access_manager.get_nodes_for_registered_pipeline(
+            registered_pipeline_id
+        ),
+        edges=data_access_manager.get_edges_for_registered_pipeline(
+            registered_pipeline_id
+        ),
+        tags=data_access_manager.tags.as_list(),
+        layers=data_access_manager.get_sorted_layers_for_registered_pipeline(
+            registered_pipeline_id
+        ),
+        pipelines=data_access_manager.registered_pipelines.as_list(),
+        selected_pipeline=registered_pipeline_id,
+        modular_pipelines=modular_pipelines_tree,
+    )
+
+

package/kedro_viz/api/rest/router.py

@@ -19,6 +19,8 @@
     GraphAPIResponse,
     NodeMetadataAPIResponse,
     get_default_response,
+    get_node_metadata_response,
+    get_selected_pipeline_response
 )
 
 router = APIRouter(
@@ -38,53 +40,12 @@ async def main():
     response_model_exclude_none=True,
 )
 async def get_single_node_metadata(node_id: str):
-    node = data_access_manager.nodes.get_node_by_id(node_id)
-    if not node:
-        return JSONResponse(status_code=404, content={"message": "Invalid node ID"})
-
-    if not node.has_metadata():
-        return JSONResponse(content={})
-
-    if isinstance(node, TaskNode):
-        return TaskNodeMetadata(node)
-
-    if isinstance(node, DataNode):
-        return DataNodeMetadata(node)
-
-    if isinstance(node, TranscodedDataNode):
-        return TranscodedDataNodeMetadata(node)
-
-    return ParametersNodeMetadata(node)
+    return get_node_metadata_response(node_id)
 
 
 @router.get(
     "/pipelines/{registered_pipeline_id}",
     response_model=GraphAPIResponse,
 )
 async def get_single_pipeline_data(registered_pipeline_id: str):
-    if not data_access_manager.registered_pipelines.has_pipeline(
-        registered_pipeline_id
-    ):
-        return JSONResponse(status_code=404, content={"message": "Invalid pipeline ID"})
-
-    modular_pipelines_tree = (
-        data_access_manager.create_modular_pipelines_tree_for_registered_pipeline(
-            registered_pipeline_id
-        )
-    )
-
-    return GraphAPIResponse(
-        nodes=data_access_manager.get_nodes_for_registered_pipeline(
-            registered_pipeline_id
-        ),
-        edges=data_access_manager.get_edges_for_registered_pipeline(
-            registered_pipeline_id
-        ),
-        tags=data_access_manager.tags.as_list(),
-        layers=data_access_manager.get_sorted_layers_for_registered_pipeline(
-            registered_pipeline_id
-        ),
-        pipelines=data_access_manager.registered_pipelines.as_list(),
-        selected_pipeline=registered_pipeline_id,
-        modular_pipelines=modular_pipelines_tree,
-    )
+    return get_selected_pipeline_response(registered_pipeline_id)

package/kedro_viz/data_access/repositories/graph.py

@@ -29,6 +29,9 @@ def as_list(self) -> List[GraphNode]:
     def as_dict(self) -> Dict[str, GraphNode]:
         return self.nodes_dict
 
+    def get_node_ids(self) -> List[str]:
+        return self.nodes_dict.keys()
+
     def get_nodes_by_ids(self, node_ids: Set[str]) -> List[GraphNode]:
         return [n for n in self.nodes_list if n.id in node_ids]
 

package/kedro_viz/data_access/repositories/registered_pipelines.py

@@ -20,6 +20,9 @@ def add_node(self, pipeline_id: str, node_id: str):
 
     def get_pipeline_by_id(self, pipeline_id: str) -> Optional[RegisteredPipeline]:
         return self.pipelines_dict.get(pipeline_id)
+
+    def get_pipeline_ids(self) -> List[str]:
+        return self.pipelines_dict.keys()
 
     def has_pipeline(self, pipeline_id: str) -> bool:
         return pipeline_id in self.pipelines_dict

package/kedro_viz/server.py

@@ -3,14 +3,18 @@
 from pathlib import Path
 from typing import Any, Dict, Optional
 
+import zipfile
+import io
+import json
+
 import uvicorn
 from fastapi.encoders import jsonable_encoder
 from kedro.io import DataCatalog
 from kedro.pipeline import Pipeline
 from watchgod import run_process
 
 from kedro_viz.api import apps
-from kedro_viz.api.rest.responses import EnhancedORJSONResponse, get_default_response
+from kedro_viz.api.rest.responses import EnhancedORJSONResponse, get_default_response, get_node_metadata_response, get_selected_pipeline_response
 from kedro_viz.constants import DEFAULT_HOST, DEFAULT_PORT
 from kedro_viz.data_access import DataAccessManager, data_access_manager
 from kedro_viz.database import create_db_engine
@@ -81,22 +85,45 @@ def run_server(
         catalog, pipelines, session_store_location = kedro_data_loader.load_data(
             path, env, extra_params
         )
+
         pipelines = (
             pipelines
             if pipeline_name is None
             else {pipeline_name: pipelines[pipeline_name]}
         )
         populate_data(data_access_manager, catalog, pipelines, session_store_location)
+
         if save_file:
-            default_response = get_default_response()
-            jsonable_default_response = jsonable_encoder(default_response)
-            encoded_default_response = EnhancedORJSONResponse.encode_to_human_readable(
-                jsonable_default_response
+
+            default_pipeline_response = get_default_response()
+            jsonable_default_pipeline_response = jsonable_encoder(
+                default_pipeline_response)
+            encoded_default_pipeline_response = EnhancedORJSONResponse.encode_to_human_readable(
+                jsonable_default_pipeline_response
             )
-            Path(save_file).write_bytes(encoded_default_response)
+
+            encoded_node_response = {}
+
+            for node in data_access_manager.nodes.get_node_ids():
+                node_response = get_node_metadata_response(node)
+                jsonable_node_response = jsonable_encoder(node_response)
+                encoded_node_response[node] = EnhancedORJSONResponse.encode_to_human_readable(
+                    jsonable_node_response
+                )
+
+            zip_buffer = io.BytesIO()
+
+            with zipfile.ZipFile(zip_buffer, 'w', zipfile.ZIP_DEFLATED) as zip_file:
+                zip_file.writestr('main', encoded_default_pipeline_response)
+                for node in encoded_node_response:
+                    zip_file.writestr(f'nodes/{node}', encoded_node_response[node])
+
+            with open(f'/{project_path}/{save_file}.zip', 'wb') as zip_file:
+                zip_file.write(zip_buffer.getvalue())
+
         app = apps.create_api_app_from_project(path, autoreload)
     else:
-        app = apps.create_api_app_from_file(load_file)
+        app = apps.create_api_app_from_file(f'{project_path}/{load_file}')
 
     if browser and is_localhost(host):
         webbrowser.open_new(f"http://{host}:{port}/")