Rename ArrowVersion -> ArrowEndpointVersion

Author: s1ck

…atascience/query_runner/arrow_version.py …e/query_runner/arrow_endpoint_version.pygraphdatascience/query_runner/arrow_version.py renamed to graphdatascience/query_runner/arrow_endpoint_version.py graphdatascience/query_runner/arrow_version.py renamed to graphdatascience/query_runner/arrow_endpoint_version.py

@@ -1,7 +1,7 @@
 from enum import Enum
 
 
-class ArrowVersion(Enum):
+class ArrowEndpointVersion(Enum):
     ALPHA = ""
     V1 = "v1/"
 
@@ -12,6 +12,6 @@ def prefix(self):
         return self._value_
 
 
-class UnsupportedArrowVersion(Exception):
+class UnsupportedArrowEndpointVersion(Exception):
     def __init__(self, server_version):
-        super().__init__(self, f"Unsupported Arrow server version: {server_version}")
+        super().__init__(self, f"Unsupported Arrow endpoint version: {server_version}")

graphdatascience/query_runner/arrow_graph_constructor.py

@@ -11,7 +11,7 @@
 from pyarrow import Table
 from tqdm.auto import tqdm
 
-from .arrow_version import ArrowVersion
+from .arrow_endpoint_version import ArrowEndpointVersion
 from .graph_constructor import GraphConstructor
 
 
@@ -22,15 +22,15 @@ def __init__(
         graph_name: str,
         flight_client: flight.FlightClient,
         concurrency: int,
-        arrow_version: ArrowVersion,
+        arrow_endpoint_version: ArrowEndpointVersion,
         undirected_relationship_types: Optional[List[str]],
         chunk_size: int = 10_000,
     ):
         self._database = database
         self._concurrency = concurrency
         self._graph_name = graph_name
         self._client = flight_client
-        self._arrow_version = arrow_version
+        self._arrow_endpoint_version = arrow_endpoint_version
         self._undirected_relationship_types = (
             [] if undirected_relationship_types is None else undirected_relationship_types
         )
@@ -124,15 +124,15 @@ def _send_dfs(self, dfs: List[DataFrame], entity_type: str) -> None:
                 raise future.exception()  # type: ignore
 
     def _versioned_action_type(self, action_type: str) -> str:
-        return self._arrow_version.prefix() + action_type
+        return self._arrow_endpoint_version.prefix() + action_type
 
     def _versioned_flight_desriptor(self, flight_descriptor: dict) -> dict:
         return (
             flight_descriptor
-            if self._arrow_version == ArrowVersion.ALPHA
+            if self._arrow_endpoint_version == ArrowEndpointVersion.ALPHA
             else {
                 "name": "PUT_MESSAGE",
-                "version": ArrowVersion.V1.name(),
+                "version": ArrowEndpointVersion.V1.name(),
                 "body": flight_descriptor,
             }
         )

graphdatascience/query_runner/arrow_query_runner.py

@@ -14,8 +14,9 @@
 
 from ..call_parameters import CallParameters
 from ..server_version.server_version import ServerVersion
+from .arrow_endpoint_version import ArrowEndpointVersion
 from .arrow_graph_constructor import ArrowGraphConstructor
-from .arrow_version import ArrowVersion, UnsupportedArrowVersion
+from .arrow_endpoint_version import ArrowEndpointVersion, UnsupportedArrowEndpointVersion
 from .graph_constructor import GraphConstructor
 from .query_runner import QueryRunner
 from graphdatascience.server_version.compatible_with import (
@@ -35,7 +36,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_version = ArrowQueryRunner._read_arrow_version(arrow_info)
+        arrow_endpoint_version = ArrowQueryRunner._read_arrow_version(arrow_info)
 
         if arrow_info["running"]:
             return ArrowQueryRunner(
@@ -46,7 +47,7 @@ def create(
                 encrypted,
                 disable_server_verification,
                 tls_root_certs,
-                arrow_version,
+                arrow_endpoint_version,
             )
         else:
             return fallback_query_runner
@@ -56,13 +57,13 @@ 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]) -> ArrowVersion:
+    def _read_arrow_version(arrow_info: Series[Any]) -> ArrowEndpointVersion:
         arrow_version = arrow_info.get("version", "alpha")
         try:
