Execution API: support fetching a single artifact to memory

6 files changed, 50 insertions, 0 deletions

diff --git a/src/buildtool/execution_api/common/execution_api.hpp b/src/buildtool/execution_api/common/execution_api.hpp
index 7a301f71..a8ba8009 100644
--- a/src/buildtool/execution_api/common/execution_api.hpp
+++ b/src/buildtool/execution_api/common/execution_api.hpp
@@ -17,6 +17,7 @@
 
 #include <map>
 #include <memory>
+#include <optional>
 #include <string>
 #include <utility>
 #include <vector>
@@ -86,6 +87,12 @@ class IExecutionApi {
         std::vector<Artifact::ObjectInfo> const& artifacts_info,
         gsl::not_null<IExecutionApi*> const& api) noexcept -> bool = 0;
 
+    /// \brief Retrieve one artifact from CAS and make it available for
+    /// furter in-memory processing
+    [[nodiscard]] virtual auto RetrieveToMemory(
+        Artifact::ObjectInfo const& artifact_info)
+        -> std::optional<std::string> = 0;
+
     /// \brief Upload blobs to CAS. Uploads only the blobs that are not yet
     /// available in CAS, unless `skip_find_missing` is specified.
     /// \param blobs                Container of blobs to upload.
diff --git a/src/buildtool/execution_api/git/git_api.hpp b/src/buildtool/execution_api/git/git_api.hpp
index 69179fe1..5b552220 100644
--- a/src/buildtool/execution_api/git/git_api.hpp
+++ b/src/buildtool/execution_api/git/git_api.hpp
@@ -168,6 +168,13 @@ class GitApi final : public IExecutionApi {
         return false;
     }
 
+    [[nodiscard]] auto RetrieveToMemory(
+        Artifact::ObjectInfo const& artifact_info)
+        -> std::optional<std::string> override {
+        return RepositoryConfig::Instance().ReadBlobFromGitCAS(
+            artifact_info.digest.hash());
+    }
+
     /// NOLINTNEXTLINE(google-default-arguments)
     [[nodiscard]] auto Upload(BlobContainer const& /*blobs*/,
                               bool /*skip_find_missing*/ = false) noexcept
diff --git a/src/buildtool/execution_api/local/TARGETS b/src/buildtool/execution_api/local/TARGETS
index b9e056b8..1ea94d69 100644
--- a/src/buildtool/execution_api/local/TARGETS
+++ b/src/buildtool/execution_api/local/TARGETS
@@ -43,6 +43,7 @@
     , "config"
     , ["src/buildtool/common", "common"]
     , ["src/buildtool/common", "bazel_types"]
+    , ["src/buildtool/file_system", "file_system_manager"]
     ]
   }
 }
diff --git a/src/buildtool/execution_api/local/local_api.hpp b/src/buildtool/execution_api/local/local_api.hpp
index d699ad4f..e0de1598 100644
--- a/src/buildtool/execution_api/local/local_api.hpp
+++ b/src/buildtool/execution_api/local/local_api.hpp
@@ -31,6 +31,7 @@
 #include "src/buildtool/execution_api/common/execution_api.hpp"
 #include "src/buildtool/execution_api/git/git_api.hpp"
 #include "src/buildtool/execution_api/local/local_action.hpp"
+#include "src/buildtool/file_system/file_system_manager.hpp"
 #include "src/buildtool/logging/logger.hpp"
 #include "src/buildtool/storage/storage.hpp"
 
@@ -237,6 +238,27 @@ class LocalApi final : public IExecutionApi {
         return api->Upload(container, /*skip_find_missing=*/true);
     }
 
+    [[nodiscard]] auto RetrieveToMemory(
+        Artifact::ObjectInfo const& artifact_info)
+        -> std::optional<std::string> override {
+        std::optional<std::filesystem::path> location{};
+        if (IsTreeObject(artifact_info.type)) {
+            location = storage_->CAS().TreePath(artifact_info.digest);
+        }
+        else {
+            location = storage_->CAS().BlobPath(
+                artifact_info.digest, IsExecutableObject(artifact_info.type));
+        }
+        if (not location) {
+            return std::nullopt;
+        }
+        auto const content = FileSystemManager::ReadFile(*location);
+        if (not content) {
+            return std::nullopt;
+        }
+        return *content;
+    }
+
     [[nodiscard]] auto Upload(BlobContainer const& blobs,
                               bool /*skip_find_missing*/) noexcept
         -> bool final {
diff --git a/src/buildtool/execution_api/remote/bazel/bazel_api.cpp b/src/buildtool/execution_api/remote/bazel/bazel_api.cpp
index c4f170a3..e2da0188 100644
--- a/src/buildtool/execution_api/remote/bazel/bazel_api.cpp
+++ b/src/buildtool/execution_api/remote/bazel/bazel_api.cpp
@@ -251,6 +251,15 @@ auto BazelApi::CreateAction(
     return api->Upload(container, /*skip_find_missing=*/true);
 }
 
+[[nodiscard]] auto BazelApi::RetrieveToMemory(
+    Artifact::ObjectInfo const& artifact_info) -> std::optional<std::string> {
+    auto blobs = network_->ReadBlobs({artifact_info.digest}).Next();
+    if (blobs.size() == 1) {
+        return blobs.at(0).data;
+    }
+    return std::nullopt;
+}
+
 [[nodiscard]] auto BazelApi::Upload(BlobContainer const& blobs,
                                     bool skip_find_missing) noexcept -> bool {
     return network_->UploadBlobs(blobs, skip_find_missing);
diff --git a/src/buildtool/execution_api/remote/bazel/bazel_api.hpp b/src/buildtool/execution_api/remote/bazel/bazel_api.hpp
index 6edfae5f..310fb347 100644
--- a/src/buildtool/execution_api/remote/bazel/bazel_api.hpp
+++ b/src/buildtool/execution_api/remote/bazel/bazel_api.hpp
@@ -81,6 +81,10 @@ class BazelApi final : public IExecutionApi {
     [[nodiscard]] auto IsAvailable(std::vector<ArtifactDigest> const& digests)
         const noexcept -> std::vector<ArtifactDigest> final;
 
+    [[nodiscard]] auto RetrieveToMemory(
+        Artifact::ObjectInfo const& artifact_info)
+        -> std::optional<std::string> final;
+
   private:
     std::shared_ptr<BazelNetwork> network_;