add store.getSize

Author: brokkoli71

src/main/java/dev/zarr/zarrjava/store/BufferedZipStore.java

@@ -296,6 +296,10 @@ public InputStream getInputStream(String[] keys, long start, long end) {
         return bufferStore.getInputStream(keys, start, end);
     }
 
+    public long getSize(String[] keys) {
+        return bufferStore.getSize(keys);
+    }
+
     @Override
     public String toString() {
         return "BufferedZipStore(" + underlyingStore.toString() + ")";

src/main/java/dev/zarr/zarrjava/store/FilesystemStore.java

@@ -170,4 +170,11 @@ public InputStream getInputStream(String[] keys, long start, long end) {
             throw new RuntimeException(e);
         }
     }
+    public long getSize(String[] keys) {
+        try {
+            return Files.size(resolveKeys(keys));
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
 }

src/main/java/dev/zarr/zarrjava/store/HttpStore.java

@@ -134,4 +134,30 @@ public void close() throws IOException {
             return null;
         }
     }
+    @Override
+    public long getSize(String[] keys) {
+        // Explicitly request "identity" encoding to prevent OkHttp from adding "gzip"
+        // and subsequently stripping the Content-Length header.
+        Request request = new Request.Builder()
+                .head()
+                .url(resolveKeys(keys))
+                .header("Accept-Encoding", "identity")
+                .build();
+
+        Call call = httpClient.newCall(request);
+        try {
+            Response response = call.execute();
+            if (!response.isSuccessful()) {
+                throw new IOException("Failed to get size: " + response.code());
+            }
+
+            String contentLength = response.header("Content-Length");
+            if (contentLength != null) {
+                return Long.parseLong(contentLength);
+            }
+            return -1;
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
 }

src/main/java/dev/zarr/zarrjava/store/MemoryStore.java

@@ -93,5 +93,13 @@ public InputStream getInputStream(String[] keys, long start, long end) {
         if (end > Integer.MAX_VALUE) throw new IllegalArgumentException("End index too large");
         return new java.io.ByteArrayInputStream(bytes, (int) start, (int)(end - start));
     }
-}
 
+    @Override
+    public long getSize(String[] keys) {
+        byte[] bytes = map.get(resolveKeys(keys));
+        if (bytes == null) {
+            throw new RuntimeException(new java.io.FileNotFoundException("Key not found: " + String.join("/", keys)));
+        }
+        return bytes.length;
+    }
+}

src/main/java/dev/zarr/zarrjava/store/S3Store.java

@@ -132,6 +132,19 @@ public InputStream getInputStream(String[] keys, long start, long end) {
     return responseInputStream;
   }
 
+  @Override
+  public long getSize(String[] keys) {
+    HeadObjectRequest req = HeadObjectRequest.builder()
+            .bucket(bucketName)
+            .key(resolveKeys(keys))
+            .build();
+    try {
+      return s3client.headObject(req).contentLength();
+    } catch (NoSuchKeyException e) {
+      throw new RuntimeException(e);
+    }
+  }
+
   @Override
   public String toString() {
     return "s3://" + bucketName + "/" + prefix;

src/main/java/dev/zarr/zarrjava/store/Store.java

@@ -49,4 +49,6 @@ default Stream<String[]> list() {
   default InputStream getInputStream(String[] keys) {
     return getInputStream(keys, 0, -1);
   }
+
+  long getSize(String[] keys);
 }

src/main/java/dev/zarr/zarrjava/store/StoreHandle.java

@@ -73,6 +73,10 @@ public Stream<String[]> list() {
     return ((Store.ListableStore) store).list(keys);
   }
 
+  public long getSize() {
+    return store.getSize(keys);
+  }
+
   @Override
   public String toString() {
     return store + "/" + String.join("/", keys);

src/test/java/dev/zarr/zarrjava/ZarrStoreTest.java

@@ -155,6 +155,14 @@ public void testStoreInputStream(StoreHandle storeHandle) throws IOException, Za
         Assertions.assertArrayEquals(expectedBuffer, buffer);
     }
 
+    @ParameterizedTest
+    @MethodSource("inputStreamStores")
+    public void testStoreGetSize(StoreHandle storeHandle) throws IOException, ZarrException {
+        long size = storeHandle.getSize();
+        long actual_size = storeHandle.read().remaining();
+        Assertions.assertEquals(actual_size, size);
+    }
+
     @Test
     public void testHttpStore() throws IOException, ZarrException {
         HttpStore httpStore = new dev.zarr.zarrjava.store.HttpStore("https://static.webknossos.org/data/zarr_v3/l4_sample");