From 22b804bd2deddf34b628a097c45abcf2d5d2cf2f Mon Sep 17 00:00:00 2001
From: Laramie Leavitt <lar@google.com>
Date: Fri, 11 Oct 2024 12:01:13 -0700
Subject: [PATCH] Add an initial tensorstore cli tool.

Currently supports:
* tscli list  --source [globs...]
* tscli copy  --source --target
* tscli search --source
* tscli print_spec --spec
* tscli print_stats --spec [boxes...]

PiperOrigin-RevId: 684912204
Change-Id: I13229643fc65d3352670dbdc5a62fdd016259ec6
---
 tensorstore/tscli/BUILD             |  84 ++++++++++++
 tensorstore/tscli/args.cc           | 198 ++++++++++++++++++++++++++++
 tensorstore/tscli/args.h            |  56 ++++++++
 tensorstore/tscli/args_test.cc      |  37 ++++++
 tensorstore/tscli/cli.h             |  37 ++++++
 tensorstore/tscli/kvstore_copy.cc   | 144 ++++++++++++++++++++
 tensorstore/tscli/kvstore_list.cc   | 144 ++++++++++++++++++++
 tensorstore/tscli/main.cc           | 145 ++++++++++++++++++++
 tensorstore/tscli/ts_print_spec.cc  |  75 +++++++++++
 tensorstore/tscli/ts_print_stats.cc | 108 +++++++++++++++
 tensorstore/tscli/ts_search.cc      | 175 ++++++++++++++++++++++++
 11 files changed, 1203 insertions(+)
 create mode 100644 tensorstore/tscli/BUILD
 create mode 100644 tensorstore/tscli/args.cc
 create mode 100644 tensorstore/tscli/args.h
 create mode 100644 tensorstore/tscli/args_test.cc
 create mode 100644 tensorstore/tscli/cli.h
 create mode 100644 tensorstore/tscli/kvstore_copy.cc
 create mode 100644 tensorstore/tscli/kvstore_list.cc
 create mode 100644 tensorstore/tscli/main.cc
 create mode 100644 tensorstore/tscli/ts_print_spec.cc
 create mode 100644 tensorstore/tscli/ts_print_stats.cc
 create mode 100644 tensorstore/tscli/ts_search.cc

