Add post_connect to sql_pool (#422)

* Add post_connect to sql_pool

Allow for a custom post_connect to run for new SQL connection, for
example registering user functions for SQLite connections.

* Adjust signature, skip standard_post_connect

---------

Co-authored-by: Gusted <postmaster@gusted.xyz>

Author: gahr

docs/web/postprocess/index.ml

@@ -1951,7 +1951,8 @@ let graphiql_expected = {|<div class="spec value" id="val-graphiql">
 |}
 
 let sql_pool_expected = {|<div class="spec value" id="val-sql_pool">
- <a href="#val-sql_pool" class="anchor"></a><code><span><span class="keyword">val</span> sql_pool : <span>?size:int <span class="arrow">-&gt;</span></span> <span>string <span class="arrow">-&gt;</span></span> <a href="#type-middleware">middleware</a></span></code>
+ <a href="#val-sql_pool" class="anchor"></a><code><span><span class="keyword">val</span> sql_pool : <span>?size:int <span class="arrow">-&gt;</span></span>
+<span>?post_connect:<span>(<span><span>(<span class="keyword">module</span> <span class="xref-unresolved">Caqti_lwt</span>.CONNECTION)</span> <span class="arrow">-&gt;</span></span> <span><span><span>(unit,&nbsp;<span class="xref-unresolved">Caqti_error</span>.t)</span> <span class="xref-unresolved">Stdlib</span>.result</span> <a href="#type-promise">promise</a></span>)</span> <span class="arrow">-&gt;</span></span> <span>string <span class="arrow">-&gt;</span></span> <a href="#type-middleware">middleware</a></span></code>
 </div>
 |}
 

src/dream.mli

@@ -1752,11 +1752,16 @@ val graphiql : ?default_query:string -> string -> handler
     {{:https://cheatsheetseries.owasp.org/cheatsheets/Database_Security_Cheat_Sheet.html}
     OWASP {i Database Security Cheat Sheet}}. *)
 
-val sql_pool : ?size:int -> string -> middleware
+val sql_pool :
+  ?size:int ->
+  ?post_connect: (Caqti_lwt.connection -> unit promise) ->
+  string ->
+  middleware
 (** Makes an SQL connection pool available to its inner handler. [?size] is the
     maximum number of concurrent connections that the pool will support. The
     default value is picked by the driver. Note that for SQLite, [?size] is
-    capped to [1]. *)
+    capped to [1]. [post_connect] is an optional callback, which is called for
+    every new connection that is opened to the database. *)
 
 val sql : request -> (Caqti_lwt.connection -> 'a promise) -> 'a promise
 (** Runs the callback with a connection from the SQL pool. See example

src/sql/sql.ml

@@ -25,12 +25,12 @@ let foreign_keys_on =
   (Caqti_type.unit ->. Caqti_type.unit) "PRAGMA foreign_keys = ON"
   [@ocaml.warning "-3"]
 
-let post_connect (module Db : Caqti_lwt.CONNECTION) =
+let standard_post_connect (module Db : Caqti_lwt.CONNECTION) =
   match Caqti_driver_info.dialect_tag Db.driver_info with
   | `Sqlite -> Db.exec foreign_keys_on ()
   | _ -> Lwt.return (Ok ())
 
-let sql_pool ?size uri =
+let sql_pool ?size ?post_connect uri =
     let pool_cell = ref None in
     fun inner_handler request ->
 
@@ -47,9 +47,15 @@ let sql_pool ?size uri =
       log.warning (fun log -> log ~request
         "Dream.sql_pool: \
         'sqlite' is not a valid scheme; did you mean 'sqlite3'?");
+    let post_connect =
+      match post_connect with
+      | None -> standard_post_connect
+      | Some f -> (fun db -> Lwt.map Result.ok (f db))
+    in
     let pool =
       let pool_config = Caqti_pool_config.create ?max_size:size () in
-      Caqti_lwt_unix.connect_pool ~pool_config ~post_connect parsed_uri in
+      Caqti_lwt_unix.connect_pool ~pool_config ~post_connect parsed_uri
+    in
     match pool with
     | Ok pool ->
       pool_cell := Some pool;