Add support for Amazon S3 bucket storage backend.

This adds support for using Amazon's S3 bucket as a storage backend via
the `AwsStore` module. It is only implemented for the `Lwt` concurrency
library at this time since the underlying AWS-S3 library does not yet
have support for `Eio`.

.github/workflows/build-and-test.yml

@@ -45,7 +45,7 @@ jobs:
         run: |
           opam install --deps-only --with-test --with-doc --yes zarr
           opam install bytesrw conf-zlib conf-zstd --yes
-          opam install lwt --yes
+          opam install lwt aws-s3-lwt --yes
           opam exec -- dune build zarr zarr-sync zarr-lwt
       
       - name: setup ocaml-5-specific

dune-project

@@ -63,6 +63,7 @@
      (and (>= 4.14.0)))
    (zarr (= :version))
    (lwt (>= 2.5.1))
+   (aws-s3-lwt (>= 4.8.1))
    (odoc :with-doc)
    (ounit2 :with-test)
    (ppx_deriving :with-test)

zarr-lwt.opam

@@ -13,6 +13,7 @@ depends: [
   "ocaml" {>= "4.14.0"}
   "zarr" {= version}
   "lwt" {>= "2.5.1"}
+  "aws-s3-lwt" {>= "4.8.1"}
   "odoc" {with-doc}
   "ounit2" {with-test}
   "ppx_deriving" {with-test}

zarr-lwt/src/dune

@@ -3,6 +3,7 @@
  (public_name zarr-lwt)
  (libraries
    zarr
+   aws-s3-lwt
    lwt
    lwt.unix)
  (ocamlopt_flags

zarr-lwt/src/storage.ml

@@ -144,3 +144,160 @@ module FilesystemStore = struct
 
   include Zarr.Storage.Make(IO)
 end
+
+module AwsStore = struct
+  open Deferred.Syntax
+
+  let fold_result :
+    ok:('a -> 'b) -> not_found:'b -> ('a, Aws_s3_lwt.S3.error) result -> 'b
+    = fun ~ok ~not_found -> function
+    | Ok v -> ok v
+    | Error Not_found -> not_found
+    | Error Failed e -> raise e
+    | Error Throttled -> failwith "Throttled"
+    | Error Forbidden -> failwith "Forbidden"
+    | Error Unknown (_, s) -> failwith @@ Printf.sprintf "Unknown: %s" s
+    | Error Redirect _ -> failwith "Redirect error"
+
+  module IO = struct
+    module Deferred = Deferred
+
+    type t =
+      {bucket : string
+      ;cred : Aws_s3_lwt.Credentials.t
+      ;endpoint : Aws_s3.Region.endpoint }
+
+    let size t key =
+      let content_size : Aws_s3_lwt.S3.content -> int = fun c -> c.size in
+      let* res =
+        Aws_s3_lwt.S3.head
+          ~bucket:t.bucket ~credentials:t.cred ~endpoint:t.endpoint ~key () in
+      Deferred.return (fold_result ~ok:content_size ~not_found:0 res)
+
+    let get t key =
+      let* res = Aws_s3_lwt.S3.get
+        ~bucket:t.bucket ~credentials:t.cred ~endpoint:t.endpoint ~key () in
+      let v = fold_result ~ok:Fun.id ~not_found:(raise @@ Zarr.Storage.Key_not_found key) res in
+      Deferred.return v 
+
+    let get_partial_values t key ranges =
+      let read_range ~t ~key acc (ofs, len) =
+        let range : Aws_s3_lwt.S3.range = match len with
+          | None -> {first = Some ofs; last = None}
+          | l -> {first = Some ofs; last = l}
+        in
+        let+ res = Aws_s3_lwt.S3.get
+          ~bucket:t.bucket ~credentials:t.cred ~endpoint:t.endpoint ~range ~key () in
+        let a = fold_result ~ok:Fun.id ~not_found:(raise @@ Zarr.Storage.Key_not_found key) res in
+        a :: acc
+      in
+      Deferred.fold_left (read_range ~t ~key) [] ranges
+
+    let set t key data =
+      let* res = Aws_s3_lwt.S3.put
+        ~bucket:t.bucket ~credentials:t.cred ~endpoint:t.endpoint ~key ~data () in
+      let _ = fold_result ~ok:Fun.id ~not_found:String.empty res in
+      Deferred.return_unit
+
+    let set_partial_values t key ?(append=false) rsv =
+      let* size = size t key in
+      let* ov = match size with
+        | 0 -> Deferred.return String.empty
+        | _ -> get t key
+      in
+      let f = if append || ov = String.empty then
+        fun acc (_, v) -> acc ^ v else
+        fun acc (rs, v) ->
+          let s = Bytes.unsafe_of_string acc in
+          Bytes.blit_string v 0 s rs String.(length v);
+          Bytes.unsafe_to_string s
+      in
+      let ov' = List.fold_left f ov rsv in
+      set t key ov'
+
+    let erase t key =
+      let+ res = Aws_s3_lwt.S3.delete
+        ~bucket:t.bucket ~credentials:t.cred ~endpoint:t.endpoint ~key () in
+      fold_result ~ok:Fun.id ~not_found:() res
+
+    (*TODO: Use delete_multi function instead. *)
+    let erase_prefix t prefix =
+      let content_key : Aws_s3_lwt.S3.content -> string = fun c -> c.key in
+      let rec delete_keys : Aws_s3_lwt.S3.content list * Aws_s3_lwt.S3.Ls.cont -> unit Deferred.t = function
+        | xs, Done -> Deferred.iter (erase t) (List.map content_key xs)
+        | xs, More k ->
+          let* res = k () in
+          let cont = fold_result ~ok:Fun.id ~not_found:(raise @@ Zarr.Storage.Key_not_found prefix) res in
+          let* () = Deferred.iter (erase t) (List.map content_key xs) in
+          delete_keys cont
+      in
+      let* res = Aws_s3_lwt.S3.ls
+        ~bucket:t.bucket ~credentials:t.cred ~endpoint:t.endpoint ~prefix () in
+      let keys = fold_result ~ok:Fun.id ~not_found:(raise @@ Zarr.Storage.Key_not_found prefix) res in
+      delete_keys keys
+
+    let list t =
+      let content_key : Aws_s3_lwt.S3.content -> string = fun c -> c.key in
+      let rec accumulate_keys :
+        string list -> Aws_s3_lwt.S3.content list * Aws_s3_lwt.S3.Ls.cont -> string list Deferred.t
+        = fun acc -> function
+        | xs, Done -> Deferred.return @@ (acc @ List.map content_key xs)
+        | xs, More k ->
+          let* res = k () in
+          let cont = fold_result ~ok:Fun.id ~not_found:([], Aws_s3_lwt.S3.Ls.Done) res in
+          let acc' = acc @ List.map content_key xs in
+          accumulate_keys acc' cont
+      in
+      let* res = Aws_s3_lwt.S3.ls
+        ~bucket:t.bucket ~credentials:t.cred ~endpoint:t.endpoint () in
+      let keys = fold_result ~ok:Fun.id ~not_found:([], Aws_s3_lwt.S3.Ls.Done) res in
+      accumulate_keys [] keys
+
+    let list_dir t prefix =
+      let module S = Set.Make(String) in
+      let add ~prefix (l, r) (c : Aws_s3_lwt.S3.content) =
+        let size = String.length prefix in
+        if not (String.contains_from c.key size '/') then c.key :: l, r else
+        l, S.add String.(sub c.key 0 @@ 1 + index_from c.key size '/') r
+      in
+      let rec accumulate_keys :
+        string list * S.t -> Aws_s3_lwt.S3.content list * Aws_s3_lwt.S3.Ls.cont -> (string list * S.t) Deferred.t
+        = fun acc -> function
+        | xs, Done -> Deferred.return @@ List.fold_left (add ~prefix) acc xs
+        | xs, More k ->
+          let* res = k () in
+          let cont = fold_result ~ok:Fun.id ~not_found:([], Aws_s3_lwt.S3.Ls.Done) res in
+          let acc' = List.fold_left (add ~prefix) acc xs in
+          accumulate_keys acc' cont
+      in
+      let* res = Aws_s3_lwt.S3.ls
+        ~bucket:t.bucket ~credentials:t.cred ~endpoint:t.endpoint ~prefix () in
+      let keys = fold_result ~ok:Fun.id ~not_found:([], Aws_s3_lwt.S3.Ls.Done) res in
+      let+ keys, prefixes = accumulate_keys ([], S.empty) keys in
+      keys, S.elements prefixes
+
+    let is_member t key =
+      let+ size = size t key in
+      if size = 0 then false else true
+
+    let rename t prefix new_prefix =
+      let add ~t ~prefix ~new_prefix acc k =
+        let l = String.length prefix in
+        let k' = new_prefix ^ String.sub k l (String.length k - l) in
+        let+ a = get t k in (k', a) :: acc
+      in
+      let* xs = list t in
+      let to_delete = List.filter (String.starts_with ~prefix) xs in
+      let* data = Deferred.fold_left (add ~t ~prefix ~new_prefix) [] to_delete in
+      let* () = Deferred.iter (fun (k, v) -> set t k v) data in
+      Deferred.iter (erase t) to_delete
+  end
+
+  let with_open ?(scheme=`Https) ?(inet=`V4) ~region ~bucket ~profile f =
+    let* res = Aws_s3_lwt.Credentials.Helper.get_credentials ~profile () in
+    let cred = Result.fold ~ok:Fun.id ~error:(fun e -> raise e) res in
+    let endpoint = Aws_s3.Region.endpoint ~inet ~scheme region in
+    f IO.{bucket; cred; endpoint}
+
+  include Zarr.Storage.Make(IO)
+end

zarr-lwt/src/storage.mli

@@ -18,3 +18,17 @@ module FilesystemStore : sig
 
       @raise Failure if [dir] is not a Zarr store path. *)
 end
+
+(** An Lwt-aware Amazon S3 bucket storage backend for a Zarr V3 hierarchy. *)
+module AwsStore : sig
+  include Zarr.Storage.STORE with type 'a Deferred.t = 'a Lwt.t
+
+  val with_open :
+    ?scheme:[ `Http | `Https ] ->
+    ?inet:[ `V4 | `V6 ] ->
+    region:Aws_s3.Region.t ->
+    bucket:string ->
+    profile:string ->
+    (t -> 'a Lwt.t) ->
+    'a Lwt.t
+end