Implement defaults and limits

DarthMax · DarthMax · commit 15387317757c · 2025-10-16T23:58:54.000+02:00
diff --git a/graphdatascience/procedure_surface/api/config_endpoints.py b/graphdatascience/procedure_surface/api/config_endpoints.py
@@ -0,0 +1,104 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Any
+
+
+class ConfigEndpoints(ABC):
+    @property
+    @abstractmethod
+    def defaults(self) -> DefaultsEndpoints:
+        pass
+
+    @property
+    @abstractmethod
+    def limits(self) -> LimitsEndpoints:
+        pass
+
+
+class DefaultsEndpoints(ABC):
+    @abstractmethod
+    def set(
+        self,
+        key: str,
+        value: Any,
+        username: str | None = None,
+    ) -> None:
+        """
+        Configure a new default configuration value.
+
+        Parameters:
+            key : str
+                The configuration key for which the default value is being set.
+            value : Any
+                The value to set as the default for the given key.
+            username : str | None, default=None
+                If set, the configuration will be set for the given user.
+
+        Returns: None
+        """
+        pass
+
+    @abstractmethod
+    def list(
+        self,
+        username: str | None = None,
+        key: str | None = None,
+    ) -> dict[str, Any]:
+        """
+        List configured default configuration values.
+
+        Parameters:
+            key : str | None (default=None)
+                List only the default value for the given key.
+            username : str | None, default=None
+                List only default values for the given user.
+
+        Returns: dict[str, Any]
+            A dictionary containing the default configuration values.
+        """
+        pass
+
+
+class LimitsEndpoints(ABC):
+    @abstractmethod
+    def set(
+        self,
+        key: str,
+        value: Any,
+        username: str | None = None,
+    ) -> None:
+        """
+        Configure a new limit for a configuration value.
+
+        Parameters:
+            key : str
+                The configuration key for which the limit is being set.
+            value : Any
+                The value to set as the limit for the given key.
+            username : str | None, default=None
+                If set, the limit will be set for the given user.
+
+        Returns: None
+        """
+        pass
+
+    @abstractmethod
+    def list(
+        self,
+        username: str | None = None,
+        key: str | None = None,
+    ) -> dict[str, Any]:
+        """
+        List configured configuration limits.
+
+        Parameters:
+            key : str | None (default=None)
+                List only the limits for the given key.
+            username : str | None, default=None
+                List only liomits for the given user.
+
+        Returns: dict[str, Any]
+            A dictionary containing the configuration limits.
+        """
+        pass
diff --git a/graphdatascience/procedure_surface/arrow/config_arrow_endpoints.py b/graphdatascience/procedure_surface/arrow/config_arrow_endpoints.py
@@ -0,0 +1,88 @@
+from __future__ import annotations
+
+from typing import Any
+
+from graphdatascience.arrow_client.authenticated_flight_client import AuthenticatedArrowClient
+from graphdatascience.arrow_client.v2.data_mapper_utils import deserialize
+from graphdatascience.procedure_surface.api.config_endpoints import (
+    ConfigEndpoints,
+    DefaultsEndpoints,
+    LimitsEndpoints,
+)
+from graphdatascience.procedure_surface.utils.config_converter import ConfigConverter
+
+
+class ConfigArrowEndpoints(ConfigEndpoints):
+    def __init__(self, arrow_client: AuthenticatedArrowClient, show_progress: bool = False):
+        self._arrow_client = arrow_client
+        self._show_progress = show_progress
+
+    @property
+    def defaults(self) -> DefaultsEndpoints:
+        return DefaultsArrowEndpoints(self._arrow_client, self._show_progress)
+
+    @property
+    def limits(self) -> LimitsEndpoints:
+        return LimitsArrowEndpoints(self._arrow_client, self._show_progress)
+
+
+class DefaultsArrowEndpoints(DefaultsEndpoints):
+    def __init__(self, arrow_client: AuthenticatedArrowClient, show_progress: bool = False):
+        self._arrow_client = arrow_client
+        self._show_progress = show_progress
+
+    def set(
+        self,
+        key: str,
+        value: Any,
+        username: str | None = None,
+    ) -> None:
+        deserialize(self._arrow_client.do_action_with_retry("v2/defaults.set", {key: value}))
+
+    def list(
+        self,
+        username: str | None = None,
+        key: str | None = None,
+    ) -> dict[str, Any]:
+        config = ConfigConverter.convert_to_gds_config(
+            key=key,
+        )
+
+        rows = deserialize(self._arrow_client.do_action_with_retry("v2/defaults.list", config))
+        result = {}
+
+        for row in rows:
+            result[row["key"]] = row["value"]
+
+        return result
+
+
+class LimitsArrowEndpoints(LimitsEndpoints):
+    def __init__(self, arrow_client: AuthenticatedArrowClient, show_progress: bool = False):
+        self._arrow_client = arrow_client
+        self._show_progress = show_progress
+
+    def set(
+        self,
+        key: str,
+        value: Any,
+        username: str | None = None,
+    ) -> None:
+        deserialize(self._arrow_client.do_action_with_retry("v2/limits.set", {key: value}))
+
+    def list(
+        self,
+        username: str | None = None,
+        key: str | None = None,
+    ) -> dict[str, Any]:
+        config = ConfigConverter.convert_to_gds_config(
+            key=key,
+        )
+
+        rows = deserialize(self._arrow_client.do_action_with_retry("v2/limits.list", config))
+        result = {}
+
+        for row in rows:
+            result[row["key"]] = row["value"]
+
+        return result
diff --git a/graphdatascience/procedure_surface/cypher/config_cypher_endpoints.py b/graphdatascience/procedure_surface/cypher/config_cypher_endpoints.py
@@ -0,0 +1,105 @@
+from __future__ import annotations
+
+from typing import Any
+
+from graphdatascience.call_parameters import CallParameters
+from graphdatascience.procedure_surface.api.config_endpoints import (
+    ConfigEndpoints,
+    DefaultsEndpoints,
+    LimitsEndpoints,
+)
+from graphdatascience.procedure_surface.utils.config_converter import ConfigConverter
+from graphdatascience.query_runner.query_runner import QueryRunner
+
+
+class ConfigCypherEndpoints(ConfigEndpoints):
+    def __init__(self, query_runner: QueryRunner):
+        self._query_runner = query_runner
+
+    @property
+    def defaults(self) -> DefaultsEndpoints:
+        return DefaultsCypherEndpoints(self._query_runner)
+
+    @property
+    def limits(self) -> LimitsEndpoints:
+        return LimitsCypherEndpoints(self._query_runner)
+
+
+class DefaultsCypherEndpoints(DefaultsEndpoints):
+    def __init__(self, query_runner: QueryRunner):
+        self._query_runner = query_runner
+
+    def set(
+        self,
+        key: str,
+        value: Any,
+        username: str | None = None,
+    ) -> None:
+        params = {
+            "key": key,
+            "value": value,
+        }
+
+        if username:
+            params["username"] = username
+
+        params = CallParameters(**params)
+
+        self._query_runner.call_procedure(endpoint="gds.config.defaults.set", params=params)
+
+    def list(
+        self,
+        username: str | None = None,
+        key: str | None = None,
+    ) -> dict[str, Any]:
+        config = ConfigConverter.convert_to_gds_config(
+            key=key,
+            username=username,
+        )
+
+        params = CallParameters(
+            config=config,
+        )
+
+        result = self._query_runner.call_procedure(endpoint="gds.config.defaults.list", params=params)
+        return {row["key"]: row["value"] for _, row in result.iterrows()}
+
+
+class LimitsCypherEndpoints(LimitsEndpoints):
+    def __init__(self, query_runner: QueryRunner):
+        self._query_runner = query_runner
+
+    def set(
+        self,
+        key: str,
+        value: Any,
+        username: str | None = None,
+    ) -> None:
+        params = {
+            "key": key,
+            "value": value,
+        }
+
+        if username:
+            params["username"] = username
+
+        params = CallParameters(**params)
+
+        self._query_runner.call_procedure(endpoint="gds.config.limits.set", params=params)
+
+    def list(
+        self,
+        username: str | None = None,
+        key: str | None = None,
+    ) -> dict[str, Any]:
+        config = ConfigConverter.convert_to_gds_config(
+            key=key,
+            username=username,
+        )
+
+        params = CallParameters(
+            config=config,
+        )
+
+        result = self._query_runner.call_procedure(endpoint="gds.config.limits.list", params=params)
+        return {row["key"]: row["value"] for _, row in result.iterrows()}
diff --git a/graphdatascience/tests/integrationV2/procedure_surface/arrow/test_config_arrow_endpoints.py b/graphdatascience/tests/integrationV2/procedure_surface/arrow/test_config_arrow_endpoints.py
@@ -0,0 +1,58 @@
+from typing import Generator
+
+import pytest
+
+from graphdatascience.arrow_client.authenticated_flight_client import AuthenticatedArrowClient
+from graphdatascience.procedure_surface.arrow.config_arrow_endpoints import ConfigArrowEndpoints
+
+
+@pytest.fixture
+def config_endpoints(arrow_client: AuthenticatedArrowClient) -> Generator[ConfigArrowEndpoints, None, None]:
+    yield ConfigArrowEndpoints(arrow_client)
+
+
+def test_defaults_set_and_list(config_endpoints: ConfigArrowEndpoints) -> None:
+    config_endpoints.defaults.set("test.key", 6)
+
+    defaults = config_endpoints.defaults.list()
+
+    assert ("test.key", 6) in defaults.items()
+
+
+def test_defaults_list_by_key(config_endpoints: ConfigArrowEndpoints) -> None:
+    config_endpoints.defaults.set("test.specific.key", "specific_value")
+
+    specific_defaults = config_endpoints.defaults.list(key="test.specific.key")
+
+    assert specific_defaults == {"test.specific.key": "specific_value"}
+
+
+def test_limits_set_and_list(config_endpoints: ConfigArrowEndpoints) -> None:
+    config_endpoints.limits.set("test.key", 6)
+
+    limits = config_endpoints.limits.list()
+
+    assert ("test.key", 6) in limits.items()
+
+
+def test_limits_list_by_key(config_endpoints: ConfigArrowEndpoints) -> None:
+    config_endpoints.limits.set("test.specific.key", 42)
+
+    specific_limits = config_endpoints.limits.list(key="test.specific.key")
+
+    assert specific_limits == {"test.specific.key": 42}
+
+
+def test_config_endpoints_properties(config_endpoints: ConfigArrowEndpoints) -> None:
+    """Test that the config endpoints have the required properties."""
+    assert hasattr(config_endpoints, "defaults")
+    assert hasattr(config_endpoints, "limits")
+
+    # Verify the properties return the correct endpoint types
+    defaults = config_endpoints.defaults
+    limits = config_endpoints.limits
+
+    assert hasattr(defaults, "set")
+    assert hasattr(defaults, "list")
+    assert hasattr(limits, "set")
+    assert hasattr(limits, "list")
diff --git a/graphdatascience/tests/integrationV2/procedure_surface/cypher/test_config_cypher_endpoints.py b/graphdatascience/tests/integrationV2/procedure_surface/cypher/test_config_cypher_endpoints.py
