Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Introduce get_kedro_project_json_data function for VSCode integration #2049

Merged
merged 10 commits into from
Sep 3, 2024
30 changes: 28 additions & 2 deletions package/kedro_viz/api/rest/responses.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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():
astrojuanlu marked this conversation as resolved.
Show resolved Hide resolved
"""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:
Expand Down
16 changes: 11 additions & 5 deletions package/kedro_viz/server.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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()

Expand Down Expand Up @@ -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")
Expand Down
23 changes: 23 additions & 0 deletions package/tests/test_api/test_rest/test_responses.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
# pylint: disable=too-many-lines
import json
import logging
import operator
from pathlib import Path
Expand All @@ -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,
Expand Down Expand Up @@ -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"
Expand Down
Loading