Count.hs

{-# LANGUAGE CPP #-}

{- | Thanks to our dynamic 'Opaleye' effect, we can write an alternative interpreter which,
as well as performing SQL operations as before, will also keep a tally of the number of
SQL operations (SELECTs, INSERTs etc) that have been performed. This is really useful for debugging.

The intended use-case is a sort of benchmark that runs several Opaleye operations for different
"sizes", counts the SQL operations, and prints the tallies to the console. This lets us detect if
some datbase operations are ineffecient.

For example, suppose our model has users with @UserId@s; those users an have multiple @Transaction@s, which
are composed of multiple @SubTransaction@s etc.
To insert a group of new users, we would need to insert the users, insert the transactions, and insert the subtransactions.
Ideally, the number of @INSERT@s should not depend on the number of @User@s or the number or size of their @Transactions@.
We would expect the number of SELECTs to remain basically constant (O(1)), while the execution time might grow linearly (O(u * t * s)).

A very naive implementation might be:

@
insertUsersNaive :: (e :> es) => 'Opaleye' e -> [User] -> Eff es ()
insertUsersNaive o users = for_ users $ \user -> do
  insertUserFlat o user
  for (transactions user) $ \transaction -> do
    insertTransactionFlat o transaction
    for (subTransactions transaction) $ \subTransaction -> do
      insertSubTransactionFlat o subTransaction
@

However, if we ran a "benchmark" that looked something like this:

@
u1, u5, u10, u50 :: [User]
u1 = [User {transactions = [Transaction [SubTransaction]]}] -- one user, one transaction, one sub-transaction
u5 = ...  -- five users, each with five transactions, each with 5 sub-transactions

benchmark ::
  (e :> es, e1 :> es, e2 :> es) =>
  'Opaleye' e ->
  State SQLOperationCounts e1 ->
  IOE e2 ->
  Eff es ()
benchmark o st ioe = for_ [(1, u1), (5, u5), (10, u10), (50, u50)] $ \(n, users) -> do
  (counts, ()) <- withCounts st $ insertUsersNaive o users
  effIO ioe . putStrLn $ "Counts at n=" <> show n <> ": " <> 'renderCountsBrief' counts

main :: IO ()
main =
  runEff $ \ioe ->
    'Conn.runWithConnectInfo' ioe connInfo $ \withConn ->
      evalState @SQLOperationCounts mempty $ \st ->
        runOpaleyeWithConnectionCounting withConn ioe st $ \o ->
          benchmark o st ioe
  where
    connInfo = ...
@

We will probably see something like:

@
Counts at n=1: INSERT: 3
Counts at n=5: INSERT: 155
Counts at n=10: INSERT: 1110
Counts at n=50: INSERT: 127550
@

This is obviously going to have a severe performance impact. Rearranging our implementation of @insertUsers@:

@
insertUsersBetter :: (e :> es) => 'Opaleye' e -> [User] -> Eff es ()
insertUsersBetter o users = do
  let transactions_ = concatMap transactions users
      subTransactions_ = concatMap subTransactions transactions_
  insertUsersFlat o users
  insertTransactionsFlat o transactions_
  insertSubTransactionsFlat o subTransactions_
@

As long as @insertTransactionsFlat@ etc are smart enough to only do one 'runInsert', then we should now get:

@
Counts at n=1: INSERT: 3
Counts at n=5: INSERT: 3
Counts at n=10: INSERT: 3
Counts at n=50: INSERT: 3
@

Note that we used 'renderCountsBrief' for simplicity. If we wanted to debug in more detail, we could have used
'renderCounts' instead:

@
Counts at n=1: INSERT: user: 1
                       transaction: 1
                       sub_transaction: 1
Counts at n=5: INSERT: user: 5
                       transaction: 25
                       sub_transaction: 125
Counts at n=10: INSERT: user: 10
                        transaction: 100
                        sub_transaction: 1000
Counts at n=50: INSERT: user: 50
                        transaction: 2500
                        sub_transaction: 125000
@
-}
module Bluefin.Opaleye.Count
  ( -- * Counting SQL operations
    opaleyeAddCounting
  , withCounts
  , module PostgreSQL.Count
  )
where

import Bluefin.Compound
import Bluefin.Eff
import Bluefin.Opaleye.Effect
import Bluefin.State
import Data.Map (Map)
import qualified Data.Map as Map
import qualified Data.Text as T
import Numeric.Natural
import qualified Opaleye as O
import qualified Opaleye.Internal.PrimQuery as O (TableIdentifier (..))
import qualified Opaleye.Internal.Table as O
import PostgreSQL.Count

------------------------------------------------------------
-- Tallying SQL operations

{- | Add counting of SQL operations to the interpreter of an 'Opaleye' effect.
Note that the effect itself is not actually interpreted. After updating our t'SQLOperationCounts' state
 based on the 'Opaleye' constructor, we then pass them
through to the upstream handler (e.g. 'Bluefin.Opaleye.runOpaleyeWithConnection' or
'Bluefin.Opaleye.runOpaleyeConnection'). See 'Bluefin.Opaleye.runOpaleyeConnectionCounting'
and 'Bluefin.Opaleye.runOpaleyeWithConnectionCounting' for interpreters that do both.

Note: this function should only be used once, otherwise the operations will be tallied
more than once. Unless you're sure, it's probably better to use
'Bluefin.Opaleye.runOpaleyeConnectionCounting' or
'Bluefin.Opaleye.runOpaleyeWithConnectionCounting'.
-}
opaleyeAddCounting ::
  forall es e1 e2 a.
  (e1 :> es, e2 :> es) =>
  State SQLOperationCounts e1 ->
  Opaleye e2 ->
  (forall e. Opaleye e -> Eff (e :& es) a) ->
  Eff es a
opaleyeAddCounting st oldEffect k =
  useImplIn
    k
    MkOpaleye
      { runSelectExplicitImpl = \ff sel -> do
          incrementSelect
          runSelectExplicit oldEffect ff sel
      , runSelectFoldExplicitImpl = \ff sel b f -> do
          incrementSelect
          runSelectFoldExplicit oldEffect ff sel b (\b' -> useImpl . f b')
      , runInsertImpl = \ins -> do
          incrementInsert $ insertTableName ins
          runInsert oldEffect ins
      , runDeleteImpl = \del -> do
          incrementDelete $ deleteTableName del
          runDeleteImpl (mapHandle oldEffect) del
      , runUpdateImpl = \upd -> do
          incrementUpdate $ updateTableName upd
          runUpdate oldEffect upd
      }
  where
    incrementSelect :: Eff (e :& es) ()
    incrementSelect = modify st $ \counts ->
      counts {sqlSelects = succ $ sqlSelects counts}

    incrementInsert :: TableName -> Eff (e :& es) ()
    incrementInsert name = modify st $ \counts ->
      counts {sqlInserts = incrementMap name $ sqlInserts counts}

    incrementUpdate :: TableName -> Eff (e :& es) ()
    incrementUpdate name = modify st $ \counts ->
      counts {sqlUpdates = incrementMap name $ sqlUpdates counts}

    incrementDelete :: TableName -> Eff (e :& es) ()
    incrementDelete name = modify st $ \counts ->
      counts {sqlDeletes = incrementMap name $ sqlDeletes counts}

    incrementMap :: TableName -> Map TableName Natural -> Map TableName Natural
    incrementMap = Map.alter (Just . maybe 1 succ)

-- | This allows us to count the number of SQL operations over the course of a sub-operation.
withCounts ::
  (e :> es) =>
  State SQLOperationCounts e ->
  Eff es a ->
  Eff es (SQLOperationCounts, a)
withCounts st eff = do
  countsBefore <- get st
  res <- eff
  countsAfter <- get st
  pure (countsAfter `subtractCounts` countsBefore, res)

------------------------------------------------------------
-- Getting table identifiers from opaleye operations

tableIdentifierToTableName :: O.TableIdentifier -> TableName
tableIdentifierToTableName (O.TableIdentifier mSchema table) =
  TableName (T.pack <$> mSchema) (T.pack table)

insertTableName :: O.Insert haskells -> TableName
insertTableName (O.Insert table _ _ _) =
  tableIdentifierToTableName . O.tableIdentifier $ table

updateTableName :: O.Update haskells -> TableName
updateTableName (O.Update table _ _ _) =
  tableIdentifierToTableName . O.tableIdentifier $ table

deleteTableName :: O.Delete haskells -> TableName
deleteTableName (O.Delete table _ _) =
  tableIdentifierToTableName . O.tableIdentifier $ table