Share logic for endpoint helpers

Author: FlorentinD

graphdatascience/arrow_client/v2/mutation_client.py

@@ -22,7 +22,7 @@ def mutate_relationship_property(
         client: AuthenticatedArrowClient,
         job_id: str,
         mutate_relationship_type: str,
-        mutate_property: str,
+        mutate_property: str | None,
     ) -> MutateResult:
         return MutationClient._mutate(
             client=client,

graphdatascience/procedure_surface/arrow/catalog/node_label_arrow_endpoints.py

@@ -9,7 +9,7 @@
     NodeLabelMutateResult,
     NodeLabelWriteResult,
 )
-from graphdatascience.procedure_surface.arrow.node_property_endpoints import NodePropertyEndpoints
+from graphdatascience.procedure_surface.arrow.node_property_endpoints import NodePropertyEndpointsHelper
 from graphdatascience.procedure_surface.utils.config_converter import ConfigConverter
 
 
@@ -20,11 +20,11 @@ def __init__(
         write_back_client: RemoteWriteBackClient | None = None,
         show_progress: bool = True,
     ):
-        self._node_property_endpoints = NodePropertyEndpoints(
+        self._node_property_endpoints = NodePropertyEndpointsHelper(
             arrow_client, write_back_client, show_progress=show_progress
         )
         self._arrow_client = arrow_client
-        self._node_property_endpoints = NodePropertyEndpoints(arrow_client, write_back_client)
+        self._node_property_endpoints = NodePropertyEndpointsHelper(arrow_client, write_back_client)
         self._show_progress = show_progress
 
     def mutate(
@@ -83,6 +83,11 @@ def write(
         )
 
         result = self._node_property_endpoints.run_job_and_write(
-            "v2/graph.nodeLabel.stream", G, config, write_concurrency, concurrency
+            "v2/graph.nodeLabel.stream",
+            G,
+            config,
+            property_overwrites={},
+            write_concurrency=write_concurrency,
+            concurrency=concurrency,
         )
         return NodeLabelWriteResult(**result)

graphdatascience/procedure_surface/arrow/catalog/node_properties_arrow_endpoints.py

@@ -14,7 +14,7 @@
     NodePropertiesWriteResult,
     NodePropertySpec,
 )
-from graphdatascience.procedure_surface.arrow.node_property_endpoints import NodePropertyEndpoints
+from graphdatascience.procedure_surface.arrow.node_property_endpoints import NodePropertyEndpointsHelper
 from graphdatascience.procedure_surface.utils.config_converter import ConfigConverter
 from graphdatascience.procedure_surface.utils.result_utils import join_db_node_properties
 
@@ -31,7 +31,7 @@ def __init__(
         self._write_back_client: RemoteWriteBackClient | None = (
             RemoteWriteBackClient(arrow_client, query_runner) if query_runner is not None else None
         )
-        self._node_property_endpoints = NodePropertyEndpoints(
+        self._node_property_endpoints = NodePropertyEndpointsHelper(
             arrow_client, self._write_back_client, show_progress=show_progress
         )
         self._show_progress = show_progress

graphdatascience/procedure_surface/arrow/centrality/articlerank_arrow_endpoints.py

@@ -12,7 +12,7 @@
     ArticleRankWriteResult,
 )
 from graphdatascience.procedure_surface.api.estimation_result import EstimationResult
-from graphdatascience.procedure_surface.arrow.node_property_endpoints import NodePropertyEndpoints
+from graphdatascience.procedure_surface.arrow.node_property_endpoints import NodePropertyEndpointsHelper
 
 
 class ArticleRankArrowEndpoints(ArticleRankEndpoints):
@@ -22,7 +22,7 @@ def __init__(
         write_back_client: RemoteWriteBackClient | None = None,
         show_progress: bool = True,
     ):
-        self._node_property_endpoints = NodePropertyEndpoints(
+        self._node_property_endpoints = NodePropertyEndpointsHelper(
             arrow_client, write_back_client, show_progress=show_progress
         )
 
@@ -60,9 +60,7 @@ def mutate(
             tolerance=tolerance,
         )
 
-        result = self._node_property_endpoints.run_job_and_mutate(
-            "v2/centrality.articleRank", G, config, mutate_property
-        )
+        result = self._node_property_endpoints.run_job_and_mutate("v2/centrality.articleRank", config, mutate_property)
 
         return ArticleRankMutateResult(**result)
 
@@ -99,9 +97,7 @@ def stats(
             tolerance=tolerance,
         )
 
-        computation_result = self._node_property_endpoints.run_job_and_get_summary(
-            "v2/centrality.articleRank", G, config
-        )
+        computation_result = self._node_property_endpoints.run_job_and_get_summary("v2/centrality.articleRank", config)
 
         return ArticleRankStatsResult(**computation_result)
 
@@ -176,7 +172,12 @@ def write(
         )
 
         result = self._node_property_endpoints.run_job_and_write(
-            "v2/centrality.articleRank", G, config, write_concurrency, concurrency, write_property
+            "v2/centrality.articleRank",
+            G,
+            config,
+            property_overwrites=write_property,
+            write_concurrency=write_concurrency,
+            concurrency=concurrency,
         )
 
         return ArticleRankWriteResult(**result)

graphdatascience/procedure_surface/arrow/centrality/articulationpoints_arrow_endpoints.py

@@ -12,7 +12,7 @@
     ArticulationPointsWriteResult,
 )
 from graphdatascience.procedure_surface.api.estimation_result import EstimationResult
-from graphdatascience.procedure_surface.arrow.node_property_endpoints import NodePropertyEndpoints
+from graphdatascience.procedure_surface.arrow.node_property_endpoints import NodePropertyEndpointsHelper
 
 
 class ArticulationPointsArrowEndpoints(ArticulationPointsEndpoints):
@@ -24,7 +24,7 @@ def __init__(
         write_back_client: RemoteWriteBackClient | None = None,
         show_progress: bool = True,
     ):
-        self._node_property_endpoints = NodePropertyEndpoints(
+        self._node_property_endpoints = NodePropertyEndpointsHelper(
             arrow_client, write_back_client, show_progress=show_progress
         )
 
@@ -52,7 +52,7 @@ def mutate(
         )
 
         result = self._node_property_endpoints.run_job_and_mutate(
-            "v2/centrality.articulationPoints", G, config, mutate_property
+            "v2/centrality.articulationPoints", config, mutate_property
         )
 
         return ArticulationPointsMutateResult(**result)
@@ -80,7 +80,7 @@ def stats(
         )
 
         computation_result = self._node_property_endpoints.run_job_and_get_summary(
-            "v2/centrality.articulationPoints", G, config
+            "v2/centrality.articulationPoints", config
         )
 
         return ArticulationPointsStatsResult(**computation_result)
@@ -127,7 +127,12 @@ def write(
         )
 
         result = self._node_property_endpoints.run_job_and_write(
-            "v2/centrality.articulationPoints", G, config, write_concurrency, concurrency, write_property
+            "v2/centrality.articulationPoints",
+            G,
+            config,
+            property_overwrites=write_property,
+            write_concurrency=write_concurrency,
+            concurrency=concurrency,
         )
 
         return ArticulationPointsWriteResult(**result)

graphdatascience/procedure_surface/arrow/centrality/betweenness_arrow_endpoints.py

@@ -12,7 +12,7 @@
     BetweennessWriteResult,
 )
 from graphdatascience.procedure_surface.api.estimation_result import EstimationResult
-from graphdatascience.procedure_surface.arrow.node_property_endpoints import NodePropertyEndpoints
+from graphdatascience.procedure_surface.arrow.node_property_endpoints import NodePropertyEndpointsHelper
 
 
 class BetweennessArrowEndpoints(BetweennessEndpoints):
@@ -24,7 +24,7 @@ def __init__(
         write_back_client: RemoteWriteBackClient | None = None,
         show_progress: bool = True,
     ):
-        self._node_property_endpoints = NodePropertyEndpoints(
+        self._node_property_endpoints = NodePropertyEndpointsHelper(
             arrow_client, write_back_client, show_progress=show_progress
         )
 
@@ -57,9 +57,7 @@ def mutate(
             username=username,
         )
 
-        result = self._node_property_endpoints.run_job_and_mutate(
-            "v2/centrality.betweenness", G, config, mutate_property
-        )
+        result = self._node_property_endpoints.run_job_and_mutate("v2/centrality.betweenness", config, mutate_property)
 
         return BetweennessMutateResult(**result)
 
@@ -91,9 +89,7 @@ def stats(
             username=username,
         )
 
-        computation_result = self._node_property_endpoints.run_job_and_get_summary(
-            "v2/centrality.betweenness", G, config
-        )
+        computation_result = self._node_property_endpoints.run_job_and_get_summary("v2/centrality.betweenness", config)
 
         return BetweennessStatsResult(**computation_result)
 
@@ -158,7 +154,12 @@ def write(
         )
 
         result = self._node_property_endpoints.run_job_and_write(
-            "v2/centrality.betweenness", G, config, write_concurrency, concurrency, write_property
+            "v2/centrality.betweenness",
+            G,
+            config,
+            property_overwrites=write_property,
+            write_concurrency=write_concurrency,
+            concurrency=concurrency,
         )
 
         return BetweennessWriteResult(**result)

graphdatascience/procedure_surface/arrow/centrality/celf_arrow_endpoints.py

@@ -12,7 +12,7 @@
     CelfWriteResult,
 )
 from graphdatascience.procedure_surface.api.estimation_result import EstimationResult
-from graphdatascience.procedure_surface.arrow.node_property_endpoints import NodePropertyEndpoints
+from graphdatascience.procedure_surface.arrow.node_property_endpoints import NodePropertyEndpointsHelper
 
 
 class CelfArrowEndpoints(CelfEndpoints):
@@ -22,7 +22,7 @@ def __init__(
         write_back_client: RemoteWriteBackClient | None = None,
         show_progress: bool = True,
     ):
-        self._node_property_endpoints = NodePropertyEndpoints(
+        self._node_property_endpoints = NodePropertyEndpointsHelper(
             arrow_client, write_back_client, show_progress=show_progress
         )
 
@@ -56,7 +56,7 @@ def mutate(
             random_seed=random_seed,
         )
 
-        result = self._node_property_endpoints.run_job_and_mutate("v2/centrality.celf", G, config, mutate_property)
+        result = self._node_property_endpoints.run_job_and_mutate("v2/centrality.celf", config, mutate_property)
 
         return CelfMutateResult(**result)
 
@@ -89,7 +89,7 @@ def stats(
             random_seed=random_seed,
         )
 
-        computation_result = self._node_property_endpoints.run_job_and_get_summary("v2/centrality.celf", G, config)
+        computation_result = self._node_property_endpoints.run_job_and_get_summary("v2/centrality.celf", config)
 
         return CelfStatsResult(**computation_result)
 

graphdatascience/procedure_surface/arrow/centrality/closeness_arrow_endpoints.py

@@ -12,7 +12,7 @@
     ClosenessWriteResult,
 )
 from graphdatascience.procedure_surface.api.estimation_result import EstimationResult
-from graphdatascience.procedure_surface.arrow.node_property_endpoints import NodePropertyEndpoints
+from graphdatascience.procedure_surface.arrow.node_property_endpoints import NodePropertyEndpointsHelper
 
 
 class ClosenessArrowEndpoints(ClosenessEndpoints):
@@ -24,7 +24,7 @@ def __init__(
         write_back_client: RemoteWriteBackClient | None = None,
         show_progress: bool = True,
     ):
-        self._node_property_endpoints = NodePropertyEndpoints(
+        self._node_property_endpoints = NodePropertyEndpointsHelper(
             arrow_client, write_back_client, show_progress=show_progress
         )
 
@@ -53,7 +53,7 @@ def mutate(
             job_id=job_id,
         )
 
-        result = self._node_property_endpoints.run_job_and_mutate("v2/centrality.closeness", G, config, mutate_property)
+        result = self._node_property_endpoints.run_job_and_mutate("v2/centrality.closeness", config, mutate_property)
 
         return ClosenessMutateResult(**result)
 
@@ -81,7 +81,7 @@ def stats(
             job_id=job_id,
         )
 
-        computation_result = self._node_property_endpoints.run_job_and_get_summary("v2/centrality.closeness", G, config)
+        computation_result = self._node_property_endpoints.run_job_and_get_summary("v2/centrality.closeness", config)
 
         return ClosenessStatsResult(**computation_result)
 
@@ -138,7 +138,12 @@ def write(
         )
 
         result = self._node_property_endpoints.run_job_and_write(
-            "v2/centrality.closeness", G, config, write_concurrency, concurrency, write_property
+            "v2/centrality.closeness",
+            G,
+            config,
+            write_concurrency=write_concurrency,
+            concurrency=concurrency,
+            property_overwrites=write_property,
         )
 
         return ClosenessWriteResult(**result)

graphdatascience/procedure_surface/arrow/centrality/closeness_harmonic_arrow_endpoints.py

@@ -12,7 +12,7 @@
     ClosenessHarmonicWriteResult,
 )
 from graphdatascience.procedure_surface.api.estimation_result import EstimationResult
-from graphdatascience.procedure_surface.arrow.node_property_endpoints import NodePropertyEndpoints
+from graphdatascience.procedure_surface.arrow.node_property_endpoints import NodePropertyEndpointsHelper
 
 
 class ClosenessHarmonicArrowEndpoints(ClosenessHarmonicEndpoints):
@@ -22,7 +22,7 @@ def __init__(
         write_back_client: RemoteWriteBackClient | None = None,
         show_progress: bool = True,
     ):
-        self._node_property_endpoints = NodePropertyEndpoints(
+        self._node_property_endpoints = NodePropertyEndpointsHelper(
             arrow_client, write_back_client, show_progress=show_progress
         )
 
@@ -49,7 +49,7 @@ def mutate(
             username=username,
         )
 
-        result = self._node_property_endpoints.run_job_and_mutate("v2/centrality.harmonic", G, config, mutate_property)
+        result = self._node_property_endpoints.run_job_and_mutate("v2/centrality.harmonic", config, mutate_property)
 
         return ClosenessHarmonicMutateResult(**result)
 
@@ -75,7 +75,7 @@ def stats(
             username=username,
         )
 
-        computation_result = self._node_property_endpoints.run_job_and_get_summary("v2/centrality.harmonic", G, config)
+        computation_result = self._node_property_endpoints.run_job_and_get_summary("v2/centrality.harmonic", config)
 
         return ClosenessHarmonicStatsResult(**computation_result)
 
@@ -128,7 +128,12 @@ def write(
         )
 
         result = self._node_property_endpoints.run_job_and_write(
-            "v2/centrality.harmonic", G, config, write_concurrency, concurrency, write_property
+            "v2/centrality.harmonic",
+            G,
+            config,
+            property_overwrites=write_property,
+            write_concurrency=write_concurrency,
+            concurrency=concurrency,
         )
 
         return ClosenessHarmonicWriteResult(**result)