Add instrumentation hook for wrapping statements

persistent-mysql/Database/Persist/MySQL.hs

@@ -152,6 +152,7 @@ open' ci logFunc = do
         , connLogFunc    = logFunc
         , connMaxParams = Nothing
         , connRepsertManySql = Just repsertManySql
+        , connStatementMiddleware = const pure
         }
 
 -- | Prepare a query.  We don't support prepared statements, but
@@ -1270,7 +1271,8 @@ mockMigration mig = do
                              connUpsertSql = undefined,
                              connPutManySql = undefined,
                              connMaxParams = Nothing,
-                             connRepsertManySql = Nothing
+                             connRepsertManySql = Nothing,
+                             connStatementMiddleware = const pure
                              }
       result = runReaderT . runWriterT . runWriterT $ mig
   resp <- result sqlbackend

persistent-postgresql/Database/Persist/Postgresql.hs

@@ -364,6 +364,7 @@ createBackend logFunc serverVersion smap conn = do
         , connLogFunc = logFunc
         , connMaxParams = Nothing
         , connRepsertManySql = upsertFunction repsertManySql serverVersion
+        , connStatementMiddleware = const pure
         }
 
 prepare' :: PG.Connection -> Text -> IO Statement
@@ -1720,7 +1721,8 @@ mockMigration mig = do
                              connLimitOffset = undefined,
                              connLogFunc = undefined,
                              connMaxParams = Nothing,
-                             connRepsertManySql = Nothing
+                             connRepsertManySql = Nothing,
+                             connStatementMiddleware = const pure
                              }
       result = runReaderT $ runWriterT $ runWriterT mig
   resp <- result sqlbackend

persistent-sqlite/Database/Persist/Sqlite.hs

@@ -288,6 +288,7 @@ wrapConnectionInfo connInfo conn logFunc = do
         , connLogFunc = logFunc
         , connMaxParams = Just 999
         , connRepsertManySql = Just repsertManySql
+        , connStatementMiddleware = const pure
         }
   where
     helper t getter = do
@@ -482,6 +483,7 @@ mockMigration mig = do
                    , connPutManySql = undefined
                    , connMaxParams = Just 999
                    , connRepsertManySql = Nothing
+                   , connStatementMiddleware = const pure
                    }
       result = runReaderT . runWriterT . runWriterT $ mig
   resp <- result sqlbackend

persistent/Database/Persist/Sql/Raw.hs

@@ -76,7 +76,7 @@ getStmtConn :: SqlBackend -> Text -> IO Statement
 getStmtConn conn sql = do
     smap <- liftIO $ readIORef $ connStmtMap conn
     case Map.lookup sql smap of
-        Just stmt -> return stmt
+        Just stmt -> connStatementMiddleware conn sql stmt
         Nothing -> do
             stmt' <- liftIO $ connPrepare conn sql
             iactive <- liftIO $ newIORef True
@@ -100,7 +100,7 @@ getStmtConn conn sql = do
                             else liftIO $ throwIO $ StatementAlreadyFinalized sql
                     }
             liftIO $ writeIORef (connStmtMap conn) $ Map.insert sql stmt smap
-            return stmt
+            connStatementMiddleware conn sql stmt
 
 -- | Execute a raw SQL statement and return its results as a
 -- list. If you do not expect a return value, use of

persistent/Database/Persist/Sql/Types/Internal.hs

@@ -190,6 +190,9 @@ data SqlBackend = SqlBackend
     -- When left as 'Nothing', we default to using 'defaultRepsertMany'.
     --
     -- @since 2.9.0
+    , connStatementMiddleware :: Text -> Statement -> IO Statement
+    -- ^ Provide facilities for injecting middleware into statements
+    -- to allow for instrumenting queries.
     }
 
 instance HasPersistBackend SqlBackend where