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
  { TableName -> Maybe Text
tableSchema :: Maybe Text
  , TableName -> Text
tableName :: Text
  }
  deriving (Int -> TableName -> ShowS
[TableName] -> ShowS
TableName -> String
(Int -> TableName -> ShowS)
-> (TableName -> String)
-> ([TableName] -> ShowS)
-> Show TableName
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> TableName -> ShowS
showsPrec :: Int -> TableName -> ShowS
$cshow :: TableName -> String
show :: TableName -> String
$cshowList :: [TableName] -> ShowS
showList :: [TableName] -> ShowS
Show, TableName -> TableName -> Bool
(TableName -> TableName -> Bool)
-> (TableName -> TableName -> Bool) -> Eq TableName
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: TableName -> TableName -> Bool
== :: TableName -> TableName -> Bool
$c/= :: TableName -> TableName -> Bool
/= :: TableName -> TableName -> Bool
Eq, Eq TableName
Eq TableName =>
(TableName -> TableName -> Ordering)
-> (TableName -> TableName -> Bool)
-> (TableName -> TableName -> Bool)
-> (TableName -> TableName -> Bool)
-> (TableName -> TableName -> Bool)
-> (TableName -> TableName -> TableName)
-> (TableName -> TableName -> TableName)
-> Ord TableName
TableName -> TableName -> Bool
TableName -> TableName -> Ordering
TableName -> TableName -> TableName
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: TableName -> TableName -> Ordering
compare :: TableName -> TableName -> Ordering
$c< :: TableName -> TableName -> Bool
< :: TableName -> TableName -> Bool
$c<= :: TableName -> TableName -> Bool
<= :: TableName -> TableName -> Bool
$c> :: TableName -> TableName -> Bool
> :: TableName -> TableName -> Bool
$c>= :: TableName -> TableName -> Bool
>= :: TableName -> TableName -> Bool
$cmax :: TableName -> TableName -> TableName
max :: TableName -> TableName -> TableName
$cmin :: TableName -> TableName -> TableName
min :: TableName -> TableName -> TableName
Ord, (forall x. TableName -> Rep TableName x)
-> (forall x. Rep TableName x -> TableName) -> Generic TableName
forall x. Rep TableName x -> TableName
forall x. TableName -> Rep TableName x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cfrom :: forall x. TableName -> Rep TableName x
from :: forall x. TableName -> Rep TableName x
$cto :: forall x. Rep TableName x -> TableName
to :: forall x. Rep TableName x -> TableName
Generic)

instance P.Pretty TableName where
  pPrint :: TableName -> Doc
pPrint = TableName -> Doc
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
  { SQLOperationCounts -> Natural
sqlSelects :: Natural
  , SQLOperationCounts -> Map TableName Natural
sqlInserts :: Map TableName Natural
  , SQLOperationCounts -> Map TableName Natural
sqlDeletes :: Map TableName Natural
  , SQLOperationCounts -> Map TableName Natural
sqlUpdates :: Map TableName Natural
  }
  deriving (Int -> SQLOperationCounts -> ShowS
[SQLOperationCounts] -> ShowS
SQLOperationCounts -> String
(Int -> SQLOperationCounts -> ShowS)
-> (SQLOperationCounts -> String)
-> ([SQLOperationCounts] -> ShowS)
-> Show SQLOperationCounts
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> SQLOperationCounts -> ShowS
showsPrec :: Int -> SQLOperationCounts -> ShowS
$cshow :: SQLOperationCounts -> String
show :: SQLOperationCounts -> String
$cshowList :: [SQLOperationCounts] -> ShowS
showList :: [SQLOperationCounts] -> ShowS
Show, SQLOperationCounts -> SQLOperationCounts -> Bool
(SQLOperationCounts -> SQLOperationCounts -> Bool)
-> (SQLOperationCounts -> SQLOperationCounts -> Bool)
-> Eq SQLOperationCounts
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: SQLOperationCounts -> SQLOperationCounts -> Bool
== :: SQLOperationCounts -> SQLOperationCounts -> Bool
$c/= :: SQLOperationCounts -> SQLOperationCounts -> Bool
/= :: SQLOperationCounts -> SQLOperationCounts -> Bool
Eq, (forall x. SQLOperationCounts -> Rep SQLOperationCounts x)
-> (forall x. Rep SQLOperationCounts x -> SQLOperationCounts)
-> Generic SQLOperationCounts
forall x. Rep SQLOperationCounts x -> SQLOperationCounts
forall x. SQLOperationCounts -> Rep SQLOperationCounts x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cfrom :: forall x. SQLOperationCounts -> Rep SQLOperationCounts x
from :: forall x. SQLOperationCounts -> Rep SQLOperationCounts x
$cto :: forall x. Rep SQLOperationCounts x -> SQLOperationCounts
to :: forall x. Rep SQLOperationCounts x -> SQLOperationCounts
Generic)

