Source file databases_core.ml

(*
 Azure cosmos database documentation: https://docs.microsoft.com/en-us/rest/api/cosmos-db/
*)

module type Account = sig
  type resource = Dbs | Colls | Docs | Users | Permissions | Offers | Pkranges

  val authorization :
    Utilities.Verb.t -> resource -> Utilities.Ms_time.t -> string -> string

  val endpoint : string
end

module Auth_credential (C : Databases_intf.Credentials) : Account = struct
  type resource = Dbs | Colls | Docs | Users | Permissions | Offers | Pkranges

  let string_of_resource = function
    | Dbs -> "dbs"
    | Colls -> "colls"
    | Docs -> "docs"
    | Users -> "users"
    | Permissions -> "permissions"
    | Offers -> "offers"
    | Pkranges -> "pkranges"

  let authorization verb resource date db_name =
    match C.credential with
    | Databases_intf.Credential.Master_key key ->
        Utility.authorization_token_using_master_key
          (Utilities.Verb.string_of_verb verb)
          (string_of_resource resource)
          db_name
          (Utilities.Ms_time.x_ms_date date)
          key
    | Databases_intf.Credential.Resource_token token ->
        Utility.authorization_token_using_resource_token token
    | Databases_intf.Credential.Resource_token_provider provider ->
        Utility.authorization_token_using_resource_token (provider ())

  let endpoint = C.endpoint
end

module Auth (Keys : Databases_intf.Auth_key) : Account = Auth_credential (struct
  let credential = Databases_intf.Credential.Master_key Keys.master_key
  let endpoint = Keys.endpoint
end)

module Response_headers = struct
  type t = {
    content_type : string option;
    date : string option;
    etag : string option;
    x_ms_activity_id : string option;
    x_ms_alt_content_path : string option;
    x_ms_continuation : string option;
    x_ms_item_count : string option;
    x_ms_request_charge : string option;
    x_ms_resource_quota : string option;
    x_ms_resource_usage : string option;
    x_ms_retry_after_ms : string option;
    x_ms_schemaversion : string option;
    x_ms_serviceversion : string option;
    x_ms_session_token : string option;
    x_ms_substatus : string option;
  }

  let empty =
    {
      content_type = None;
      date = None;
      etag = None;
      x_ms_activity_id = None;
      x_ms_alt_content_path = None;
      x_ms_continuation = None;
      x_ms_item_count = None;
      x_ms_request_charge = None;
      x_ms_resource_quota = None;
      x_ms_resource_usage = None;
      x_ms_retry_after_ms = None;
      x_ms_schemaversion = None;
      x_ms_serviceversion = None;
      x_ms_session_token = None;
      x_ms_substatus = None;
    }

  let update t value_tuple =
    let key, value = value_tuple in
    match key with
    | "Content-Type" -> { t with content_type = Some value }
    | "Date" -> { t with date = Some value }
    | "etag" -> { t with etag = Some value }
    | "x-ms-activity-id" -> { t with x_ms_activity_id = Some value }
    | "x-ms-alt-content-path" -> { t with x_ms_alt_content_path = Some value }
    | "x-ms-continuation" -> { t with x_ms_continuation = Some value }
    | "x-ms-item-count" -> { t with x_ms_item_count = Some value }
    | "x-ms-request-charge" -> { t with x_ms_request_charge = Some value }
    | "x-ms-resource-quota" -> { t with x_ms_resource_quota = Some value }
    | "x-ms-resource-usage" -> { t with x_ms_resource_usage = Some value }
    | "x-ms-retry-after-ms" -> { t with x_ms_retry_after_ms = Some value }
    | "x-ms-schemaversion" -> { t with x_ms_schemaversion = Some value }
    | "x-ms-serviceversion" -> { t with x_ms_serviceversion = Some value }
    | "x-ms-session-token" -> { t with x_ms_session_token = Some value }
    | "x-ms-substatus" -> { t with x_ms_substatus = Some value }
    | _ -> t

  let get_header resp =
    resp |> Cohttp.Response.headers |> Cohttp.Header.to_list
    |> List.fold_left update empty

  let content_type t = t.content_type
  let date t = t.date
  let etag t = t.etag
  let x_ms_activity_id t = t.x_ms_activity_id
  let x_ms_alt_content_path t = t.x_ms_alt_content_path
  let x_ms_continuation t = t.x_ms_continuation
  let x_ms_item_count t = t.x_ms_item_count
  let x_ms_request_charge t = t.x_ms_request_charge
  let x_ms_resource_quota t = t.x_ms_resource_quota
  let x_ms_resource_usage t = t.x_ms_resource_usage
  let x_ms_retry_after_ms t = t.x_ms_retry_after_ms
  let x_ms_schemaversion t = t.x_ms_schemaversion
  let x_ms_serviceversion t = t.x_ms_serviceversion
  let x_ms_session_token t = t.x_ms_session_token
  let x_ms_substatus t = t.x_ms_substatus

  let string_of t =
    [
      ("Content-Type", t.content_type);
      ("Date", t.date);
      ("etag", t.etag);
      ("x-ms-activity-id", t.x_ms_activity_id);
      ("x-ms-alt-content-path", t.x_ms_alt_content_path);
      ("x-ms-continuation", t.x_ms_continuation);
      ("x-ms-item-count", t.x_ms_item_count);
      ("x-ms-request-charge", t.x_ms_request_charge);
      ("x-ms-resource-quota", t.x_ms_resource_quota);
      ("x-ms-resource-usage", t.x_ms_resource_usage);
      ("x-ms-retry-after-ms", t.x_ms_retry_after_ms);
      ("x-ms-schemaversion", t.x_ms_schemaversion);
      ("x-ms-serviceversion", t.x_ms_serviceversion);
      ("x-ms-session-token", t.x_ms_session_token);
      ("x-ms-substatus", t.x_ms_substatus);
    ]
    |> List.filter_map (fun (name, value) ->
        Option.map (fun v -> name ^ ": " ^ v) value)
    |> String.concat "; "
end

type batch_validation_error =
  | Too_many_operations of int
  | Mixed_patch_operations
  | Empty_batch

module Cosmos_error = struct
  type t =
    | Timeout_error
    | Connection_error
    | Http_error of string
    | Azure_error of int * Response_headers.t

  let string_of = function
    | Timeout_error -> "Timeout"
    | Connection_error -> "Connection refused"
    | Http_error message -> "HTTP error: " ^ message
    | Azure_error (code, headers) -> (
        let status = "Azure error: HTTP status " ^ string_of_int code in
        match Response_headers.string_of headers with
        | "" -> status
        | s -> status ^ " (" ^ s ^ ")")
end

module Make_account
    (IO : Databases_intf.IO)
    (Http : Databases_intf.Http_client with type 'a io := 'a IO.t)
    (Account : Account) =