-            arrow_version = ArrowVersion[arrow_version.upper()]
+            arrow_version = ArrowEndpointVersion[arrow_version.upper()]
             return arrow_version
         except KeyError:
-            raise UnsupportedArrowVersion(arrow_version)
+            raise UnsupportedArrowEndpointVersion(arrow_version)
 
     def __init__(
         self,
@@ -73,11 +74,11 @@ def __init__(
         encrypted: bool = False,
         disable_server_verification: bool = False,
         tls_root_certs: Optional[bytes] = None,
-        arrow_version: ArrowVersion = ArrowVersion.ALPHA,
+        arrow_endpoint_version: ArrowEndpointVersion = ArrowEndpointVersion.ALPHA,
     ):
         self._fallback_query_runner = fallback_query_runner
         self._server_version = server_version
-        self._arrow_version = arrow_version
+        self._arrow_endpoint_version = arrow_endpoint_version
 
         host, port_string = uri.split(":")
 
@@ -286,10 +287,10 @@ def _run_arrow_property_get(self, graph_name: str, procedure_name: str, configur
             "configuration": configuration,
         }
 
-        if self._arrow_version == ArrowVersion.V1:
+        if self._arrow_endpoint_version == ArrowEndpointVersion.V1:
             payload = {
                 "name": "GET_MESSAGE",
-                "version": ArrowVersion.V1.name(),
+                "version": ArrowEndpointVersion.V1.name(),
                 "body": payload,
             }
 
@@ -322,7 +323,7 @@ def create_graph_constructor(
             )
 
         return ArrowGraphConstructor(
-            database, graph_name, self._flight_client, concurrency, self._arrow_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:

…science/tests/unit/test_arrow_version.py …ests/unit/test_arrow_endpoint_version.pygraphdatascience/tests/unit/test_arrow_version.py renamed to graphdatascience/tests/unit/test_arrow_endpoint_version.py graphdatascience/tests/unit/test_arrow_version.py renamed to graphdatascience/tests/unit/test_arrow_endpoint_version.py

@@ -1,13 +1,13 @@
 import pytest
 
 from graphdatascience.graph_data_science import ArrowQueryRunner
-from graphdatascience.query_runner.arrow_version import (
-    ArrowVersion,
-    UnsupportedArrowVersion,
+from graphdatascience.query_runner.arrow_endpoint_version import (
+    ArrowEndpointVersion,
+    UnsupportedArrowEndpointVersion,
 )
 
 
-@pytest.mark.parametrize("arrow_version", [ArrowVersion.V1, ArrowVersion.ALPHA])
+@pytest.mark.parametrize("arrow_version", [ArrowEndpointVersion.V1, ArrowEndpointVersion.ALPHA])
 def test_arrow_version_parsing(arrow_version):
     arrow_info = dict()
     arrow_info["version"] = arrow_version.name()
@@ -19,11 +19,11 @@ def test_arrow_version_parsing(arrow_version):
 def test_arrow_version_parsing_fails(arrow_version):
     arrow_info = dict()
     arrow_info["version"] = arrow_version
-    with pytest.raises(UnsupportedArrowVersion) as e:
+    with pytest.raises(UnsupportedArrowEndpointVersion) as e:
         ArrowQueryRunner._read_arrow_version(arrow_info)
     assert arrow_version in str(e.value)
 
 
 def test_arrow_version_prefix():
-    assert ArrowVersion.ALPHA.prefix() == ""
-    assert ArrowVersion.V1.prefix() == "v1/"
+    assert ArrowEndpointVersion.ALPHA.prefix() == ""
+    assert ArrowEndpointVersion.V1.prefix() == "v1/"