instance Semigroup SQLOperationCounts where
  SQLOperationCounts Natural
s1 Map TableName Natural
i1 Map TableName Natural
d1 Map TableName Natural
u1 <> :: SQLOperationCounts -> SQLOperationCounts -> SQLOperationCounts
<> SQLOperationCounts Natural
s2 Map TableName Natural
i2 Map TableName Natural
d2 Map TableName Natural
u2 =
    Natural
-> Map TableName Natural
-> Map TableName Natural
-> Map TableName Natural
-> SQLOperationCounts
SQLOperationCounts
      (Natural
s1 Natural -> Natural -> Natural
forall a. Num a => a -> a -> a
+ Natural
s2)
      (Map TableName Natural
i1 Map TableName Natural
-> Map TableName Natural -> Map TableName Natural
`addNatMaps` Map TableName Natural
i2)
      (Map TableName Natural
d1 Map TableName Natural
-> Map TableName Natural -> Map TableName Natural
`addNatMaps` Map TableName Natural
d2)
      (Map TableName Natural
u1 Map TableName Natural
-> Map TableName Natural -> Map TableName Natural
`addNatMaps` Map TableName Natural
u2)
    where
      addNatMaps :: Map TableName Natural
-> Map TableName Natural -> Map TableName Natural
addNatMaps = (Natural -> Natural -> Natural)
-> Map TableName Natural
-> Map TableName Natural
-> Map TableName Natural
forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
Map.unionWith Natural -> Natural -> Natural
forall a. Num a => a -> a -> a
(+)

instance Monoid SQLOperationCounts where
  mempty :: SQLOperationCounts
mempty = Natural
-> Map TableName Natural
-> Map TableName Natural
-> Map TableName Natural
-> SQLOperationCounts
SQLOperationCounts Natural
0 Map TableName Natural
forall a. Monoid a => a
mempty Map TableName Natural
forall a. Monoid a => a
mempty Map TableName Natural
forall a. Monoid a => a
mempty

subtractNat :: Natural -> Natural -> Natural
Natural
a subtractNat :: Natural -> Natural -> Natural
`subtractNat` Natural
b = if Natural
a Natural -> Natural -> Bool
forall a. Ord a => a -> a -> Bool
> Natural
b then Natural
a Natural -> Natural -> Natural
forall a. Num a => a -> a -> a
- Natural
b else Natural
0

subtractNatMaps :: (Ord k) => Map k Natural -> Map k Natural -> Map k Natural
subtractNatMaps :: forall k. Ord k => Map k Natural -> Map k Natural -> Map k Natural
subtractNatMaps Map k Natural
c1 Map k Natural
c2 =
  let f :: k -> Natural -> Map k Natural -> Map k Natural
f k
op Natural
count = (Natural -> Natural) -> k -> Map k Natural -> Map k Natural
forall k a. Ord k => (a -> a) -> k -> Map k a -> Map k a
Map.adjust (Natural -> Natural -> Natural
`subtractNat` Natural
count) k
op
  in  (k -> Natural -> Map k Natural -> Map k Natural)
-> Map k Natural -> Map k Natural -> Map k Natural
forall k a b. (k -> a -> b -> b) -> b -> Map k a -> b
Map.foldrWithKey k -> Natural -> Map k Natural -> Map k Natural
forall {k}. Ord k => k -> Natural -> Map k Natural -> Map k Natural
f Map k Natural
c1 Map k Natural
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 -> SQLOperationCounts -> SQLOperationCounts
subtractCounts (SQLOperationCounts Natural
s1 Map TableName Natural
i1 Map TableName Natural
d1 Map TableName Natural
u1) (SQLOperationCounts Natural
s2 Map TableName Natural
i2 Map TableName Natural
d2 Map TableName Natural
u2) =
  Natural
