Merge pull request #139 from seung-lab/l2cache_check

add table map endpoint and l2 cache check

caveclient/endpoints.py

@@ -262,6 +262,7 @@
 l2cache_endpoints_v1 = {
     "l2cache_data": l2cache_v1 + "/table/{table_id}/attributes",
     "l2cache_meta": l2cache_v1 + "/attribute_metadata",
+    "l2cache_table_mapping": l2cache_v1 + "/table_mapping",
 }
 
 l2cache_api_versions = {1: l2cache_endpoints_v1}

caveclient/l2cache.py

@@ -4,7 +4,10 @@
     l2cache_endpoints_common,
 )
 from .auth import AuthClient
+from requests.exceptions import HTTPError
+from warnings import warn
 import json
+from urllib.parse import urlparse
 
 server_key = "l2cache_server_address"
 
@@ -135,6 +138,48 @@ def attributes(self):
             self._available_attributes = list(self.cache_metadata().keys())
         return self._available_attributes
 
+    def table_mapping(self):
+        """Retrieves table mappings for l2 cache.
+        Parameters
+        ----------
+
+        Returns
+        -------
+        dict
+            keys are pcg table names, values are dicts with fields `l2cache_id` and `cv_path`.
+        """
+        endpoint_mapping = self.default_url_mapping
+        url = self._endpoints["l2cache_table_mapping"].format_map(endpoint_mapping)
+        response = self.session.get(url)
+        return handle_response(response)
+
+    def has_cache(self, datastack_name=None):
+        """Checks if the l2 cache is available for the dataset
+
+        Parameters
+        ----------
+        datastack_name : str, optional
+            The name of the datastack to check, by default None (if None, uses the client's datastack)
+
+        Returns
+        -------
+        bool
+            True if the l2 cache is available, False otherwise
+        """
+        seg_source = self.fc.info.segmentation_source(datastack_name=datastack_name)
+        if urlparse(seg_source).scheme != "graphene":
+            return False
+        table_name = self.fc.chunkedgraph.table_name
+        try:
+            table_mapping = self.table_mapping()
+        except HTTPError as e:
+            if e.response.status_code == 404:
+                warn(f"L2cache deployment '{self.server_address}/l2cache' does not have a l2 cache table mapping. Assuming no cache.")
+                return False
+            else:
+                raise e
+        return table_name in table_mapping
+
 
 client_mapping = {
     1: L2CacheClientLegacy,