Introduce query builder DSL (#4812)

* Introduce query builder DSL

* Add haddocks

* Add CHANGELOG entry

* Typo

Author: pcapriotti

changelog.d/5-internal/query-builder

@@ -0,0 +1 @@
+Add postgres dynamic query builder

libs/wire-subsystems/src/Wire/PaginationState.hs

@@ -0,0 +1,8 @@
+module Wire.PaginationState (PaginationState (..)) where
+
+import Data.Time.Clock
+import Imports
+
+data PaginationState id
+  = PaginationSortByName (Maybe (Text, id))
+  | PaginationSortByCreatedAt (Maybe (UTCTime, id))

libs/wire-subsystems/src/Wire/Postgres.hs

@@ -1,8 +1,55 @@
-module Wire.Postgres where
+module Wire.Postgres
+  ( -- | This module provides a composable DSL for constructing postgres
+    -- statements. Queries are assembled from smaller 'QueryFragment's that
+    -- carry both their SQL text and parameter encoders.
+    --
+    -- Typical usage involves combining fragments with monoidal operators and
+    -- building a final 'Statement' using 'buildStatement'.
+    --
+    -- Example:
+    --
+    -- > let q =
+    -- >       literal "select * from users"
+    -- >       <> where_ [like "name" "alice"]
+    -- >       <> orderBy [("created_at", Desc)]
+    -- >       <> limit (10 :: Int)
+    -- > in buildStatement q userDecoder
+    --
+    -- Note that the encoders are specialised to the specific values passed when
+    -- constructing the fragments, so they don't require further values. The
+    -- resulting statement can be run with something like @runStatement ()@.
 
-import Hasql.Pipeline
+    -- * Runners
+    runStatement,
+    runTransaction,
+    runPipeline,
+
+    -- * Query builder
+    QueryFragment,
+    literal,
+    where_,
+    like,
+    Clause,
+    mkClause,
+    clause,
+    clause1,
+    orderBy,
+    limit,
+    buildStatement,
+  )
+where
+
+import Control.Monad.Trans.State
+import Data.Functor.Contravariant
+import Data.Id
+import Data.Text qualified as T
+import Data.Text.Encoding qualified as T
+import Data.Time.Clock
+import Hasql.Decoders qualified as Dec
+import Hasql.Encoders qualified as Enc
+import Hasql.Pipeline (Pipeline)
 import Hasql.Pool
-import Hasql.Session qualified as Session
+import Hasql.Session
 import Hasql.Statement
 import Hasql.Transaction (Transaction)
 import Hasql.Transaction.Sessions
@@ -11,6 +58,7 @@ import Imports
 import Polysemy
 import Polysemy.Error (Error, throw)
 import Polysemy.Input
+import Wire.API.Pagination
 
 runStatement ::
   ( Member (Input Pool) r,
@@ -22,7 +70,7 @@ runStatement ::
   Sem r b
 runStatement a stmt = do
   pool <- input
-  liftIO (use pool (Session.statement a stmt)) >>= either throw pure
+  liftIO (use pool (statement a stmt)) >>= either throw pure
 
 runTransaction ::
   (Member (Input Pool) r, Member (Embed IO) r, Member (Error UsageError) r) =>
@@ -43,4 +91,160 @@ runPipeline ::
   Sem r a
 runPipeline p = do
   pool <- input
-  liftIO (use pool $ Session.pipeline p) >>= either throw pure
+  liftIO (use pool $ pipeline p) >>= either throw pure
+
+class PostgresValue a where
+  postgresType :: Text
+  postgresValue :: a -> Enc.Value ()
+
+instance PostgresValue (Id a) where
+  postgresType = "uuid"
+  postgresValue u = const (toUUID u) >$< Enc.uuid
+
+instance PostgresValue Text where
+  postgresType = "text"
+  postgresValue x = const x >$< Enc.text
+
+instance PostgresValue UTCTime where
+  postgresType = "timestamptz"
+  postgresValue t = const t >$< Enc.timestamptz
+
+instance PostgresValue Int32 where
+  postgresType = "int"
+  postgresValue n = const n >$< Enc.int4
+
+--------------------------------------------------------------------------------
+-- Query builder DSL
+
+data QueryFragment = QueryFragment
+  { query :: State Int Text,
+    encoder :: Enc.Params ()
+  }
+
+joinFragments :: Text -> QueryFragment -> QueryFragment -> QueryFragment
+joinFragments sep f1 f2 =
+  QueryFragment
+    { query = separate <$> f1.query <*> f2.query,
+      encoder = f1.encoder <> f2.encoder
+    }
+  where
+    separate "" q = q
+    separate q "" = q
+    separate q1 q2 = q1 <> sep <> q2
+
+instance Semigroup QueryFragment where
+  (<>) = joinFragments " "
+
+instance Monoid QueryFragment where
+  mempty =
+    QueryFragment
+      { query = pure "",
+        encoder = mempty
+      }
+
+literal :: Text -> QueryFragment
+literal q =
+  QueryFragment
+    { query = pure q,
+      encoder = mempty
+    }
+
+-- | Construct a WHERE clause from a list of fragments.
+where_ :: [QueryFragment] -> QueryFragment
+where_ frags = literal "where" <> foldr (joinFragments " and ") mempty frags
+
+like :: Text -> Text -> QueryFragment
+like field pat =
+  QueryFragment
+    { query = do
+        i <- nextIndex
+        pure $ field <> " ilike ($" <> T.pack (show i) <> " :: text)",
+      encoder = const (fuzzy pat) >$< Enc.param (Enc.nonNullable Enc.text)
+    }
+
+-- | A portion of a WHERE clause with multiple values. The monoidal operation
+-- of this type can be used to combine values into one clause. For example:
+--
+-- > clause "=" (mkClause "foo" 3 <> mkClause "bar" 4)
+--
+-- generates a pattern that will end up being expanded as @"(foo, bar) = (3, 4)"@.
+data Clause = Clause
+  { fields :: [Text],
+    types :: [Text],
+    encoder :: Enc.Params ()
+  }
+
+instance Semigroup Clause where
+  cl1 <> cl2 =
+    Clause
+      { fields = cl1.fields <> cl2.fields,
+        types = cl1.types <> cl2.types,
+        encoder = cl1.encoder <> cl2.encoder
+      }
+
+instance Monoid Clause where
+  mempty =
+    Clause
+      { fields = mempty,
+        types = mempty,
+        encoder = mempty
+      }
+
+mkClause :: forall a. (PostgresValue a) => Text -> a -> Clause
+mkClause field value =
+  Clause
+    { fields = [field],
+      types = [postgresType @a],
+      encoder = Enc.param (Enc.nonNullable (postgresValue value))
+    }
+
+-- | Convert a 'Clause' to a 'QueryFragment'.
+clause :: Text -> Clause -> QueryFragment
+clause op cl =
+  QueryFragment
+    { query = do
+        types <-
+          fmap wrap $
+            for cl.types $ \ty -> do
+              i <- nextIndex
+              pure $ "$" <> T.pack (show i) <> " :: " <> ty <> ""
+        let fields = wrap cl.fields
+        pure $ fields <> " " <> op <> " " <> types,
+      encoder = cl.encoder
+    }
+  where
+    wrap :: [Text] -> Text
+    wrap xs = "(" <> T.intercalate ", " xs <> ")"
+
+-- | Fragment for a clause with a single value.
+clause1 :: forall a. (PostgresValue a) => Text -> Text -> a -> QueryFragment
+clause1 field op value = clause op (mkClause field value)
+
+orderBy :: [(Text, SortOrder)] -> QueryFragment
+orderBy os =
+  literal $
+    "order by "
+      <> T.intercalate ", " (map (\(field, o) -> field <> " " <> sortOrderClause o) os)
+
+limit :: forall a. (PostgresValue a) => a -> QueryFragment
+limit n =
+  QueryFragment
+    { query = do
+        i <- nextIndex
+        pure $ "limit ($" <> T.pack (show i) <> " :: " <> postgresType @a <> ")",
+      encoder = Enc.param (Enc.nonNullable (postgresValue n))
+    }
+
+buildStatement :: QueryFragment -> Dec.Result b -> Statement () b
+buildStatement frag dec =
+  Statement
+    (T.encodeUtf8 (evalState frag.query 1))
+    frag.encoder
+    dec
+    True
+
+nextIndex :: State Int Int
+nextIndex = get <* modify succ
+
+fuzzy :: Text -> Text
+fuzzy x = "%" <> x <> "%"

libs/wire-subsystems/src/Wire/UserGroupStore.hs

@@ -4,30 +4,30 @@ module Wire.UserGroupStore where
 
 import Data.Id
 import Data.Json.Util
+import Data.Time.Clock
 import Data.Vector
 import Imports
 import Polysemy
 import Wire.API.Pagination
 import Wire.API.User.Profile
 import Wire.API.UserGroup
 import Wire.API.UserGroup.Pagination
+import Wire.PaginationState
 
 data UserGroupPageRequest = UserGroupPageRequest
   { team :: TeamId,
     searchString :: Maybe Text,
-    paginationState :: PaginationState,
+    paginationState :: PaginationState UserGroupId,
     sortOrder :: SortOrder,
     pageSize :: PageSize,
     includeMemberCount :: Bool,
     includeChannels :: Bool
   }
 
-data PaginationState = PaginationSortByName (Maybe (UserGroupName, UserGroupId)) | PaginationSortByCreatedAt (Maybe (UTCTimeMillis, UserGroupId))
+userGroupCreatedAtPaginationState :: UserGroup_ f -> (UTCTime, UserGroupId)
+userGroupCreatedAtPaginationState ug = (fromUTCTimeMillis ug.createdAt, ug.id_)
 
-userGroupCreatedAtPaginationState :: UserGroup_ f -> (UTCTimeMillis, UserGroupId)
-userGroupCreatedAtPaginationState ug = (ug.createdAt, ug.id_)
-
-toSortBy :: PaginationState -> SortBy
+toSortBy :: PaginationState UserGroupId -> SortBy
 toSortBy = \case
   PaginationSortByName _ -> SortByName
   PaginationSortByCreatedAt _ -> SortByCreatedAt