Alternate context preserving solution

This fixes the logging context in a different way, since the other solution
was flaky.

- Use Scheduler.preserve_context when it's reasonable
- Move the query logging out of the background thread since there's no
  reason to put it there anyway

src/client.ml

@@ -79,10 +79,9 @@ let format_query query params =
 let execute' ?params ~query ~formatted_query ({ month_offset; _ } as t) ~f =
   sequencer_enqueue t
   @@ fun conn ->
-  let context = Scheduler.current_execution_context () in
+  Logger.debug !"Executing query: %s" formatted_query;
   In_thread.run
   @@ fun () ->
-  Logger.debug ~context !"Executing query: %s" formatted_query;
   Mssql_error.with_wrap ~query ?params ~formatted_query [%here] (fun () ->
       Dblib.cancel conn;
       Dblib.sqlexec conn formatted_query;
@@ -97,7 +96,7 @@ let execute' ?params ~query ~formatted_query ({ month_offset; _ } as t) ~f =
             Iter.from_fun (fun () ->
                 try
                   let row = Dblib.nextrow conn in
-                  let row = Row.create_exn ~context ~month_offset ~colnames row in
+                  let row = Row.create_exn ~month_offset ~colnames row in
                   Some row
                 with
                 | Caml.Not_found -> None)
@@ -119,6 +118,7 @@ let execute_multi_result ?params conn query =
 (* Execute [f iter] for the first result set iterator and throw an exception if there is more than
    one result set *)
 let execute_f' ?params ~f conn query =
+  let f = Scheduler.preserve_execution_context' f |> Staged.unstage in
   execute_multi_result' ?params conn query ~f:(fun result_sets ->
       let result =
         let input =

src/db_field.ml

@@ -26,7 +26,7 @@ type t =
 
     If [str] contains a character that can't be represented in [~dst], we
     skip that character and move on. *)
-let recode ?context ~src ~dst input =
+let recode ~src ~dst input =
   (*  Note that we originally used //TRANSLIT and //IGNORE to do this in iconv,
       but iconv is inconsistent between platforms so we do the conversion one
       char at a time. *)
@@ -62,15 +62,11 @@ let recode ?context ~src ~dst input =
     String.sub output ~pos:0 ~len:!enc_i
   with
   | exn ->
-    Logger.info
-      ?context
-      !"Recoding error, falling back to ascii filter %{sexp: exn} %s"
-      exn
-      input;
+    Logger.info !"Recoding error, falling back to ascii filter %{sexp: exn} %s" exn input;
     String.filter input ~f:(fun c -> Char.to_int c < 128)
 ;;
 
-let of_data ~context ~month_offset data =
+let of_data ~month_offset data =
   match data with
   | Dblib.BIT b -> Some (Bool b)
   | INT i | SMALL i | TINY i -> Some (Int i)
@@ -79,7 +75,7 @@ let of_data ~context ~month_offset data =
   | FLOAT f | MONEY f -> Some (Float f)
   | DECIMAL s | NUMERIC s -> Some (Bignum (Bignum.of_string s))
   | BINARY s -> Some (String s)
-  | STRING s -> Some (String (recode ~context ~src:"CP1252" ~dst:"UTF-8" s))
+  | STRING s -> Some (String (recode ~src:"CP1252" ~dst:"UTF-8" s))
   | DATETIME (y, mo, day, hr, min, sec, ms, _zone) ->
     (* FIXME: Timezones don't match in FreeTDS 0.91 and 1.0, so for now we
        just assume everything in UTC. *)

src/db_field.mli

@@ -16,12 +16,7 @@ type t =
   | Array of t list
 [@@deriving sexp]
 
-val of_data
-  :  context:Async_kernel.Execution_context.t
-  -> month_offset:int
-  -> Dblib.data
-  -> t option
-
+val of_data : month_offset:int -> Dblib.data -> t option
 val to_string : t option -> string
 val to_string_escaped : t option -> string
 val bignum : ?column:string -> t -> Bignum.t

src/logger.ml

@@ -1,24 +1,11 @@
 open Core
-open Async
 
 let src = Logs.Src.create "mssql"
 let lib_tag = Logs.Tag.def "lib" Format.pp_print_string
 
-let msg ?(tags = Logs.Tag.empty) ?context level fmt =
+let msg ?(tags = Logs.Tag.empty) level fmt =
   Async_helper.safely_run_in_async
   @@ fun () ->
-  let in_context f =
-    match context with
-    | Some context ->
-      (* Once we're on Async v0.13, use Scheduler.enqueue, since that will apparently pass exceptions
-        through correctly *)
-      Scheduler.within_context context f
-      |> Result.ok
-      |> Option.value_exn ~here:[%here] ~message:"Unknown exception in logging"
-    | None -> f ()
-  in
-  in_context
-  @@ fun () ->
   ksprintf
     (fun msg ->
       Logs.msg ~src level (fun m ->
@@ -27,6 +14,6 @@ let msg ?(tags = Logs.Tag.empty) ?context level fmt =
     fmt
 ;;
 
-let debug ?tags ?context fmt = msg ?tags ?context Logs.Debug fmt
-let info ?tags ?context fmt = msg ?tags ?context Logs.Info fmt
-let error ?tags ?context fmt = msg ?tags ?context Logs.Error fmt
+let debug ?tags fmt = msg ?tags Logs.Debug fmt
+let info ?tags fmt = msg ?tags Logs.Info fmt
+let error ?tags fmt = msg ?tags Logs.Error fmt

src/row.ml

@@ -2,8 +2,8 @@ open Core
 
 type t = Db_field.t option String.Map.t [@@deriving sexp_of]
 
-let create_exn ~context ~month_offset ~colnames row =
-  List.map row ~f:(Db_field.of_data ~context ~month_offset)
+let create_exn ~month_offset ~colnames row =
+  List.map row ~f:(Db_field.of_data ~month_offset)
   |> List.zip_exn colnames
   |> String.Map.of_alist_exn
 ;;

src/row.mli

@@ -6,12 +6,7 @@ open Freetds
 
 type t [@@deriving sexp_of]
 
-val create_exn
-  :  context:Async_kernel.Execution_context.t
-  -> month_offset:int
-  -> colnames:string list
-  -> Dblib.data list
-  -> t
+val create_exn : month_offset:int -> colnames:string list -> Dblib.data list -> t
 
 (** [to_alist t] converts t to a list of (colname, value) pairs *)
 val to_alist : t -> (string * string) list