Merge pull request #429 from nodestream-proj/pipeline-specific-schemas

Updating Cypheresque to allow for filtering by pipeline(s)

Author: rreddy2

nodestream/cli/commands/print_schema.py

@@ -1,13 +1,16 @@
-from cleo.helpers import option
+from typing import ClassVar
+
+from cleo.helpers import argument, option
 
 from ..operations import InitializeProject, PrintProjectSchema
 from .nodestream_command import NodestreamCommand
-from .shared_options import PROJECT_FILE_OPTION
+from .shared_options import MANY_PIPELINES_ARGUMENT, PROJECT_FILE_OPTION
 
 
 class PrintSchema(NodestreamCommand):
     name = "print schema"
     description = "Print the schema for the current project"
+    arguments: ClassVar[list[argument]] = [MANY_PIPELINES_ARGUMENT]
     options = [
         PROJECT_FILE_OPTION,
         option(
@@ -37,5 +40,6 @@ async def handle_async(self):
                 format_string=self.option("format"),
                 type_overrides_file=self.option("overrides"),
                 output_file=self.option("out"),
+                pipeline_names=self.argument("pipelines"),
             )
         )

nodestream/cli/operations/print_project_schema.py

@@ -1,5 +1,5 @@
 from pathlib import Path
-from typing import Optional
+from typing import List, Optional
 
 from ...project import Project
 from ...schema.printers import SchemaPrinter
@@ -14,17 +14,25 @@ def __init__(
         format_string: str,
         type_overrides_file: Optional[str] = None,
         output_file: Optional[str] = None,
+        pipeline_names: Optional[List[str]] = None,
     ) -> None:
         self.project = project
         self.format_string = format_string
         self.output_file = output_file
         self.type_overrides_file = type_overrides_file
+        self.pipeline_names = pipeline_names or []
 
     async def perform(self, command: NodestreamCommand):
         type_overrides_file = (
             Path(self.type_overrides_file) if self.type_overrides_file else None
         )
-        schema = self.project.get_schema(type_overrides_file=type_overrides_file)
+        if self.pipeline_names:
+            schema = self.project.get_pipelines_schema(
+                pipeline_names=self.pipeline_names,
+                type_overrides_file=type_overrides_file,
+            )
+        else:
+            schema = self.project.get_schema(type_overrides_file=type_overrides_file)
 
         # Import all schema printers so that they can register themselves
         SchemaPrinter.import_all()

nodestream/project/project.py

@@ -13,7 +13,12 @@
 from ..pipeline import Step
 from ..pipeline.object_storage import ObjectStore
 from ..pluggable import Pluggable
-from ..schema import ExpandsSchema, ExpandsSchemaFromChildren, Schema
+from ..schema import (
+    ExpandsSchema,
+    ExpandsSchemaFromChildren,
+    Schema,
+    SchemaExpansionCoordinator,
+)
 from .pipeline_definition import PipelineDefinition
 from .pipeline_scope import PipelineScope
 from .plugin import PluginConfiguration
@@ -245,6 +250,16 @@ def get_pipeline_by_name(self, pipeline_name: str) -> PipelineDefinition:
             if pipeline is not None:
                 return pipeline
 