struct
  let ( let* ) = IO.bind
  let host = Utility.adjust_host Account.endpoint
  let timeout_error = IO.return (Error Cosmos_error.Timeout_error)
  let connection_error = IO.return (Error Cosmos_error.Connection_error)

  let wrap_timeout timeout command =
    match timeout with
    | None ->
        let* result = command in
        IO.return (Some result)
    | Some t -> IO.with_timeout t command

  let headers resource verb db_name =
    let ms_date = Utilities.Ms_time.create_now () in
    let header = Cohttp.Header.init () in
    let header =
      Cohttp.Header.add header "authorization"
        (Account.authorization verb resource ms_date db_name)
    in
    let header = Cohttp.Header.add header "x-ms-version" "2018-12-31" in
    let header =
      Cohttp.Header.add header "x-ms-date" (Utilities.Ms_time.x_ms_date ms_date)
    in
    header

  let json_headers resource verb db_name =
    let header = headers resource verb db_name in
    let header = Cohttp.Header.add header "content_type" "application/json" in
    header

  let add_header name value header = Cohttp.Header.add header name value

  let get_code resp =
    resp |> Cohttp.Response.status |> Cohttp.Code.code_of_status

  let azure_error resp =
    Cosmos_error.Azure_error (get_code resp, Response_headers.get_header resp)

  let result_or_error expected_code resp =
    let code = get_code resp in
    if code = expected_code then Result.ok expected_code
    else Result.error (azure_error resp)

  let result_or_error_with_result expected_code f resp body =
    let code = get_code resp in
    if code = expected_code then IO.return (Ok (expected_code, f body))
    else IO.return (Error (azure_error resp))

  let with_204_do = result_or_error 204

  let handle_http_error = function
    | Http.Connection_refused -> connection_error
    | Http.Other_error exn ->
        IO.return (Error (Cosmos_error.Http_error (Printexc.to_string exn)))

  let make_uri path = Uri.make ~scheme:"https" ~host ~port:443 ~path ()

  let header_path_of_path path =
    if String.length path > 0 && path.[0] = '/' then
      String.sub path 1 (String.length path - 1)
    else path

  let handle_response response f =
    match response with
    | None -> timeout_error
    | Some (Error e) -> handle_http_error e
    | Some (Ok (resp, body)) -> f resp body

  let with_throttle_retry ?timeout ~max_retries f =
    let rec retry_loop attempt () =
      IO.catch
        (fun () ->
          let* result = f () |> wrap_timeout timeout in
          match result with
          | None -> timeout_error
          | Some result -> (
              match result with
              | Ok (resp, body) ->
                  let code = get_code resp in
                  if code = 429 then
                    let response_header = Response_headers.get_header resp in
                    let milliseconds =
                      Response_headers.x_ms_retry_after_ms response_header
                      |> Option.value ~default:"0" |> int_of_string_opt
                      |> Option.value ~default:0 |> Int.to_float
                    in
                    let* () = IO.sleep (milliseconds /. 1000.) in
                    if attempt > 0 then retry_loop (attempt - 1) ()
                    else timeout_error
                  else IO.return (Ok (code, resp, body))
              | Error Http.Connection_refused ->
                  if attempt > 0 then
                    let sleep_time = Random.int 5 |> float_of_int in
                    let* () = IO.sleep sleep_time in
                    retry_loop (attempt - 1) ()
                  else connection_error
              | Error (Http.Other_error exn) -> raise exn))
        (fun _exn ->
          if attempt > 0 then
            let sleep_time = Random.int 5 |> float_of_int in
            let* () = IO.sleep sleep_time in
            retry_loop (attempt - 1) ()
          else connection_error)
    in
    retry_loop max_retries ()

  let list_databases ?timeout () =
    let uri = make_uri "dbs" in
    let* response =
      Http.get ~headers:(headers Account.Dbs Utilities.Verb.Get "") uri
      |> wrap_timeout timeout
    in
    handle_response response (fun resp body ->
        let code = get_code resp in
        let value = Json_converter_j.list_databases_of_string body in
        IO.return @@ Result.ok (code, value))

  let create ?timeout name =
    let body =
      ({ id = name } : Json_converter_j.create_database)
      |> Json_converter_j.string_of_create_database
    in
    let uri = make_uri "dbs" in
    let hdrs = json_headers Account.Dbs Utilities.Verb.Post "" in
    let* response = Http.post ~headers:hdrs ~body uri |> wrap_timeout timeout in
    handle_response response (fun resp body ->
        let result body = Some (Json_converter_j.database_of_string body) in
        result_or_error_with_result 201 result resp body)

  let get ?timeout name =
    let path = "dbs/" ^ name in
    let uri = make_uri path in
    let* response =
      Http.get ~headers:(headers Account.Dbs Utilities.Verb.Get path) uri
      |> wrap_timeout timeout
    in
    handle_response response (fun resp body ->
        let result body = Some (Json_converter_j.database_of_string body) in
        result_or_error_with_result 200 result resp body)

  let create_if_not_exists ?timeout name =
    let* exists = get ?timeout name in
    match exists with
    | Ok (code, result) -> IO.return (Ok (code, result))
    | Error (Cosmos_error.Azure_error (404, _)) -> create ?timeout name
    | Error x -> IO.return (Error x)

  let delete ?timeout name =
    let path = "dbs/" ^ name in
    let uri = make_uri path in
    let* response =
      Http.delete ~headers:(headers Account.Dbs Utilities.Verb.Delete path) uri
      |> wrap_timeout timeout
    in
    handle_response response (fun resp _body -> IO.return (with_204_do resp))

  module Collection = struct
    let string_of_partition_key s = Printf.sprintf "[%S]" s
    let path_of_collections dbname = Printf.sprintf "dbs/%s/colls" dbname

    let path_of_collection dbname coll_name =
      Printf.sprintf "dbs/%s/colls/%s" dbname coll_name

    let path_of_docs dbname coll_name =
      Printf.sprintf "/dbs/%s/colls/%s/docs" dbname coll_name

    let path_of_pkranges dbname coll_name =
      Printf.sprintf "/dbs/%s/colls/%s/pkranges" dbname coll_name

    let path_of_doc dbname coll_name doc_id =
      Printf.sprintf "/dbs/%s/colls/%s/docs/%s" dbname coll_name doc_id

    let list ?timeout dbname =
      let path = path_of_collections dbname in
      let uri = make_uri path in
      let* response =
        Http.get
          ~headers:(headers Account.Colls Utilities.Verb.Get ("dbs/" ^ dbname))
          uri
        |> wrap_timeout timeout
      in
      handle_response response (fun resp body ->
          let value body = Json_converter_j.list_collections_of_string body in
          result_or_error_with_result 200 value resp body)

    let create ?(indexing_policy = None) ?offer_throughput ~partition_key
        ?timeout dbname coll_name =
      let body =
        ({ id = coll_name; indexing_policy; partition_key }
          : Json_converter_j.create_collection)
        |> Json_converter_j.string_of_create_collection
      in
      let path = path_of_collections dbname in
      let uri = make_uri path in
      let hdrs =
        json_headers Account.Colls Utilities.Verb.Post ("dbs/" ^ dbname)
        |> Utilities.apply_to_header_if_some "x-ms-offer-throughput"
             string_of_int offer_throughput
      in
      let* response =
        Http.post ~headers:hdrs ~body uri |> wrap_timeout timeout
      in
      handle_response response (fun resp body ->
          let value body = Some (Json_converter_j.collection_of_string body) in
          result_or_error_with_result 201 value resp body)

    let get ?timeout name coll_name =
      let path = path_of_collection name coll_name in
      let uri = make_uri path in
      let* response =
        Http.get
          ~headers:
            (headers Account.Colls Utilities.Verb.Get (header_path_of_path path))
          uri
        |> wrap_timeout timeout
      in
      handle_response response (fun resp body ->
          let value body = Some (Json_converter_j.collection_of_string body) in
          result_or_error_with_result 200 value resp body)

    let create_if_not_exists ?(indexing_policy = None) ?offer_throughput
        ~partition_key ?timeout dbname coll_name =
      let* exists = get ?timeout dbname coll_name in
      match exists with
      | Ok result -> IO.return (Result.ok result)
      | Error (Cosmos_error.Azure_error (404, _)) ->
          create ?timeout ~indexing_policy ?offer_throughput ~partition_key
            dbname coll_name
      | Error x -> IO.return (Error x)

    let delete ?timeout name coll_name =
      let path = path_of_collection name coll_name in
      let uri = make_uri path in
      let* response =
        Http.delete
          ~headers:
            (headers Account.Colls Utilities.Verb.Delete
               (header_path_of_path path))
          uri
        |> wrap_timeout timeout
      in
      handle_response response (fun resp _body -> IO.return (with_204_do resp))

    module Partition_key_range = struct
      let list ?max_item_count ?continuation ?timeout dbname coll_name =
        let path = path_of_pkranges dbname coll_name in
        let uri = make_uri path in
        let hdrs =
          headers Account.Pkranges Utilities.Verb.Get
            (path_of_collection dbname coll_name)
          |> Utilities.apply_to_header_if_some "x-ms-max-item-count"
               string_of_int max_item_count
          |> Utilities.apply_to_header_if_some "x-ms-continuation" Fun.id
               continuation
        in
        let do_get () = Http.get ~headers:hdrs uri in
        let* response = with_throttle_retry ?timeout ~max_retries:10 do_get in
        match response with
        | Error e -> IO.return (Error e)
        | Ok (code, resp, body) ->
            if code = 200 then
              IO.return
                (Ok
                   ( code,
                     Response_headers.get_header resp,
                     Json_converter_j.list_partition_key_ranges_of_string body
                   ))
            else IO.return (Error (azure_error resp))

      let ids ?timeout dbname coll_name =
        let* result = list ?timeout dbname coll_name in
        match result with
        | Error e -> IO.return (Error e)
        | Ok (code, _, ranges) ->
            IO.return
              (Ok
                 ( code,
                   List.map
                     (fun (range : Json_converter_t.partition_key_range) ->
                       range.id)
                     ranges.partition_key_ranges ))
    end

    module Document = struct
      type indexing_directive = Include | Exclude

      let string_of_indexing_directive = function
        | Include -> "Include"
        | Exclude -> "Exclude"

      let create ?is_upsert ?indexing_directive ~partition_key ?timeout dbname
          coll_name content =
        let path = path_of_docs dbname coll_name in
        let uri = make_uri path in
        let hdrs =
          json_headers Account.Docs Utilities.Verb.Post
            ("dbs/" ^ dbname ^ "/colls/" ^ coll_name)
          |> Utilities.apply_to_header_if_some "x-ms-documentdb-is-upsert"
               string_of_bool is_upsert
          |> Utilities.apply_to_header_if_some "x-ms-indexing-directive"
               string_of_indexing_directive indexing_directive
          |> add_header "x-ms-documentdb-partitionkey"
               (string_of_partition_key partition_key)
        in
        let () = Random.self_init () in
        let do_post () = Http.post ~headers:hdrs ~body:content uri in
        let extract_value code body_str =
          match code with
          | 200 -> Some (Json_converter_j.collection_of_string body_str)
          | _ -> None
        in
        let* retry_result = with_throttle_retry ~max_retries:10 do_post in
        match retry_result with
        | Error e -> IO.return (Error e)
        | Ok (code, _resp, body_str) -> (
            let value = extract_value code body_str in
            match timeout with
            | None -> IO.return (Result.ok (code, value))
            | Some t -> (
                let* timed_result = IO.with_timeout t (IO.return ()) in
                match timed_result with
                | None -> timeout_error
                | Some () -> IO.return (Result.ok (code, value))))

      let create_multiple ?is_upsert ?indexing_directive ?timeout
          ?(chunk_size = 100) dbname coll_name content_list =
        let rec take_first n acc = function
          | [] -> IO.return @@ acc
          | x ->
              let first, rest = Utilities.take_first n x in
              let* r =
                IO.parallel_map
                  (fun (partition_key, content) ->
                    create ?is_upsert ?indexing_directive ~partition_key
                      ?timeout dbname coll_name content)
                  first
              in
              take_first n (r @ acc) rest
        in
        take_first chunk_size [] content_list

      type list_result_meta_data = {
        rid : string;
        self : string;
        etag : string;
        ts : int;
        attachments : string;
      }

      let convert_to_list_result_meta_data json =
        let open Yojson.Basic.Util in
        try
          let rid = json |> member "_rid" |> to_string in
          let self = json |> member "_self" |> to_string in
          let etag = json |> member "_etag" |> to_string in
          let ts = json |> member "_ts" |> to_int in
          let attachments = json |> member "_attachments" |> to_string in
          Some { rid; self; etag; ts; attachments }
        with Type_error _ -> None

      type list_result = {
        rid : string;
        documents : (string * list_result_meta_data option) list;
        count : int;
      }

      let convert_to_list_result value =
        let open Yojson.Basic.Util in
        let json = Yojson.Basic.from_string value in
        let rid = json |> member "_rid" |> to_string in
        let count = json |> member "_count" |> to_int in
        let docs = json |> member "Documents" |> to_list in
        let string_docs =
          List.map
            (fun x ->
              (Yojson.Basic.to_string x, convert_to_list_result_meta_data x))
            docs
        in
        { rid; documents = string_docs; count }

      let list ?max_item_count ?continuation ?consistency_level ?session_token
          ?a_im ?if_none_match ?partition_key_range_id ?timeout dbname coll_name
          =
        let apply_a_im_to_header_if_some name values headers =
          match values with
          | None -> headers
          | Some false -> headers
          | Some true -> Cohttp.Header.add headers name "Incremental feed"
        in
        let path = path_of_docs dbname coll_name in
        let uri = make_uri path in
        let hdrs =
          json_headers Account.Docs Utilities.Verb.Get
            ("dbs/" ^ dbname ^ "/colls/" ^ coll_name)
          |> Utilities.apply_to_header_if_some "x-ms-max-item-count"
               string_of_int max_item_count
          |> Utilities.apply_to_header_if_some "x-ms-continuation" Fun.id
               continuation
          |> Utilities.apply_to_header_if_some "x-ms-consistency-level" Fun.id
               consistency_level
          |> Utilities.apply_to_header_if_some "x-ms-session-token" Fun.id
               session_token
          |> apply_a_im_to_header_if_some "A-IM" a_im
          |> Utilities.apply_to_header_if_some "If-None-Match" Fun.id
               if_none_match
          |> Utilities.apply_to_header_if_some
               "x-ms-documentdb-partitionkeyrangeid" Fun.id
               partition_key_range_id
        in
        let* get_action = Http.get ~headers:hdrs uri |> wrap_timeout timeout in
        handle_response get_action (fun resp body ->
            let code = get_code resp in
            let response_header = Response_headers.get_header resp in
            if code = 200 then
              let result = convert_to_list_result body in
              IO.return (Ok (200, response_header, result))
            else
              IO.return
                (Error (Cosmos_error.Azure_error (code, response_header))))

      type consistency_level = Strong | Bounded | Session | Eventual

      let string_of_consistency_level = function
        | Strong -> "Strong"
        | Bounded -> "Bounded"
        | Session -> "Session"
        | Eventual -> "Eventual"

      let get ?if_none_match ~partition_key ?consistency_level ?session_token
          ?timeout dbname coll_name doc_id =
        let path = path_of_doc dbname coll_name doc_id in
        let headers =
          json_headers Account.Docs Utilities.Verb.Get
            (header_path_of_path path)
          |> Utilities.apply_to_header_if_some "If-None-Match" Fun.id
               if_none_match
          |> add_header "x-ms-documentdb-partitionkey"
               (string_of_partition_key partition_key)
          |> Utilities.apply_to_header_if_some "x-ms-consistency-level"
               string_of_consistency_level consistency_level
          |> Utilities.apply_to_header_if_some "x-ms-session-token" Fun.id
               session_token
        in
        let uri = make_uri path in
        let* response = Http.get ~headers uri |> wrap_timeout timeout in
        handle_response response (fun resp body ->
            let code = get_code resp in
            IO.return (Result.ok (code, body)))

      let replace ?indexing_directive ~partition_key ?if_match ?timeout dbname
          coll_name doc_id content =
        let path = path_of_doc dbname coll_name doc_id in
        let headers =
          json_headers Account.Docs Utilities.Verb.Put
            (header_path_of_path path)
          |> Utilities.apply_to_header_if_some "x-ms-indexing-directive"
               string_of_indexing_directive indexing_directive
          |> add_header "x-ms-documentdb-partitionkey"
               (string_of_partition_key partition_key)
          |> Utilities.apply_to_header_if_some "If-Match" Fun.id if_match
        in
        let uri = make_uri path in
        let* response =
          Http.put ~headers ~body:content uri |> wrap_timeout timeout
        in
        handle_response response (fun resp body ->
            let code = get_code resp in
            IO.return (Ok (code, body)))

      let delete ~partition_key ?timeout dbname coll_name doc_id =
        let path = path_of_doc dbname coll_name doc_id in
        let hdrs =
          header_path_of_path path
          |> headers Account.Docs Utilities.Verb.Delete
          |> add_header "x-ms-documentdb-partitionkey"
               (string_of_partition_key partition_key)
        in
        let uri = make_uri path in
        let do_delete () = Http.delete ~headers:hdrs uri in
        let* retry_result = with_throttle_retry ~max_retries:3 do_delete in
        match retry_result with
        | Error e -> IO.return (Error e)
        | Ok (code, _resp, _body) -> (
            match timeout with
            | None -> IO.return (Ok code)
            | Some t -> (
                let* timed_result = IO.with_timeout t (IO.return ()) in
                match timed_result with
                | None -> timeout_error
                | Some () -> IO.return (Ok code)))

      let delete_multiple ~partition_key ?timeout ?(chunk_size = 100) dbname
          coll_name doc_ids =
        let rec take_first n acc = function
          | [] -> IO.return @@ acc
          | x ->
              let first, rest = Utilities.take_first n x in
              let* r =
                IO.parallel_map
                  (delete ~partition_key ?timeout dbname coll_name)
                  first
              in
              take_first n (r @ acc) rest
        in
        take_first chunk_size [] doc_ids

      let query ?max_item_count ?continuation ?consistency_level ?session_token
          ?is_partition ?partition_key ?timeout dbname coll_name query =
        let path = path_of_docs dbname coll_name in
        let make_headers s =
          let h = headers Account.Docs Utilities.Verb.Post s in
          Cohttp.Header.add h "x-ms-documentdb-isquery" (string_of_bool true)
          |> Utilities.apply_to_header_if_some "x-ms-max-item-count"
               string_of_int max_item_count
          |> Utilities.apply_to_header_if_some "x-ms-continuation" Fun.id
               continuation
          |> Utilities.apply_to_header_if_some "x-ms-consistency-level" Fun.id
               consistency_level
          |> Utilities.apply_to_header_if_some "x-ms-session-token" Fun.id
               session_token
          |> Utilities.apply_to_header_if_some
               "x-ms-documentdb-query-enablecrosspartition" string_of_bool
               is_partition
          |> Utilities.apply_to_header_if_some "x-ms-documentdb-partitionkey"
               string_of_partition_key partition_key
          |> add_header "content-type" "application/query+json"
        in
        let headers = make_headers ("dbs/" ^ dbname ^ "/colls/" ^ coll_name) in
        let body = Json_converter_j.string_of_query query in
        let uri = make_uri path in
        let* response = Http.post ~headers ~body uri |> wrap_timeout timeout in
        handle_response response (fun resp body ->
            let code = get_code resp in
            let response_header = Response_headers.get_header resp in
            if code = 200 then
              let result = convert_to_list_result body in
              IO.return (Ok (200, response_header, result))
            else
              IO.return
                (Error (Cosmos_error.Azure_error (code, response_header))))
    end

    module Change_feed = struct
      module Mode = struct
        type t = Latest_version

        let string_of = function Latest_version -> "Incremental feed"
      end

      module Start_from = struct
        type t =
          | Beginning
          | Now
          | Point_in_time of float
          | Continuation of string

        let string_of = function
          | Beginning -> "Beginning"
          | Now -> "Now"
          | Point_in_time time ->
              Utilities.Ms_time.x_ms_date (Utilities.Ms_time.create time)
          | Continuation etag -> etag
      end

      module Scope = struct
        type t =
          | Container
          | Partition_key of string
          | Partition_key_range of string

        let string_of = function
          | Container -> "Container"
          | Partition_key key -> Printf.sprintf "Partition_key(%s)" key
          | Partition_key_range id ->
              Printf.sprintf "Partition_key_range(%s)" id
      end

      type page = {
        rid : string;
        documents : (string * Document.list_result_meta_data option) list;
        count : int;
        continuation : string;
        has_more_pages : bool;
        session_token : string option;
      }

      type drain_result = {
        pages : page list;
        checkpoint : string;
        caught_up : bool;
      }

      let apply_start_from start_from headers =
        match start_from with
        | Start_from.Beginning -> headers
        | Start_from.Now -> add_header "If-None-Match" "*" headers
        | Start_from.Continuation etag ->
            add_header "If-None-Match" etag headers
        | Start_from.Point_in_time time ->
            add_header "If-Modified-Since"
              (Utilities.Ms_time.x_ms_date (Utilities.Ms_time.create time))
              headers

      let apply_scope scope headers =
        match scope with
        | Scope.Container -> headers
        | Scope.Partition_key key ->
            add_header "x-ms-documentdb-partitionkey"
              (string_of_partition_key key)
              headers
        | Scope.Partition_key_range id ->
            add_header "x-ms-documentdb-partitionkeyrangeid" id headers

      let read ?(mode = Mode.Latest_version)
          ?(start_from = Start_from.Beginning) ?(scope = Scope.Container)
          ?max_item_count ?session_token ?timeout dbname coll_name =
        let path = path_of_docs dbname coll_name in
        let hdrs =
          headers Account.Docs Utilities.Verb.Get
            (path_of_collection dbname coll_name)
          |> add_header "A-IM" (Mode.string_of mode)
          |> Utilities.apply_to_header_if_some "x-ms-max-item-count"
               string_of_int max_item_count
          |> Utilities.apply_to_header_if_some "x-ms-session-token" Fun.id
               session_token
          |> apply_start_from start_from
          |> apply_scope scope
        in
        let uri = make_uri path in
        let do_get () = Http.get ~headers:hdrs uri in
        let* result = with_throttle_retry ?timeout ~max_retries:10 do_get in
        match result with
        | Error e -> IO.return (Error e)
        | Ok (code, resp, body) -> (
            let response_headers = Response_headers.get_header resp in
            if code <> 200 && code <> 304 then
              IO.return
                (Error (Cosmos_error.Azure_error (code, response_headers)))
            else
              match Response_headers.etag response_headers with
              | None ->
                  IO.return
                    (Error (Cosmos_error.Azure_error (code, response_headers)))
              | Some continuation ->
                  if code = 304 then
                    IO.return (Ok (304, response_headers, None))
                  else
                    let result = Document.convert_to_list_result body in
                    let page =
                      {
                        rid = result.rid;
                        documents = result.documents;
                        count = result.count;
                        continuation;
                        has_more_pages =
                          Option.is_some
                            (Response_headers.x_ms_continuation response_headers);
                        session_token =
                          Response_headers.x_ms_session_token response_headers;
                      }
                    in
                    IO.return (Ok (200, response_headers, Some page)))

      let is_partition_split = function
        | Cosmos_error.Azure_error (410, headers) -> (
            match Response_headers.x_ms_substatus headers with
            | Some ("1002" | "1007") -> true
            | _ -> false)
        | _ -> false

      let drain ?mode ?start_from ?scope ?max_item_count ?(max_pages = 100)
          ?timeout dbname coll_name =
        let rec loop pages checkpoint start_from remaining =
          if remaining = 0 then
            IO.return
              (Ok { pages = List.rev pages; checkpoint; caught_up = false })
          else
            let* result =
              read ?mode ?start_from ?scope ?max_item_count ?timeout dbname
                coll_name
            in
            match result with
            | Error e -> IO.return (Error e)
            | Ok (304, headers, None) -> (
                match Response_headers.etag headers with
                | Some checkpoint ->
                    IO.return
                      (Ok
                         {
                           pages = List.rev pages;
                           checkpoint;
                           caught_up = true;
                         })
                | None ->
                    IO.return (Error (Cosmos_error.Azure_error (304, headers))))
            | Ok (200, _, Some page) ->
                loop (page :: pages) page.continuation
                  (Some (Start_from.Continuation page.continuation))
                  (remaining - 1)
            | Ok (code, headers, _) ->
                IO.return (Error (Cosmos_error.Azure_error (code, headers)))
        in
        loop [] "" start_from max_pages

      let fold ?mode ?start_from ?scope ?max_item_count ?(poll_interval = 1.0)
          ?(max_polls = 1) ?timeout dbname coll_name ~init ~f =
        let rec poll acc checkpoint start_from remaining =
          if remaining = 0 then IO.return (Ok (acc, checkpoint))
          else
            let* result =
              drain ?mode ?start_from ?scope ?max_item_count ?timeout dbname
                coll_name
            in
            match result with
            | Error e -> IO.return (Error e)
            | Ok drained -> (
                let rec apply_pages acc checkpoint = function
                  | [] -> IO.return (Ok (acc, checkpoint))
                  | page :: pages -> (
                      let* result = f acc page in
                      match result with
                      | Error _ -> IO.return (Error (acc, checkpoint))
                      | Ok acc -> apply_pages acc page.continuation pages)
                in
                let* applied = apply_pages acc checkpoint drained.pages in
                match applied with
                | Error (acc, checkpoint) -> IO.return (Ok (acc, checkpoint))
                | Ok (acc, _) ->
                    if remaining = 1 then
                      IO.return (Ok (acc, drained.checkpoint))
                    else
                      let* () =
                        if drained.caught_up then IO.sleep poll_interval
                        else IO.return ()
                      in
                      poll acc drained.checkpoint
                        (Some (Start_from.Continuation drained.checkpoint))
                        (remaining - 1))
        in
        poll init "" start_from max_polls
    end

    module Batch = struct
      type operation =
        | Create of {
            if_match : string option;
            if_none_match : string option;
            body : string;
          }
        | Upsert of {
            if_match : string option;
            if_none_match : string option;
            body : string;
          }
        | Read of {
            id : string;
            if_match : string option;
            if_none_match : string option;
          }
        | Delete of {
            id : string;
            if_match : string option;
            if_none_match : string option;
          }
        | Replace of {
            id : string;
            if_match : string option;
            if_none_match : string option;
            body : string;
          }
        | Patch of {
            id : string;
            if_match : string option;
            patch_op : patch_operation;
          }

      and patch_operation =
        | Add of { path : string; value : string }
        | Set of { path : string; value : string }
        | ReplacePath of { path : string; value : string }
        | Remove of { path : string }
        | Increment of { path : string; value : int }

      type operation_result = {
        status_code : int;
        request_charge : float;
        etag : string option;
        resource_body : string option;
      }

      type batch_result = {
        outcomes : operation_result list;
        total_request_charge : float;
      }

      type validation_error = batch_validation_error =
        | Too_many_operations of int
        | Mixed_patch_operations
        | Empty_batch

      let string_of_operation_type = function
        | Create _ -> "Create"
        | Upsert _ -> "Upsert"
        | Read _ -> "Read"
        | Delete _ -> "Delete"
        | Replace _ -> "Replace"
        | Patch _ -> "Patch"

      let construct_batch_request_body batch_ops =
        let json_ops =
          List.map
            (fun (op_type, op) ->
              let base_fields = [ ("operationType", op_type) ] in
              let optional_fields =
                let add_string_field name opt acc =
                  Option.fold opt ~none:acc ~some:(fun v ->
                      (name, Printf.sprintf "\"%s\"" v) :: acc)
                in
                let add_field name opt acc =
                  Option.fold opt ~none:acc ~some:(fun v -> (name, v) :: acc)
                in
                []
                |> add_string_field "ifMatch" op.Json_converter_t.ifMatch
                |> add_string_field "ifNoneMatch"
                     op.Json_converter_t.ifNoneMatch
                |> add_string_field "id" op.Json_converter_t.id
                |> add_field "resourceBody" op.Json_converter_t.resourceBody
                |> add_string_field "from" op.Json_converter_t.from
                |> add_field "value" op.Json_converter_t.value
              in
              let all_fields = base_fields @ optional_fields in
              let field_strs =
                List.map
                  (fun (k, v) -> Printf.sprintf "\"%s\":%s" k v)
                  all_fields
              in
              Printf.sprintf "{%s}" (String.concat "," field_strs))
            batch_ops
        in
        Printf.sprintf "[%s]" (String.concat "," json_ops)

      let make_batch_op partition_key ~operation_type ~if_match ~if_none_match
          ~id ~resource_body ~from ~value () =
        ( Json_converter_j.string_of_batch_operation_type operation_type,
          {
            Json_converter_t.operationType = operation_type;
            partitionKey = string_of_partition_key partition_key;
            ifMatch = if_match;
            ifNoneMatch = if_none_match;
            id;
            resourceBody = resource_body;
            from;
            value;
          } )

      let operation_to_batch_op partition_key op =
        let open Json_converter_t in
        match op with
        | Create { if_match; if_none_match; body } ->
            make_batch_op partition_key ~operation_type:`Create ~if_match
              ~if_none_match ~id:None ~resource_body:(Some body) ~from:None
              ~value:None ()
        | Upsert { if_match; if_none_match; body } ->
            make_batch_op partition_key ~operation_type:`Upsert ~if_match
              ~if_none_match ~id:None ~resource_body:(Some body) ~from:None
              ~value:None ()
        | Read { id; if_match; if_none_match } ->
            make_batch_op partition_key ~operation_type:`Read ~if_match
              ~if_none_match ~id:(Some id) ~resource_body:None ~from:None
              ~value:None ()
        | Delete { id; if_match; if_none_match } ->
            make_batch_op partition_key ~operation_type:`Delete ~if_match
              ~if_none_match ~id:(Some id) ~resource_body:None ~from:None
              ~value:None ()
        | Replace { id; if_match; if_none_match; body } ->
            make_batch_op partition_key ~operation_type:`Replace ~if_match
              ~if_none_match ~id:(Some id) ~resource_body:(Some body) ~from:None
              ~value:None ()
        | Patch { id; if_match; patch_op } ->
            let value_str =
              match patch_op with
              | Add { path; value } ->
                  Printf.sprintf {|{"path": "%s", "op": "add", "value": %s}|}
                    path value
              | Set { path; value } ->
                  Printf.sprintf {|{"path": "%s", "op": "set", "value": %s}|}
                    path value
              | ReplacePath { path; value } ->
                  Printf.sprintf
                    {|{"path": "%s", "op": "replace", "value": %s}|} path value
              | Remove { path } ->
                  Printf.sprintf {|{"path": "%s", "op": "remove"}|} path
              | Increment { path; value } ->
                  Printf.sprintf {|{"path": "%s", "op": "incr", "value": %d}|}
                    path value
            in
            make_batch_op partition_key ~operation_type:`Patch ~if_match
              ~if_none_match:None ~id:(Some id) ~resource_body:None
              ~from:(Some "") ~value:(Some value_str) ()

      let has_patch_operation ops =
        List.exists (function Patch _ -> true | _ -> false) ops

      let has_non_patch_operation ops =
        List.exists
          (function
            | Create _ | Upsert _ | Delete _ | Replace _ -> true | _ -> false)
          ops

      let validate ops =
        let count = List.length ops in
        if count = 0 then Error Empty_batch
        else if count > 100 then Error (Too_many_operations count)
        else if has_patch_operation ops && has_non_patch_operation ops then
          Error Mixed_patch_operations
        else Ok ()

      let execute ?timeout ?(atomic = true) ~partition_key dbname coll_name
          operations =
        let do_execute () =
          let path = path_of_docs dbname coll_name in
          let uri = make_uri path in
          let batch_ops =
            List.map (operation_to_batch_op partition_key) operations
          in
          let body = construct_batch_request_body batch_ops in
          let hdrs =
            json_headers Account.Docs Utilities.Verb.Post
              (path_of_collection dbname coll_name)
            |> add_header "x-ms-cosmos-is-batch-request" "True"
            |> add_header "x-ms-documentdb-partitionkey"
                 (string_of_partition_key partition_key)
            |> add_header "x-ms-cosmos-batch-atomic" (string_of_bool atomic)
          in
          let* response =
            Http.post ~headers:hdrs ~body uri |> wrap_timeout timeout
          in
          handle_response response (fun resp body ->
              let code = get_code resp in
              let response_header = Response_headers.get_header resp in
              if code = 200 || code = 207 then
                let outcomes =
                  let parse_result json =
                    let open Yojson.Safe.Util in
                    {
                      status_code = json |> member "statusCode" |> to_int;
                      request_charge =
                        json |> member "requestCharge" |> to_number;
                      etag = json |> member "eTag" |> to_string_option;
                      resource_body =
                        (match json |> member "resourceBody" with
                        | `Null -> None
                        | v -> Some (Yojson.Safe.to_string v));
                    }
                  in
                  Yojson.Safe.from_string body
                  |> Yojson.Safe.Util.to_list |> List.map parse_result
                in
                let total_charge =
                  List.fold_left
                    (fun acc r -> acc +. r.request_charge)
                    0.0 outcomes
                in
                IO.return (Ok { outcomes; total_request_charge = total_charge })
              else
                IO.return
                  (Error (Cosmos_error.Azure_error (code, response_header))))
        in
        do_execute ()
    end

    module Batch_builder = struct
      type t = { ops : Batch.operation list; count : int }

      let empty = { ops = []; count = 0 }

      let add_create ?if_match ?if_none_match ~body t =
        {
          ops = Batch.Create { if_match; if_none_match; body } :: t.ops;
          count = t.count + 1;
        }

      let add_upsert ?if_match ?if_none_match ~body t =
        {
          ops = Batch.Upsert { if_match; if_none_match; body } :: t.ops;
          count = t.count + 1;
        }

      let add_read ?if_match ?if_none_match ~id t =
        {
          ops = Batch.Read { id; if_match; if_none_match } :: t.ops;
          count = t.count + 1;
        }

      let add_delete ?if_match ?if_none_match ~id t =
        {
          ops = Batch.Delete { id; if_match; if_none_match } :: t.ops;
          count = t.count + 1;
        }

      let add_replace ?if_match ?if_none_match ~id ~body t =
        {
          ops = Batch.Replace { id; if_match; if_none_match; body } :: t.ops;
          count = t.count + 1;
        }

      let add_patch ?if_match ~id ~patch_op t =
        {
          ops = Batch.Patch { id; if_match; patch_op } :: t.ops;
          count = t.count + 1;
        }

      let to_operations t = List.rev t.ops

      let build t =
        let ops = to_operations t in
        Batch.validate ops |> Result.map (fun () -> ops)

      let length t = t.count
    end
  end

  module User = struct
    let resource = Account.Users
    let headers = headers resource

    let create ?timeout dbname user_name =
      let body =
        ({ id = user_name } : Json_converter_j.create_user)
        |> Json_converter_j.string_of_create_user
      in
      let uri = make_uri ("/dbs/" ^ dbname ^ "/users") in
      let hdrs = json_headers resource Utilities.Verb.Post ("dbs/" ^ dbname) in
      let* response =
        Http.post ~headers:hdrs ~body uri |> wrap_timeout timeout
      in
      handle_response response (fun resp body ->
          let value body = Json_converter_j.user_of_string body in
          result_or_error_with_result 201 value resp body)

    let list ?timeout dbname =
      let path = "/dbs/" ^ dbname ^ "/users" in
      let uri = make_uri path in
      let* response =
        Http.get ~headers:(headers Utilities.Verb.Get ("dbs/" ^ dbname)) uri
        |> wrap_timeout timeout
      in
      handle_response response (fun resp body ->
          let value body = Json_converter_j.list_users_of_string body in
          result_or_error_with_result 200 value resp body)

    let get ?timeout dbname user_name =
      let path = "/dbs/" ^ dbname ^ "/users/" ^ user_name in
      let uri = make_uri path in
      let* response =
        Http.get
          ~headers:(headers Utilities.Verb.Get (header_path_of_path path))
          uri
        |> wrap_timeout timeout
      in
      handle_response response (fun resp body ->
          let value body = Json_converter_j.user_of_string body in
          result_or_error_with_result 200 value resp body)

    let replace ?timeout dbname user_name new_user_name =
      let body =
        ({ id = new_user_name } : Json_converter_j.create_user)
        |> Json_converter_j.string_of_create_user
      in
      let path = "/dbs/" ^ dbname ^ "/users/" ^ user_name in
      let uri = make_uri path in
      let* response =
        Http.put
          ~headers:(headers Utilities.Verb.Put (header_path_of_path path))
          ~body uri
        |> wrap_timeout timeout
      in
      handle_response response (fun resp body ->
          let value body = Json_converter_j.user_of_string body in
          result_or_error_with_result 200 value resp body)

    let delete ?timeout dbname user_name =
      let path = "/dbs/" ^ dbname ^ "/users/" ^ user_name in
      let uri = make_uri path in
      let* response =
        Http.delete
          ~headers:(headers Utilities.Verb.Delete (header_path_of_path path))
          uri
        |> wrap_timeout timeout
      in
      handle_response response (fun resp _body ->
          IO.return (result_or_error 204 resp))
  end

  module Permission = struct
    let resource = Account.Permissions
    let headers = headers resource

    type permission_mode = Read | All

    let string_of_permission_mode = function Read -> "Read" | All -> "All"

    let create ?timeout ?expiry_seconds ~dbname ~user_name ~coll_name
        permission_mode ~permission_name =
      let permission_mode = string_of_permission_mode permission_mode in
      let resource_string = Printf.sprintf "dbs/%s/colls/%s" dbname coll_name in
      let body =
        ({ id = permission_name; permission_mode; resource = resource_string }
          : Json_converter_j.create_permission)
        |> Json_converter_j.string_of_create_permission
      in
      let path =
        Printf.sprintf "/dbs/%s/users/%s/permissions" dbname user_name
      in
      let uri = make_uri path in
      let hdrs =
        json_headers resource Utilities.Verb.Post
          (Printf.sprintf "dbs/%s/users/%s" dbname user_name)
        |> Utilities.apply_to_header_if_some "x-ms-documentdb-expiry-seconds"
             string_of_int expiry_seconds
      in
      let* response =
        Http.post ~headers:hdrs ~body uri |> wrap_timeout timeout
      in
      handle_response response (fun resp body ->
          let value body = Json_converter_j.permission_of_string body in
          result_or_error_with_result 201 value resp body)

    let list ?timeout ~dbname ~user_name () =
      let path =
        Printf.sprintf "/dbs/%s/users/%s/permissions" dbname user_name
      in
      let uri = make_uri path in
      let* response =
        Http.get
          ~headers:
            (headers Utilities.Verb.Get
               (Printf.sprintf "dbs/%s/users/%s" dbname user_name))
          uri
        |> wrap_timeout timeout
      in
      handle_response response (fun resp body ->
          let value body = Json_converter_j.list_permissions_of_string body in
          result_or_error_with_result 200 value resp body)

    let get ?timeout ?expiry_seconds ~dbname ~user_name ~permission_name () =
      let path =
        Printf.sprintf "/dbs/%s/users/%s/permissions/%s" dbname user_name
          permission_name
      in
      let uri = make_uri path in
      let* response =
        Http.get
          ~headers:
            (headers Utilities.Verb.Get (header_path_of_path path)
            |> Utilities.apply_to_header_if_some
                 "x-ms-documentdb-expiry-seconds" string_of_int expiry_seconds)
          uri
        |> wrap_timeout timeout
      in
      handle_response response (fun resp body ->
          let value body = Json_converter_j.permission_of_string body in
          result_or_error_with_result 200 value resp body)

    let replace ?timeout ?expiry_seconds ~dbname ~user_name ~coll_name
        permission_mode ~permission_name =
      let permission_mode = string_of_permission_mode permission_mode in
      let resource_string = Printf.sprintf "dbs/%s/colls/%s" dbname coll_name in
      let body =
        ({ id = permission_name; permission_mode; resource = resource_string }
          : Json_converter_j.create_permission)
        |> Json_converter_j.string_of_create_permission
      in
      let path =
        Printf.sprintf "/dbs/%s/users/%s/permissions/%s" dbname user_name
          permission_name
      in
      let uri = make_uri path in
      let* response =
        Http.put
          ~headers:
            (headers Utilities.Verb.Put (header_path_of_path path)
            |> Utilities.apply_to_header_if_some
                 "x-ms-documentdb-expiry-seconds" string_of_int expiry_seconds)
          ~body uri
        |> wrap_timeout timeout
      in
      handle_response response (fun resp body ->
          let value body = Json_converter_j.permission_of_string body in
          result_or_error_with_result 200 value resp body)

    let delete ?timeout ~dbname ~user_name ~permission_name () =
      let path =
        Printf.sprintf "/dbs/%s/users/%s/permissions/%s" dbname user_name
          permission_name
      in
      let uri = make_uri path in
      let* response =
        Http.delete
          ~headers:(headers Utilities.Verb.Delete (header_path_of_path path))
          uri
        |> wrap_timeout timeout
      in
      handle_response response (fun resp _body ->
          IO.return (result_or_error 204 resp))
  end

  let get_database = get

  module Offer = struct
    let resource = Account.Offers
    let headers = headers resource
    let path_of_offers = "offers"
    let path_of_offer offer_rid = Printf.sprintf "offers/%s" offer_rid
    let auth_path_of_offer = String.lowercase_ascii

    module Throughput = struct
      type t = Manual of int | Autoscale of { max_throughput : int }

      let to_content = function
        | Manual offer_throughput ->
            Json_converter_t.
              {
                offer_throughput = Some offer_throughput;
                offer_is_ru_per_minute_throughput_enabled = None;
                offer_autopilot_settings = None;
              }
        | Autoscale { max_throughput } ->
            Json_converter_t.
              {
                offer_throughput = None;
                offer_is_ru_per_minute_throughput_enabled = None;
                offer_autopilot_settings = Some { max_throughput };
              }

      let of_content content =
        Option.fold content.Json_converter_t.offer_autopilot_settings
          ~none:
            (Option.fold content.Json_converter_t.offer_throughput ~none:None
               ~some:(fun value -> Some (Manual value)))
          ~some:(fun settings ->
            Some
              (Autoscale
                 { max_throughput = settings.Json_converter_t.max_throughput }))

      let string_of = function
        | Manual throughput -> Printf.sprintf "Manual %d" throughput
        | Autoscale { max_throughput } ->
            Printf.sprintf "Autoscale %d" max_throughput
    end

    let list ?timeout () =
      let uri = make_uri path_of_offers in
      let* response =
        Http.get ~headers:(headers Utilities.Verb.Get "") uri
        |> wrap_timeout timeout
      in
      handle_response response (fun resp body ->
          let value body = Json_converter_j.list_offers_of_string body in
          result_or_error_with_result 200 value resp body)

    let get ?timeout offer_rid =
      let path = path_of_offer offer_rid in
      let uri = make_uri path in
      let* response =
        Http.get
          ~headers:(headers Utilities.Verb.Get (auth_path_of_offer offer_rid))
          uri
        |> wrap_timeout timeout
      in
      handle_response response (fun resp body ->
          let value body = Json_converter_j.offer_of_string body in
          result_or_error_with_result 200 value resp body)

    let query ?max_item_count ?continuation ?timeout query =
      let hdrs =
        headers Utilities.Verb.Post ""
        |> add_header "x-ms-documentdb-isquery" (string_of_bool true)
        |> Utilities.apply_to_header_if_some "x-ms-max-item-count" string_of_int
             max_item_count
        |> Utilities.apply_to_header_if_some "x-ms-continuation" Fun.id
             continuation
        |> add_header "content-type" "application/query+json"
      in
      let body = Json_converter_j.string_of_query query in
      let uri = make_uri path_of_offers in
      let* response =
        Http.post ~headers:hdrs ~body uri |> wrap_timeout timeout
      in
      handle_response response (fun resp body ->
          let code = get_code resp in
          let response_headers = Response_headers.get_header resp in
          if code = 200 then
            let result = Json_converter_j.list_offers_of_string body in
            IO.return (Ok (code, response_headers, result))
          else
            IO.return
              (Error (Cosmos_error.Azure_error (code, response_headers))))

    let replace ?migrate ?timeout (offer : Json_converter_t.offer) throughput =
      let offer =
        {
          offer with
          Json_converter_t.content = Throughput.to_content throughput;
        }
      in
      let body = Json_converter_j.string_of_offer offer in
      let path = path_of_offer offer.Json_converter_t.rid in
      let uri = make_uri path in
      let hdrs =
        json_headers resource Utilities.Verb.Put
          (auth_path_of_offer offer.Json_converter_t.rid)
        |> Utilities.apply_to_header_if_some
             "x-ms-cosmos-migrate-offer-to-autopilot" string_of_bool
             (match migrate with Some `To_autoscale -> Some true | _ -> None)
        |> Utilities.apply_to_header_if_some
             "x-ms-cosmos-migrate-offer-to-manual-throughput" string_of_bool
             (match migrate with Some `To_manual -> Some true | _ -> None)
      in
      let do_put () = Http.put ~headers:hdrs ~body uri in
      let* retry_result =
        match timeout with
        | None ->
            let* result = with_throttle_retry ~max_retries:3 do_put in
            IO.return (Some result)
        | Some timeout ->
            IO.with_timeout timeout (with_throttle_retry ~max_retries:3 do_put)
      in
      match retry_result with
      | None -> timeout_error
      | Some (Error error) -> IO.return (Error error)
      | Some (Ok (_code, resp, body)) ->
          let value body = Json_converter_j.offer_of_string body in
          result_or_error_with_result 200 value resp body

    let query_for_resource_id ?timeout resource_id =
      let query_body =
        Json_converter_t.
          {
            query = "SELECT * FROM c WHERE c.offerResourceId = @rid";
            parameters = [ { name = "@rid"; value = resource_id } ];
          }
      in
      let* result = query ?timeout query_body in
      match result with
      | Error error -> IO.return (Error error)
      | Ok (code, response_headers, { Json_converter_t.offers; _ }) ->
          IO.return (Ok (code, response_headers, List.nth_opt offers 0))

    let get_for_collection_with_headers ?timeout dbname coll_name =
      let* collection = Collection.get ?timeout dbname coll_name in
      match collection with
      | Error error -> IO.return (Error error)
      | Ok (code, None) -> IO.return (Ok (code, Response_headers.empty, None))
      | Ok (_, Some { Json_converter_t.rid; _ }) ->
          query_for_resource_id ?timeout rid

    let get_for_collection ?timeout dbname coll_name =
      let* result = get_for_collection_with_headers ?timeout dbname coll_name in
      match result with
      | Error error -> IO.return (Error error)
      | Ok (code, _response_headers, offer) -> IO.return (Ok (code, offer))

    let get_for_database ?timeout dbname =
      let* database = get_database ?timeout dbname in
      match database with
      | Error error -> IO.return (Error error)
      | Ok (code, None) -> IO.return (Ok (code, None))
      | Ok (_, Some { Json_converter_t._rid; _ }) -> (
          let* result = query_for_resource_id ?timeout _rid in
          match result with
          | Error error -> IO.return (Error error)
          | Ok (code, _response_headers, offer) -> IO.return (Ok (code, offer)))

    let get_throughput ?timeout dbname coll_name =
      let* result = get_for_collection ?timeout dbname coll_name in
      match result with
      | Error error -> IO.return (Error error)
      | Ok (code, offer) ->
          let throughput =
            Option.fold offer ~none:None ~some:(fun offer ->
                Throughput.of_content offer.Json_converter_t.content)
          in
          IO.return (Ok (code, throughput))

    let set_throughput ?migrate ?timeout dbname coll_name throughput =
      let* result = get_for_collection_with_headers ?timeout dbname coll_name in
      match result with
      | Error error -> IO.return (Error error)
      | Ok (_code, response_headers, None) ->
          IO.return (Error (Cosmos_error.Azure_error (404, response_headers)))
      | Ok (_, _, Some offer) -> replace ?migrate ?timeout offer throughput
  end
end

module Make_credential
    (IO : Databases_intf.IO)
    (Http : Databases_intf.Http_client with type 'a io := 'a IO.t)
    (C : Databases_intf.Credentials) =
  Make_account (IO) (Http) (Auth_credential (C))

module Make
    (IO : Databases_intf.IO)
    (Http : Databases_intf.Http_client with type 'a io := 'a IO.t)
    (Auth_key : Databases_intf.Auth_key) =
  Make_account (IO) (Http) (Auth (Auth_key))