Separate out beta endpoints

adamnsch · adamnsch · commit d04017365337 · 2023-02-22T10:29:15.000+01:00
diff --git a/graphdatascience/alpha_endpoints.py b/graphdatascience/alpha_endpoints.py
@@ -1,6 +1,5 @@
-from .algo.algo_endpoints import AlgoEndpoints
 from .algo.single_mode_algo_endpoints import SingleModeAlgoEndpoints
-from .call_builder import CallBuilder
+from .call_builder import IndirectAlphaCallBuilder
 from .graph.graph_endpoints import GraphAlphaEndpoints
 from .model.model_endpoints import ModelAlphaEndpoints
 from .pipeline.pipeline_endpoints import PipelineAlphaEndpoints
@@ -10,14 +9,8 @@
 from .system.system_endpoints import SystemAlphaEndpoints
 from .topological_lp.topological_lp_endpoints import TopologicalLPEndpoints
 
-"""
-This class should inherit endpoint classes that only contain endpoints that can be called directly from
-the `gds` namespace. Example of such endpoints are: "graph" and "list".
-"""
-
 
 class AlphaEndpoints(
-    AlgoEndpoints,
     GraphAlphaEndpoints,
     PipelineAlphaEndpoints,
     TopologicalLPEndpoints,
@@ -29,5 +22,5 @@ class AlphaEndpoints(
     def __init__(self, query_runner: QueryRunner, namespace: str, server_version: ServerVersion):
         super().__init__(query_runner, namespace, server_version)
 
-    def __getattr__(self, attr: str) -> CallBuilder:
-        return CallBuilder(self._query_runner, f"{self._namespace}.{attr}", self._server_version)
+    def __getattr__(self, attr: str) -> IndirectAlphaCallBuilder:
+        return IndirectAlphaCallBuilder(self._query_runner, f"{self._namespace}.{attr}", self._server_version)
diff --git a/graphdatascience/beta_endpoints.py b/graphdatascience/beta_endpoints.py
@@ -0,0 +1,15 @@
+from .call_builder import IndirectBetaCallBuilder
+from .graph.graph_endpoints import GraphBetaEndpoints
+from .model.model_endpoints import ModelBetaEndpoints
+from .pipeline.pipeline_endpoints import PipelineBetaEndpoints
+from .query_runner.query_runner import QueryRunner
+from .server_version.server_version import ServerVersion
+from .system.system_endpoints import SystemBetaEndpoints
+
+
+class BetaEndpoints(GraphBetaEndpoints, PipelineBetaEndpoints, 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)
diff --git a/graphdatascience/call_builder.py b/graphdatascience/call_builder.py
@@ -1,8 +1,21 @@
+from .algo.algo_endpoints import AlgoEndpoints
 from .error.uncallable_namespace import UncallableNamespace
-from .indirect_endpoints import IndirectEndpoints
+from .utils.util_endpoints import IndirectUtilAlphaEndpoints
 
 
-class CallBuilder(IndirectEndpoints, UncallableNamespace):
-    def __getattr__(self, attr: str) -> "CallBuilder":
+class IndirectCallBuilder(AlgoEndpoints, UncallableNamespace):
+    def __getattr__(self, attr: str) -> "IndirectCallBuilder":
         namespace = f"{self._namespace}.{attr}"
-        return CallBuilder(self._query_runner, namespace, self._server_version)
+        return IndirectCallBuilder(self._query_runner, namespace, self._server_version)
+
+
+class IndirectBetaCallBuilder(AlgoEndpoints, UncallableNamespace):
+    def __getattr__(self, attr: str) -> "IndirectBetaCallBuilder":
+        namespace = f"{self._namespace}.{attr}"
+        return IndirectBetaCallBuilder(self._query_runner, namespace, self._server_version)
+
+
+class IndirectAlphaCallBuilder(AlgoEndpoints, IndirectUtilAlphaEndpoints, UncallableNamespace):
+    def __getattr__(self, attr: str) -> "IndirectAlphaCallBuilder":
+        namespace = f"{self._namespace}.{attr}"
+        return IndirectAlphaCallBuilder(self._query_runner, namespace, self._server_version)
diff --git a/graphdatascience/graph/graph_beta_proc_runner.py b/graphdatascience/graph/graph_beta_proc_runner.py
@@ -0,0 +1,47 @@
+from typing import Any, List, Tuple, Union
+
+from pandas import Series
+
+from ..caller_base import CallerBase
+from ..error.illegal_attr_checker import IllegalAttrChecker
+from ..error.uncallable_namespace import UncallableNamespace
+from .graph_export_runner import GraphExportCsvEndpoints
+from .graph_object import Graph
+from .graph_project_runner import GraphProjectBetaRunner
+from graphdatascience.graph.graph_entity_ops_runner import GraphRelationshipsBetaRunner
+
+Strings = Union[str, List[str]]
+
+
+class GraphBetaProcRunner(CallerBase, UncallableNamespace, IllegalAttrChecker):
+    @property
+    def project(self) -> GraphProjectBetaRunner:
+        self._namespace += ".project"
+        return GraphProjectBetaRunner(self._query_runner, self._namespace, self._server_version)
+
+    @property
+    def export(self) -> GraphExportCsvEndpoints:
+        self._namespace += ".export"
+        return GraphExportCsvEndpoints(self._query_runner, self._namespace, self._server_version)
+
+    @property
+    def relationships(self) -> GraphRelationshipsBetaRunner:
+        self._namespace += ".relationships"
+        return GraphRelationshipsBetaRunner(self._query_runner, self._namespace, self._server_version)
+
+    def generate(
+        self, graph_name: str, node_count: int, average_degree: int, **config: Any
+    ) -> Tuple[Graph, "Series[Any]"]:
+        self._namespace += ".generate"
+
+        query = f"CALL {self._namespace}($graph_name, $node_count, $average_degree, $config)"
+        params = {
+            "graph_name": graph_name,
+            "node_count": node_count,
+            "average_degree": average_degree,
+            "config": config,
+        }
+
+        result = self._query_runner.run_query(query, params).squeeze()
+
+        return Graph(graph_name, self._query_runner, self._server_version), result
diff --git a/graphdatascience/graph/graph_endpoints.py b/graphdatascience/graph/graph_endpoints.py
@@ -1,5 +1,6 @@
 from ..caller_base import CallerBase
 from .graph_alpha_proc_runner import GraphAlphaProcRunner
+from .graph_beta_proc_runner import GraphBetaProcRunner
 from .graph_proc_runner import GraphProcRunner
 
 
@@ -13,3 +14,9 @@ class GraphAlphaEndpoints(CallerBase):
     @property
     def graph(self) -> GraphAlphaProcRunner:
         return GraphAlphaProcRunner(self._query_runner, f"{self._namespace}.graph", self._server_version)
+
+
+class GraphBetaEndpoints(CallerBase):
+    @property
+    def graph(self) -> GraphBetaProcRunner:
+        return GraphBetaProcRunner(self._query_runner, f"{self._namespace}.graph", self._server_version)
diff --git a/graphdatascience/graph/graph_entity_ops_runner.py b/graphdatascience/graph/graph_entity_ops_runner.py
@@ -228,6 +228,8 @@ def drop(
 
         return self._query_runner.run_query(query, params).squeeze()  # type: ignore
 
+
+class GraphRelationshipsBetaRunner(GraphEntityOpsBaseRunner):
     @compatible_with("stream", min_inclusive=ServerVersion(2, 2, 0))
     @graph_type_check
     def stream(self, G: Graph, relationship_types: List[str] = ["*"], **config: Any) -> TopologyDataFrame:
diff --git a/graphdatascience/graph/graph_export_runner.py b/graphdatascience/graph/graph_export_runner.py
@@ -2,7 +2,9 @@
 
 from pandas import Series
 
+from ..caller_base import CallerBase
 from ..error.illegal_attr_checker import IllegalAttrChecker
+from ..error.uncallable_namespace import UncallableNamespace
 from .graph_object import Graph
 from .graph_type_check import graph_type_check
 
@@ -26,7 +28,15 @@ def estimate(self, G: Graph, **config: Any) -> "Series[Any]":
         return self._export_call(G, config)
 
 
-class GraphExportRunner(IllegalAttrChecker):
+class GraphExportCsvEndpoints(CallerBase, UncallableNamespace, IllegalAttrChecker):
+    @property
+    def csv(self) -> GraphExportCsvRunner:
+        self._namespace += ".csv"
+
+        return GraphExportCsvRunner(self._query_runner, self._namespace, self._server_version)
+
+
+class GraphExportRunner(CallerBase, IllegalAttrChecker):
     def __call__(self, G: Graph, **config: Any) -> "Series[Any]":
         return self._export_call(G, config)
 
diff --git a/graphdatascience/graph/graph_project_runner.py b/graphdatascience/graph/graph_project_runner.py
@@ -40,6 +40,8 @@ def estimate(self, node_spec: Any, relationship_spec: Any, **config: Any) -> "Se
     def cypher(self) -> "GraphProjectRunner":
         return GraphProjectRunner(self._query_runner, self._namespace + ".cypher", self._server_version)
 
+
+class GraphProjectBetaRunner(CallerBase, IllegalAttrChecker):
     @from_graph_type_check
     def subgraph(
         self,
diff --git a/graphdatascience/graph_data_science.py b/graphdatascience/graph_data_science.py
@@ -5,7 +5,8 @@
 from pandas import DataFrame, Series
 
 from .alpha_endpoints import AlphaEndpoints
-from .call_builder import CallBuilder
+from .beta_endpoints import BetaEndpoints
+from .call_builder import IndirectCallBuilder
 from .direct_endpoints import DirectEndpoints
 from .error.unable_to_connect import UnableToConnectError
 from .error.uncallable_namespace import UncallableNamespace
@@ -130,12 +131,12 @@ def __init__(
     def alpha(self) -> AlphaEndpoints:
         return AlphaEndpoints(self._query_runner, "gds.alpha", self._server_version)
 
-    # TODO
-    # def beta(self) -> None:
-    #     pass
+    @property
+    def beta(self) -> BetaEndpoints:
+        return BetaEndpoints(self._query_runner, "gds.beta", self._server_version)
 
-    def __getattr__(self, attr: str) -> CallBuilder:
-        return CallBuilder(self._query_runner, f"gds.{attr}", self._server_version)
+    def __getattr__(self, attr: str) -> IndirectCallBuilder:
+        return IndirectCallBuilder(self._query_runner, f"gds.{attr}", self._server_version)
 
     def set_database(self, database: str) -> None:
         self._query_runner.set_database(database)
diff --git a/graphdatascience/indirect_endpoints.py b/graphdatascience/indirect_endpoints.py
@@ -1,11 +1,6 @@
 from .algo.algo_endpoints import AlgoEndpoints
-from .graph.graph_endpoints import GraphEndpoints
-from .model.model_endpoints import ModelEndpoints
-from .pipeline.pipeline_endpoints import PipelineEndpoints
 from .query_runner.query_runner import QueryRunner
 from .server_version.server_version import ServerVersion
-from .system.system_endpoints import IndirectSystemEndpoints
-from .utils.util_endpoints import IndirectUtilEndpoints
 
 """
 This class should inherit endpoint classes that only contain endpoints that needs more of a prefix
@@ -14,13 +9,6 @@
 """
 
 
-class IndirectEndpoints(
-    AlgoEndpoints,
-    GraphEndpoints,
-    ModelEndpoints,
-    PipelineEndpoints,
-    IndirectSystemEndpoints,
-    IndirectUtilEndpoints,
-):
+class IndirectProductEndpoints(AlgoEndpoints):
     def __init__(self, query_runner: QueryRunner, namespace: str, server_version: ServerVersion):
         super().__init__(query_runner, namespace, server_version)
diff --git a/graphdatascience/model/model_beta_proc_runner.py b/graphdatascience/model/model_beta_proc_runner.py
@@ -0,0 +1,38 @@
+from typing import Any, Optional
+
+from pandas import DataFrame, Series
+
+from ..caller_base import CallerBase
+from ..error.illegal_attr_checker import IllegalAttrChecker
+from ..error.uncallable_namespace import UncallableNamespace
+from .model import Model
+
+
+class ModelBetaProcRunner(CallerBase, UncallableNamespace, IllegalAttrChecker):
+    def list(self, model: Optional[Model] = None) -> DataFrame:
+        self._namespace += ".list"
+
+        if model:
+            query = f"CALL {self._namespace}($model_name)"
+            params = {"model_name": model.name()}
+        else:
+            query = f"CALL {self._namespace}()"
+            params = {}
+
+        return self._query_runner.run_query(query, params)
+
+    def exists(self, model_name: str) -> "Series[Any]":
+        self._namespace += ".exists"
+
+        query = f"CALL {self._namespace}($model_name)"
+        params = {"model_name": model_name}
+
+        return self._query_runner.run_query(query, params).squeeze()  # type: ignore
+
+    def drop(self, model: Model) -> "Series[Any]":
+        self._namespace += ".drop"
+
+        query = f"CALL {self._namespace}($model_name)"
+        params = {"model_name": model.name()}
+
+        return self._query_runner.run_query(query, params).squeeze()  # type: ignore
diff --git a/graphdatascience/model/model_endpoints.py b/graphdatascience/model/model_endpoints.py
@@ -1,5 +1,6 @@
 from ..caller_base import CallerBase
 from .model_alpha_proc_runner import ModelAlphaProcRunner
+from .model_beta_proc_runner import ModelBetaProcRunner
 from .model_proc_runner import ModelProcRunner
 
 
@@ -9,6 +10,12 @@ def model(self) -> ModelProcRunner:
         return ModelProcRunner(self._query_runner, f"{self._namespace}.model", self._server_version)
 
 
+class ModelBetaEndpoints(CallerBase):
+    @property
+    def model(self) -> ModelBetaProcRunner:
+        return ModelBetaProcRunner(self._query_runner, f"{self._namespace}.model", self._server_version)
+
+
 class ModelAlphaEndpoints(CallerBase):
     @property
     def model(self) -> ModelAlphaProcRunner:
diff --git a/graphdatascience/model/model_proc_runner.py b/graphdatascience/model/model_proc_runner.py
@@ -1,41 +1,9 @@
-from typing import Any, Optional
-
-from pandas import DataFrame, Series
-
 from ..error.client_only_endpoint import client_only_endpoint
 from .model import Model
 from .model_resolver import ModelResolver
 
 
 class ModelProcRunner(ModelResolver):
-    def list(self, model: Optional[Model] = None) -> DataFrame:
-        self._namespace += ".list"
-
-        if model:
-            query = f"CALL {self._namespace}($model_name)"
-            params = {"model_name": model.name()}
-        else:
-            query = f"CALL {self._namespace}()"
-            params = {}
-
-        return self._query_runner.run_query(query, params)
-
-    def exists(self, model_name: str) -> "Series[Any]":
-        self._namespace += ".exists"
-
-        query = f"CALL {self._namespace}($model_name)"
-        params = {"model_name": model_name}
-
-        return self._query_runner.run_query(query, params).squeeze()  # type: ignore
-
-    def drop(self, model: Model) -> "Series[Any]":
-        self._namespace += ".drop"
-
-        query = f"CALL {self._namespace}($model_name)"
-        params = {"model_name": model.name()}
-
-        return self._query_runner.run_query(query, params).squeeze()  # type: ignore
-
     @client_only_endpoint("gds.model")
     def get(self, model_name: str) -> Model:
         query = "CALL gds.beta.model.list($model_name)"
diff --git a/graphdatascience/pipeline/pipeline_beta_proc_runner.py b/graphdatascience/pipeline/pipeline_beta_proc_runner.py
@@ -0,0 +1,49 @@
+from typing import Any, Optional
+
+from pandas import DataFrame, Series
+
+from ..caller_base import CallerBase
+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 .training_pipeline import TrainingPipeline
+
+
+class PipelineBetaProcRunner(CallerBase, 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: Optional[TrainingPipeline[PipelineModel]] = None) -> DataFrame:
+        self._namespace += ".list"
+
+        if pipeline:
+            query = f"CALL {self._namespace}($pipeline_name)"
+            params = {"pipeline_name": pipeline.name()}
+        else:
+            query = f"CALL {self._namespace}()"
+            params = {}
+
+        return self._query_runner.run_query(query, params)
+
+    def exists(self, pipeline_name: str) -> "Series[Any]":
+        self._namespace += ".exists"
+
+        query = f"CALL {self._namespace}($pipeline_name)"
+        params = {"pipeline_name": pipeline_name}
+
+        return self._query_runner.run_query(query, params).squeeze()  # type: ignore
+
+    def drop(self, pipeline: TrainingPipeline[PipelineModel]) -> "Series[Any]":
+        self._namespace += ".drop"
+
+        query = f"CALL {self._namespace}($pipeline_name)"
+        params = {"pipeline_name": pipeline.name()}
+
+        return self._query_runner.run_query(query, params).squeeze()  # type: ignore
diff --git a/graphdatascience/pipeline/pipeline_endpoints.py b/graphdatascience/pipeline/pipeline_endpoints.py
diff --git a/graphdatascience/pipeline/pipeline_proc_runner.py b/graphdatascience/pipeline/pipeline_proc_runner.py
diff --git a/graphdatascience/system/system_endpoints.py b/graphdatascience/system/system_endpoints.py
diff --git a/graphdatascience/utils/util_endpoints.py b/graphdatascience/utils/util_endpoints.py
