Add KvStore URL pipeline support

This is in line with zarr-developers/zeps#48 and
the syntax supported by Neuroglancer.

Currently, zip is supported.  OCDBT support will be added in a
subsequent commit.

PiperOrigin-RevId: 755691199
Change-Id: Ia6cb84c12a986a7dd0ba65e41454fbe6d415aed0

Author: jbms

tensorstore/internal/BUILD

@@ -1612,6 +1612,7 @@ tensorstore_cc_library(
     hdrs = ["uri_utils.h"],
     deps = [
         ":ascii_set",
+        "@abseil-cpp//absl/status",
         "@abseil-cpp//absl/strings",
     ],
 )

tensorstore/internal/driver_kind_registry.cc

@@ -35,6 +35,8 @@ struct DriverKindRegistry {
   absl::Mutex mutex;
   absl::flat_hash_map<std::string, DriverKind> driver_kinds
       ABSL_GUARDED_BY(mutex);
+  absl::flat_hash_map<std::string, UrlSchemeKind> scheme_kinds
+      ABSL_GUARDED_BY(mutex);
 };
 
 DriverKindRegistry& GetDriverKindRegistry() {
@@ -84,5 +86,44 @@ std::ostream& operator<<(std::ostream& os, DriverKind x) {
   return os << DriverKindToStringView(x);
 }
 
+void RegisterUrlSchemeKind(std::string_view scheme, UrlSchemeKind scheme_kind) {
+  auto& registry = GetDriverKindRegistry();
+  absl::MutexLock lock(&registry.mutex);
+  if (auto result = registry.scheme_kinds.emplace(scheme, scheme_kind);
+      !result.second) {
+    ABSL_LOG(FATAL) << scheme << " already registered as "
+                    << result.first->second;
+  }
+}
+
+std::optional<UrlSchemeKind> GetUrlSchemeKind(std::string_view scheme) {
+  auto& registry = GetDriverKindRegistry();
+  absl::MutexLock lock(&registry.mutex);
+  auto it = registry.scheme_kinds.find(scheme);
+  if (it == registry.scheme_kinds.end()) return std::nullopt;
+  return it->second;
+}
+
+std::string_view UrlSchemeKindToStringView(UrlSchemeKind x) {
+  switch (x) {
+    case UrlSchemeKind::kKvStoreRoot:
+      return "root kvstore";
+    case UrlSchemeKind::kKvStoreAdapter:
+      return "kvstore adapter";
+    case UrlSchemeKind::kTensorStoreRoot:
+      return "root TensorStore";
+    case UrlSchemeKind::kTensorStoreKvStoreAdapter:
+      return "kvstore-based TensorStore";
+    case UrlSchemeKind::kTensorStoreAdapter:
+      return "TensorStore adapter";
+    default:
+      ABSL_UNREACHABLE();
+  }
+}
+
+std::ostream& operator<<(std::ostream& os, UrlSchemeKind x) {
+  return os << UrlSchemeKindToStringView(x);
+}
+
 }  // namespace internal
 }  // namespace tensorstore

tensorstore/internal/driver_kind_registry.h

@@ -54,6 +54,53 @@ void RegisterDriverKind(std::string_view id, DriverKind kind,
 // registered.
 std::optional<DriverKind> GetDriverKind(std::string_view id);
 
+// Indicates the kind of a registered URL scheme in the tensorstore library.
+//
+// The kvstore URL scheme handlers are registered in
+// tensorstore/kvstore/registry.h and the TensorStore URL scheme
+// handlers are registered in tensorstore/driver/registry.h.
+enum class UrlSchemeKind {
+  // KvStore driver that directly accesses storage, e.g. "file"
+  //
+  // This must be first in a URL pipeline.
+  kKvStoreRoot,
+  // KvStore driver that adapts a base kvstore, e.g. "zip" or "ocdbt".
+  //
+  // This must immediately follow a `kKvStoreBase` component in a URL pipeline.
+  kKvStoreAdapter,
+  // TensorStore driver that directly accesses storage.
+  //
+  // This must be first in a URL pipeline.
+  kTensorStoreRoot,
+  // TensorStore driver that adapts a base kvstore, e.g. "zarr3"
+  //
+  // This must immediately follow a `kKvStoreBase` component in a URL pipeline.
+  kTensorStoreKvStoreAdapter,
+  // TensorStore driver that adapts a base TensorStore, e.g. "cast"
+  //
+  // This must immediately follow a `kTensorStoreBase` or
+  // `kTensorStoreKvStoreAdapter` component in a URL pipeline.
+  kTensorStoreAdapter,
+};
+
+// Returns the name of a URL scheme kind for use in error messages.
+std::string_view UrlSchemeKindToStringView(UrlSchemeKind x);
+
+template <typename Sink>
+void AbslStringify(Sink& sink, UrlSchemeKind x) {
+  return sink.Append(UrlSchemeKindToStringView(x));
+}
+std::ostream& operator<<(std::ostream& os, UrlSchemeKind x);
+
+// Registers a URL scheme.  This is called automatically by
+// `internal_kvstore::UrlSchemeRegistration` and
+// `internal::UrlSchemeRegistration`.
+void RegisterUrlSchemeKind(std::string_view scheme, UrlSchemeKind scheme_kind);
+
+// Returns the kind of a registered URL scheme, or `std::nullopt` if not
+// registered.
+std::optional<UrlSchemeKind> GetUrlSchemeKind(std::string_view scheme);
+
 }  // namespace internal
 }  // namespace tensorstore
 

