Generalize config extraction for arrow endpoints

Author: DarthMax

graphdatascience/procedure_surface/arrow/arrow_config_converter.py

@@ -0,0 +1,38 @@
+from typing import Optional, Any, Dict
+
+from graphdatascience import Graph
+
+
+class ArrowConfigConverter:
+
+    @staticmethod
+    def build_configuration(G: Graph, **kwargs: Optional[Any]) -> dict[str, Any]:
+        config: dict[str, Any] = {
+            "graphName": G.name(),
+        }
+
+        # Process kwargs
+        processed_kwargs = ArrowConfigConverter._process_dict_values(kwargs)
+        config.update(processed_kwargs)
+
+        return config
+
+    @staticmethod
+    def _convert_to_camel_case(name: str) -> str:
+        """Convert a snake_case string to camelCase."""
+        parts = name.split('_')
+        return ''.join([word.capitalize() if i > 0 else word.lower() for i, word in enumerate(parts)])
+
+    @staticmethod
+    def _process_dict_values(input_dict: Dict[str, Any]) -> Dict[str, Any]:
+        """Process dictionary values, converting keys to camelCase and handling nested dictionaries."""
+        result = {}
+        for key, value in input_dict.items():
+            if value is not None:
+                camel_key = ArrowConfigConverter._convert_to_camel_case(key)
+                # Recursively process nested dictionaries
+                if isinstance(value, dict):
+                    result[camel_key] = ArrowConfigConverter._process_dict_values(value)
+                else:
+                    result[camel_key] = value
+        return result

graphdatascience/procedure_surface/arrow/wcc_arrow_endpoints.py

@@ -2,6 +2,7 @@
 
 from pandas import DataFrame
 
+from .arrow_config_converter import ArrowConfigConverter
 from ...arrow_client.authenticated_arrow_client import AuthenticatedArrowClient
 from ...arrow_client.v2.job_client import JobClient
 from ...arrow_client.v2.mutation_client import MutationClient
@@ -33,19 +34,18 @@ def mutate(
         consecutive_ids: Optional[bool] = None,
         relationship_weight_property: Optional[str] = None,
     ) -> WccMutateResult:
-        config = self._build_configuration(
+        config = ArrowConfigConverter.build_configuration(
             G,
-            concurrency,
-            consecutive_ids,
-            job_id,
-            log_progress,
-            None,
-            node_labels,
-            relationship_types,
-            relationship_weight_property,
-            seed_property,
-            sudo,
-            threshold,
+            concurrency = concurrency,
+            consecutive_ids = consecutive_ids,
+            job_id = job_id,
+            log_progress = log_progress,
+            node_labels = node_labels,
+            relationship_types = relationship_types,
+            relationship_weight_property = relationship_weight_property,
+            seed_property = seed_property,
+            sudo = sudo,
+            threshold = threshold,
         )
 
         job_id = JobClient.run_job_and_wait(self._arrow_client, WCC_ENDPOINT, config)
@@ -79,19 +79,18 @@ def stats(
         consecutive_ids: Optional[bool] = None,
         relationship_weight_property: Optional[str] = None,
     ) -> WccStatsResult:
-        config = self._build_configuration(
+        config = ArrowConfigConverter.build_configuration(
             G,
-            concurrency,
-            consecutive_ids,
-            job_id,
-            log_progress,
-            None,
-            node_labels,
-            relationship_types,
-            relationship_weight_property,
-            seed_property,
-            sudo,
-            threshold,
+            concurrency = concurrency,
+            consecutive_ids = consecutive_ids,
+            job_id = job_id,
+            log_progress = log_progress,
+            node_labels = node_labels,
+            relationship_types = relationship_types,
+            relationship_weight_property = relationship_weight_property,
+            seed_property = seed_property,
+            sudo = sudo,
+            threshold = threshold,
         )
 
         job_id = JobClient.run_job_and_wait(self._arrow_client, WCC_ENDPOINT, config)
@@ -122,19 +121,19 @@ def stream(
         consecutive_ids: Optional[bool] = None,
         relationship_weight_property: Optional[str] = None,
     ) -> DataFrame:
-        config = self._build_configuration(
+        config = ArrowConfigConverter.build_configuration(
             G,
-            concurrency,
-            consecutive_ids,
-            job_id,
-            log_progress,
-            min_component_size,
-            node_labels,
-            relationship_types,
-            relationship_weight_property,
-            seed_property,
-            sudo,
-            threshold,
+            concurrency = concurrency,
+            consecutive_ids = consecutive_ids,
+            job_id = job_id,
+            log_progress = log_progress,
+            min_component_size = min_component_size,
+            node_labels = node_labels,
+            relationship_types = relationship_types,
+            relationship_weight_property = relationship_weight_property,
+            seed_property = seed_property,
+            sudo = sudo,
+            threshold = threshold,
         )
 
         job_id = JobClient.run_job_and_wait(self._arrow_client, WCC_ENDPOINT, config)
@@ -158,19 +157,20 @@ def write(
         relationship_weight_property: Optional[str] = None,
         write_concurrency: Optional[int] = None,
     ) -> WccWriteResult:
-        config = self._build_configuration(
+        
+        config = ArrowConfigConverter.build_configuration(
             G,
-            concurrency,
-            consecutive_ids,
-            job_id,
-            log_progress,
-            min_component_size,
-            node_labels,
-            relationship_types,
-            relationship_weight_property,
-            seed_property,
-            sudo,
-            threshold,
+            concurrency = concurrency,
+            consecutive_ids = consecutive_ids,
+            job_id = job_id,
+            log_progress = log_progress,
+            min_component_size = min_component_size,
+            node_labels = node_labels,
+            relationship_types = relationship_types,
+            relationship_weight_property = relationship_weight_property,
+            seed_property = seed_property,
+            sudo = sudo,
+            threshold = threshold,
         )
 
         job_id = JobClient.run_job_and_wait(self._arrow_client, WCC_ENDPOINT, config)
@@ -192,48 +192,4 @@ def write(
             computation_result["postProcessingMillis"],
             computation_result["nodePropertiesWritten"],
             computation_result["configuration"],
-        )
-
-    @staticmethod
-    def _build_configuration(
-        G: Graph,
-        concurrency: Optional[int],
-        consecutive_ids: Optional[bool],
-        job_id: Optional[str],
-        log_progress: Optional[bool],
-        min_component_size: Optional[int],
-        node_labels: Optional[List[str]],
-        relationship_types: Optional[List[str]],
-        relationship_weight_property: Optional[str],
-        seed_property: Optional[str],
-        sudo: Optional[bool],
-        threshold: Optional[float],
-    ) -> dict[str, Any]:
-        config: dict[str, Any] = {
-            "graphName": G.name(),
-        }
-
-        if min_component_size is not None:
-            config["minComponentSize"] = min_component_size
-        if threshold is not None:
-            config["threshold"] = threshold
-        if relationship_types is not None:
-            config["relationshipTypes"] = relationship_types
-        if node_labels is not None:
-            config["nodeLabels"] = node_labels
-        if sudo is not None:
-            config["sudo"] = sudo
-        if log_progress is not None:
-            config["logProgress"] = log_progress
-        if concurrency is not None:
-            config["concurrency"] = concurrency
-        if job_id is not None:
-            config["jobId"] = job_id
-        if seed_property is not None:
-            config["seedProperty"] = seed_property
-        if consecutive_ids is not None:
-            config["consecutiveIds"] = consecutive_ids
-        if relationship_weight_property is not None:
-            config["relationshipWeightProperty"] = relationship_weight_property
-
-        return config
+        )