commit 68482572c61ad0d0e1fb79bd0e7164f0a30cfeab · anil.recoil.org/slop

+18 -46
stack/zulip/lib/zulip_bot/lib/bot_storage.ml
···

       11
       11
        
       

     

       12
       12
        
       (** {1 JSON Codecs for Bot Storage} *)

     

       13
       13
        
       

     

       14
       14
       +
       (* String map for storage values *)

     

       15
       15
       +
       module String_map = Map.Make (String)

     

       16
       16
       +
       

     

       14
       17
        
       (* Storage response type - {"storage": {...}} *)

     

       15
       18
        
       type storage_response = {

     

       16
       16
       -
         storage : (string * string) list;

     

       19
       19
       +
         storage : string String_map.t;

     

       17
       20
        
         unknown : Jsont.json;  (** Unknown/extra JSON fields preserved during parsing *)

     

       18
       21
        
       }

     

       19
       22
        
       

     

       20
       20
       -
       (* Custom codec for storage_response that handles the dictionary *)

     

       23
       23
       +
       (* Codec for storage response using Jsont.Object with keep_unknown *)

     

       21
       24
        
       let storage_response_jsont : storage_response Jsont.t =

     

       22
       22
       -
         let of_string s =

     

       23
       23
       -
           match Jsont_bytesrw.decode_string' Jsont.json s with

     

       24
       24
       -
           | Error _ -> Error "Failed to decode JSON"

     

       25
       25
       -
           | Ok json ->

     

       26
       26
       -
               match json with

     

       27
       27
       -
               | Jsont.Object (fields, _) ->

     

       28
       28
       -
                   let assoc = List.map (fun ((k, _), v) -> (k, v)) fields in

     

       29
       29
       -
                   (match List.assoc_opt "storage" assoc with

     

       30
       30
       -
                    | Some (Jsont.Object (storage_fields, _)) ->

     

       31
       31
       -
                        let storage = List.filter_map (fun ((k, _), v) ->

     

       32
       32
       -
                          match v with

     

       33
       33
       -
                          | Jsont.String (s, _) -> Some (k, s)

     

       34
       34
       -
                          | _ -> None

     

       35
       35
       -
                        ) storage_fields in

     

       36
       36
       -
                        (* Keep unknown fields *)

     

       37
       37
       -
                        let unknown_fields = List.filter (fun (k, _) -> k <> "storage") assoc in

     

       38
       38
       -
                        let unknown_mems = List.map (fun (k, v) -> ((k, Jsont.Meta.none), v)) unknown_fields in

     

       39
       39
       -
                        let unknown = Jsont.Object (unknown_mems, Jsont.Meta.none) in

     

       40
       40
       -
                        Ok { storage; unknown }

     

       41
       41
       -
                    | Some _ -> Error "Expected 'storage' field to be an object"

     

       42
       42
       -
                    | None -> Ok { storage = []; unknown = Jsont.Object ([], Jsont.Meta.none) })

     

       43
       43
       -
               | _ -> Error "Expected JSON object for storage response"

     

       44
       44
       -
         in

     

       45
       45
       -
         let to_string { storage; unknown } =

     

       46
       46
       -
           (* Create storage object *)

     

       47
       47
       -
           let storage_fields = List.map (fun (k, v) ->

     

       48
       48
       -
             ((k, Jsont.Meta.none), Jsont.String (v, Jsont.Meta.none))

     

       49
       49
       -
           ) storage in

     

       50
       50
       -
           let storage_obj = Jsont.Object (storage_fields, Jsont.Meta.none) in

     

       51
       51
       -
       

     

       52
       52
       -
           (* Merge with unknown fields *)

     

       53
       53
       -
           let storage_mem = (("storage", Jsont.Meta.none), storage_obj) in

     

       54
       54
       -
           let unknown_mems = match unknown with

     

       55
       55
       -
             | Jsont.Object (fields, _) -> fields

     

       56
       56
       -
             | _ -> []

     

       57
       57
       -
           in

     

       58
       58
       -
           let json = Jsont.Object (storage_mem :: unknown_mems, Jsont.Meta.none) in

     

       59
       59
       -
           match Jsont_bytesrw.encode_string' Jsont.json json with

     

       60
       60
       -
           | Ok s -> s

     

       61
       61
       -
           | Error e -> failwith ("Failed to encode storage response: " ^ Jsont.Error.to_string e)

     

       25
       25
       +
         let make storage unknown = { storage; unknown } in

     

       26
       26
       +
         let storage_map_jsont =

     

       27
       27
       +
           Jsont.Object.map ~kind:"StorageMap" Fun.id

     

       28
       28
       +
           |> Jsont.Object.keep_unknown (Jsont.Object.Mems.string_map Jsont.string) ~enc:Fun.id

     

       29
       29
       +
           |> Jsont.Object.finish

     

       62
       30
        
         in

     

       63
       63
       -
         Jsont.of_of_string ~kind:"StorageResponse" of_string ~enc:to_string

     

       31
       31
       +
         Jsont.Object.map ~kind:"StorageResponse" make

     

       32
       32
       +
         |> Jsont.Object.mem "storage" storage_map_jsont ~enc:(fun r -> r.storage)

     

       33
       33
       +
           ~dec_absent:String_map.empty

     

       34
       34
       +
         |> Jsont.Object.keep_unknown Jsont.json_mems ~enc:(fun r -> r.unknown)

     

       35
       35
       +
         |> Jsont.Object.finish

     

       64
       36
        
       

     

       65
       37
        
       let create client ~bot_email =

     

       66
       38
        
         Log.info (fun m -> m "Creating bot storage for %s" bot_email);

     
···

       74
       46
        
         | Ok json ->

     

       75
       47
        
             (match Zulip.Encode.from_json storage_response_jsont json with

     

       76
       48
        
              | Ok response ->

     

       77
       77
       -
                  List.iter (fun (k, v) ->

     

       49
       49
       +
                  String_map.iter (fun k v ->

     

       78
       50
        
                    Log.debug (fun m -> m "Loaded key from server: %s" k);

     

       79
       51
        
                    Hashtbl.add cache k v

     

       80
       52
        
                  ) response.storage

     
···

       121
       93
        
             | Ok json ->

     

       122
       94
        
                 (match Zulip.Encode.from_json storage_response_jsont json with

     

       123
       95
        
                  | Ok response ->

     

       124
       124
       -
                      (match List.assoc_opt key response.storage with

     

       96
       96
       +
                      (match String_map.find_opt key response.storage with

     

       125
       97
        
                       | Some value ->

     

       126
       98
        
                           (* Cache the value *)

     

       127
       99
        
                           Log.debug (fun m -> m "Retrieved key from API: %s" key);

     
···

       195
       167
        
         | Ok json ->

     

       196
       168
        
             (match Zulip.Encode.from_json storage_response_jsont json with

     

       197
       169
        
              | Ok response ->

     

       198
       198
       -
                  let api_keys = List.map fst response.storage in

     

       170
       170
       +
                  let api_keys = String_map.fold (fun k _ acc -> k :: acc) response.storage [] in

     

       199
       171
        
                  (* Merge with cache keys *)

     

       200
       172
        
                  let cache_keys = Hashtbl.fold (fun k _ acc -> k :: acc) t.cache [] in

     

       201
       173
        
                  let all_keys = List.sort_uniq String.compare (api_keys @ cache_keys) in