Add support for reading version lists

Author: s1ck

graphdatascience/query_runner/arrow_endpoint_version.py

@@ -1,17 +1,35 @@
+from __future__ import annotations
+
 from enum import Enum
 
 
 class ArrowEndpointVersion(Enum):
     ALPHA = ""
     V1 = "v1/"
 
-    def name(self):
+    def name(self) -> str:
         return self._name_.lower()
 
-    def prefix(self):
+    def prefix(self) -> str:
         return self._value_
 
+    @staticmethod
+    def from_arrow_info(arrow_info: Series[Any]) -> ArrowEndpointVersion:
+        supported_arrow_versions = arrow_info.get("versions", [])
+        # Fallback for pre 2.6.0 servers that do not support versions
+        if len(supported_arrow_versions) == 0:
+            return ArrowEndpointVersion.ALPHA
+
+        # If the server supports versioned endpoints, we try v1 first
+        if ArrowEndpointVersion.V1.name() in supported_arrow_versions:
+            return ArrowEndpointVersion.V1
+
+        if ArrowEndpointVersion.ALPHA.name() in supported_arrow_versions:
+            return ArrowEndpointVersion.ALPHA
+
+        raise UnsupportedArrowEndpointVersion(supported_arrow_versions)
+
 
 class UnsupportedArrowEndpointVersion(Exception):
     def __init__(self, server_version):
-        super().__init__(self, f"Unsupported Arrow endpoint version: {server_version}")
+        super().__init__(self, f"Unsupported Arrow endpoint versions: {server_version}")

graphdatascience/query_runner/arrow_query_runner.py

@@ -16,7 +16,6 @@
 from ..server_version.server_version import ServerVersion
 from .arrow_endpoint_version import ArrowEndpointVersion
 from .arrow_graph_constructor import ArrowGraphConstructor
-from .arrow_endpoint_version import ArrowEndpointVersion, UnsupportedArrowEndpointVersion
 from .graph_constructor import GraphConstructor
 from .query_runner import QueryRunner
 from graphdatascience.server_version.compatible_with import (
@@ -36,7 +35,7 @@ def create(
         arrow_info = ArrowQueryRunner._get_arrow_debug_info(fallback_query_runner)
         server_version = fallback_query_runner.server_version()
         listen_address: str = arrow_info.get("advertisedListenAddress", arrow_info["listenAddress"])  # type: ignore
-        arrow_endpoint_version = ArrowQueryRunner._read_arrow_version(arrow_info)
+        arrow_endpoint_version = ArrowEndpointVersion.from_arrow_info(arrow_info)
 
         if arrow_info["running"]:
             return ArrowQueryRunner(
@@ -56,15 +55,6 @@ def create(
     def _get_arrow_debug_info(query_runner) -> Series[Any]:
         return query_runner.call_procedure(endpoint="gds.debug.arrow", custom_error=False).squeeze()
 
-    @staticmethod
-    def _read_arrow_version(arrow_info: Series[Any]) -> ArrowEndpointVersion:
-        arrow_version = arrow_info.get("version", "alpha")
-        try:
-            arrow_version = ArrowEndpointVersion[arrow_version.upper()]
-            return arrow_version
-        except KeyError:
-            raise UnsupportedArrowEndpointVersion(arrow_version)
-
     def __init__(
         self,
         uri: str,
@@ -323,7 +313,12 @@ def create_graph_constructor(
             )
 
         return ArrowGraphConstructor(
-            database, graph_name, self._flight_client, concurrency, self._arrow_endpoint_version, undirected_relationship_types
+            database,
+            graph_name,
+            self._flight_client,
+            concurrency,
+            self._arrow_endpoint_version,
+            undirected_relationship_types,
         )
 
     def _sanitize_arrow_table(self, arrow_table: Table) -> Table:

graphdatascience/tests/unit/test_arrow_endpoint_version.py

@@ -7,23 +7,35 @@
 )
 
 
-@pytest.mark.parametrize("arrow_version", [ArrowEndpointVersion.V1, ArrowEndpointVersion.ALPHA])
-def test_arrow_version_parsing(arrow_version):
+@pytest.mark.parametrize(
+    "arrow_versions",
+    [
+        (ArrowEndpointVersion.ALPHA, []),
+        (ArrowEndpointVersion.ALPHA, ["alpha"]),
+        (ArrowEndpointVersion.V1, ["v1"]),
+        (ArrowEndpointVersion.V1, ["alpha", "v1"]),
+        (ArrowEndpointVersion.V1, ["v1", "v2"]),
+        (ArrowEndpointVersion.ALPHA, ["alpha"]),
+        (ArrowEndpointVersion.ALPHA, ["alpha", "v2"]),
+    ],
+)
+def test_from_arrow_info_multiple_versions(arrow_versions):
     arrow_info = dict()
-    arrow_info["version"] = arrow_version.name()
-    actual = ArrowQueryRunner._read_arrow_version(arrow_info)
-    assert actual == arrow_version
+    arrow_info["versions"] = arrow_versions[1]
+    actual = ArrowEndpointVersion.from_arrow_info(arrow_info)
+    assert actual == arrow_versions[0]
 
 
 @pytest.mark.parametrize("arrow_version", ["v2", "unsupported"])
-def test_arrow_version_parsing_fails(arrow_version):
+def test_from_arrow_info_fails(arrow_version):
     arrow_info = dict()
-    arrow_info["version"] = arrow_version
+    arrow_info["versions"] = [arrow_version]
     with pytest.raises(UnsupportedArrowEndpointVersion) as e:
-        ArrowQueryRunner._read_arrow_version(arrow_info)
+        ArrowEndpointVersion.from_arrow_info(arrow_info)
+    assert "Unsupported" in str(e.value)
     assert arrow_version in str(e.value)
 
 
-def test_arrow_version_prefix():
+def test_prefix():
     assert ArrowEndpointVersion.ALPHA.prefix() == ""
     assert ArrowEndpointVersion.V1.prefix() == "v1/"