Make logging pass the original Async context around

This should make tagging systems which use the Async context
work correctly.

src/client.ml

@@ -79,9 +79,10 @@ 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
   In_thread.run
   @@ fun () ->
-  Logger.debug !"Executing query: %s" formatted_query;
+  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;
@@ -96,7 +97,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 ~month_offset ~colnames row in
+                  let row = Row.create_exn ~context ~month_offset ~colnames row in
                   Some row
                 with
                 | Caml.Not_found -> None)

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 ~src ~dst input =
+let recode ?context ~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,11 +62,15 @@ let recode ~src ~dst input =
     String.sub output ~pos:0 ~len:!enc_i
   with
   | exn ->
-    Logger.info !"Recoding error, falling back to ascii filter %{sexp: exn} %s" exn input;
+    Logger.info
+      ?context
+      !"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 ~month_offset data =
+let of_data ~context ~month_offset data =
   match data with
   | Dblib.BIT b -> Some (Bool b)
   | INT i | SMALL i | TINY i -> Some (Int i)
@@ -75,7 +79,7 @@ let of_data ~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 ~src:"CP1252" ~dst:"UTF-8" s))
+  | STRING s -> Some (String (recode ~context ~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,7 +16,12 @@ type t =
   | Array of t list
 [@@deriving sexp]
 
-val of_data : month_offset:int -> Dblib.data -> t option
+val of_data
+  :  context:Async_kernel.Execution_context.t
+  -> 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,11 +1,24 @@
 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) level fmt =
+let msg ?(tags = Logs.Tag.empty) ?context 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 ->
@@ -14,6 +27,6 @@ let msg ?(tags = Logs.Tag.empty) level fmt =
     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
+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

src/row.ml

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

src/row.mli

@@ -6,7 +6,12 @@ open Freetds
 
 type t [@@deriving sexp_of]
 
-val create_exn : month_offset:int -> colnames:string list -> Dblib.data list -> t
+val create_exn
+  :  context:Async_kernel.Execution_context.t
+  -> 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