diff --git a/tensorstore/tscli/BUILD b/tensorstore/tscli/BUILD
new file mode 100644
index 00000000..59054a2f
--- /dev/null
+++ b/tensorstore/tscli/BUILD
@@ -0,0 +1,84 @@
+load("//bazel:tensorstore.bzl", "tensorstore_cc_binary", "tensorstore_cc_library", "tensorstore_cc_test")
+
+package(default_visibility = ["//visibility:public"])
+
+licenses(["notice"])
+
+tensorstore_cc_test(
+    name = "args_test",
+    srcs = ["args_test.cc"],
+    deps = [
+        ":tsclilib",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+tensorstore_cc_library(
+    name = "tsclilib",
+    srcs = [
+        "args.cc",
+        "kvstore_copy.cc",
+        "kvstore_list.cc",
+        "ts_print_spec.cc",
+        "ts_print_stats.cc",
+        "ts_search.cc",
+    ],
+    hdrs = [
+        "args.h",
+        "cli.h",
+    ],
+    deps = [
+        "//tensorstore",
+        "//tensorstore:array_storage_statistics",
+        "//tensorstore:box",
+        "//tensorstore:context",
+        "//tensorstore:open",
+        "//tensorstore:open_mode",
+        "//tensorstore:spec",
+        "//tensorstore/index_space:dim_expression",
+        "//tensorstore/internal:path",
+        "//tensorstore/internal/json_binding",
+        "//tensorstore/kvstore",
+        "//tensorstore/kvstore:generation",
+        "//tensorstore/kvstore:key_range",
+        "//tensorstore/util:executor",
+        "//tensorstore/util:future",
+        "//tensorstore/util:json_absl_flag",
+        "//tensorstore/util:quote_string",
+        "//tensorstore/util:result",
+        "//tensorstore/util:span",
+        "//tensorstore/util:status",
+        "@com_github_nlohmann_json//:json",
+        "@com_google_absl//absl/container:flat_hash_set",
+        "@com_google_absl//absl/flags:parse",
+        "@com_google_absl//absl/functional:function_ref",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:string_view",
+        "@com_google_absl//absl/synchronization",
+        "@com_google_re2//:re2",
+    ],
+)
+
+tensorstore_cc_binary(
+    name = "tscli",
+    srcs = [
+        "main.cc",
+    ],
+    deps = [
+        ":tsclilib",
+        "//tensorstore:context",
+        "//tensorstore/driver:all_drivers",
+        "//tensorstore/internal/metrics:collect",
+        "//tensorstore/internal/metrics:registry",
+        "//tensorstore/kvstore:all_drivers",
+        "//tensorstore/util:json_absl_flag",
+        "@com_google_absl//absl/base:log_severity",
+        "@com_google_absl//absl/flags:flag",
+        "@com_google_absl//absl/flags:parse",
+        "@com_google_absl//absl/log:globals",
+        "@com_google_absl//absl/log:initialize",
+        "@com_google_absl//absl/status",
+    ],
+)
diff --git a/tensorstore/tscli/args.cc b/tensorstore/tscli/args.cc
new file mode 100644
index 00000000..c00bf2d5
--- /dev/null
+++ b/tensorstore/tscli/args.cc
@@ -0,0 +1,198 @@
+// Copyright 2024 The TensorStore Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "tensorstore/tscli/args.h"
+
+#include <cassert>
+#include <cstdint>
+#include <string>
+#include <string_view>
+#include <vector>
+
+#include "absl/container/flat_hash_set.h"
+#include "absl/status/status.h"
+#include "absl/strings/ascii.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+#include "absl/strings/strip.h"
+#include "tensorstore/util/span.h"
+#include "tensorstore/util/status.h"
+
+namespace tensorstore {
+namespace cli {
+
+absl::Status TryParseOptions(CommandFlags& flags,
+                             tensorstore::span<Option> options) {
+  absl::flat_hash_set<uintptr_t> handled;
+
+  auto it = flags.argv.begin() + 1;
+  while (it != flags.argv.end()) {
+    bool parsed = false;
+    std::string_view it_str = *it;
+
+    // Try to parse as a long option.
+    for (auto& opt : options) {
+      if (opt.longname.empty()) continue;
+      std::string_view arg = it_str;
+      if (absl::ConsumePrefix(&arg, opt.longname)) {
+        if (arg.empty()) {
+          parsed = true;
+          std::string_view value;
+          if (it + 1 != flags.argv.end()) {
+            value = *(it + 1);
+            handled.insert(reinterpret_cast<uintptr_t>(*it));
+            it++;
+          }
+          handled.insert(reinterpret_cast<uintptr_t>(*it));
+          TENSORSTORE_RETURN_IF_ERROR(opt.parse(value));
+        } else if (absl::ConsumePrefix(&arg, "=")) {
+          parsed = true;
+          TENSORSTORE_RETURN_IF_ERROR(opt.parse(arg));
+          handled.insert(reinterpret_cast<uintptr_t>(*it));
+        }
+      }
+      if (parsed) break;
+    }
+    it++;
+  }
+
+  // Erase any additionally used values from positional args.
+  auto i = flags.positional_args.begin();
+  for (auto j = flags.positional_args.begin(); j != flags.positional_args.end();
+       ++j) {
+    if (handled.contains(reinterpret_cast<uintptr_t>(j->data()))) {
+      continue;
+    }
+    *i++ = *j;
+  }
+  flags.positional_args.erase(i, flags.positional_args.end());
+  return absl::OkStatus();
+}
+
+std::string GlobToRegex(std::string_view glob) {
+  std::string re;
+  re.reserve(glob.size() * 2);
+  re.append("^");
+
+  while (!glob.empty()) {
+    char c = glob[0];
+    glob.remove_prefix(1);
+    switch (c) {
+      case '*': {
+        bool is_star_star = false;
+        while (!glob.empty() && glob[0] == '*') {
+          glob.remove_prefix(1);
+          is_star_star = true;
+        }
+        // TODO: Handle **? / **/, etc.
+        if (is_star_star) {
+          absl::StrAppend(&re, ".*");
+        } else {
+          absl::StrAppend(&re, "[^/]*");
+        }
+        break;
+      }
+      case '?':
+        absl::StrAppend(&re, "[^/]");
+        break;
+      case '[': {
+        if (glob.size() < 2 || glob.find(']', 1) == std::string_view::npos) {
+          // Literal [
+          absl::StrAppend(&re, "\\[");
+          break;
+        }
+        re.push_back('[');
+        bool is_exclude = false;
+        if (glob[0] == '!' || glob[0] == '^') {
+          is_exclude = true;
+          re.push_back('^');
+          re.push_back('/');
+          glob.remove_prefix(1);
+        }
+        // Copy the characters.
+        while (glob[0] != ']') {
+          if (glob[0] == '[' && glob[1] == ':') {
+            // Escape '[' to avoid character classes.
+            absl::StrAppend(&re, "\\[");
+            glob.remove_prefix(1);
+          } else if (glob[1] != '-' || glob[2] == ']') {
+            // Not a range, so copy the character unless it is '/'.
+            if (glob[0] != '/') re.push_back(glob[0]);
+            glob.remove_prefix(1);
+          } else if (!is_exclude && glob[0] <= '/' && '/' <= glob[2]) {
+            // Make sure that the included range does not contain '/'.
+            //
+            // NOTE: "/-/" is dropped entirely, which it should,
+            // because by definition there is no matching pathname.
+            if (glob[0] < '/') {
+              re.push_back(glob[0]);
+              re.push_back('-');
+              re.push_back('/' - 1);
+            }
+            if ('/' < glob[2]) {
+              re.push_back('/' + 1);
+              re.push_back('-');
+              re.push_back(glob[2]);
+            }
+            glob.remove_prefix(3);
+          } else {
+            // Range will not match '/', so copy it blindly
+            re.push_back(glob[0]);
+            re.push_back('-');
+            re.push_back(glob[2]);
+            glob.remove_prefix(3);
+          }
+        }
+        re.push_back(']');
+        glob.remove_prefix(1);
+        break;
+      }
+      case '.':
+      case '+':
+      case '{':
+      case '}':
+      case '(':
+      case ')':
+      case '|':
+      case '^':
+      case '$': {
+        // Escape special characters.
+        re.push_back('\\');
+        re.push_back(c);
+        break;
+      }
+      case '\\':
+        if (glob.empty()) {
+          re.push_back('\\');
+          re.push_back('\\');
+        } else if (!absl::ascii_isalnum(glob[0])) {
+          re.push_back('\\');
+          re.push_back(glob[0]);
+          glob.remove_prefix(1);
+        } else {
+          // ignore.
+        }
+        break;
+
+      default:
+        re.push_back(c);
+        break;
+    }
+  }
+  re.push_back('$');
+  return re;
+}
+
+}  // namespace cli
+}  // namespace tensorstore
diff --git a/tensorstore/tscli/args.h b/tensorstore/tscli/args.h
new file mode 100644
index 00000000..85b4eaaf
--- /dev/null
+++ b/tensorstore/tscli/args.h
@@ -0,0 +1,56 @@
+// Copyright 2024 The TensorStore Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef TENSORSTORE_TSCLI_ARGS_H_
+#define TENSORSTORE_TSCLI_ARGS_H_
+
+#include <string>
+#include <string_view>
+#include <vector>
+
+#include "absl/flags/parse.h"
+#include "absl/functional/function_ref.h"
+#include "absl/status/status.h"
+#include "absl/strings/string_view.h"
+#include "tensorstore/util/span.h"
+
+namespace tensorstore {
+namespace cli {
+
+/// Flag representation after calling absl::ParseAbseilFlagsOnly.
+struct CommandFlags {
+  std::vector<char*> argv;
+  std::vector<absl::UnrecognizedFlag> unrecognized_flags;
+  std::vector<std::string_view> positional_args;
+};
+
+/// Representation of a cli option.
+struct Option {
+  // Long option name. Should begin with a "--" prefix.
+  std::string_view longname;
+  // Parsing function.  An error status indicates an invalid argument.
+  absl::FunctionRef<absl::Status(std::string_view)> parse;
+};
+
+// Try to parse cli options.
+absl::Status TryParseOptions(CommandFlags& flags,
+                             tensorstore::span<Option> options);
+
+// Convert a glob pattern to a regular expression.
+std::string GlobToRegex(std::string_view glob);
+
+}  // namespace cli
+}  // namespace tensorstore
+
+#endif  // TENSORSTORE_TSCLI_ARGS_H_
diff --git a/tensorstore/tscli/args_test.cc b/tensorstore/tscli/args_test.cc
new file mode 100644
index 00000000..3e7c329c
--- /dev/null
+++ b/tensorstore/tscli/args_test.cc
@@ -0,0 +1,37 @@
+// Copyright 2024 The TensorStore Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "tensorstore/tscli/args.h"
+
+#include <gtest/gtest.h>
+
+using ::tensorstore::cli::GlobToRegex;
+
+namespace {
+
+TEST(GlobToRegex, Basic) {
+  EXPECT_EQ(GlobToRegex("a*b"), "^a[^/]*b$");
+  EXPECT_EQ(GlobToRegex("a**b"), "^a.*b$");
+  EXPECT_EQ(GlobToRegex("a?b"), "^a[^/]b$");
+
+  EXPECT_EQ(GlobToRegex("a[b"), "^a\\[b$");
+  EXPECT_EQ(GlobToRegex("a[A-Z]b"), "^a[A-Z]b$");
+  EXPECT_EQ(GlobToRegex("a[!A-Z]b"), "^a[^/A-Z]b$");
+  EXPECT_EQ(GlobToRegex("a[A-]b"), "^a[A-]b$");
+
+  EXPECT_EQ(GlobToRegex("a.+{}()|^$b"), "^a\\.\\+\\{\\}\\(\\)\\|\\^\\$b$");
+  EXPECT_EQ(GlobToRegex("a\\-b\\"), "^a\\-b\\\\$");
+}
+
+}  // namespace
diff --git a/tensorstore/tscli/cli.h b/tensorstore/tscli/cli.h
new file mode 100644
index 00000000..9cf26ae5
--- /dev/null
+++ b/tensorstore/tscli/cli.h
@@ -0,0 +1,37 @@
+// Copyright 2024 The TensorStore Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef TENSORSTORE_TSCLI_CLI_H_
+#define TENSORSTORE_TSCLI_CLI_H_
+
+#include "absl/status/status.h"
+#include "tensorstore/context.h"
+#include "tensorstore/tscli/args.h"
+
+namespace tensorstore {
+namespace cli {
+
+// Kvstore commands
+absl::Status RunKvstoreCopy(Context::Spec context_spec, CommandFlags flags);
+absl::Status RunKvstoreList(Context::Spec context_spec, CommandFlags flags);
+
+// Tensorstore commands
+absl::Status RunTsSearch(Context::Spec context_spec, CommandFlags flags);
+absl::Status RunTsPrintSpec(Context::Spec context_spec, CommandFlags flags);
+absl::Status RunTsPrintStorageStatistics(Context::Spec context_spec,
+                                         CommandFlags flags);
+}  // namespace cli
+}  // namespace tensorstore
+
+#endif  // TENSORSTORE_TSCLI_CLI_H_
diff --git a/tensorstore/tscli/kvstore_copy.cc b/tensorstore/tscli/kvstore_copy.cc
new file mode 100644
index 00000000..36e158f2
--- /dev/null
+++ b/tensorstore/tscli/kvstore_copy.cc
@@ -0,0 +1,144 @@
+// Copyright 2022 The TensorStore Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <iostream>
+#include <optional>
+#include <string>
+#include <string_view>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/synchronization/mutex.h"
+#include "tensorstore/context.h"
+#include "tensorstore/internal/json_binding/std_optional.h"  // IWYU pragma: keep
+#include "tensorstore/kvstore/generation.h"
+#include "tensorstore/kvstore/kvstore.h"
+#include "tensorstore/kvstore/operations.h"
+#include "tensorstore/kvstore/read_result.h"
+#include "tensorstore/kvstore/spec.h"
+#include "tensorstore/tscli/args.h"
+#include "tensorstore/tscli/cli.h"
+#include "tensorstore/util/executor.h"
+#include "tensorstore/util/future.h"
+#include "tensorstore/util/json_absl_flag.h"
+#include "tensorstore/util/quote_string.h"
+#include "tensorstore/util/result.h"
+#include "tensorstore/util/span.h"
+#include "tensorstore/util/status.h"
+
+namespace tensorstore {
+namespace cli {
+
+absl::Status KvstoreCopy(Context::Spec context_spec,
+                         tensorstore::kvstore::Spec source_spec,
+                         tensorstore::kvstore::Spec target_spec) {
+  static absl::Mutex log_mutex;
+  tensorstore::Context context(context_spec);
+
+  TENSORSTORE_ASSIGN_OR_RETURN(auto source,
+                               kvstore::Open(source_spec, context).result());
+  TENSORSTORE_ASSIGN_OR_RETURN(auto target,
+                               kvstore::Open(target_spec, context).result());
+
+  TENSORSTORE_ASSIGN_OR_RETURN(auto list_entries,
+                               kvstore::ListFuture(source).result());
+
+  std::vector<Future<const void>> write_futures;
+
+  for (const auto& entry : list_entries) {
+    std::string key = entry.key;
+    write_futures.push_back(MapFutureValue(
+        InlineExecutor{},
+        [&](const Result<kvstore::ReadResult>& read_result) -> Future<void> {
+          if (!read_result.ok()) {
+            absl::MutexLock lock(&log_mutex);
+            std::cout << "Error reading: " << tensorstore::QuoteString(key)
+                      << ": " << read_result.status() << std::endl;
+          }
+          if (!read_result->has_value()) {
+            return absl::OkStatus();
+          }
+          {
+            absl::MutexLock lock(&log_mutex);
+            std::cout << "Read: " << tensorstore::QuoteString(key) << std::endl;
+          }
+          return MapFuture(
+              InlineExecutor{},
+              [&](const Result<TimestampedStorageGeneration>& stamp)
+                  -> Result<void> {
+                if (!stamp.ok()) {
+                  absl::MutexLock lock(&log_mutex);
+                  std::cout
+                      << "Error writing: " << tensorstore::QuoteString(key)
+                      << ": " << stamp.status() << std::endl;
+                  return stamp.status();
+                }
+                {
+                  absl::MutexLock lock(&log_mutex);
+                  std::cout << "Wrote: " << tensorstore::QuoteString(key)
+                            << std::endl;
+                }
+                return absl::OkStatus();
+              },
+              kvstore::Write(target, key, read_result->value));
+        },
+        kvstore::Read(source, key)));
+  }
+
+  absl::Status status = absl::OkStatus();
+  for (const auto& future : write_futures) {
+    const auto& result = future.result();
+    status.Update(result.status());
+  }
+  return status;
+}
+
+absl::Status RunKvstoreCopy(Context::Spec context_spec, CommandFlags flags) {
+  tensorstore::JsonAbslFlag<std::optional<tensorstore::kvstore::Spec>> source;
+  tensorstore::JsonAbslFlag<std::optional<tensorstore::kvstore::Spec>> target;
+
+  std::vector<Option> options({
+      Option{"--source",
+             [&](std::string_view value) {
+               std::string error;
+               if (!AbslParseFlag(value, &source, &error)) {
+                 return absl::InvalidArgumentError(error);
+               }
+               return absl::OkStatus();
+             }},
+      Option{"--target",
+             [&](std::string_view value) {
+               std::string error;
+               if (!AbslParseFlag(value, &target, &error)) {
+                 return absl::InvalidArgumentError(error);
+               }
+               return absl::OkStatus();
+             }},
+  });
+
+  TENSORSTORE_RETURN_IF_ERROR(TryParseOptions(flags, options));
+
+  if (!source.value) {
+    return absl::InvalidArgumentError("Must specify --source");
+  }
+  if (!target.value) {
+    return absl::InvalidArgumentError("Must specify --target");
+  }
+
+  // TODO: Use positional args as optional keys.
+  return KvstoreCopy(context_spec, *source.value, *target.value);
+}
+
+}  // namespace cli
+}  // namespace tensorstore
diff --git a/tensorstore/tscli/kvstore_list.cc b/tensorstore/tscli/kvstore_list.cc
new file mode 100644
index 00000000..56a6c00d
--- /dev/null
+++ b/tensorstore/tscli/kvstore_list.cc
@@ -0,0 +1,144 @@
+// Copyright 2024 The TensorStore Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <iostream>
+#include <optional>
+#include <string>
+#include <string_view>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/strings/match.h"
+#include "absl/strings/str_cat.h"
+#include "re2/re2.h"
+#include "tensorstore/context.h"
+#include "tensorstore/internal/json_binding/std_optional.h"  // IWYU pragma: keep
+#include "tensorstore/kvstore/key_range.h"
+#include "tensorstore/kvstore/kvstore.h"
+#include "tensorstore/kvstore/operations.h"
+#include "tensorstore/kvstore/spec.h"
+#include "tensorstore/tscli/args.h"
+#include "tensorstore/tscli/cli.h"
+#include "tensorstore/util/json_absl_flag.h"
+#include "tensorstore/util/result.h"
+#include "tensorstore/util/span.h"
+#include "tensorstore/util/status.h"
+
+namespace tensorstore {
+namespace cli {
+namespace {
+
+std::string_view NonGlobPrefix(std::string_view glob) {
+  auto first_glob_char = glob.find_first_of("*?[]");
+  if (first_glob_char == std::string_view::npos) {
+    return glob;
+  }
+  return glob.substr(0, first_glob_char);
+}
+
+}  // namespace
+absl::Status KvstoreList(Context::Spec context_spec,
+                         tensorstore::kvstore::Spec source_spec,
+                         tensorstore::span<std::string_view> args) {
+  tensorstore::Context context(context_spec);
+  TENSORSTORE_ASSIGN_OR_RETURN(auto source,
+                               kvstore::Open(source_spec, context).result());
+
+  // Derive the common prefix and a regex that matches all the args
+  // to minimize the number of list operations on the kvstore.
+  std::string common_prefix;
+  std::string re_string;
+  for (const std::string_view glob : args) {
+    if (glob.empty()) {
+      common_prefix.clear();
+      re_string.clear();
+      break;
+    }
+
+    auto non_glob_prefix = NonGlobPrefix(glob);
+    auto glob_range = non_glob_prefix.empty()
+                          ? KeyRange()
+                          : KeyRange::Prefix(std::string(non_glob_prefix));
+
+    if (re_string.empty()) {
+      common_prefix = std::string(non_glob_prefix);
+      absl::StrAppend(&re_string, "(", GlobToRegex(glob), ")");
+      continue;
+    }
+    if (!common_prefix.empty()) {
+      common_prefix = std::string(
+          absl::FindLongestCommonPrefix(common_prefix, non_glob_prefix));
+    }
+    absl::StrAppend(&re_string, "|(", GlobToRegex(glob), ")");
+  }
+
+  kvstore::ListOptions options;
+  options.range =
+      common_prefix.empty() ? KeyRange() : KeyRange::Prefix(common_prefix);
+  TENSORSTORE_ASSIGN_OR_RETURN(auto list_entries,
+                               kvstore::ListFuture(source, options).result());
+
+  RE2 re2(re_string);
+  for (const auto& entry : list_entries) {
+    if (re_string.empty() || RE2::FullMatch(entry.key, re2)) {
+      std::cout << entry.key << std::endl;
+    }
+  }
+  return absl::OkStatus();
+}
+
+absl::Status RunKvstoreList(Context::Spec context_spec, CommandFlags flags) {
+  tensorstore::JsonAbslFlag<std::optional<tensorstore::kvstore::Spec>> source;
+  std::vector<Option> options({
+      Option{"--source",
+             [&](std::string_view value) {
+               std::string error;
+               if (!AbslParseFlag(value, &source, &error)) {
+                 return absl::InvalidArgumentError(error);
+               }
+               return absl::OkStatus();
+             }},
+  });
+
+  TENSORSTORE_RETURN_IF_ERROR(TryParseOptions(flags, options));
+
+  if (source.value) {
+    if (flags.positional_args.empty()) flags.positional_args.push_back("");
+    return KvstoreList(context_spec, *source.value, flags.positional_args);
+  } else if (flags.positional_args.empty()) {
+    return absl::InvalidArgumentError(
+        "list: Must include --source or a sequence of specs");
+  }
+
+  absl::Status status;
+  std::vector<std::string_view> args({""});
+  for (const auto spec : flags.positional_args) {
+    auto from_url = kvstore::Spec::FromUrl(spec);
+    if (from_url.ok()) {
+      status.Update(KvstoreList(context_spec, from_url.value(), args));
+      continue;
+    }
+    tensorstore::JsonAbslFlag<tensorstore::kvstore::Spec> arg_spec;
+    std::string error;
+    if (AbslParseFlag(spec, &arg_spec, &error)) {
+      status.Update(KvstoreList(context_spec, arg_spec.value, args));
+      continue;
+    }
+    std::cerr << "Invalid spec: " << spec << ": " << error << std::endl;
+  }
+  return status;
+}
+
+}  // namespace cli
+}  // namespace tensorstore
diff --git a/tensorstore/tscli/main.cc b/tensorstore/tscli/main.cc
new file mode 100644
index 00000000..9885d38f
--- /dev/null
+++ b/tensorstore/tscli/main.cc
@@ -0,0 +1,145 @@
+// Copyright 2024 The TensorStore Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <algorithm>
+#include <iostream>
+#include <optional>
+#include <string>
+#include <string_view>
+#include <utility>
+#include <vector>
+
+#include "absl/base/log_severity.h"  // IWYU pragma: keep
+#include "absl/flags/flag.h"
+#include "absl/flags/parse.h"
+#include "absl/log/globals.h"     // IWYU pragma: keep
+#include "absl/log/initialize.h"  // IWYU pragma: keep
+#include "absl/status/status.h"
+#include "tensorstore/context.h"
+#include "absl/flags/parse.h"
+#include "tensorstore/internal/metrics/collect.h"
+#include "tensorstore/internal/metrics/registry.h"
+#include "tensorstore/tscli/args.h"
+#include "tensorstore/tscli/cli.h"
+#include "tensorstore/util/json_absl_flag.h"
+
+using ::tensorstore::internal_metrics::FormatCollectedMetric;
+using ::tensorstore::internal_metrics::GetMetricRegistry;
+
+ABSL_FLAG(std::optional<std::string>, metrics_prefix, std::nullopt,
+          "Prefix for metrics to output.");
+
+ABSL_FLAG(tensorstore::JsonAbslFlag<tensorstore::Context::Spec>, context_spec,
+          {},
+          "Context spec for writing data.  This can be used to control the "
+          "number of concurrent write operations of the underlying key-value "
+          "store.  See examples at the start of the source file.");
+
+namespace {
+
+const char kUsage[] =
+    "Usage:\n"
+    "  tscli [global options..] <command> [command args]\n\n"
+    "Commands:\n"
+    "  copy --source <source-kvstore-spec> --target <target-kvstore-spec>\n"
+    "  list --source <source-kvstore-spec> [glob...] \n"
+    "  list [spec...]\n"
+    "  print_spec --spec <tensorstore-spec>\n"
+    "  print_stats --spec <tensorstore-spec> [box...]\n"
+    "  search --source <source-kvstore-spec>\n"
+    "  search [spec...]\n"
+    "  help\n";
+
+void PrintShortHelp(std::string_view program) {
+  std::cerr << program << "\n" << kUsage << std::endl;
+}
+
+void PrintHelp(std::string_view program) {
+  std::cerr << program << "\n" << kUsage << std::endl;
+}
+
+void DumpMetrics(std::string_view prefix) {
+  std::vector<std::string> lines;
+  for (const auto& metric : GetMetricRegistry().CollectWithPrefix(prefix)) {
+    FormatCollectedMetric(metric, [&lines](bool has_value, std::string line) {
+      if (has_value) lines.emplace_back(std::move(line));
+    });
+  }
+
+  // `lines` is unordered, which isn't great for benchmark comparison.
+  std::sort(std::begin(lines), std::end(lines));
+  std::cerr << std::endl;
+  for (const auto& l : lines) {
+    std::cerr << l << std::endl;
+  }
+  std::cerr << std::endl;
+}
+
+}  // namespace
+
+int main(int argc, char** argv) {
+  absl::InitializeLog();
+  absl::SetStderrThreshold(absl::LogSeverityAtLeast::kInfo);
+
+  ::tensorstore::cli::CommandFlags flags;
+  flags.argv = std::vector<char*>(argv, argv + argc);
+
+  std::vector<char*> positional_args;
+  absl::ParseAbseilFlagsOnly(argc, argv, positional_args,
+                             flags.unrecognized_flags);
+
+  if (positional_args.size() < 2) {
+    PrintShortHelp(argv[0]);
+    return 1;
+  }
+
+  std::string_view command = positional_args[1];
+  flags.positional_args.assign(positional_args.begin() + 2,
+                               positional_args.end());
+
+  if (command == "help") {
+    PrintHelp(argv[0]);
+    return 0;
+  }
+
+  absl::Status status;
+  if (command == "copy") {
+    status = ::tensorstore::cli::RunKvstoreCopy(
+        absl::GetFlag(FLAGS_context_spec).value, std::move(flags));
+  } else if (command == "ls" || command == "list") {
+    status = ::tensorstore::cli::RunKvstoreList(
+        absl::GetFlag(FLAGS_context_spec).value, std::move(flags));
+  } else if (command == "search") {
+    status = ::tensorstore::cli::RunTsSearch(
+        absl::GetFlag(FLAGS_context_spec).value, std::move(flags));
+  } else if (command == "print_spec") {
+    status = ::tensorstore::cli::RunTsPrintSpec(
+        absl::GetFlag(FLAGS_context_spec).value, std::move(flags));
+  } else if (command == "print_stats") {
+    status = ::tensorstore::cli::RunTsPrintStorageStatistics(
+        absl::GetFlag(FLAGS_context_spec).value, std::move(flags));
+  }
+
+  if (absl::GetFlag(FLAGS_metrics_prefix).has_value()) {
+    DumpMetrics(*absl::GetFlag(FLAGS_metrics_prefix));
+  }
+
+  if (status.ok()) return 0;
+
+  std::cerr << status << std::endl;
+  if (absl::IsInvalidArgument(status)) {
+    PrintHelp(argv[0]);
+  }
+  return 1;
+}
diff --git a/tensorstore/tscli/ts_print_spec.cc b/tensorstore/tscli/ts_print_spec.cc
new file mode 100644
index 00000000..b3f3e1b2
--- /dev/null
+++ b/tensorstore/tscli/ts_print_spec.cc
@@ -0,0 +1,75 @@
+// Copyright 2024 The TensorStore Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <iostream>
+#include <optional>
+#include <string>
+#include <string_view>
+#include <vector>
+
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "tensorstore/context.h"
+#include "tensorstore/internal/json_binding/json_binding.h"  // IWYU pragma: keep
+#include "tensorstore/internal/json_binding/std_optional.h"  // IWYU pragma: keep
+#include "tensorstore/open.h"
+#include "tensorstore/open_mode.h"
+#include "tensorstore/spec.h"
+#include "tensorstore/tensorstore.h"
+#include "tensorstore/tscli/args.h"
+#include "tensorstore/tscli/cli.h"
+#include "tensorstore/util/json_absl_flag.h"
+#include "tensorstore/util/result.h"
+#include "tensorstore/util/span.h"
+#include "tensorstore/util/status.h"
+
+namespace tensorstore {
+namespace cli {
+
+absl::Status TsPrintSpec(Context::Spec context_spec, tensorstore::Spec spec) {
+  tensorstore::Context context(context_spec);
+  TENSORSTORE_ASSIGN_OR_RETURN(
+      auto ts,
+      tensorstore::Open(spec, context, tensorstore::ReadWriteMode::read,
+                        tensorstore::OpenMode::open)
+          .result());
+  TENSORSTORE_ASSIGN_OR_RETURN(auto actual_spec, ts.spec());
+  TENSORSTORE_ASSIGN_OR_RETURN(auto json_spec, actual_spec.ToJson());
+  std::cout << json_spec.dump() << std::endl;
+  return absl::OkStatus();
+}
+
+absl::Status RunTsPrintSpec(Context::Spec context_spec, CommandFlags flags) {
+  tensorstore::JsonAbslFlag<std::optional<tensorstore::Spec>> spec;
+  std::vector<Option> options({
+      Option{"--spec",
+             [&](std::string_view value) {
+               std::string error;
+               if (!AbslParseFlag(value, &spec, &error)) {
+                 return absl::InvalidArgumentError(error);
+               }
+               return absl::OkStatus();
+             }},
+  });
+
+  TENSORSTORE_RETURN_IF_ERROR(TryParseOptions(flags, options));
+
+  if (!spec.value) {
+    return absl::InvalidArgumentError("print_spec: Must include --spec");
+  }
+  return TsPrintSpec(context_spec, *spec.value);
+}
+
+}  // namespace cli
+}  // namespace tensorstore
diff --git a/tensorstore/tscli/ts_print_stats.cc b/tensorstore/tscli/ts_print_stats.cc
new file mode 100644
index 00000000..1b68ccf5
--- /dev/null
+++ b/tensorstore/tscli/ts_print_stats.cc
@@ -0,0 +1,108 @@
+// Copyright 2024 The TensorStore Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <iostream>
+#include <optional>
+#include <string>
+#include <string_view>
+#include <vector>
+
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "tensorstore/array_storage_statistics.h"
+#include "tensorstore/box.h"
+#include "tensorstore/context.h"
+#include "tensorstore/index_space/dim_expression.h"
+#include "tensorstore/internal/json_binding/box.h"  // IWYU pragma: keep
+#include "tensorstore/internal/json_binding/json_binding.h"  // IWYU pragma: keep
+#include "tensorstore/internal/json_binding/std_optional.h"  // IWYU pragma: keep
+#include "tensorstore/open.h"
+#include "tensorstore/open_mode.h"
+#include "tensorstore/spec.h"
+#include "tensorstore/tensorstore.h"
+#include "tensorstore/tscli/args.h"
+#include "tensorstore/tscli/cli.h"
+#include "tensorstore/util/json_absl_flag.h"
+#include "tensorstore/util/result.h"
+#include "tensorstore/util/span.h"
+#include "tensorstore/util/status.h"
+
+namespace tensorstore {
+namespace cli {
+
+absl::Status TsPrintStorageStatistics(
+    Context::Spec context_spec, tensorstore::Spec spec,
+    tensorstore::span<std::string_view> args) {
+  tensorstore::Context context(context_spec);
+  TENSORSTORE_ASSIGN_OR_RETURN(
+      auto ts,
+      tensorstore::Open(spec, context, tensorstore::ReadWriteMode::read,
+                        tensorstore::OpenMode::open)
+          .result());
+
+  for (const auto& arg : args) {
+    tensorstore::JsonAbslFlag<Box<>> box_flag;
+    std::string error;
+    if (!AbslParseFlag(arg, &box_flag, &error)) {
+      std::cerr << "Invalid box: " << arg << ": " << error << std::endl;
+      continue;
+    }
+    auto future = GetStorageStatistics(
+        ts | tensorstore::AllDims().BoxSlice(box_flag.value),
+        tensorstore::ArrayStorageStatistics::query_not_stored,
+        tensorstore::ArrayStorageStatistics::query_fully_stored);
+    future.Wait();
+    if (!future.status().ok()) {
+      std::cerr << "Error getting storage statistics for " << arg << ": "
+                << future.status() << std::endl;
+      continue;
+    }
+    const char* txt = "partial";
+    if (future.value().not_stored) {
+      txt = "not_stored";
+    } else if (future.value().fully_stored) {
+      txt = "fully_stored";
+    }
+    std::cout << "Box: " << arg << " " << txt << std::endl;
+  }
+
+  return absl::OkStatus();
+}
+
+absl::Status RunTsPrintStorageStatistics(Context::Spec context_spec,
+                                         CommandFlags flags) {
+  tensorstore::JsonAbslFlag<std::optional<tensorstore::Spec>> spec;
+  std::vector<Option> options({
+      Option{"--spec",
+             [&](std::string_view value) {
+               std::string error;
+               if (!AbslParseFlag(value, &spec, &error)) {
+                 return absl::InvalidArgumentError(error);
+               }
+               return absl::OkStatus();
+             }},
+  });
+
+  TENSORSTORE_RETURN_IF_ERROR(TryParseOptions(flags, options));
+
+  if (!spec.value) {
+    return absl::InvalidArgumentError("print_spec: Must include --spec");
+  }
+
+  return TsPrintStorageStatistics(context_spec, *spec.value,
+                                  flags.positional_args);
+}
+
+}  // namespace cli
+}  // namespace tensorstore
diff --git a/tensorstore/tscli/ts_search.cc b/tensorstore/tscli/ts_search.cc
new file mode 100644
index 00000000..67b815cf
--- /dev/null
+++ b/tensorstore/tscli/ts_search.cc
@@ -0,0 +1,175 @@
+// Copyright 2024 The TensorStore Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <iostream>
+#include <optional>
+#include <string>
+#include <string_view>
+#include <utility>
+#include <vector>
+
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "absl/strings/match.h"
+#include <nlohmann/json_fwd.hpp>
+#include <nlohmann/json.hpp>
+#include "tensorstore/context.h"
+#include "tensorstore/internal/json_binding/box.h"  // IWYU pragma: keep
+#include "tensorstore/internal/json_binding/json_binding.h"  // IWYU pragma: keep
+#include "tensorstore/internal/json_binding/std_optional.h"  // IWYU pragma: keep
+#include "tensorstore/internal/path.h"
+#include "tensorstore/kvstore/kvstore.h"
+#include "tensorstore/kvstore/operations.h"
+#include "tensorstore/kvstore/spec.h"
+#include "tensorstore/open.h"
+#include "tensorstore/open_mode.h"
+#include "tensorstore/spec.h"
+#include "tensorstore/tscli/args.h"
+#include "tensorstore/tscli/cli.h"
+#include "tensorstore/util/json_absl_flag.h"
+#include "tensorstore/util/result.h"
+#include "tensorstore/util/span.h"
+#include "tensorstore/util/status.h"
+
+namespace tensorstore {
+namespace cli {
+namespace {
+
+absl::Status TryOpenTensorstore(Context context, kvstore::Spec base_spec,
+                                std::string_view driver_name,
+                                std::string_view key) {
+  auto [dirname, basename] = internal::PathDirnameBasename(key);
+  if (!dirname.empty()) {
+    base_spec.AppendPathComponent(dirname);
+  }
+  if (!absl::EndsWith(base_spec.path, "/")) {
+    base_spec.AppendSuffix("/");
+  }
+
+  // Adjust the base_spec to be an ocdbt spec.
+  TENSORSTORE_ASSIGN_OR_RETURN(auto base_json, base_spec.ToJson());
+  ::nlohmann::json json{
+      {"driver", std::string(driver_name)},
+      {"kvstore", std::move(base_json)},
+  };
+  TENSORSTORE_ASSIGN_OR_RETURN(auto parsed_spec,
+                               tensorstore::Spec::FromJson(json));
+  TENSORSTORE_ASSIGN_OR_RETURN(
+      auto ts,
+      tensorstore::Open(parsed_spec, context, tensorstore::ReadWriteMode::read,
+                        tensorstore::OpenMode::open)
+          .result());
+
+  std::cout << json.dump() << std::endl;
+
+  return absl::OkStatus();
+}
+
+}  // namespace
+
+absl::Status TsSearch(Context::Spec context_spec,
+                      tensorstore::kvstore::Spec source_spec) {
+  tensorstore::Context context(context_spec);
+
+  TENSORSTORE_ASSIGN_OR_RETURN(auto source,
+                               kvstore::Open(source_spec, context).result());
+
+  TENSORSTORE_ASSIGN_OR_RETURN(auto list_entries,
+                               kvstore::ListFuture(source, {}).result());
+
+  for (const auto& entry : list_entries) {
+    auto [dirname, basename] = internal::PathDirnameBasename(entry.key);
+
+    if (basename == "info.json") {
+      TryOpenTensorstore(context, source_spec, "neuroglancer_precomputed",
+                         entry.key)
+          .IgnoreError();
+    } else if (basename == "metadata.json") {
+      TryOpenTensorstore(context, source_spec, "n5", entry.key).IgnoreError();
+    } else if (basename == "zarr.json") {
+      TryOpenTensorstore(context, source_spec, "zarr3", entry.key)
+          .IgnoreError();
+    } else if (absl::EndsWith(basename, "zarray")) {
+      TryOpenTensorstore(context, source_spec, "zarr", entry.key).IgnoreError();
+    } else if (absl::StrContains(basename, "manifest.ocdbt")) {
+      // Found an ocdbt manifest file. try to recurse into it.
+      kvstore::Spec base_spec = source_spec;
+      base_spec.AppendPathComponent(dirname);
+      if (!absl::EndsWith(base_spec.path, "/")) {
+        base_spec.AppendSuffix("/");
+      }
+      auto base_json = base_spec.ToJson();
+      if (!base_json.ok()) continue;
+      if (auto* obj = base_json->get_ptr<nlohmann::json::object_t*>();
+          obj != nullptr) {
+        if (auto it = obj->find("driver");
+            it != obj->end() && it->second.get<std::string>() == "ocdbt") {
+          // Already an ocdbt driver;
+          continue;
+        }
+      }
+      auto ocdbt_spec = kvstore::Spec::FromJson({
+          {"driver", "ocdbt"},
+          {"base", std::move(base_json.value())},
+      });
+      if (!ocdbt_spec.ok()) continue;
+      TsSearch(context_spec, ocdbt_spec.value()).IgnoreError();
+    }
+  }
+
+  return absl::OkStatus();
+}
+
+absl::Status RunTsSearch(Context::Spec context_spec, CommandFlags flags) {
+  tensorstore::JsonAbslFlag<std::optional<tensorstore::kvstore::Spec>> source;
+  std::vector<Option> options({
+      Option{"--source",
+             [&](std::string_view value) {
+               std::string error;
+               if (!AbslParseFlag(value, &source, &error)) {
+                 return absl::InvalidArgumentError(error);
+               }
+               return absl::OkStatus();
+             }},
+  });
+
+  TENSORSTORE_RETURN_IF_ERROR(TryParseOptions(flags, options));
+
+  if (source.value) {
+    return TsSearch(context_spec, *source.value);
+  }
+  if (flags.positional_args.empty()) {
+    return absl::InvalidArgumentError(
+        "search: Must include --source or a sequence of specs");
+  }
+  absl::Status status;
+  for (const std::string_view spec : flags.positional_args) {
+    auto from_url = kvstore::Spec::FromUrl(spec);
+    if (from_url.ok()) {
+      status.Update(TsSearch(context_spec, from_url.value()));
+      continue;
+    }
+    tensorstore::JsonAbslFlag<tensorstore::kvstore::Spec> arg_spec;
+    std::string error;
+    if (AbslParseFlag(spec, &arg_spec, &error)) {
+      status.Update(TsSearch(context_spec, arg_spec.value));
+      continue;
+    }
+    std::cerr << "Invalid spec: " << spec << ": " << error << std::endl;
+  }
+  return status;
+}
+
+}  // namespace cli
+}  // namespace tensorstore