Fix bug with cloudstore detection

mdio/dataset.h

@@ -101,29 +101,8 @@ inline Future<void> write_zmetadata(
 inline Future<tensorstore::KvStore> dataset_kvs_store(
     const std::string& dataset_path,
     tensorstore::Context context = tensorstore::Context::Default()) {
-  ::nlohmann::json kvstore;
-
-  // Use shared utility to infer driver from path prefix
-  std::string driver = zarr::InferDriverFromPath(dataset_path);
-  kvstore["driver"] = driver;
-
-  if (driver == "file") {
-    // Local file system - just normalize with trailing slash
-    kvstore["path"] = zarr::NormalizePathWithSlash(dataset_path);
-    return tensorstore::kvstore::Open(kvstore, context);
-  }
-
-  // Cloud storage (GCS or S3) - extract bucket and path
-  auto [bucket, path] = zarr::ExtractCloudPath(dataset_path);
-  if (bucket.empty()) {
-    return absl::InvalidArgumentError(
-        "gcs/s3 drivers requires [s3/gs]://[bucket]/[path_to_file]");
-  }
-
-  kvstore["bucket"] = bucket;
-  kvstore["path"] = zarr::NormalizePathWithSlash(path);
-
-  return tensorstore::kvstore::Open(kvstore, context);
+  MDIO_ASSIGN_OR_RETURN(auto loc, zarr::ResolveKvStoreLocation(dataset_path));
+  return tensorstore::kvstore::Open(zarr::BuildKvStoreSpec(loc), context);
 }
 
 /**

mdio/dataset_factory.h

@@ -422,28 +422,9 @@ inline void transform_shape(
  */
 inline absl::Status transform_metadata(const std::string& path,
                                        nlohmann::json& variable /*NOLINT*/) {
-  // Use shared utilities for driver inference and path handling
-  std::string driver = mdio::zarr::InferDriverFromPath(path);
   std::string var_name = variable["kvstore"]["path"].get<std::string>();
-
-  variable["kvstore"]["driver"] = driver;
-
-  if (driver == "file") {
-    // Local filesystem - normalize path with trailing slash
-    variable["kvstore"]["path"] =
-        mdio::zarr::NormalizePathWithSlash(path) + var_name;
-  } else {
-    // Cloud storage (GCS or S3) - extract bucket and path
-    auto [bucket, cloud_path] = mdio::zarr::ExtractCloudPath(path);
-    if (bucket.empty()) {
-      return absl::InvalidArgumentError(
-          "Cloud path requires [gs/s3]://[bucket]/[path to file] name");
-    }
-    variable["kvstore"]["bucket"] = bucket;
-    variable["kvstore"]["path"] =
-        mdio::zarr::NormalizePathWithSlash(cloud_path) + var_name;
-  }
-
+  MDIO_ASSIGN_OR_RETURN(auto loc, mdio::zarr::ResolveKvStoreLocation(path));
+  variable["kvstore"] = mdio::zarr::BuildKvStoreSpec(loc, var_name);
   return absl::OkStatus();
 }
 

mdio/zarr/zarr_driver.h

@@ -223,10 +223,8 @@ inline std::string GetConsolidatedMetadataFileName() { return ".zmetadata"; }
  * @return The driver name ("file", "gcs", or "s3").
  */
 inline std::string InferDriverFromPath(const std::string& path) {
-  if (path.length() > 5) {
-    if (path.substr(0, 5) == "gs://") return "gcs";
-    if (path.substr(0, 5) == "s3://") return "s3";
-  }
+  if (path.compare(0, 5, "gs://") == 0) return "gcs";
+  if (path.compare(0, 5, "s3://") == 0) return "s3";
   return "file";
 }
 
@@ -244,15 +242,21 @@ inline std::string NormalizePath(const std::string& path) {
 }
 
 /**
- * @brief Normalizes a path by ensuring it has a trailing slash.
+ * @brief Normalizes a path so it ends in exactly one trailing slash.
+ *
+ * Repeated trailing slashes are collapsed. Cloud object keys are not path
+ * normalized by the store, so "a//var" and "a/var" are distinct objects.
+ *
  * @param path The input path.
- * @return The normalized path with a trailing slash.
+ * @return The normalized path with a single trailing slash.
  */
 inline std::string NormalizePathWithSlash(const std::string& path) {
-  std::string result = path;
-  if (!result.empty() && result.back() != '/') {
-    result.push_back('/');
+  std::string result = NormalizePath(path);
+  if (result.empty()) {
+    // Preserve an empty path, and collapse a root-only path back to "/".
+    return path.empty() ? path : "/";
   }
+  result.push_back('/');
   return result;
 }
 
@@ -273,6 +277,76 @@ inline std::pair<std::string, std::string> ExtractCloudPath(
           without_scheme.substr(bucket_end + 1)};
 }
 
+/**
+ * @brief Resolved TensorStore kvstore location for a dataset path.
+ *
+ * Cloud URLs (`gs://`, `s3://`) are split into `bucket` plus a bucket-relative
+ * `path`. File paths keep the original path. `path` has a trailing slash when
+ * non-empty so a child name can be appended.
+ */
+struct KvStoreLocation {
+  std::string driver;
+  std::string bucket;
+  std::string path;
+};
+
+/**
+ * @brief Splits a dataset path into TensorStore kvstore fields.
+ *
+ * This is the single place that maps a user path onto `driver` / `bucket` /
+ * `path`.
+ *
+ * @param dataset_path Local path or `gs://` / `s3://` URI.
+ * @return The resolved location, or InvalidArgument for a missing cloud bucket.
+ */
+inline Result<KvStoreLocation> ResolveKvStoreLocation(
+    const std::string& dataset_path) {
+  std::string driver = InferDriverFromPath(dataset_path);
+  if (driver == "file") {
+    return KvStoreLocation{driver, "", NormalizePathWithSlash(dataset_path)};
+  }
+
+  auto [bucket, path] = ExtractCloudPath(dataset_path);
+  if (bucket.empty()) {
+    return absl::InvalidArgumentError(
+        "Cloud path requires [gs/s3]://[bucket]/[path to file]");
+  }
+  return KvStoreLocation{driver, bucket, NormalizePathWithSlash(path)};
+}
+
+/**
+ * @brief Builds a TensorStore kvstore JSON spec from a resolved location.
+ *
+ * For non-file drivers, `bucket` is always set. `child` is appended to `path`
+ * (typically a variable / array name).
+ *
+ * @param loc The resolved dataset location.
+ * @param child Optional child key appended to `path`.
+ * @return The kvstore JSON object.
+ */
+inline nlohmann::json BuildKvStoreSpec(const KvStoreLocation& loc,
+                                       const std::string& child = "") {
+  nlohmann::json kvstore = {{"driver", loc.driver}, {"path", loc.path + child}};
+  if (loc.driver != "file") {
+    kvstore["bucket"] = loc.bucket;
+  }
+  return kvstore;
+}
+
+/**
+ * @brief Builds a TensorStore Variable open spec (zarr/zarr3 + kvstore).
+ * @param zarr_driver TensorStore array driver (`zarr` or `zarr3`).
+ * @param loc The resolved dataset location.
+ * @param var_name Variable / array name appended to the kvstore path.
+ * @return The Variable spec JSON object.
+ */
+inline nlohmann::json BuildVariableSpec(const std::string& zarr_driver,
+                                        const KvStoreLocation& loc,
+                                        const std::string& var_name) {
+  return {{"driver", zarr_driver},
+          {"kvstore", BuildKvStoreSpec(loc, var_name)}};
+}
+
 // ============================================================================
 // Shared JSON Utilities
 // ============================================================================

mdio/zarr/zarr_test.cc

@@ -266,10 +266,13 @@ TEST(ZarrDriver, NormalizePathWithSlash_AddsTrailingSlash) {
 TEST(ZarrDriver, NormalizePathWithSlash_AlreadyHasSlash) {
   EXPECT_EQ(mdio::zarr::NormalizePathWithSlash("/path/to/data/"),
             "/path/to/data/");
+  EXPECT_EQ(mdio::zarr::NormalizePathWithSlash("/path/to/data///"),
+            "/path/to/data/");
 }
 
 TEST(ZarrDriver, NormalizePathWithSlash_EmptyPath) {
   EXPECT_EQ(mdio::zarr::NormalizePathWithSlash(""), "");
+  EXPECT_EQ(mdio::zarr::NormalizePathWithSlash("///"), "/");
 }
 
 TEST(ZarrDriver, ExtractCloudPath_GCS) {
@@ -298,6 +301,69 @@ TEST(ZarrDriver, ExtractCloudPath_ShortUrl) {
   EXPECT_EQ(path, "");
 }
 
+TEST(ZarrDriver, ResolveKvStoreLocation_LocalFile) {
+  auto loc = mdio::zarr::ResolveKvStoreLocation("/path/to/data");
+  ASSERT_TRUE(loc.ok()) << loc.status();
+  EXPECT_EQ(loc->driver, "file");
+  EXPECT_EQ(loc->bucket, "");
+  EXPECT_EQ(loc->path, "/path/to/data/");
+}
+
+TEST(ZarrDriver, ResolveKvStoreLocation_GCS) {
+  auto loc = mdio::zarr::ResolveKvStoreLocation("gs://bucket/path");
+  ASSERT_TRUE(loc.ok()) << loc.status();
+  EXPECT_EQ(loc->driver, "gcs");
+  EXPECT_EQ(loc->bucket, "bucket");
+  EXPECT_EQ(loc->path, "path/");
+}
+
+TEST(ZarrDriver, ResolveKvStoreLocation_S3) {
+  auto loc =
+      mdio::zarr::ResolveKvStoreLocation("s3://my-bucket/nested/path/volume");
+  ASSERT_TRUE(loc.ok()) << loc.status();
+  EXPECT_EQ(loc->driver, "s3");
+  EXPECT_EQ(loc->bucket, "my-bucket");
+  EXPECT_EQ(loc->path, "nested/path/volume/");
+}
+
+TEST(ZarrDriver, ResolveKvStoreLocation_CollapsesTrailingSlashes) {
+  auto loc = mdio::zarr::ResolveKvStoreLocation("s3://my-bucket/nested/path//");
+  ASSERT_TRUE(loc.ok()) << loc.status();
+  EXPECT_EQ(loc->path, "nested/path/");
+}
+
+TEST(ZarrDriver, ResolveKvStoreLocation_RejectsMissingBucket) {
+  EXPECT_FALSE(mdio::zarr::ResolveKvStoreLocation("gs://").ok());
+  EXPECT_FALSE(mdio::zarr::ResolveKvStoreLocation("s3:///path").ok());
+}
+
+TEST(ZarrDriver, BuildKvStoreSpec_LocalFile) {
+  auto loc = mdio::zarr::ResolveKvStoreLocation("/path/to/dataset");
+  ASSERT_TRUE(loc.ok()) << loc.status();
+  auto spec = mdio::zarr::BuildKvStoreSpec(loc.value(), "myvar");
+  EXPECT_EQ(spec["driver"], "file");
+  EXPECT_EQ(spec["path"], "/path/to/dataset/myvar");
+  EXPECT_FALSE(spec.contains("bucket"));
+}
+
+TEST(ZarrDriver, BuildKvStoreSpec_GCS) {
+  auto loc = mdio::zarr::ResolveKvStoreLocation("gs://bucket/path");
+  ASSERT_TRUE(loc.ok()) << loc.status();
+  auto spec = mdio::zarr::BuildKvStoreSpec(loc.value(), "myvar");
+  EXPECT_EQ(spec["driver"], "gcs");
+  EXPECT_EQ(spec["bucket"], "bucket");
+  EXPECT_EQ(spec["path"], "path/myvar");
+}
+
+TEST(ZarrDriver, BuildKvStoreSpec_S3) {
+  auto loc = mdio::zarr::ResolveKvStoreLocation("s3://my-bucket/nested/path");
+  ASSERT_TRUE(loc.ok()) << loc.status();
+  auto spec = mdio::zarr::BuildKvStoreSpec(loc.value(), "amplitude");
+  EXPECT_EQ(spec["driver"], "s3");
+  EXPECT_EQ(spec["bucket"], "my-bucket");
+  EXPECT_EQ(spec["path"], "nested/path/amplitude");
+}
+
 // =============================================================================
 // JSON Utility Tests
 // =============================================================================
@@ -895,20 +961,36 @@ TEST(ZarrV3, ExtractChildArrayCandidates_Empty) {
 }
 
 TEST(ZarrV3, BuildVariableSpec_LocalFile) {
-  auto spec =
-      mdio::zarr::v3::BuildVariableSpec("file", "/path/to/dataset", "myvar");
+  auto loc = mdio::zarr::ResolveKvStoreLocation("/path/to/dataset");
+  ASSERT_TRUE(loc.ok()) << loc.status();
+  auto spec = mdio::zarr::BuildVariableSpec("zarr3", loc.value(), "myvar");
 
   EXPECT_EQ(spec["driver"], "zarr3");
   EXPECT_EQ(spec["kvstore"]["driver"], "file");
   EXPECT_EQ(spec["kvstore"]["path"], "/path/to/dataset/myvar");
+  EXPECT_FALSE(spec["kvstore"].contains("bucket"));
 }
 
 TEST(ZarrV3, BuildVariableSpec_GCS) {
-  auto spec = mdio::zarr::v3::BuildVariableSpec("gcs", "bucket/path", "myvar");
+  auto loc = mdio::zarr::ResolveKvStoreLocation("gs://bucket/path");
+  ASSERT_TRUE(loc.ok()) << loc.status();
+  auto spec = mdio::zarr::BuildVariableSpec("zarr3", loc.value(), "myvar");
 
   EXPECT_EQ(spec["driver"], "zarr3");
   EXPECT_EQ(spec["kvstore"]["driver"], "gcs");
-  EXPECT_EQ(spec["kvstore"]["path"], "bucket/path/myvar");
+  EXPECT_EQ(spec["kvstore"]["bucket"], "bucket");
+  EXPECT_EQ(spec["kvstore"]["path"], "path/myvar");
+}
+
+TEST(ZarrV3, BuildVariableSpec_S3) {
+  auto loc = mdio::zarr::ResolveKvStoreLocation("s3://my-bucket/nested/path");
+  ASSERT_TRUE(loc.ok()) << loc.status();
+  auto spec = mdio::zarr::BuildVariableSpec("zarr3", loc.value(), "amplitude");
+
+  EXPECT_EQ(spec["driver"], "zarr3");
+  EXPECT_EQ(spec["kvstore"]["driver"], "s3");
+  EXPECT_EQ(spec["kvstore"]["bucket"], "my-bucket");
+  EXPECT_EQ(spec["kvstore"]["path"], "nested/path/amplitude");
 }
 
 }  // namespace ZarrV3Tests

mdio/zarr/zarr_v2.h

@@ -326,22 +326,12 @@ struct V2MetadataState {
   PromiseType promise;
   tensorstore::KvStore kvs;
   std::string dataset_path;
-  std::string normalized_path;
-  std::string driver;
-  std::string bucket;
-  std::string cloud_path;
+  KvStoreLocation location;
 
-  explicit V2MetadataState(PromiseType p, const std::string& path)
+  V2MetadataState(PromiseType p, std::string path, KvStoreLocation loc)
       : promise(std::move(p)),
-        dataset_path(path),
-        normalized_path(NormalizePathWithSlash(path)),
-        driver(InferDriverFromPath(path)) {
-    if (driver != "file") {
-      auto cloud_parts = ExtractCloudPath(path);
-      bucket = cloud_parts.first;
-      cloud_path = NormalizePathWithSlash(cloud_parts.second);
-    }
-  }
+        dataset_path(std::move(path)),
+        location(std::move(loc)) {}
 
   /// Completes with an error status.
   void Fail(absl::Status status) { promise.SetResult(std::move(status)); }
@@ -352,16 +342,8 @@ struct V2MetadataState {
   }
 
   /// Builds a variable spec for the given variable name.
-  nlohmann::json BuildVariableSpec(const std::string& var_name) const {
-    nlohmann::json spec = {{"driver", std::string(kDriverName)},
-                           {"kvstore", {{"driver", driver}}}};
-    if (driver != "file") {
-      spec["kvstore"]["bucket"] = bucket;
-      spec["kvstore"]["path"] = cloud_path + var_name;
-    } else {
-      spec["kvstore"]["path"] = normalized_path + var_name;
-    }
-    return spec;
+  nlohmann::json MakeVariableSpec(const std::string& var_name) const {
+    return BuildVariableSpec(std::string(kDriverName), location, var_name);
   }
 };
 
@@ -413,7 +395,7 @@ inline void OnZmetadataRead(
     if (element.value().contains("dtype") &&
         IsMetadataOnlyDataType(element.value()["dtype"])) {
       std::string var_name = ExtractVariableName(key);
-      auto spec = state->BuildVariableSpec(var_name);
+      auto spec = state->MakeVariableSpec(var_name);
       spec["_mdio_header_only"] = true;
       spec["_mdio_zarray"] = element.value();
       const std::string zattrs_key = var_name + "/.zattrs";
@@ -426,7 +408,7 @@ inline void OnZmetadataRead(
       continue;
     }
     std::string var_name = ExtractVariableName(key);
-    json_vars.push_back(state->BuildVariableSpec(var_name));
+    json_vars.push_back(state->MakeVariableSpec(var_name));
   }
 
   if (json_vars.empty()) {
@@ -467,11 +449,13 @@ inline void OnV2KvStoreReady(
 inline Future<std::tuple<::nlohmann::json, std::vector<::nlohmann::json>>>
 ReadConsolidatedMetadata(const std::string& dataset_path,
                          tensorstore::Future<tensorstore::KvStore> kvs_future) {
+  MDIO_ASSIGN_OR_RETURN(auto location, ResolveKvStoreLocation(dataset_path));
+
   auto pair = tensorstore::PromiseFuturePair<
       std::tuple<::nlohmann::json, std::vector<::nlohmann::json>>>::Make();
 
   auto state = std::make_shared<internal::V2MetadataState>(
-      std::move(pair.promise), dataset_path);
+      std::move(pair.promise), dataset_path, std::move(location));
 
   kvs_future.ExecuteWhenReady(
       [state](tensorstore::ReadyFuture<tensorstore::KvStore> ready) {