graphdatascience 1.21a1__tar.gz → 1.22__tar.gz

{graphdatascience-1.21a1/src/graphdatascience.egg-info → graphdatascience-1.22}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: graphdatascience
-Version: 1.21a1
+Version: 1.22
 Summary: A Python client for the Neo4j Graph Data Science (GDS) library
 Author-email: Neo4j <team-gds@neo4j.org>
 License-Expression: Apache-2.0

{graphdatascience-1.21a1 → graphdatascience-1.22}/pyproject.toml

@@ -59,6 +59,7 @@ dev = [
     { include-group = "dev-base" },
     { include-group = "test" },
     { include-group = "docs-ci" },
+    { include-group = "notebook-base"}
 ]
 
 dev-base = [
@@ -77,6 +78,7 @@ dev-base = [
 ]
 test = [
     "pytest == 9.0.2",
+    "pytest-timeout",
     "requests_mock == 1.12.1",
     "pytest_mock == 3.15.1",
     "testcontainers == 4.14.0",

{graphdatascience-1.21a1 → graphdatascience-1.22}/src/graphdatascience/arrow_client/authenticated_flight_client.py

@@ -10,7 +10,7 @@ from typing import Any, Iterator, Type
 import certifi
 from pyarrow import Schema, flight
 from pyarrow import __version__ as arrow_version
-from pyarrow._flight import (
+from pyarrow.flight import (
     Action,
     ActionType,
     FlightInternalError,
@@ -118,6 +118,14 @@ class AuthenticatedArrowClient:
 
         self._flight_client: flight.FlightClient = self._instantiate_flight_client()
 
+    def _reconnect(self) -> None:
+        try:
+            self._flight_client.close()
+        except Exception:
+            pass
+
+        self._flight_client = self._instantiate_flight_client()
+
     def connection_info(self) -> ConnectionInfo:
         """
         Returns the host and port of the GDS Arrow server.
@@ -156,10 +164,14 @@ class AuthenticatedArrowClient:
 
         @self._retry_config.decorator(operation_name="Request token", logger=self._logger)
         def auth_with_retry() -> None:
-            client = self._flight_client
-            if self._auth:
-                auth_pair = self._auth.auth_pair()
-                client.authenticate_basic_token(auth_pair[0], auth_pair[1])
+            try:
+                client = self._flight_client
+                if self._auth:
+                    auth_pair = self._auth.auth_pair()
+                    client.authenticate_basic_token(auth_pair[0], auth_pair[1])
+            except (FlightTimedOutError, FlightUnavailableError, FlightInternalError):
+                self._reconnect()
+                raise
 
         if self._auth:
             auth_with_retry()
@@ -178,9 +190,13 @@ class AuthenticatedArrowClient:
     def do_action_with_retry(self, endpoint: str, payload: bytes | dict[str, Any]) -> list[Result]:
         @self._retry_config.decorator(operation_name="Send action", logger=self._logger)
         def run_with_retry() -> list[Result]:
-            # the Flight response error code is only checked on iterator consumption
-            # we eagerly collect iterator here to trigger retry in case of an error
-            return list(self.do_action(endpoint, payload))
+            try:
+                # the Flight response error code is only checked on iterator consumption
+                # we eagerly collect iterator here to trigger retry in case of an error
+                return list(self.do_action(endpoint, payload))
+            except (FlightTimedOutError, FlightUnavailableError, FlightInternalError):
+                self._reconnect()
+                raise
 
         return run_with_retry()
 
@@ -192,7 +208,11 @@ class AuthenticatedArrowClient:
     ) -> tuple[flight.FlightStreamWriter, flight.FlightMetadataReader]:
         @self._retry_config.decorator(operation_name="Do put", logger=self._logger)
         def run_with_retry() -> tuple[flight.FlightStreamWriter, flight.FlightMetadataReader]:
-            return self._flight_client.do_put(descriptor, schema)  # type: ignore
+            try:
+                return self._flight_client.do_put(descriptor, schema)  # type: ignore
+            except (FlightTimedOutError, FlightUnavailableError, FlightInternalError):
+                self._reconnect()
+                raise
 
         return run_with_retry()
 

{graphdatascience-1.21a1 → graphdatascience-1.22}/src/graphdatascience/arrow_client/middleware/auth_middleware.py

@@ -4,7 +4,7 @@ import base64
 import time
 from typing import Any
 
-from pyarrow._flight import ClientMiddleware, ClientMiddlewareFactory
+from pyarrow.flight import ClientMiddleware, ClientMiddlewareFactory
 
 from graphdatascience.arrow_client.arrow_authentication import ArrowAuthentication
 

{graphdatascience-1.21a1 → graphdatascience-1.22}/src/graphdatascience/arrow_client/middleware/user_agent_middleware.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 from typing import Any
 
-from pyarrow._flight import ClientMiddleware, ClientMiddlewareFactory
+from pyarrow.flight import ClientMiddleware, ClientMiddlewareFactory
 
 
 class UserAgentFactory(ClientMiddlewareFactory):  # type: ignore

{graphdatascience-1.21a1/src/graphdatascience/arrow_client/v2 → graphdatascience-1.22/src/graphdatascience/arrow_client/v1}/data_mapper_utils.py

@@ -1,7 +1,7 @@
 import json
 from typing import Any
 
-from pyarrow._flight import Result
+from pyarrow.flight import Result
 
 
 def deserialize_single(input_stream: list[Result]) -> dict[str, Any]:

{graphdatascience-1.21a1/src/graphdatascience/arrow_client/v1 → graphdatascience-1.22/src/graphdatascience/arrow_client/v2}/data_mapper_utils.py

@@ -1,7 +1,7 @@
 import json
 from typing import Any
 
-from pyarrow._flight import Result
+from pyarrow.flight import Result
 
 
 def deserialize_single(input_stream: list[Result]) -> dict[str, Any]:

{graphdatascience-1.21a1 → graphdatascience-1.22}/src/graphdatascience/arrow_client/v2/gds_arrow_client.py

@@ -123,7 +123,8 @@ class GdsArrowClient:
             job_id=job_id,
         )
 
-        return JobClient.run_job_and_wait(self._flight_client, "v2/graph.nodeLabel.stream", config, log_progress)
+        compute_job = JobClient.run_job_and_wait(self._flight_client, "v2/graph.nodeLabel.stream", config, log_progress)
+        return JobClient.start_export_result(self._flight_client, graph_name, compute_job)
 
     def get_relationships(
         self,
@@ -175,14 +176,12 @@ class GdsArrowClient:
 
         return JobClient.run_job(self._flight_client, endpoint, config)
 
-    def stream_job(self, graph_name: str, job_id: str) -> pandas.DataFrame:
+    def stream_job(self, job_id: str) -> pandas.DataFrame:
         """
         Streams the results of a previously started job.
 
         Parameters
         ----------
-        graph_name
-           The graph name.
         job_id
             Identifier for the computation.
 
@@ -191,7 +190,7 @@ class GdsArrowClient:
         DataFrame
             A pandas DataFrame containing the results of the job.
         """
-        return JobClient().stream_results(self._flight_client, graph_name, job_id)
+        return JobClient().get_stream(self._flight_client, job_id)
 
     def create_graph(
         self,
@@ -415,7 +414,7 @@ class GdsArrowClient:
         job_id : str
             The job id of the process
         """
-        self._flight_client.do_action_with_retry("v2/jobs.cancel", {"jobId": job_id})
+        JobClient.cancel_job(self._flight_client, job_id)
 
     def job_status(self, job_id: str) -> JobStatus:
         """
@@ -472,32 +471,21 @@ class GdsArrowClient:
 
         return self._flight_client.request_token()
 
+    def flight_client(self) -> AuthenticatedArrowClient:
+        """
+        Returns the underlying flight client.
+        """
+        return self._flight_client
+
     def _upload_data(
         self,
         endpoint: str,
         job_id: str,
-        data: pyarrow.Table | list[pyarrow.RecordBatch] | list[pandas.DataFrame],
+        data: pyarrow.Table | list[pyarrow.RecordBatch] | list[pandas.DataFrame] | pandas.DataFrame,
         batch_size: int = 10000,
         progress_callback: ProgressCallback = lambda num_rows: None,
         termination_flag: TerminationFlag | None = None,
     ) -> None:
-        if isinstance(data, pandas.DataFrame):
-            data = [data]
-
-        def map_list_entry_to_batches(e: pyarrow.RecordBatch | pandas.DataFrame) -> list[pyarrow.RecordBatch]:
-            return [e] if isinstance(e, pyarrow.RecordBatch) else pyarrow.Table.from_pandas(e).to_batches(batch_size)
-
-        batches: list[pyarrow.RecordBatch] = []
-        match data:
-            case list():
-                for entry in data:
-                    partial_batches = map_list_entry_to_batches(entry)
-                    batches.extend(partial_batches)
-            case pyarrow.Table():
-                batches = data.to_batches(batch_size)
-            case _:
-                raise ValueError(f"Unsupported data type: {type(data)}")
-
         flight_descriptor = {
             "name": endpoint,
             "version": ArrowEndpointVersion.V2.version(),
@@ -507,11 +495,45 @@ class GdsArrowClient:
         }
         upload_descriptor = flight.FlightDescriptor.for_command(json.dumps(flight_descriptor).encode("utf-8"))
 
+        for batches in self._batch_groups(data, batch_size):
+            self._upload_batches(upload_descriptor, batches, job_id, progress_callback, termination_flag)
+
+    @staticmethod
+    def _batch_groups(
+        data: pyarrow.Table | list[pyarrow.RecordBatch] | list[pandas.DataFrame] | pandas.DataFrame,
+        batch_size: int,
+    ) -> list[list[pyarrow.RecordBatch]]:
+        if isinstance(data, pandas.DataFrame):
+            return [pyarrow.Table.from_pandas(data).to_batches(batch_size)]
+
+        if isinstance(data, pyarrow.Table):
+            return [data.to_batches(batch_size)]
+
+        if isinstance(data, list):
+            if all(isinstance(entry, pandas.DataFrame) for entry in data):
+                return [pyarrow.Table.from_pandas(entry).to_batches(batch_size) for entry in data]
+
+            if all(isinstance(entry, pyarrow.RecordBatch) for entry in data):
+                return [data]
+
+        raise ValueError(f"Unsupported data type: {type(data)}")
+
+    def _upload_batches(
+        self,
+        upload_descriptor: flight.FlightDescriptor,
+        batches: list[pyarrow.RecordBatch],
+        job_id: str,
+        progress_callback: ProgressCallback,
+        termination_flag: TerminationFlag | None,
+    ) -> None:
+        if not batches:
+            return
+
         put_stream, ack_stream = self._flight_client.do_put_with_retry(upload_descriptor, batches[0].schema)
 
         @self._flight_client._retry_config.decorator(operation_name="Upload batch", logger=self._logger)
-        def upload_batch(p: RecordBatch) -> None:
-            put_stream.write_batch(p)
+        def upload_batch(batch: RecordBatch) -> None:
+            put_stream.write_batch(batch)
 
         with put_stream:
             for partition in batches:

{graphdatascience-1.21a1 → graphdatascience-1.22}/src/graphdatascience/arrow_client/v2/job_client.py

@@ -2,7 +2,7 @@ import json
 from typing import Any
 
 from pandas import ArrowDtype, DataFrame
-from pyarrow._flight import Ticket
+from pyarrow.flight import Ticket
 from tenacity import Retrying, retry_if_result
 
 from graphdatascience.arrow_client.authenticated_flight_client import AuthenticatedArrowClient
@@ -13,6 +13,7 @@ from graphdatascience.query_runner.termination_flag import TerminationFlag
 from graphdatascience.retry_utils.retry_utils import job_wait_strategy
 
 JOB_STATUS_ENDPOINT = "v2/jobs.status"
+JOBS_CANCEL_ENDPOINT = "v2/jobs.cancel"
 RESULTS_SUMMARY_ENDPOINT = "v2/results.summary"
 
 
@@ -73,6 +74,10 @@ class JobClient:
                     if progress_bar:
                         progress_bar.update(job_status.status, job_status.progress_percent(), job_status.sub_tasks())
 
+    @staticmethod
+    def cancel_job(client: AuthenticatedArrowClient, job_id: str) -> None:
+        client.do_action_with_retry(JOBS_CANCEL_ENDPOINT, JobIdConfig(jobId=job_id).dump_camel())
+
     @staticmethod
     def get_job_status(client: AuthenticatedArrowClient, job_id: str) -> JobStatus:
         arrow_res = client.do_action_with_retry(JOB_STATUS_ENDPOINT, JobIdConfig(jobId=job_id).dump_camel())
@@ -86,14 +91,22 @@ class JobClient:
 
     @staticmethod
     def stream_results(client: AuthenticatedArrowClient, graph_name: str, job_id: str) -> DataFrame:
+        export_job_id = JobClient.start_export_result(client, graph_name, job_id)
+
+        return JobClient.get_stream(client, export_job_id)
+
+    @staticmethod
+    def start_export_result(client: AuthenticatedArrowClient, graph_name: str, job_id: str) -> str:
         payload = {
             "graphName": graph_name,
             "jobId": job_id,
         }
 
         res = client.do_action_with_retry("v2/results.stream", payload)
-        export_job_id = JobIdConfig(**deserialize_single(res)).job_id
+        return JobIdConfig(**deserialize_single(res)).job_id
 
+    @staticmethod
+    def get_stream(client: AuthenticatedArrowClient, export_job_id: str) -> DataFrame:
         stream_payload = {"version": "v2", "name": export_job_id, "body": {}}
 
         ticket = Ticket(json.dumps(stream_payload).encode("utf-8"))

{graphdatascience-1.21a1 → graphdatascience-1.22}/src/graphdatascience/arrow_client/v2/mutation_client.py

@@ -1,5 +1,6 @@
 import math
 import time
+from typing import Any
 
 from graphdatascience.arrow_client.authenticated_flight_client import AuthenticatedArrowClient
 from graphdatascience.arrow_client.v2.api_types import MutateResult
@@ -23,17 +24,33 @@ class MutationClient:
             process_job_id=process_job_id,
         )
 
+    @staticmethod
+    def mutate_node_properties(
+        client: AuthenticatedArrowClient,
+        result_store_job_id: str,
+        mutate_properties: dict[str, str] | None = None,
+        process_job_id: str | None = None,
+    ) -> MutateResult:
+        return MutationClient._mutate(
+            client=client,
+            result_store_job_id=result_store_job_id,
+            mutate_properties=mutate_properties,
+            process_job_id=process_job_id,
+        )
+
     @staticmethod
     def mutate_relationship_property(
         client: AuthenticatedArrowClient,
         job_id: str,
         mutate_relationship_type: str,
         mutate_property: str | None,
+        mutate_properties: dict[str, str] | None = None,
     ) -> MutateResult:
         return MutationClient._mutate(
             client=client,
             result_store_job_id=job_id,
             mutate_property=mutate_property,
+            mutate_properties=mutate_properties,
             mutate_relationship_type=mutate_relationship_type,
         )
 
@@ -43,9 +60,10 @@ class MutationClient:
         result_store_job_id: str,
         process_job_id: str | None = None,
         mutate_property: str | None = None,
+        mutate_properties: dict[str, str] | None = None,
         mutate_relationship_type: str | None = None,
     ) -> MutateResult:
-        mutate_config = {
+        mutate_config: dict[str, Any] = {
             "jobId": result_store_job_id,
         }
         if process_job_id:
@@ -54,6 +72,8 @@ class MutationClient:
             mutate_config["mutateRelationshipType"] = mutate_relationship_type
         if mutate_property:
             mutate_config["mutateProperty"] = mutate_property
+        if mutate_properties:
+            mutate_config["mutateProperties"] = mutate_properties
 
         start_time = time.time()
         mutate_arrow_res = client.do_action_with_retry(MutationClient.MUTATE_ENDPOINT, mutate_config)

{graphdatascience-1.21a1 → graphdatascience-1.22}/src/graphdatascience/endpoints.py

@@ -13,6 +13,9 @@ from .pipeline.pipeline_endpoints import (
     PipelineAlphaEndpoints,
     PipelineBetaEndpoints,
     PipelineEndpoints,
+    SessionPipelineAlphaEndpoints,
+    SessionPipelineBetaEndpoints,
+    SessionPipelineEndpoints,
 )
 from .query_runner.query_runner import QueryRunner
 from .server_version.server_version import ServerVersion
@@ -43,6 +46,17 @@ class DirectEndpoints(
         super().__init__(query_runner, namespace, server_version)
 
 
+class DirectSessionEndpoints(
+    SingleModeAlgoEndpoints,
+    DirectSystemEndpoints,
+    DirectUtilEndpoints,
+    SessionPipelineEndpoints,
+    ConfigEndpoints,
+):
+    def __init__(self, query_runner: QueryRunner, namespace: str, server_version: ServerVersion):
+        super().__init__(query_runner, namespace, server_version)
+
+
 """
 This class should inherit endpoint classes that only expose calls of the `gds.alpha` namespace.
 Example of such endpoints: "gds.alpha.listProgress".
@@ -67,7 +81,7 @@ class AlphaEndpoints(
 
 class AlphaRemoteEndpoints(
     GraphAlphaEndpoints,
-    PipelineAlphaEndpoints,
+    SessionPipelineAlphaEndpoints,
     ModelAlphaEndpoints,
     SingleModeAlphaAlgoEndpoints,
     SystemAlphaEndpoints,
@@ -80,6 +94,14 @@ class AlphaRemoteEndpoints(
         return IndirectAlphaCallBuilder(self._query_runner, f"{self._namespace}.{attr}", self._server_version)
 
 
+class BetaSessionEndpoints(GraphBetaEndpoints, SessionPipelineBetaEndpoints, ModelBetaEndpoints, SystemBetaEndpoints):
+    def __init__(self, query_runner: QueryRunner, namespace: str, server_version: ServerVersion):
+        super().__init__(query_runner, namespace, server_version)
+
+    def __getattr__(self, attr: str) -> IndirectBetaCallBuilder:
+        return IndirectBetaCallBuilder(self._query_runner, f"{self._namespace}.{attr}", self._server_version)
+
+
 """
 This class should inherit endpoint classes that only expose calls of the `gds.beta` namespace.
 Example of such endpoints:  "gds.beta.listProgress".

{graphdatascience-1.21a1 → graphdatascience-1.22}/src/graphdatascience/error/client_only_endpoint.py

@@ -7,6 +7,10 @@ from ..caller_base import CallerBase
 F = TypeVar("F", bound=Callable[..., Any])
 
 
+def deprecated_endpoint_message(old_endpoint: str, new_endpoint: str) -> str:
+    return f"Deprecated `{old_endpoint}` in favor of `{new_endpoint}`"
+
+
 def client_only_endpoint(expected_namespace_prefix: str) -> Callable[[F], F]:
     def decorator(func: F) -> F:
         wraps(func)
@@ -35,7 +39,7 @@ def client_deprecated(
 
         @wraps(func)
         def wrapper(self: CallerBase, *args: Any, **kwargs: Any) -> Any:
-            warnings.warn(f"Deprecated `{old_endpoint}` in favor of `{new_endpoint}`", DeprecationWarning)
+            warnings.warn(deprecated_endpoint_message(old_endpoint, new_endpoint), DeprecationWarning)
             return func(self, *args, **kwargs)
 
         return cast(F, wrapper)

graphdatascience-1.22/src/graphdatascience/graph/graph_remote_proc_runner.py

@@ -0,0 +1,76 @@
+from typing import Any
+from uuid import uuid4
+
+import pandas as pd
+
+from graphdatascience import Graph, GraphCreateResult, QueryRunner, ServerVersion
+from graphdatascience.arrow_client.v2.gds_arrow_client import GdsArrowClient
+from graphdatascience.graph.base_graph_proc_runner import BaseGraphProcRunner
+from graphdatascience.procedure_surface.arrow.error_handler import handle_flight_error
+from graphdatascience.query_runner.protocol.project_protocols import ProjectProtocol
+from graphdatascience.query_runner.protocol.projection_runner import ProjectionRunner
+from graphdatascience.query_runner.termination_flag import TerminationFlag
+from graphdatascience.session.dbms.protocol_resolver import ProtocolVersionResolver
+
+
+class GraphRemoteProcRunner(BaseGraphProcRunner):
+    def __init__(
+        self,
+        query_runner: QueryRunner,
+        arrow_client: GdsArrowClient,
+        db_query_runner: QueryRunner | None,
+        namespace: str,
+        server_version: ServerVersion,
+    ):
+        super().__init__(query_runner, namespace, server_version)
+        self._arrow_client = arrow_client
+        self._db_query_runner = db_query_runner
+
+    def project(
+        self,
+        graph_name: str,
+        query: str,
+        job_id: str | None = None,
+        query_parameters: dict[str, Any] | None = None,
+        concurrency: int = 4,
+        undirected_relationship_types: list[str] | None = None,
+        inverse_indexed_relationship_types: list[str] | None = None,
+        batch_size: int | None = None,
+        logging: bool = True,
+    ) -> GraphCreateResult:
+        if not self._db_query_runner:
+            raise Exception("Projections from a database are not supported for standalone sessions")
+
+        resolved_protocol_version = ProtocolVersionResolver(self._query_runner).resolve()
+
+        project_protocol = ProjectProtocol.select(
+            resolved_protocol_version,
+            self._arrow_client.flight_client(),
+            self._db_query_runner,
+            TerminationFlag.create(),
+        )
+
+        try:
+            result = ProjectionRunner(
+                project_protocol,
+                self._arrow_client.flight_client(),
+                TerminationFlag.create(),
+            ).run_cypher_projection(
+                graph_name,
+                query,
+                job_id or str(uuid4()),
+                query_parameters,
+                concurrency,
+                undirected_relationship_types,
+                inverse_indexed_relationship_types,
+                batch_size,
+                logging,
+            )
+
+            # explicitly specify empty dtype to avoid pandas 1.x FutureWarning
+            result_series = pd.Series(result) if result else pd.Series(result, dtype=object)
+            return GraphCreateResult(Graph(graph_name, self._query_runner), result_series)
+
+        except Exception as e:
+            handle_flight_error(e)
+            raise e  # above should already raise

{graphdatascience-1.21a1 → graphdatascience-1.22}/src/graphdatascience/model/v2/model.py

@@ -12,9 +12,6 @@ class Model(ABC):
         self._name = name
         self._model_api = model_api
 
-    # TODO estimate mode, predict modes on here?
-    # implement Cypher and Arrow info_provider and stuff
-
     def name(self) -> str:
         """
         Get the name of the model.
@@ -52,7 +49,7 @@ class Model(ABC):
         return self._model_api.drop(self._name, failIfMissing)
 
     def __str__(self) -> str:
-        return f"{self.__class__.__name__}(name={self.name()}, type={self.details().type})"
+        return f"{self.__class__.__name__}(name={self.name()}, type={self.details().model_type})"
 
     def __repr__(self) -> str:
         return f"{self.__class__.__name__}({self.details().model_dump()})"

{graphdatascience-1.21a1 → graphdatascience-1.22}/src/graphdatascience/model/v2/model_details.py

@@ -1,13 +1,13 @@
 import datetime
 from typing import Any
 
-from pydantic import BaseModel, Field
+from pydantic import BaseModel
 from pydantic.alias_generators import to_camel
 
 
 class ModelDetails(BaseModel, alias_generator=to_camel):
-    name: str = Field(alias="modelName")
-    type: str = Field(alias="modelType")
+    model_name: str
+    model_type: str
     train_config: dict[str, Any]
     graph_schema: dict[str, Any]
     loaded: bool

graphdatascience-1.22/src/graphdatascience/pipeline/pipeline_alpha_proc_runner.py

@@ -0,0 +1,25 @@
+import warnings
+
+from ..error.client_only_endpoint import deprecated_endpoint_message
+from ..error.illegal_attr_checker import IllegalAttrChecker
+from ..error.uncallable_namespace import UncallableNamespace
+from .nr_pipeline_create_runner import NRPipelineCreateRunner
+
+
+class PipelineAlphaProcRunner(UncallableNamespace, IllegalAttrChecker):
+    @property
+    def nodeRegression(self) -> NRPipelineCreateRunner:
+        return NRPipelineCreateRunner(self._query_runner, f"{self._namespace}.nodeRegression", self._server_version)
+
+
+class SessionPipelineAlphaProcRunner(PipelineAlphaProcRunner):
+    @property
+    def nodeRegression(self) -> NRPipelineCreateRunner:
+        warnings.warn(
+            deprecated_endpoint_message(
+                "gds.alpha.pipeline.nodeRegression.create",
+                "gds.v2.pipeline.node_regression.create",
+            ),
+            DeprecationWarning,
+        )
+        return super().nodeRegression

graphdatascience-1.22/src/graphdatascience/pipeline/pipeline_beta_proc_runner.py

@@ -0,0 +1,79 @@
+from __future__ import annotations
+
+import warnings
+from typing import Any
+
+from pandas import DataFrame, Series
+
+from ..error.client_only_endpoint import deprecated_endpoint_message
+from ..error.illegal_attr_checker import IllegalAttrChecker
+from ..error.uncallable_namespace import UncallableNamespace
+from ..model.pipeline_model import PipelineModel
+from .lp_pipeline_create_runner import LPPipelineCreateRunner
+from .nc_pipeline_create_runner import NCPipelineCreateRunner
+from .pipeline_proc_runner import PipelineProcRunner
+from .training_pipeline import TrainingPipeline
+
+
+class PipelineBetaProcRunner(UncallableNamespace, IllegalAttrChecker):
+    @property
+    def linkPrediction(self) -> LPPipelineCreateRunner:
+        return LPPipelineCreateRunner(self._query_runner, f"{self._namespace}.linkPrediction", self._server_version)
+
+    @property
+    def nodeClassification(self) -> NCPipelineCreateRunner:
+        return NCPipelineCreateRunner(self._query_runner, f"{self._namespace}.nodeClassification", self._server_version)
+
+    def list(self, pipeline: TrainingPipeline[PipelineModel] | None = None) -> DataFrame:
+        return PipelineProcRunner(self._query_runner, self._namespace, self._server_version).list(pipeline)
+
+    def exists(self, pipeline_name: str) -> Series[Any]:
+        return PipelineProcRunner(self._query_runner, self._namespace, self._server_version).exists(pipeline_name)
+
+    def drop(self, pipeline: TrainingPipeline[PipelineModel]) -> Series[Any]:
+        return PipelineProcRunner(self._query_runner, self._namespace, self._server_version).drop(pipeline)
+
+
+class SessionPipelineBetaProcRunner(PipelineBetaProcRunner):
+    @property
+    def linkPrediction(self) -> LPPipelineCreateRunner:
+        warnings.warn(
+            deprecated_endpoint_message(
+                "gds.beta.pipeline.linkPrediction.create",
+                "gds.v2.pipeline.link_prediction.create",
+            ),
+            DeprecationWarning,
+        )
+        return super().linkPrediction
+
+    @property
+    def nodeClassification(self) -> NCPipelineCreateRunner:
+        warnings.warn(
+            deprecated_endpoint_message(
+                "gds.beta.pipeline.nodeClassification.create",
+                "gds.v2.pipeline.node_classification.create",
+            ),
+            DeprecationWarning,
+        )
+        return super().nodeClassification
+
+    def list(self, pipeline: TrainingPipeline[PipelineModel] | None = None) -> DataFrame:
+        warnings.warn(
+            deprecated_endpoint_message("gds.beta.pipeline.list", "gds.v2.pipeline.list"),
+            DeprecationWarning,
+        )
+        return super().list(pipeline)
+
+    def exists(self, pipeline_name: str) -> Series[Any]:
+        warnings.warn(
+            deprecated_endpoint_message("gds.beta.pipeline.exists", "gds.v2.pipeline.exists"),
+            DeprecationWarning,
+        )
+        return super().exists(pipeline_name)
+
+    def drop(self, pipeline: TrainingPipeline[PipelineModel]) -> Series[Any]:
+        warnings.warn(
+            deprecated_endpoint_message("gds.beta.pipeline.drop", "gds.v2.pipeline.drop"),
+            DeprecationWarning,
+        )
+        return super().drop(pipeline)