Read Arrow Server version from debug procedure

Author: s1ck

graphdatascience/query_runner/arrow_query_runner.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import base64
 import json
 import time
@@ -14,6 +16,7 @@
 from ..server_version.server_version import ServerVersion
 from .arrow_graph_constructor import ArrowGraphConstructor
 from .graph_constructor import GraphConstructor
+from .arrow_version import ArrowVersion, UnsupportedArrowVersion
 from .query_runner import QueryRunner
 from graphdatascience.server_version.compatible_with import (
     IncompatibleServerVersionError,
@@ -32,14 +35,16 @@ def create(
         server_version = fallback_query_runner.server_version()
 
         yield_fields = (
-            ["running", "listenAddress"]
+            ["running", "listenAddress", "version"]
             if server_version >= ServerVersion(2, 2, 1)
-            else ["running", "advertisedListenAddress"]
+            else ["running", "advertisedListenAddress", "version"]
         )
-        arrow_info: "Series[Any]" = fallback_query_runner.call_procedure(
+        arrow_info: Series[Any] = fallback_query_runner.call_procedure(
             endpoint="gds.debug.arrow", yields=yield_fields, custom_error=False
         ).squeeze()
         listen_address: str = arrow_info.get("advertisedListenAddress", arrow_info["listenAddress"])  # type: ignore
+        arrow_version = ArrowQueryRunner.read_arrow_version(arrow_info)
+
         if arrow_info["running"]:
             return ArrowQueryRunner(
                 listen_address,
@@ -49,10 +54,20 @@ def create(
                 encrypted,
                 disable_server_verification,
                 tls_root_certs,
+                arrow_version,
             )
         else:
             return fallback_query_runner
 
+    @staticmethod
+    def read_arrow_version(arrow_info: Series[Any]) -> ArrowVersion:
+        arrow_version = arrow_info.get("version", "alpha")
+        try:
+            arrow_version = ArrowVersion[arrow_version.upper()]
+            return arrow_version
+        except KeyError:
+            raise UnsupportedArrowVersion(arrow_version)
+
     def __init__(
         self,
         uri: str,
@@ -62,9 +77,11 @@ def __init__(
         encrypted: bool = False,
         disable_server_verification: bool = False,
         tls_root_certs: Optional[bytes] = None,
+        arrow_version: ArrowVersion = ArrowVersion.ALPHA,
     ):
         self._fallback_query_runner = fallback_query_runner
         self._server_version = server_version
+        self._arrow_version = arrow_version
 
         host, port_string = uri.split(":")
 

graphdatascience/query_runner/arrow_version.py

@@ -0,0 +1,17 @@
+from enum import Enum
+
+
+class ArrowVersion(Enum):
+    ALPHA = ""
+    V1 = "v1/"
+
+    def name(self):
+        return self._name_.lower()
+
+    def prefix(self):
+        return self._value_
+
+
+class UnsupportedArrowVersion(Exception):
+    def __init__(self, server_version):
+        super().__init__(self, f"Unsupported Arrow server version: {server_version}")

graphdatascience/tests/unit/test_arrow_version.py

@@ -0,0 +1,21 @@
+import pytest
+
+from graphdatascience.graph_data_science import ArrowQueryRunner
+from graphdatascience.query_runner.arrow_version import ArrowVersion, UnsupportedArrowVersion
+
+
+@pytest.mark.parametrize("arrow_version", [ArrowVersion.V1, ArrowVersion.ALPHA])
+def test_arrow_version_parsing(arrow_version):
+    arrow_info = dict()
+    arrow_info["version"] = arrow_version.name()
+    actual = ArrowQueryRunner.read_arrow_version(arrow_info)
+    assert actual == arrow_version
+
+
+@pytest.mark.parametrize("arrow_version", ["v2", "unsupported"])
+def test_arrow_version_parsing_fails(arrow_version):
+    arrow_info = dict()
+    arrow_info["version"] = arrow_version
+    with pytest.raises(UnsupportedArrowVersion) as e:
+        ArrowQueryRunner.read_arrow_version(arrow_info)
+    assert arrow_version in str(e.value)