Distinct of a count?

[TomMD]

I'm probably burned out after a long week, but am having serious trouble seeing how the library allows for filtering based on distinct fields of a count.

For example (full code below). Say you have two tables of people and pets (foreign key to people). You want all people who own more than two types of pets. Thus, you must group by people, count the number of distinct pets and restrict the people to just those rows. Writing the actual query has eluded me, is there something I've entirely overlooked?

#!/usr/bin/env cabal
{- cabal:
    build-depends: base, selda, selda-sqlite
-}
{-# LANGUAGE DeriveGeneric, OverloadedStrings, OverloadedLabels #-}
import Database.Selda
import Database.Selda.SQLite

data Pet = Dog | Horse | Dragon
  deriving (Show, Read, Bounded, Enum)
instance SqlType Pet

data Person = Person
  { name :: Text
  } deriving Generic
instance SqlRow Person
data Has = Has
    { ownership_id :: ID Has
    , owner :: ID Person
    , pet :: Pet
    } deriving (Generic)
instance SqlRow Has

people :: Table Person
people = table "people" [#name :- primary]

pets :: Table Has
pets = table "pets" [#ownership_id :- autoPrimary]

main = withSQLite "people.sqlite" $ do
  createTable people
  createTable pets
  Just mainId <- insertWithPK people [ Person "Main" ]
  Just ellyId <- insertWithPK people [ Person "Elizabeth" ]
  insert_ pets
    [ Has def mainId Dragon
    , Has def ellyId Dragon
    , Has def ellyId Horse
    , Has def ellyId Horse -- Change to dog and expect elly to appear in the result
    ]

  peopleWithMoreThanTwoKindsOfPets <- query $ do
    person <- select people
    restrict (person ! #name `isIn` do
        (owner :*: nr) <- aggregate $ fmap count $ distinct $ do
            os <- select pets
            owner <- groupBy (os ! owner)
            nr <- count <$> distinct (os ! pet) -- Obviously not right
            pure (owner :*: nr)
        restrict (nr .> 2)
        pure owner)
  liftIO $ print peopleWithMoreThanTwoKindsOfPets

[TomMD]

Figured it out, though I feel it is a bit unintuitive. Probably just tired.

  peopleWithMoreThanTwoPets <- query $ do
    person <- select people
    restrict (person ! #name `isIn` do
        (c:*:o) <- aggregate (do
            ownerPet <- distinct (select pets)
            owner <- groupBy (ownerPet ! #owner)
            pure (count (ownerPet ! #pet) :*: owner)
            )
        restrict (c .> 2)
        pure o)
    pure person