-> Map TableName Natural
-> Map TableName Natural
-> Map TableName Natural
-> SQLOperationCounts
SQLOperationCounts
    (Natural
s1 Natural -> Natural -> Natural
`subtractNat` Natural
s2)
    (Map TableName Natural
i1 Map TableName Natural
-> Map TableName Natural -> Map TableName Natural
forall k. Ord k => Map k Natural -> Map k Natural -> Map k Natural
`subtractNatMaps` Map TableName Natural
i2)
    (Map TableName Natural
d1 Map TableName Natural
-> Map TableName Natural -> Map TableName Natural
forall k. Ord k => Map k Natural -> Map k Natural -> Map k Natural
`subtractNatMaps` Map TableName Natural
d2)
    (Map TableName Natural
u1 Map TableName Natural
-> Map TableName Natural -> Map TableName Natural
forall k. Ord k => Map k Natural -> Map k Natural -> Map k Natural
`subtractNatMaps` Map TableName Natural
u2)

------------------------------------------------------------
-- Pretty rendering and printing counts

instance P.Pretty SQLOperationCounts where
  pPrint :: SQLOperationCounts -> Doc
pPrint = SQLOperationCounts -> Doc
prettyCounts

{- | Print an t'SQLOperationCounts' to stdout using 'prettyCounts'.
For less verbose output, see 'printCountsBrief'.
-}
printCounts :: (MonadIO m) => SQLOperationCounts -> m ()
printCounts :: forall (m :: * -> *). MonadIO m => SQLOperationCounts -> m ()
printCounts = IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ())
-> (SQLOperationCounts -> IO ()) -> SQLOperationCounts -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> IO ()
putStrLn (String -> IO ())
-> (SQLOperationCounts -> String) -> SQLOperationCounts -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SQLOperationCounts -> String
renderCounts

{- | Print an t'SQLOperationCounts' to stdout using 'prettyCountsBrief'.
For more verbose output, see 'printCounts'.
-}
printCountsBrief :: (MonadIO m) => SQLOperationCounts -> m ()
printCountsBrief :: forall (m :: * -> *). MonadIO m => SQLOperationCounts -> m ()
printCountsBrief = IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ())
-> (SQLOperationCounts -> IO ()) -> SQLOperationCounts -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> IO ()
putStrLn (String -> IO ())
-> (SQLOperationCounts -> String) -> SQLOperationCounts -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SQLOperationCounts -> String
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 :: SQLOperationCounts -> String
renderCounts = Doc -> String
P.render (Doc -> String)
-> (SQLOperationCounts -> Doc) -> SQLOperationCounts -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SQLOperationCounts -> Doc
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 :: SQLOperationCounts -> String
renderCountsBrief = Doc -> String
P.render (Doc -> String)
-> (SQLOperationCounts -> Doc) -> SQLOperationCounts -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SQLOperationCounts -> Doc
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 :: SQLOperationCounts -> Doc
prettyCounts = (Map TableName Natural -> Maybe Doc) -> SQLOperationCounts -> Doc
prettyCountsWith ((Map TableName Natural -> Maybe Doc) -> SQLOperationCounts -> Doc)
-> (Map TableName Natural -> Maybe Doc)
-> SQLOperationCounts
-> Doc
forall a b. (a -> b) -> a -> b
$ \Map TableName Natural
mp ->
  let counts :: [(TableName, Natural)]
counts = Map TableName Natural -> [(TableName, Natural)]
forall k a. Map k a -> [(k, a)]
Map.toList Map TableName Natural
mp
      renderPair :: (TableName, Natural) -> Maybe Doc
renderPair (TableName
name, Natural
count) = Doc -> Doc -> Doc
prefix (TableName -> Doc
renderTableName TableName
name) (Doc -> Doc) -> Maybe Doc -> Maybe Doc
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Natural -> Maybe Doc
renderNat Natural
count
  in  (NonEmpty Doc -> Doc) -> Maybe (NonEmpty Doc) -> Maybe Doc
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ([Doc] -> Doc
P.vcat ([Doc] -> Doc) -> (NonEmpty Doc -> [Doc]) -> NonEmpty Doc -> Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. NonEmpty Doc -> [Doc]
forall a. NonEmpty a -> [a]
NE.toList) (Maybe (NonEmpty Doc) -> Maybe Doc)
-> ([Doc] -> Maybe (NonEmpty Doc)) -> [Doc] -> Maybe Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Doc] -> Maybe (NonEmpty Doc)
forall a. [a] -> Maybe (NonEmpty a)
NE.nonEmpty ([Doc] -> Maybe Doc) -> [Doc] -> Maybe Doc
forall a b. (a -> b) -> a -> b
$ ((TableName, Natural) -> Maybe Doc)
-> [(TableName, Natural)] -> [Doc]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (TableName, Natural) -> Maybe Doc
renderPair [(TableName, Natural)]
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 :: SQLOperationCounts -> Doc
prettyCountsBrief = (Map TableName Natural -> Maybe Doc) -> SQLOperationCounts -> Doc
prettyCountsWith ((Map TableName Natural -> Maybe Doc) -> SQLOperationCounts -> Doc)
-> (Map TableName Natural -> Maybe Doc)
-> SQLOperationCounts
-> Doc
forall a b. (a -> b) -> a -> b
$ \Map TableName Natural
mp ->
  let total :: Natural