+    def get_pipeline_by_names(
+        self, pipeline_names: List[str]
+    ) -> Iterable[PipelineDefinition]:
+        """Returns pipeline objects in the project."""
+        for scope in self.scopes_by_name.values():
+            for pipeline_name in pipeline_names:
+                pipeline = scope.pipelines_by_name.get(pipeline_name, None)
+                if pipeline is not None:
+                    yield pipeline
+
     def delete_pipeline(
         self,
         scope_name: Optional[str],
@@ -291,6 +306,51 @@ def get_schema(self, type_overrides_file: Optional[Path] = None) -> Schema:
             schema.merge(overrides_schema)
         return schema
 
+    def get_pipelines_schema(
+        self, pipeline_names: List[str], type_overrides_file: Optional[Path] = None
+    ) -> Schema:
+        """Returns a `GraphSchema` representing only the specified pipelines.
+
+        This method generates a schema from only the specified pipelines,
+        allowing you to see what schema specific pipelines contribute without interference
+        from other pipelines in the project.
+
+        Args:
+            pipeline_names (List[str]): List of pipeline names to include in schema generation.
+            type_overrides_file (Optional[Path], optional): A path to a YAML file containing type overrides. Defaults to None.
+
+        Returns:
+            Schema: The schema representing only the specified pipelines.
+
+        Raises:
+            ValueError: If none of the specified pipelines are found.
+        """
+        pipeline_definitions = list(self.get_pipeline_by_names(pipeline_names))
+
+        if not pipeline_definitions:
+            available_pipelines = [
+                name
+                for scope in self.scopes_by_name.values()
+                for name in scope.pipelines_by_name.keys()
+            ]
+            raise ValueError(
+                f"None of the specified pipelines {pipeline_names} were found. Available pipelines: {available_pipelines}"
+            )
+
+        coordinator = SchemaExpansionCoordinator(schema := Schema())
+
+        for pipeline in pipeline_definitions:
+            pipeline.initialize_for_introspection()
+            pipeline.expand_schema(coordinator=coordinator)
+
+        schema = coordinator.schema
+
+        if type_overrides_file is not None:
+            overrides_schema = Schema.read_from_file(type_overrides_file)
+            schema.merge(overrides_schema)
+
+        return schema
+
     def get_child_expanders(self) -> Iterable[ExpandsSchema]:
         return self.scopes_by_name.values()
 

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "nodestream"
-version = "0.14.13"
+version = "0.14.14"
 description = "A Fast, Declarative ETL for Graph Databases."
 license = "GPL-3.0-only"
 authors = [

tests/unit/cli/operations/test_print_project_schema.py

@@ -1,3 +1,5 @@
+from pathlib import Path
+
 import pytest
 
 from nodestream.cli.operations import PrintProjectSchema
@@ -42,3 +44,53 @@ async def test_print_project_schema_prints_schema_to_file(
     _, stdout, fileout = await ran_print_project_schema_operation("some/path")
     assert not stdout.called
     assert fileout.called
+
+
+@pytest.mark.asyncio
+async def test_print_project_schema_with_pipeline_filtering(
+    mocker, project_with_default_scope
+):
+    std_out = mocker.patch(
+        "nodestream.schema.printers.SchemaPrinter.print_schema_to_stdout"
+    )
+    project_with_default_scope.get_schema = mocker.Mock(return_value="full_schema")
+    project_with_default_scope.get_pipelines_schema = mocker.Mock(
+        return_value="filtered_schema"
+    )
+
+    operation = PrintProjectSchema(
+        project=project_with_default_scope,
+        format_string="plain",
+        pipeline_names=["test_pipeline"],
+    )
+    await operation.perform(mocker.Mock())
+
+    # Should call get_pipelines_schema, not get_schema
+    project_with_default_scope.get_pipelines_schema.assert_called_once_with(
+        pipeline_names=["test_pipeline"], type_overrides_file=None
+    )
+    project_with_default_scope.get_schema.assert_not_called()
+    std_out.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_print_project_schema_with_pipeline_filtering_and_overrides(
+    mocker, project_with_default_scope
+):
+    project_with_default_scope.get_pipelines_schema = mocker.Mock(
+        return_value="filtered_schema"
+    )
+
+    operation = PrintProjectSchema(
+        project=project_with_default_scope,
+        format_string="plain",
+        type_overrides_file="overrides.yaml",
+        pipeline_names=["test_pipeline", "another_pipeline"],
+    )
+    await operation.perform(mocker.Mock())
+
+    # Should call get_pipelines_schema with correct arguments
+    project_with_default_scope.get_pipelines_schema.assert_called_once_with(
+        pipeline_names=["test_pipeline", "another_pipeline"],
+        type_overrides_file=Path("overrides.yaml"),
+    )

tests/unit/project/test_project.py

@@ -378,3 +378,93 @@ def test_project_get_from_storage(project, mocker):
         result,
         same_instance(project.storage_configuration.initialize_by_name.return_value),
     )
+
+
+def test_get_pipelines_schema_multiple_pipelines(project, mocker):
+    mock_schema = mocker.Mock(spec=Schema)
+    mock_coordinator = mocker.Mock()
+    mock_coordinator.schema = mock_schema
+
+    mocker.patch(
+        "nodestream.project.project.SchemaExpansionCoordinator",
+        return_value=mock_coordinator,
+    )
+
+    mock_pipeline1 = mocker.Mock()
+    mock_pipeline1.initialize_for_introspection = mocker.Mock()
+    mock_pipeline1.expand_schema = mocker.Mock()
+
+    mock_pipeline2 = mocker.Mock()
+    mock_pipeline2.initialize_for_introspection = mocker.Mock()
+    mock_pipeline2.expand_schema = mocker.Mock()
+
+    mocker.patch.object(
+        project, "get_pipeline_by_names", return_value=[mock_pipeline1, mock_pipeline2]
+    )
+
+    result = project.get_pipelines_schema(["test", "test2"])
+
+    mock_pipeline1.initialize_for_introspection.assert_called_once()
+    mock_pipeline1.expand_schema.assert_called_once_with(coordinator=mock_coordinator)
+    mock_pipeline2.initialize_for_introspection.assert_called_once()
+    mock_pipeline2.expand_schema.assert_called_once_with(coordinator=mock_coordinator)
+    assert_that(result, same_instance(mock_schema))
+
+
+def test_get_pipelines_schema_with_type_overrides(project, mocker):
+    mock_base_schema = mocker.Mock(spec=Schema)
+    mock_overrides_schema = mocker.Mock(spec=Schema)
+    mock_coordinator = mocker.Mock()
+    mock_coordinator.schema = mock_base_schema
+
+    mocker.patch(
+        "nodestream.project.project.SchemaExpansionCoordinator",
+        return_value=mock_coordinator,
+    )
+    Schema.read_from_file = mocker.Mock(return_value=mock_overrides_schema)
+
+    mock_pipeline = mocker.Mock()
+    mock_pipeline.initialize_for_introspection = mocker.Mock()
+    mock_pipeline.expand_schema = mocker.Mock()
+    mocker.patch.object(project, "get_pipeline_by_names", return_value=[mock_pipeline])
+
+    overrides_path = Path("some/overrides.yaml")
+    result = project.get_pipelines_schema(["test"], overrides_path)
+
+    Schema.read_from_file.assert_called_once_with(overrides_path)
+    mock_base_schema.merge.assert_called_once_with(mock_overrides_schema)
+    assert_that(result, same_instance(mock_base_schema))
+
+
+def test_get_pipelines_schema_nonexistent_pipeline_raises_error(project, mocker):
+    mocker.patch.object(project, "get_pipeline_by_names", return_value=[])
+
+    with pytest.raises(ValueError) as exc_info:
+        project.get_pipelines_schema(["nonexistent"])
+
+    error_message = str(exc_info.value)
+    assert_that(
+        error_message,
+        equal_to(
+            "None of the specified pipelines ['nonexistent'] were found. Available pipelines: ['test', 'test2']"
+        ),
+    )
+
+
+def test_get_pipeline_by_names(project):
+    pipelines = list(project.get_pipeline_by_names(["test", "test2"]))
+
+    assert_that(pipelines, has_length(2))
+    assert_that([p.name for p in pipelines], contains_inanyorder("test", "test2"))
+
+
+def test_get_pipeline_by_names_nonexistent_pipeline(project):
+    pipelines = list(project.get_pipeline_by_names(["nonexistent"]))
+    assert_that(pipelines, has_length(0))
+
+
+def test_get_pipeline_by_names_mixed_existing_nonexisting(project):
+    pipelines = list(project.get_pipeline_by_names(["test", "nonexistent", "test2"]))
+
+    assert_that(pipelines, has_length(2))
+    assert_that([p.name for p in pipelines], contains_inanyorder("test", "test2"))