From e09ba1fcc90e83adae9234a346fd4c34eb6aa126 Mon Sep 17 00:00:00 2001 From: Jitendra Gundaniya <38945204+jitu5@users.noreply.github.com> Date: Tue, 3 Sep 2024 10:55:33 +0100 Subject: [PATCH] Introduce `get_kedro_project_json_data` function for VSCode integration (#2049) * get_kedro_project_json_data function added Signed-off-by: Jitendra Gundaniya * lint fix Signed-off-by: Jitendra Gundaniya * Lint fix Signed-off-by: Jitendra Gundaniya * Test fix Signed-off-by: Jitendra Gundaniya * lint fix Signed-off-by: Jitendra Gundaniya * Deferred dependencies import into run_server (#2060) * Load data functions move out of server Signed-off-by: Jitendra Gundaniya * Refactor Signed-off-by: Jitendra Gundaniya * Lint fix Signed-off-by: Jitendra Gundaniya * lint fix Signed-off-by: Jitendra Gundaniya --------- Signed-off-by: Jitendra Gundaniya --------- Signed-off-by: Jitendra Gundaniya Co-authored-by: Nok Lam Chan --- package/kedro_viz/api/rest/responses.py | 30 +++++++++++++++++-- package/kedro_viz/server.py | 16 ++++++---- .../test_api/test_rest/test_responses.py | 23 ++++++++++++++ 3 files changed, 62 insertions(+), 7 deletions(-) diff --git a/package/kedro_viz/api/rest/responses.py b/package/kedro_viz/api/rest/responses.py index 8da71d0354..aee13b0b6e 100644 --- a/package/kedro_viz/api/rest/responses.py +++ b/package/kedro_viz/api/rest/responses.py @@ -2,6 +2,7 @@ # pylint: disable=missing-class-docstring,invalid-name import abc +import json import logging from importlib.metadata import PackageNotFoundError from typing import Any, Dict, List, Optional, Union @@ -400,17 +401,42 @@ def get_package_compatibilities_response( return package_requirements_response -def write_api_response_to_fs(file_path: str, response: Any, remote_fs: Any): - """Encodes, enhances responses and writes it to a file""" +def get_encoded_response(response: Any) -> bytes: + """Encodes and enhances the default response using human-readable format.""" jsonable_response = jsonable_encoder(response) encoded_response = EnhancedORJSONResponse.encode_to_human_readable( jsonable_response ) + return encoded_response + + +def write_api_response_to_fs(file_path: str, response: Any, remote_fs: Any): + """Get encoded responses and writes it to a file""" + encoded_response = get_encoded_response(response) + with remote_fs.open(file_path, "wb") as file: file.write(encoded_response) +def get_kedro_project_json_data(): + """Decodes the default response and returns the Kedro project JSON data. + This will be used in VSCode extension to get current Kedro project data.""" + encoded_response = get_encoded_response(get_default_response()) + + try: + response_str = encoded_response.decode("utf-8") + json_data = json.loads(response_str) + except UnicodeDecodeError as exc: # pragma: no cover + json_data = None + logger.error("Failed to decode response string. Error: %s", str(exc)) + except json.JSONDecodeError as exc: # pragma: no cover + json_data = None + logger.error("Failed to parse JSON data. Error: %s", str(exc)) + + return json_data + + def save_api_main_response_to_fs(main_path: str, remote_fs: Any): """Saves API /main response to a directory.""" try: diff --git a/package/kedro_viz/server.py b/package/kedro_viz/server.py index 384a3545dc..b950cc8b77 100644 --- a/package/kedro_viz/server.py +++ b/package/kedro_viz/server.py @@ -1,18 +1,13 @@ """`kedro_viz.server` provides utilities to launch a webserver for Kedro pipeline visualisation.""" -import multiprocessing from pathlib import Path from typing import Any, Dict, Optional -import fsspec -import uvicorn from kedro.framework.session.store import BaseSessionStore from kedro.io import DataCatalog from kedro.pipeline import Pipeline -from watchgod import RegExpWatcher, run_process -from kedro_viz.api import apps from kedro_viz.api.rest.responses import save_api_responses_to_fs from kedro_viz.constants import DEFAULT_HOST, DEFAULT_PORT from kedro_viz.data_access import DataAccessManager, data_access_manager @@ -78,6 +73,7 @@ def load_and_populate_data( populate_data(data_access_manager, catalog, pipelines, session_store, stats_dict) +# pylint: disable=too-many-locals def run_server( host: str = DEFAULT_HOST, port: int = DEFAULT_PORT, @@ -113,6 +109,13 @@ def run_server( take precedence over) the parameters retrieved from the project configuration. """ + # Importing below dependencies inside `run_server` to avoid ImportError + # when calling `load_and_populate_data` from VSCode + + import fsspec # pylint: disable=C0415 + import uvicorn # pylint: disable=C0415 + + from kedro_viz.api import apps # pylint: disable=C0415 path = Path(project_path) if project_path else Path.cwd() @@ -142,6 +145,9 @@ def run_server( if __name__ == "__main__": # pragma: no cover import argparse + import multiprocessing + + from watchgod import RegExpWatcher, run_process parser = argparse.ArgumentParser(description="Launch a development viz server") parser.add_argument("project_path", help="Path to a Kedro project") diff --git a/package/tests/test_api/test_rest/test_responses.py b/package/tests/test_api/test_rest/test_responses.py index 16aabd6aad..8873cfc3c9 100644 --- a/package/tests/test_api/test_rest/test_responses.py +++ b/package/tests/test_api/test_rest/test_responses.py @@ -1,4 +1,5 @@ # pylint: disable=too-many-lines +import json import logging import operator from pathlib import Path @@ -13,6 +14,7 @@ from kedro_viz.api.rest.responses import ( EnhancedORJSONResponse, PackageCompatibilityAPIResponse, + get_kedro_project_json_data, get_package_compatibilities_response, save_api_main_response_to_fs, save_api_node_response_to_fs, @@ -964,6 +966,27 @@ def test_write_api_response_to_fs( mockremote_fs.open.assert_called_once_with(file_path, "wb") mock_encode_to_human_readable.assert_called_once() + def test_get_kedro_project_json_data(self, mocker): + expected_json_data = {"key": "value"} + encoded_response = json.dumps(expected_json_data).encode("utf-8") + + mock_get_default_response = mocker.patch( + "kedro_viz.api.rest.responses.get_default_response", + return_value={"key": "value"}, + ) + mock_get_encoded_response = mocker.patch( + "kedro_viz.api.rest.responses.get_encoded_response", + return_value=encoded_response, + ) + + json_data = get_kedro_project_json_data() + + mock_get_default_response.assert_called_once() + mock_get_encoded_response.assert_called_once_with( + mock_get_default_response.return_value + ) + assert json_data == expected_json_data + def test_save_api_main_response_to_fs(self, mocker): expected_default_response = {"test": "json"} main_path = "/main"