total = [Natural] -> Natural
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([Natural] -> Natural) -> [Natural] -> Natural
forall a b. (a -> b) -> a -> b
$ Map TableName Natural -> [Natural]
forall k a. Map k a -> [a]
Map.elems Map TableName Natural
mp
  in  Natural -> Maybe Doc
renderNat Natural
total

prettyCountsWith :: (Map TableName Natural -> Maybe P.Doc) -> SQLOperationCounts -> P.Doc
prettyCountsWith :: (Map TableName Natural -> Maybe Doc) -> SQLOperationCounts -> Doc
prettyCountsWith Map TableName Natural -> Maybe Doc
renderMap (SQLOperationCounts Natural
selects Map TableName Natural
inserts Map TableName Natural
deletes Map TableName Natural
updates) =
  let parts :: [Doc]
parts =
        [Maybe Doc] -> [Doc]
forall a. [Maybe a] -> [a]
catMaybes
          [ Doc -> Doc -> Doc
prefix Doc
"SELECT" (Doc -> Doc) -> Maybe Doc -> Maybe Doc
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Natural -> Maybe Doc
renderNat Natural
selects
          , Doc -> Doc -> Doc
prefix Doc
"INSERT" (Doc -> Doc) -> Maybe Doc -> Maybe Doc
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Map TableName Natural -> Maybe Doc
renderMap Map TableName Natural
inserts
          , Doc -> Doc -> Doc
prefix Doc
"UPDATE" (Doc -> Doc) -> Maybe Doc -> Maybe Doc
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Map TableName Natural -> Maybe Doc
renderMap Map TableName Natural
updates
          , Doc -> Doc -> Doc
prefix Doc
"DELETE" (Doc -> Doc) -> Maybe Doc -> Maybe Doc
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Map TableName Natural -> Maybe Doc
renderMap Map TableName Natural
deletes
          ]
  in  case [Doc]
parts of
        [] -> Doc
"None"
        [Doc]
_ -> [Doc] -> Doc
P.vcat [Doc]
parts

prefix :: P.Doc -> P.Doc -> P.Doc
prefix :: Doc -> Doc -> Doc
prefix Doc
t Doc
n = Doc
t Doc -> Doc -> Doc
P.<> Doc
":" Doc -> Doc -> Doc
P.<+> Doc
n

renderNat :: Natural -> Maybe P.Doc
renderNat :: Natural -> Maybe Doc
renderNat = \case
  Natural
0 -> Maybe Doc
forall a. Maybe a
Nothing
  Natural
n -> Doc -> Maybe Doc
forall a. a -> Maybe a
Just (Doc -> Maybe Doc) -> Doc -> Maybe Doc
forall a b. (a -> b) -> a -> b
$ forall a. Pretty a => a -> Doc
P.pPrint @Integer (Integer -> Doc) -> Integer -> Doc
forall a b. (a -> b) -> a -> b
$ Natural -> Integer
forall a. Integral a => a -> Integer
toInteger Natural
n

renderTableName :: TableName -> P.Doc
renderTableName :: TableName -> Doc
renderTableName (TableName Maybe Text
mSchema Text
table) =
  case Maybe Text
mSchema of
    Maybe Text
Nothing -> Text -> Doc
renderText Text
table
    Just Text
schema -> Text -> Doc
renderText Text
schema Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc
"." Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Text -> Doc
renderText Text
table

renderText :: T.Text -> P.Doc
renderText :: Text -> Doc
renderText = String -> Doc
P.text (String -> Doc) -> (Text -> String) -> Text -> Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> String
T.unpack