module PostgreSQL.Count
( -- * Counting SQL operations
SQLOperationCounts (..)
, TableName (..)
, subtractCounts
-- * Pretty-printing
, printCounts
, printCountsBrief
, renderCounts
, renderCountsBrief
, prettyCounts
, prettyCountsBrief
)
where
import Control.Monad.IO.Class
import qualified Data.List.NonEmpty as NE
import Data.Map (Map)
import qualified Data.Map as Map
import Data.Maybe (catMaybes, mapMaybe)
import Data.Text (Text)
import qualified Data.Text as T
import GHC.Generics
import Numeric.Natural
import qualified Text.PrettyPrint as P
import qualified Text.PrettyPrint.HughesPJClass as P
-- | The name of a table, optionally qualified with a schema.
data TableName = TableName
{ tableSchema :: Maybe Text
, tableName :: Text
}
deriving (Show, Eq, Ord, Generic)
instance P.Pretty TableName where
pPrint = renderTableName
------------------------------------------------------------
-- Tallying SQL operations
{- | This tracks the number of SQL operations that have been performed, along with which table they were
performed on (where possible).
@INSERT@, @DELETE@ and @UPDATE@ operations act on one table only, so we can tally the number
of each that are performed on each table (indexed by a t'TableName').
@SELECT@ operations can act on multiple tables, so we just track the total number of selects.
If required, t'SQLOperationCounts' can be constructed using 'Monoid', added using 'Semigroup',
and subtracted using `subtractCounts`.
We use non-negative 'Natural's as a tally since a negative number of operations makes no sense.
-}
data SQLOperationCounts = SQLOperationCounts
{ sqlSelects :: Natural
, sqlInserts :: Map TableName Natural
, sqlDeletes :: Map TableName Natural
, sqlUpdates :: Map TableName Natural
}
deriving (Show, Eq, Generic)
instance Semigroup SQLOperationCounts where
SQLOperationCounts s1 i1 d1 u1 <> SQLOperationCounts s2 i2 d2 u2 =
SQLOperationCounts
(s1 + s2)
(i1 `addNatMaps` i2)
(d1 `addNatMaps` d2)
(u1 `addNatMaps` u2)
where
addNatMaps = Map.unionWith (+)
instance Monoid SQLOperationCounts where
mempty = SQLOperationCounts 0 mempty mempty mempty
subtractNat :: Natural -> Natural -> Natural
a `subtractNat` b = if a > b then a - b else 0
subtractNatMaps :: (Ord k) => Map k Natural -> Map k Natural -> Map k Natural
subtractNatMaps c1 c2 =
let f op count = Map.adjust (`subtractNat` count) op
in Map.foldrWithKey f c1 c2
{- | Subtract one set of counts from another. Note that the results are all still 'Natural's,
so will all be non-negative.
-}
subtractCounts :: SQLOperationCounts -> SQLOperationCounts -> SQLOperationCounts
subtractCounts (SQLOperationCounts s1 i1 d1 u1) (SQLOperationCounts s2 i2 d2 u2) =
SQLOperationCounts
(s1 `subtractNat` s2)
(i1 `subtractNatMaps` i2)
(d1 `subtractNatMaps` d2)
(u1 `subtractNatMaps` u2)
------------------------------------------------------------
-- Pretty rendering and printing counts
instance P.Pretty SQLOperationCounts where
pPrint = prettyCounts
{- | Print an t'SQLOperationCounts' to stdout using 'prettyCounts'.
For less verbose output, see 'printCountsBrief'.
-}
printCounts :: (MonadIO m) => SQLOperationCounts -> m ()
printCounts = liftIO . putStrLn . renderCounts
{- | Print an t'SQLOperationCounts' to stdout using 'prettyCountsBrief'.
For more verbose output, see 'printCounts'.
-}
printCountsBrief :: (MonadIO m) => SQLOperationCounts -> m ()
printCountsBrief = liftIO . putStrLn . renderCountsBrief
{- | Render an t'SQLOperationCounts' using 'prettyCounts'.
For less verbose output, see 'renderCountsBrief'.
For more control over how the 'P.Doc' gets rendered, use 'P.renderStyle' with a custom 'P.style'.
-}
renderCounts :: SQLOperationCounts -> String
renderCounts = P.render . prettyCounts
{- | Render an t'SQLOperationCounts' using 'prettyCountsBrief'.
For more verbose output, see 'renderCounts'.
For more control over how the 'P.Doc' gets rendered, use 'P.renderStyle' with a custom 'P.style'.
-}
renderCountsBrief :: SQLOperationCounts -> String
renderCountsBrief = P.render . prettyCountsBrief
{- | Pretty-print an t'SQLOperationCounts' using "Text.PrettyPrint".
For each 'Map', we'll print one line for each table. For less verbose output,
see 'prettyCountsBrief'.
This is also the implementation of 'P.pPrint' for t'SQLOperationCounts'.
-}
prettyCounts :: SQLOperationCounts -> P.Doc
prettyCounts = prettyCountsWith $ \mp ->
let counts = Map.toList mp
renderPair (name, count) = prefix (renderTableName name) <$> renderNat count
in fmap (P.vcat . NE.toList) . NE.nonEmpty $ mapMaybe renderPair counts
{- | Pretty-print an t'SQLOperationCounts' using "Text.PrettyPrint".
For each 'Map', we'll print just the sum of the counts. For more verbose output,
see 'prettyCounts'.
-}
prettyCountsBrief :: SQLOperationCounts -> P.Doc
prettyCountsBrief = prettyCountsWith $ \mp ->
let total = sum $ Map.elems mp
in renderNat total
prettyCountsWith :: (Map TableName Natural -> Maybe P.Doc) -> SQLOperationCounts -> P.Doc
prettyCountsWith renderMap (SQLOperationCounts selects inserts deletes updates) =
let parts =
catMaybes
[ prefix "SELECT" <$> renderNat selects
, prefix "INSERT" <$> renderMap inserts
, prefix "UPDATE" <$> renderMap updates
, prefix "DELETE" <$> renderMap deletes
]
in case parts of
[] -> "None"
_ -> P.vcat parts
prefix :: P.Doc -> P.Doc -> P.Doc
prefix t n = t P.<> ":" P.<+> n
renderNat :: Natural -> Maybe P.Doc
renderNat = \case
0 -> Nothing
n -> Just $ P.pPrint @Integer $ toInteger n
renderTableName :: TableName -> P.Doc
renderTableName (TableName mSchema table) =
case mSchema of
Nothing -> renderText table
Just schema -> renderText schema <> "." <> renderText table
renderText :: T.Text -> P.Doc
renderText = P.text . T.unpack