tensorstore/internal/uri_utils.cc

@@ -23,6 +23,7 @@
 #include <string>
 #include <string_view>
 
+#include "absl/status/status.h"
 #include "absl/strings/ascii.h"
 #include "absl/strings/match.h"
 #include "tensorstore/internal/ascii_set.h"
@@ -87,17 +88,18 @@ void PercentDecodeAppend(std::string_view src, std::string& dest) {
   }
 }
 
-ParsedGenericUri ParseGenericUri(std::string_view uri) {
-  static constexpr std::string_view kSchemeSep("://");
+namespace {
+ParsedGenericUri ParseGenericUriImpl(std::string_view uri,
+                                     std::string_view scheme_delimiter) {
   ParsedGenericUri result;
-  const auto scheme_start = uri.find(kSchemeSep);
+  const auto scheme_start = uri.find(scheme_delimiter);
   std::string_view uri_suffix;
   if (scheme_start == std::string_view::npos) {
     // No scheme
     uri_suffix = uri;
   } else {
     result.scheme = uri.substr(0, scheme_start);
-    uri_suffix = uri.substr(scheme_start + kSchemeSep.size());
+    uri_suffix = uri.substr(scheme_start + scheme_delimiter.size());
   }
   const auto fragment_start = uri_suffix.find('#');
   const auto query_start = uri_suffix.substr(0, fragment_start).find('?');
@@ -127,6 +129,32 @@ ParsedGenericUri ParseGenericUri(std::string_view uri) {
   }
   return result;
 }
+}  // namespace
+
+ParsedGenericUri ParseGenericUri(std::string_view uri) {
+  return ParseGenericUriImpl(uri, "://");
+}
+
+ParsedGenericUri ParseGenericUriWithoutSlashSlash(std::string_view uri) {
+  return ParseGenericUriImpl(uri, ":");
+}
+
+absl::Status EnsureNoQueryOrFragment(const ParsedGenericUri& parsed_uri) {
+  if (!parsed_uri.query.empty()) {
+    return absl::InvalidArgumentError("Query string not supported");
+  }
+  if (!parsed_uri.fragment.empty()) {
+    return absl::InvalidArgumentError("Fragment identifier not supported");
+  }
+  return absl::OkStatus();
+}
+
+absl::Status EnsureNoPathOrQueryOrFragment(const ParsedGenericUri& parsed_uri) {
+  if (!parsed_uri.authority_and_path.empty()) {
+    return absl::InvalidArgumentError("Path not supported");
+  }
+  return EnsureNoQueryOrFragment(parsed_uri);
+}
 
 std::optional<HostPort> SplitHostPort(std::string_view host_port) {
   if (host_port.empty()) return std::nullopt;

tensorstore/internal/uri_utils.h

@@ -22,6 +22,7 @@
 #include <string>
 #include <string_view>
 
+#include "absl/status/status.h"
 #include "tensorstore/internal/ascii_set.h"
 
 namespace tensorstore {
@@ -120,9 +121,20 @@ struct ParsedGenericUri {
 };
 
 /// Parses a "generic" URI of the form
-/// `<scheme>://<authority-and-path>?<query>#<fragment>` where the `?<query>`
-/// and `#<fragment>` portions are optional.
+/// `<scheme><scheme_delimiter>//<authority-and-path>?<query>#<fragment>`
+/// where the `?<query>` and `#<fragment>` portions are optional.
+///
+/// `<scheme_delimiter>`  is:
+/// - "://" for `ParseGenericUri`
+/// - ":" for `ParseGenericUriWithoutSlashSlash`
 ParsedGenericUri ParseGenericUri(std::string_view uri);
+ParsedGenericUri ParseGenericUriWithoutSlashSlash(std::string_view uri);
+
+// Returns an error if there is a query or fragment.
+absl::Status EnsureNoQueryOrFragment(const ParsedGenericUri& parsed_uri);
+
+// Returns an error if there is a path, query or fragment.
+absl::Status EnsureNoPathOrQueryOrFragment(const ParsedGenericUri& parsed_uri);
 
 struct HostPort {
   std::string_view host;

tensorstore/kvstore/file/file_key_value_store.cc

@@ -910,12 +910,7 @@ Future<kvstore::DriverPtr> FileKeyValueStoreSpec::DoOpen() const {
 Result<kvstore::Spec> ParseFileUrl(std::string_view url) {
   auto parsed = internal::ParseGenericUri(url);
   assert(parsed.scheme == internal_file_kvstore::FileKeyValueStoreSpec::id);
-  if (!parsed.query.empty()) {
-    return absl::InvalidArgumentError("Query string not supported");
-  }
-  if (!parsed.fragment.empty()) {
-    return absl::InvalidArgumentError("Fragment identifier not supported");
-  }
+  TENSORSTORE_RETURN_IF_ERROR(internal::EnsureNoQueryOrFragment(parsed));
   std::string path = internal::PercentDecode(parsed.authority_and_path);
   auto driver_spec = internal::MakeIntrusivePtr<FileKeyValueStoreSpec>();
   driver_spec->data_.file_io_concurrency =

tensorstore/kvstore/file/file_key_value_store_test.cc

@@ -401,6 +401,7 @@ TEST(FileKeyValueStoreTest, SpecRoundtrip) {
   std::string root = tempdir.path() + "/root";
   tensorstore::internal::KeyValueStoreSpecRoundtripOptions options;
   options.full_spec = {{"driver", "file"}, {"path", root}};
+  options.url = "file://" + root;
   tensorstore::internal::TestKeyValueStoreSpecRoundtrip(options);
 }
 
@@ -419,6 +420,7 @@ TEST(FileKeyValueStoreTest, SpecRoundtripSync) {
            {"file_io_locking", {{"mode", "lockfile"}}},
        }},
   };
+  options.url = "file://" + root;
   options.spec_request_options.Set(tensorstore::retain_context).IgnoreError();
   tensorstore::internal::TestKeyValueStoreSpecRoundtrip(options);
 }

tensorstore/kvstore/gcs_grpc/gcs_grpc.cc

@@ -1026,12 +1026,7 @@ Future<kvstore::DriverPtr> GcsGrpcKeyValueStoreSpec::DoOpen() const {
 Result<kvstore::Spec> ParseGcsGrpcUrl(std::string_view url) {
   auto parsed = internal::ParseGenericUri(url);
   assert(parsed.scheme == kUriScheme);
-  if (!parsed.query.empty()) {
-    return absl::InvalidArgumentError("Query string not supported");
-  }
-  if (!parsed.fragment.empty()) {
-    return absl::InvalidArgumentError("Fragment identifier not supported");
-  }
+  TENSORSTORE_RETURN_IF_ERROR(internal::EnsureNoQueryOrFragment(parsed));
   if (!IsValidBucketName(parsed.authority)) {
     return absl::InvalidArgumentError(tensorstore::StrCat(
         "Invalid GCS bucket name: ", QuoteString(parsed.authority)));

tensorstore/kvstore/gcs_http/gcs_key_value_store.cc

@@ -1272,12 +1272,7 @@ Future<const void> GcsKeyValueStore::DeleteRange(KeyRange range) {
 Result<kvstore::Spec> ParseGcsUrl(std::string_view url) {
   auto parsed = internal::ParseGenericUri(url);
   assert(parsed.scheme == kUriScheme);
-  if (!parsed.query.empty()) {
-    return absl::InvalidArgumentError("Query string not supported");
-  }
-  if (!parsed.fragment.empty()) {
-    return absl::InvalidArgumentError("Fragment identifier not supported");
-  }
+  TENSORSTORE_RETURN_IF_ERROR(internal::EnsureNoQueryOrFragment(parsed));
   if (!IsValidBucketName(parsed.authority)) {
     return absl::InvalidArgumentError(absl::StrCat(
         "Invalid GCS bucket name: ", QuoteString(parsed.authority)));

tensorstore/kvstore/gcs_http/gcs_key_value_store_test.cc

@@ -426,7 +426,7 @@ TEST(GcsKeyValueStoreTest, SpecRoundtrip) {
 
   tensorstore::internal::KeyValueStoreSpecRoundtripOptions options;
   options.full_spec = {{"driver", kDriver}, {"bucket", "my-bucket"}};
-
+  options.url = "gs://my-bucket/";
   tensorstore::internal::TestKeyValueStoreSpecRoundtrip(options);
 }